Add index-parents when splitting loops along store/load

Author: chriselrod

src/broadcast.jl

@@ -404,8 +404,8 @@ end
   loopsyms = [gensym!(ls, "n") for n ∈ 1:N]
   add_broadcast_loops!(ls, loopsyms, :dest)
   elementbytes = sizeof(T)
-  add_broadcast!(ls, :dest, :bc, loopsyms, BC, elementbytes)
-  storeop = add_simple_store!(ls, :dest, ArrayReference(:dest, loopsyms), elementbytes)
+  add_broadcast!(ls, :destination, :bc, loopsyms, BC, elementbytes)
+  storeop = add_simple_store!(ls, :destination, ArrayReference(:dest, loopsyms), elementbytes)
   doaddref!(ls, storeop)
   resize!(ls.loop_order, num_loops(ls)) # num_loops may be greater than N, eg Product
   # return ls
@@ -426,8 +426,8 @@ end
   pushprepreamble!(ls, Expr(:(=), :dest, Expr(:call, :parent, :dest′)))
   add_broadcast_loops!(ls, loopsyms, :dest′)
   elementbytes = sizeof(T)
-  add_broadcast!(ls, :dest, :bc, loopsyms, BC, elementbytes)
-  storeop = add_simple_store!(ls, :dest, ArrayReference(:dest, reverse(loopsyms)), elementbytes)
+  add_broadcast!(ls, :destination, :bc, loopsyms, BC, elementbytes)
+  storeop = add_simple_store!(ls, :destination, ArrayReference(:dest, reverse(loopsyms)), elementbytes)
   doaddref!(ls, storeop)
   resize!(ls.loop_order, num_loops(ls)) # num_loops may be greater than N, eg Product
   Expr(:block, Expr(:meta,:inline), setup_call(ls, :(Base.Broadcast.materialize!(dest′, bc)), LineNumberNode(0), inline, false, u₁, u₂, v, threads%Int, warncheckarg), :dest′)

src/codegen/split_loops.jl

@@ -13,13 +13,18 @@ function add_operation!(ls_new::LoopSet, included::Vector{Int}, ls::LoopSet, op:
         if isstore(opc) && identifier(opc) ∉ ids
           # @show opp opc op
           # replace opp with a load from opc
+          parentsopc = parents(opc)
+          parentsnew = length(parentsopc) > 1 ? Operation[] : NOPARENTS
           opnew = Operation(
             length(operations(ls_new)), name(opp), opc.elementbytes, instruction(:getindex), memload,
-            loopdependencies(opc), reduceddependencies(opc), NOPARENTS, opc.ref, reducedchildren(opc)
+            loopdependencies(opc), reduceddependencies(opc), parentsnew, opc.ref, reducedchildren(opc)
           )
           addsetv!(ls_new.includedactualarrays, vptr(opc.ref))
           push!(operations(ls_new), opnew)
           push!(vparents, opnew)
+          for i ∈ 2:length(parentsopc)
+            push!(parentsnew, add_operation!(ls_new, included, ls, parentsopc[i], ids, issecond))
+          end
           included[identifier(opp)] = identifier(opnew)
           found = true
           break