Add additional check that a reduction is in fact a reduction before determining whether it can be eliminated.

Author: chriselrod

Project.toml

@@ -1,7 +1,7 @@
 name = "LoopVectorization"
 uuid = "bdcacae8-1622-11e9-2a5c-532679323890"
 authors = ["Chris Elrod <[email protected]>"]
-version = "0.12.71"
+version = "0.12.72"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

src/codegen/operation_evaluation_order.jl

@@ -114,15 +114,26 @@ function load_short_static_reduction_first!(ls::LoopSet, u₁loop::Symbol, u₂l
   for op ∈ operations(ls)
     iscompute(op) || continue
     length(reduceddependencies(op)) == 0 && continue
+    parents_op = parents(op)
+    length(parents_op) == 2 || continue
+    found = false
+    parent₁deps = loopdependencies(parents_op[1])
+    parent₂deps = loopdependencies(parents_op[2])
+    for reduced_dep ∈ reduceddependencies(op)
+      if (reduced_dep ∈ parent₁deps) || (reduced_dep ∈ parent₂deps)
+        found = true
+        break
+      end
+    end
+    found || continue
     if (instruction(op).instr === :reduced_add)
       vecloop = getloop(ls, vectorized)
       if isstaticloop(vecloop) && (length(vecloop) ≤ 16) && nounrollreduction(op, u₁loop, u₂loop, vectorized)
         opsub = parents(op)[2]
         length(children(opsub)) == 1 || continue
         opsearch = parents(op)[1]
         opcheck = search_for_reductinit!(opsearch, opsub, name(opsearch), loopdependencies(op))
-        opcheck === opsearch || replace_reduct_init!(ls, op, opsub, opcheck)
-        
+        opcheck === opsearch || replace_reduct_init!(ls, op, opsub, opcheck)        
       end
     elseif (instruction(op).instr === :add_fast) && (instruction(first(parents(op))).instr === :identity)
       vecloop = getloop(ls, vectorized)

src/reconstruct_loopset.jl

@@ -726,7 +726,7 @@ Execute an `@turbo` block. The block's code is represented via the arguments:
     )
   else
     # Main.BODY[] = avx_body(ls, var"#UNROLL#")
-    # @show avx_body(ls, var"#UNROLL#")
+    # return @show avx_body(ls, var"#UNROLL#")
     avx_body(ls, var"#UNROLL#")
   end
   # @show var"#UNROLL#", var"#OPS#", var"#ARF#", var"#AM#", var"#LPSYM#", var"#LB#"

test/gemm.jl

@@ -558,6 +558,17 @@
         end
         return C
     end
+  function dense!(f::F, C, A, B) where {F}
+    K = ArrayInterface.size(A, StaticInt(2))
+    Kp1 = K + StaticInt(1)
+    @turbo for n ∈ indices((B,C),2), m ∈ indices((A,C),1)
+      Cmn = zero(eltype(C))
+      for k ∈ 1:K
+        Cmn += A[m,k] * B[k,n]
+      end
+      C[m,n] = f(Cmn + A[m,Kp1])
+    end
+  end
 
     # TODO: add fast=false option to `@turbo`
     # function gemm_accurate!(C, A, B)
@@ -712,6 +723,9 @@
             At = copy(A');
             Bt = copy(B');
             C2 = similar(C);
+            A2 = rand(R, M, K+1)
+            dense!(VectorizationBase.relu, C, A2, B);
+            @test C ≈ VectorizationBase.relu.(@view(A2[:,begin:end-1]) * B .+ @view(A2[:,end]))
             @testset "avx $T dynamc gemm" begin
                 AmulB!(C2, A, B)
                 AmulBavx1!(C, A, B)