static reduction cleanup

Author: chriselrod

src/codegen/lower_compute.jl

@@ -234,7 +234,6 @@ end
             push!(q.args, :($gf(vargs, $k, false)))
         end
         return Expr(:block, Expr(:meta, :inline), q)
-        # return Expr(:block, Expr(:meta, :inline), :(@show($q)))
     end
     if Sreduced
         M = N
@@ -273,7 +272,6 @@ end
         push!(t.args, :($gf(dd, $m, false)))
     end
     push!(q.args, :(VecUnroll($t)))
-    # push!(q.args, :(@show(VecUnroll($t))))
     q
 end
 
@@ -366,20 +364,11 @@ function getu₁forreduct(ls::LoopSet, op::Operation, u₁::Int)
 end
 isidentityop(op::Operation) = iscompute(op) && (instruction(op).instr === :identity) && (length(parents(op)) == 1)
 function reduce_parent!(q::Expr, ls::LoopSet, op::Operation, opp::Operation, parent::Symbol)
-  # if instruction(op).instr === :log_fast
-  #   @show op opp isvectorized(op) isvectorized(opp) dependent_outer_reducts(ls, op)
-  # end
   isvectorized(op) && return parent
-  # if dependent_outer_reducts(ls, op)
-    
-  #   return parent
-  # end
-  # @show op opp isvectorized(opp)
   if isvectorized(opp)
     oppt = opp
   elseif isidentityop(opp)
     oppt = parents(opp)[1]
-    # @show oppt
     isvectorized(oppt) || return parent
   else
     return parent
@@ -510,13 +499,6 @@ function lower_compute!(
       modsuffix = suffix % ls.ureduct
       Symbol(mangledvar(op), modsuffix)
     end
-    # @show op, u₁unrolledsym, u₂unrolledsym
-    # end
-    # dopartialmap = u₁ > 1
-
-    # Symbol(mvar, modsuffix)
-    # elseif u₁unrolledsym
-    #     Symbol(mvar, u)
   elseif u₁unrolledsym
     if isreduct #(isanouterreduction(ls, op))
       # isouterreduct = true
@@ -532,7 +514,6 @@ function lower_compute!(
   end
   selfopname = varsym
   selfdep = 0
-  # @show u₁unrolledsym
   for n ∈ 1:nparents
     opp = parents_op[n]
     if isloopvalue(opp)
@@ -544,7 +525,6 @@ function lower_compute!(
         (parents_u₁syms[n] != u₁unrolledsym) || (parents_u₂syms[n] != u₂unrolledsym)
 
         selfopname, uₚ = parent_op_name!(q, ls, parents_op, n, modsuffix, suffix_, parents_u₁syms, parents_u₂syms, u₁, u₂max, u₂unrolledsym, op, tiledouterreduction)
-        # @show selfopname, uₚ, tiledouterreduction opp op
         push!(instrcall.args, selfopname)
       else
         push!(instrcall.args, varsym)
@@ -574,7 +554,6 @@ function lower_compute!(
     end
   end
   selfdepreduce = ifelse(((!u₁unrolledsym) & isu₁unrolled(op)) & (u₁ > 1), selfdep, 0)
-  # @show selfdepreduce, selfdep, maskreduct, op
   if maskreduct
     ifelsefunc = if us.u₁ == 1
       :ifelse # don't need to be fancy

src/codegen/lowering.jl

@@ -292,19 +292,19 @@ function lower_unrolled_dynamic(ls::LoopSet, us::UnrollSpecification, n::Int, in
     end
   end
   q = if unsigned(Ureduct) < unsigned(UF) # unsigned(-1) == typemax(UInt); 
-    add_cleanup = true
+    add_cleanup = !loopisstatic# true
     if isone(Ureduct)
       UF_cleanup = 1
       if nisvectorized
         blockhead = :while
       else
         blockhead = if UF == 2
-        if loopisstatic
-          add_cleanup = UFt == 1
-          :block
-        else
-          :if
-        end
+          if loopisstatic
+            # add_cleanup = UFt == 1
+            :block
+          else
+            :if
+          end
         else
           :while
         end
@@ -319,6 +319,12 @@ function lower_unrolled_dynamic(ls::LoopSet, us::UnrollSpecification, n::Int, in
     end
     _q = if dynamicbounded
       initialize_outer_reductions!(q, ls, Ureduct); q
+    elseif loopisstatic
+      if length(loop) < UF*W
+        Expr(:block)
+      else
+        Expr(:block, add_upper_outer_reductions(ls, q, Ureduct, UF, loop, nisvectorized))
+      end
     else
       Expr(:block, add_upper_outer_reductions(ls, q, Ureduct, UF, loop, nisvectorized))
     end
@@ -587,6 +593,12 @@ function add_upper_outer_reductions(ls::LoopSet, loopq::Expr, Ulow::Int, Uhigh::
   end
   push!(ifq.args, t)
   ifqfull = Expr(:let, ifqlet, ifq)
+  if isstaticloop(unrolledloop)
+    W = Core.ifelse(reductisvectorized, ls.vector_width, 1)
+    if Uhigh*W*gethint(step(unrolledloop)) ≤ length(unrolledloop)
+      return Expr(:(=), mvartl, ifqfull)
+    end
+  end
   ncomparison = if reductisvectorized
     add_upper_comp_check(unrolledloop, mulexpr(VECTORWIDTHSYMBOL, Uhigh, step(unrolledloop)))
   elseif isknown(step(unrolledloop))
@@ -848,44 +860,49 @@ end
 #     cld(u₂, cld(u₂, unroll))
 # end
 function calc_Ureduct!(ls::LoopSet, us::UnrollSpecification)
-    @unpack u₁loopnum, u₁, u₂, vloopnum = us
-    ur = if iszero(length(ls.outer_reductions))
-        -1
-    elseif u₂ == -1
-        if u₁loopnum == num_loops(ls)
-            loopisstatic = isstaticloop(getloop(ls, u₁loopnum))
-            loopisstatic &= ((vloopnum != u₁loopnum) | (!iszero(ls.vector_width)))
-            # loopisstatic ? u₁ : min(u₁, 4) # much worse than the other two options, don't use this one
-            if Sys.CPU_NAME === "znver1"
-                loopisstatic ? u₁ : 1
-            else
-                loopisstatic ? u₁ : (u₁ ≥ 4 ? 2 : 1)
-            end
-        else
-            -1
-        end
+  @unpack u₁loopnum, u₁, u₂, vloopnum = us
+  ur = if iszero(length(ls.outer_reductions))
+    -1
+  elseif u₂ == -1
+    if u₁loopnum == num_loops(ls)
+      u₁loop = getloop(ls, u₁loopnum)
+      loopisstatic = isstaticloop(u₁loop)
+      loopisstatic &= ((vloopnum != u₁loopnum) | (!iszero(ls.vector_width)))
+      # loopisstatic ? u₁ : min(u₁, 4) # much worse than the other two options, don't use this one
+      if loopisstatic
+        W = Core.ifelse(vloopnum == u₁loopnum, ls.vector_width, 1)
+        UFt = cld(length(u₁loop) % (W*u₁), W)
+        Core.ifelse(UFt == 0, u₁, UFt)
+        # rem = length(u₁loop) - 
+        # max(1, cld(rem, u₁))
+      else
+        Core.ifelse(Sys.CPU_NAME === "znver1", 1, Core.ifelse(u₁ ≥ 4, 2, 1))
+      end
     else
-        u₁ui = u₂ui = -1
-        u₁loopsym = getloop(ls,    u₁loopnum).itersymbol
-        u₂loopsym = getloop(ls, us.u₂loopnum).itersymbol
-        vloopsym  = getloop(ls,     vloopnum).itersymbol
-        for or ∈ ls.outer_reductions
-            op = ls.operations[or]
-            u₁u, u₂u = isunrolled_sym(op, u₁loopsym, u₂loopsym, vloopsym, us)
-            if u₁ui == -1
-                u₁ui = Int(u₁u)
-                u₂ui = Int(u₁u)
-            elseif !((u₁ui == Int(u₁u)) & (u₂ui == Int(u₁u)))
-                throw(ArgumentError("Doesn't currenly handle differently unrolled reductions yet, please file an issue with an example."))
-            end
-        end
-        if u₁ui % Bool
-            u₁
-        else
-            u₂
-        end
+      -1
     end
-    ls.ureduct = ur
+  else
+    u₁ui = u₂ui = -1
+    u₁loopsym = getloop(ls,    u₁loopnum).itersymbol
+    u₂loopsym = getloop(ls, us.u₂loopnum).itersymbol
+    vloopsym  = getloop(ls,     vloopnum).itersymbol
+    for or ∈ ls.outer_reductions
+      op = ls.operations[or]
+      u₁u, u₂u = isunrolled_sym(op, u₁loopsym, u₂loopsym, vloopsym, us)
+      if u₁ui == -1
+        u₁ui = Int(u₁u)
+        u₂ui = Int(u₁u)
+      elseif !((u₁ui == Int(u₁u)) & (u₂ui == Int(u₁u)))
+        throw(ArgumentError("Doesn't currenly handle differently unrolled reductions yet, please file an issue with an example."))
+      end
+    end
+    if u₁ui % Bool
+      u₁
+    else
+      u₂
+    end
+  end
+  ls.ureduct = ur
 end
 ureduct(ls::LoopSet) = ls.ureduct
 function lower_unrollspec(ls::LoopSet)

test/staticsize.jl

@@ -78,6 +78,21 @@ function update_turbo!(B⁻¹yₖ, B⁻¹, yₖ, sₖᵀyₖ⁻¹)
   yₖᵀB⁻¹yₖ
 end
 
+function maxabs(x)
+  s = -Inf
+  @turbo for i ∈ eachindex(x)
+    s = max(s, abs(x[i]))
+  end
+  s
+end
+function sum_turbo(x)
+  s = zero(eltype(x))
+  @turbo for i ∈ eachindex(x)
+    s += x[i]
+  end
+  s
+end
+
 @testset "Statically Sized Arrays" begin
   @show @__LINE__
   for n1 ∈ 1:MAXTESTSIZE, n3 ∈ 1:MAXTESTSIZE
@@ -97,6 +112,13 @@ end
       @test update_turbo!(By0, output1, y, 0.124) ≈ update!(By1, output1, y, 0.124)
       @test By0 ≈ By1
     end
+    
+  end
+  for i in 1:65
+    x = StrideArray(undef, StaticInt(i))
+    x .= randn.();
+    @test maxabs(x) == maximum(abs, x)
+    @test sum_turbo(x) ≈ sum(x)
   end
 end