Tweaks/improvements to lowering of statically sized loops.

Author: chriselrod

src/determinestrategy.jl

@@ -378,11 +378,7 @@ function solve_unroll(X, R, u₁max, u₂max, u₁L, u₂L, u₁step, u₂step)
     u₁, u₂, cost
 end
 function maybedemotesize(U::Int, N::Int)
-    # U > 1 || return 1
-    Um1 = U - 1
-    urep = num_iterations(N, U)
-    um1rep = num_iterations(N, Um1)
-    um1rep > urep ? U : Um1
+    num_iterations(N, num_iterations(N, U))
 end
 function maybedemotesize(u₂::Int, N::Int, U::Int, Uloop::Loop, maxu₂base::Int)
     u₂ > 1 || return 1

src/lower_compute.jl

@@ -19,6 +19,23 @@ function load_constrained(op, u₁loop, u₂loop, forprefetch = false)
     dependsonu₂ && push!(unrolleddeps, u₂loop)
     any(opp -> isload(opp) && all(in(loopdependencies(opp)), unrolleddeps), parents(op))
 end
+function check_if_remfirst(ls, ua)
+    usorig = ls.unrollspecification[]
+    @unpack u₁, u₁loopsym, u₂loopsym, u₂max = ua
+    u₁loop = getloop(ls, u₁loopsym)
+    u₂loop = getloop(ls, u₂loopsym)
+    if isstaticloop(u₁loop) && (usorig.u₁ != u₁)
+        return true
+    end
+    if isstaticloop(u₂loop) && (usorig.u₂ != u₂max)
+        return true
+    end
+    false
+end
+function sub_fmas(ls::LoopSet, op::Operation, ua::UnrollArgs)
+    @unpack u₁, u₁loopsym, u₂loopsym, u₂max = ua
+    !(load_constrained(op, u₁loopsym, u₂loopsym) || check_if_remfirst(ls, ua))
+end
 
 struct FalseCollection end
 Base.getindex(::FalseCollection, i...) = false
@@ -106,7 +123,7 @@ function add_loopvalue!(instrcall::Expr, loopval, ua::UnrollArgs, u::Int)
 end
 
 function lower_compute!(
-    q::Expr, op::Operation, ua::UnrollArgs, mask::Union{Nothing,Symbol,Unsigned} = nothing,
+    q::Expr, op::Operation, ls::LoopSet, ua::UnrollArgs, mask::Union{Nothing,Symbol,Unsigned} = nothing,
 )
     @unpack u₁, u₁loopsym, u₂loopsym, vectorized, suffix = ua
     var = name(op)
@@ -176,16 +193,16 @@ function lower_compute!(
         instrfid = findfirst(isequal(instr.instr), (:vfmadd_fast, :vfnmadd_fast, :vfmsub_fast, :vfnmsub_fast))
         # want to instcombine when parent load's deps are superset
         # also make sure opp is unrolled
-        if instrfid !== nothing && (opunrolled && u₁ > 1) && !load_constrained(op, u₁loopsym, u₂loopsym)
-            specific_fmas = Base.libllvm_version > v"11.0.0" ? (:vfmadd, :vfnmadd, :vfmsub, :vfnmsub) : (:vfmadd231, :vfnmadd231, :vfmsub231, :vfnmsub231)
+        if !isnothing(instrfid) && (opunrolled && u₁ > 1) && sub_fmas(ls, op, ua)
+            specific_fmas = Base.libllvm_version >= v"11.0.0" ? (:vfmadd, :vfnmadd, :vfmsub, :vfnmsub) : (:vfmadd231, :vfnmadd231, :vfmsub231, :vfnmsub231)
             # specific_fmas = (:vfmadd231, :vfnmadd231, :vfmsub231, :vfnmsub231)
             instr = Instruction(specific_fmas[instrfid])
         end
     end
     # @show instr.instr
     reduceddeps = reduceddependencies(op)
     vecinreduceddeps = isreduct && vectorized ∈ reduceddeps
-    maskreduct = mask !== nothing && vecinreduceddeps #any(opp -> opp.variable === var, parents_op)
+    maskreduct = !isnothing(mask) && vecinreduceddeps #any(opp -> opp.variable === var, parents_op)
     # if vecinreduceddeps && vectorized ∉ loopdependencies(op) # screen parent opps for those needing a reduction to scalar
     #     # parents_op = reduce_vectorized_parents!(q, op, parents_op, U, u₁loopsym, u₂loopsym, vectorized, suffix)
     #     isreducingidentity!(q, op, parents_op, U, u₁loopsym, u₂loopsym, vectorized, suffix) && return

src/lowering.jl

@@ -21,7 +21,7 @@ function lower!(
     elseif isload(op)
         lower_load!(q, op, ls, unrollargs, mask)
     elseif iscompute(op)
-        lower_compute!(q, op, unrollargs, mask)
+        lower_compute!(q, op, ls, unrollargs, mask)
     elseif isstore(op)
         lower_store!(q, ls, op, unrollargs, mask)
     # elseif isloopvalue(op)
@@ -45,7 +45,7 @@ function lower!(
         elseif isload(op)
             lower_load!(q, op, ls, unrollargs, mask)
         elseif iscompute(op)
-            lower_compute!(q, op, unrollargs, mask)
+            lower_compute!(q, op, ls, unrollargs, mask)
         end
     end
 end
@@ -317,13 +317,26 @@ function lower_unrolled_dynamic(ls::LoopSet, us::UnrollSpecification, n::Int, in
         1
     end
     remfirst = loopisstatic & (UFt > 0) & !(unsigned(Ureduct) < unsigned(UF))
+#    @show remfirst, loopsym
     tc = terminatecondition(ls, us, n, inclmask, remfirst ? 1 : UF)
     usorig = ls.unrollspecification[]
 
     # tc = (usorig.u₁ == us.u₁) && (usorig.u₂ == us.u₂) && !loopisstatic && !inclmask && !ls.loadelimination[] ? expect(tc) : tc
 
     body = lower_block(ls, us, n, inclmask, UF)
-    q = Expr(:while, tc, body)
+    
+    q = if loopisstatic
+        iters = length(loop) ÷ UF
+        if iters ≤ 4
+            q = Expr(:block)
+            foreach(_ -> push!(q.args, body), 1:iters)
+            q
+        else
+            Expr(:while, tc, body)
+        end
+    else
+        Expr(:while, tc, body)
+    end
     remblock = init_remblock(loop, ls.lssm[], n)#loopsym)
     q = if unsigned(Ureduct) < unsigned(UF) # unsigned(-1) == typemax(UInt); is logic relying on twos-complement bad?
         UF_cleanup = UF - Ureduct
@@ -579,7 +592,7 @@ function setup_preamble!(ls::LoopSet, us::UnrollSpecification)
     lower_licm_constants!(ls)
     isone(num_loops(ls)) || pushpreamble!(ls, definemask(getloop(ls, vectorized)))#, u₁ > 1 && u₁loopnum == vectorizedloopnum))
     for op ∈ operations(ls)
-        (iszero(length(loopdependencies(op))) && iscompute(op)) && lower_compute!(ls.preamble, op, UnrollArgs(u₁, u₁loopsym, u₂loopsym, vectorized, u₂, nothing), nothing)
+        (iszero(length(loopdependencies(op))) && iscompute(op)) && lower_compute!(ls.preamble, op, ls, UnrollArgs(u₁, u₁loopsym, u₂loopsym, vectorized, u₂, nothing), nothing)
     end
     # define_remaining_ops!( ls, vectorized, W, u₁loop, u₂loop, u₁ )
 end