Move operations after inner loop when possible.

Author: chriselrod

src/determinestrategy.jl

@@ -270,6 +270,12 @@ function solve_tilesize(
     solve_tilesize(cost_vec, reg_pressure, maxU, maxT)
 end
 
+function set_for_each_parent!(adal::Vector{T}, op::Operation, val::T) where {T}
+    @inbounds for opp ∈ parents(op)
+        adal[identifier(opp)] = val
+    end
+end
+
 # Just tile outer two loops?
 # But optimal order within tile must still be determined
 # as well as size of the tiles.
@@ -280,7 +286,13 @@ function evaluate_cost_tile(
     @assert N ≥ 2 "Cannot tile merely $N loops!"
     tiled = order[1]
     unrolled = order[2]
-    included_vars = fill(false, length(operations(ls)))
+    ops = operations(ls)
+    nops = length(ops)
+    included_vars = fill(false, nops)
+    unrolledtiled = fill(false, 2, nops)
+    descendentsininnerloop = fill(false, nops)
+    innerloop = last(order)
+    iters = fill(-99.9, nops)
     nested_loop_syms = Symbol[]# Set{Symbol}()
     iter = 1.0
     # Need to check if fusion is possible
@@ -306,7 +318,7 @@ function evaluate_cost_tile(
             iter *= Float64(length(ls, itersym))
         end
         # check which vars we can define at this level of loop nest
-        for (id, op) ∈ enumerate(operations(ls))
+        for (id, op) ∈ enumerate(ops)
             # isconstant(op) && continue
             # @assert id == identifier(op)+1 # testing, for now
             # won't define if already defined...
@@ -318,27 +330,37 @@ function evaluate_cost_tile(
             rd = reduceddependencies(op)
             hasintersection(rd, nested_loop_syms[1:end-length(rd)]) && return 0,0,Inf
             included_vars[id] = true
-            rt, lat, rp = cost(op, vectorized, Wshift, size_T)
+            unrolledtiled[1,id] = unrolled ∈ loopdependencies(op)
+            unrolledtiled[2,id] = tiled ∈ loopdependencies(op)
+            iters[id] = iter
+            innerloop ∈ loopdependencies(op) && set_for_each_parent!(descendentsininnerloop, op, true)
+        end
+    end
+    for (id, op) ∈ enumerate(ops)
+        iters[id] == -99.9 && continue
+        descendentsininnerloop[id] || continue
+        isunrolled = unrolledtiled[1,id]
+        istiled = unrolledtiled[2,id]
+        rt, lat, rp = cost(op, vectorized, Wshift, size_T)
             # @show instruction(op), rt, lat, rp, iter
-            rt *= iter
-            isunrolled = unrolled ∈ loopdependencies(op)
-            istiled = tiled ∈ loopdependencies(op)
+        rt *= iters[id]
             # @show isunrolled, istiled
-            if isunrolled && istiled # no cost decrease; cost must be repeated
-                cost_vec[1] += rt
-                reg_pressure[1] += rp
-            elseif isunrolled # cost decreased by tiling
-                cost_vec[2] += rt
-                reg_pressure[2] += rp
-            elseif istiled # cost decreased by unrolling
-                cost_vec[3] += rt
-                reg_pressure[3] += rp
-            else# neither unrolled or tiled
-                cost_vec[4] += rt
-                reg_pressure[4] += rp
-            end
+        if isunrolled && istiled # no cost decrease; cost must be repeated
+            cost_vec[1] += rt
+            reg_pressure[1] += rp
+        elseif isunrolled # cost decreased by tiling
+            cost_vec[2] += rt
+            reg_pressure[2] += rp
+        elseif istiled # cost decreased by unrolling
+            cost_vec[3] += rt
+            reg_pressure[3] += rp
+        else# neither unrolled or tiled
+            cost_vec[4] += rt
+            reg_pressure[4] += rp
         end
     end
+    # @show order, vectorized cost_vec reg_pressure
+    # @show solve_tilesize(ls, unrolled, tiled, cost_vec, reg_pressure)
     solve_tilesize(ls, unrolled, tiled, cost_vec, reg_pressure)
 end
 

src/graphs.jl

@@ -603,15 +603,17 @@ function Base.push!(ls::LoopSet, ex::Expr, elementbytes::Int = 8)
     end
 end
 
-function place_after_loop(op::Operation)
-    if isload(op) || length(reduceddependencies(op)) == 0
+function place_after_loop!(adal::Vector{Bool}, op::Operation)
+    pal = if isload(op) || length(reduceddependencies(op)) == 0
         1
     elseif length(reduceddependencies(op)) > 1
         2
     else
         rd = first(reduceddependencies(op))
         any(d -> d === rd, loopdependencies(op)) ? 1 : 2
     end
+    pal == 1 && set_for_each_parent!(adal, op, false)
+    pal
 end
 
 function fillorder!(ls::LoopSet, order::Vector{Symbol}, loopistiled::Bool)
@@ -620,7 +622,6 @@ function fillorder!(ls::LoopSet, order::Vector{Symbol}, loopistiled::Bool)
     # @show 1, ro, order
     # copyto!(ro, order)
     # @show 2, ro, order
-    empty!(lo)
     nloops = length(order)
     if loopistiled
         tiled    = order[1]
@@ -629,23 +630,36 @@ function fillorder!(ls::LoopSet, order::Vector{Symbol}, loopistiled::Bool)
         tiled = Symbol("##UNDEFINED##")
         unrolled = first(order)
     end
-    included_vars = fill(false, length(operations(ls)))
+    ops = operations(ls)
+    nops = length(ops)
+    included_vars = fill(false, nops)
+    all_descendents_after_loop = fill(true, nops)
+    positions = fill((-1,-1,-1,-1,-1), nops)#Vector{NTuple{5,Int}}(undef, nops)
     # to go inside out, we just have to include all those not-yet included depending on the current sym
     for _n ∈ 1:nloops
         n = 1 + nloops - _n
         ro[_n] = loopsym = order[n]
         #loopsym = order[n]
-        for (id,op) ∈ enumerate(operations(ls))
+        for (id,op) ∈ enumerate(ops)
             included_vars[id] && continue
             loopsym ∈ loopdependencies(op) || continue
             included_vars[id] = true
             isunrolled = (unrolled ∈ loopdependencies(op)) + 1
             istiled = (loopistiled ? (tiled ∈ loopdependencies(op)) : false) + 1
             optype = Int(op.node_type) + 1
-            after_loop = place_after_loop(op)
-            push!(lo[optype,isunrolled,istiled,after_loop,_n], op)
+            after_loop = place_after_loop!(all_descendents_after_loop, op)
+            positions[id] = (optype,isunrolled,istiled,after_loop,_n)
         end
-    end    
+    end
+    empty!(lo)
+    for id ∈ 1:nops
+        optype,isunrolled,istiled,after_loop,_n = positions[id]
+        optype == -1 && continue#@show ops[id]
+        if all_descendents_after_loop[id]
+            after_loop = 2
+        end
+        push!(lo[optype,isunrolled,istiled,after_loop,_n], ops[id])
+    end
     # 3, ro, order
 end