Add a few many loop reduction tests.

Author: chriselrod

src/codegen/lower_store.jl

@@ -149,7 +149,6 @@ function lower_store!(
     reductfunc::Symbol = storeinstr_preprend(op, ua.vloop.itersymbol), inds_calc_by_ptr_offset = indices_calculated_by_pointer_offsets(ls, op.ref)
 )
     @unpack u₁, u₁loopsym, u₂loopsym, vloopsym, vloop, u₂max, suffix = ua
-
     omop = offsetloadcollection(ls)
     batchid, opind = omop.batchedcollectionmap[identifier(op)]
     if ((batchid ≠ 0) && isvectorized(op)) && (!rejectinterleave(op))
@@ -161,10 +160,13 @@ function lower_store!(
     # trueexpr = Expr(:call, lv(:True));
     rs = staticexpr(reg_size(ls));
     opp = first(parents(op))
-    if ((opp.instruction.instr === reductfunc) || (opp.instruction.instr === :identity)) && isone(length(parents(opp)))
-        opp = only(parents(opp))
+    if (((opp.instruction.instr === reductfunc) || (opp.instruction.instr === :identity)) && isone(length(parents(opp))))
+        oppp = only(parents(opp))
+        if isu₂unrolled(op) == isu₂unrolled(oppp)
+            opp = oppp
+        end
     end
-    # __u₂max = ls.unrollspecification.u₂
+  # __u₂max = ls.unrollspecification.u₂
     isu₁, isu₂ = isunrolled_sym(opp, u₁loopsym, u₂loopsym, vloopsym, ls)#, __u₂max)
     # @show isu₁, isu₂, u₁loopsym, u₂loopsym
     # @show isu₁, isu₂, opp, u₁loopsym, u₂loopsym, vloopsym

test/grouptests.jl

@@ -49,6 +49,8 @@ const START_TIME = time()
     @time include("multiassignments.jl")
 
     @time include("reduction_untangling.jl")
+
+    @time include("manyloopreductions.jl")
   end
 
   @time if LOOPVECTORIZATION_TEST == "all" || LOOPVECTORIZATION_TEST == "part2"

test/manyloopreductions.jl

@@ -0,0 +1,57 @@
+
+@testset "Many Loop Reductions" begin
+  A = rand((2:6)...);
+  N = ndims(A)
+  T = eltype(A)
+  let dims = (3,5)
+
+    sᵢ = size(A)
+    sₒ = ntuple(Val(N)) do d
+      ifelse(d ∈ dims, 1, sᵢ[d])
+    end
+    Tₒ = Base.promote_op(+, T, Int)
+    B = similar(A, Tₒ, sₒ);
+
+    Bᵥ = view(B, Colon(), Colon(), firstindex(B, 3), Colon(), firstindex(B, 5))
+    @turbo for i_1 = indices((A, B), 1)
+      for i_2 = indices((A, B), 2)
+        for i_4 = indices((A, B), 4)
+          Σ = zero(eltype(Bᵥ))
+          for i_3 = axes(A, 3)
+            for i_5 = axes(A, 5)
+              Σ += A[i_1, i_2, i_3, i_4, i_5]
+            end
+          end
+          Bᵥ[i_1, i_2, i_4] = Σ
+        end
+      end
+    end
+    @test B ≈ sum(A, dims = dims)
+  end
+  let dims = (1,2,5)
+
+    sᵢ = size(A)
+    sₒ = ntuple(Val(N)) do d
+      ifelse(d ∈ dims, 1, sᵢ[d])
+    end
+    Tₒ = Base.promote_op(+, T, Int)
+    B = similar(A, Tₒ, sₒ);
+
+    Bᵥ = view(B, firstindex(B, 1), firstindex(B, 2), Colon(), Colon(), firstindex(B, 5))
+    @turbo for i_3 = indices((A, B), 3)
+      for i_4 = indices((A, B), 4)
+        Σ = zero(eltype(Bᵥ))
+        for i_1 = axes(A, 1)
+          for i_2 = axes(A, 2)
+            for i_5 = axes(A, 5)
+              Σ += A[i_1, i_2, i_3, i_4, i_5]
+            end
+          end
+        end
+        Bᵥ[i_3, i_4] = Σ
+      end
+    end
+    @test B ≈ sum(A, dims = dims)
+  end
+end
+