Fix second example from #295

Author: chriselrod

Project.toml

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

src/codegen/lower_compute.jl

@@ -176,21 +176,19 @@ function ifelselastexpr(hasf::Bool, M::Int, vargtypes, K::Int, S::Int, maskearly
     q
 end
 @generated function ifelselast(f::F, m::AbstractMask{W}, ::StaticInt{M}, ::StaticInt{S}, vargs::Vararg{Any,K}) where {F,W,K,M,S}
-    1+1
     ifelselastexpr(true, M, vargs, K, S, false)
 end
 @generated function ifelselast(m::AbstractMask{W}, ::StaticInt{M}, ::StaticInt{S}, varg_1::V1, varg_2::V2) where {W,V1,V2,M,S}
-    1+1
     ifelselastexpr(false, M, (V1,V2), 2, S, false)
 end
 @generated function ifelsepartial(f::F, m::AbstractMask{W}, ::StaticInt{M}, ::StaticInt{S}, vargs::Vararg{Any,K}) where {F,W,K,M,S}
-    1+1
     ifelselastexpr(true, M, vargs, K, S, true)
 end
 @generated function ifelsepartial(m::AbstractMask{W}, ::StaticInt{M}, ::StaticInt{S}, varg_1::V1, varg_2::V2) where {W,V1,V2,M,S}
-    1+1
     ifelselastexpr(false, M, (V1,V2), 2, S, true)
 end
+# @inline ifelselast(f::F, m::AbstractMask{W}, ::StaticInt{M}, ::StaticInt{S}, vargs::Vararg{NativeTypes,K}) where {F,W,K,M,S} = f(vargs...)
+# @inline ifelsepartial(f::F, m::AbstractMask{W}, ::StaticInt{M}, ::StaticInt{S}, vargs::Vararg{NativeTypes,K}) where {F,W,K,M,S} = f(vargs...)
 @generated function subset_vec_unroll(vu::VecUnroll{N}, ::StaticInt{S}) where {N,S}
     (1 ≤ S ≤ N + 1) || throw(ArgumentError("`vu` isa `VecUnroll` of `$(N+1)` elements, but trying to subset $S of them."))
     t = Expr(:tuple)

src/modeling/determinestrategy.jl

@@ -1318,6 +1318,26 @@ function choose_tile(ls::LoopSet, sld::Vector{Vector{Symbol}} = store_load_deps(
   shouldinline = (looplengthprod(ls) < 4097.0) #|| any(op -> iscompute(op) && iszero(length(loopdependencies(op))), operations(ls))
   best_order, bestu₁, bestu₂, best_vec, u₁, u₂, lowest_cost, shouldinline
 end
+function mismatchedstorereductions(ls::LoopSet)
+  reduceddeps = Vector{Symbol}[]
+  nreduceddeps = 0
+  for op ∈ operations(ls)
+    isstore(op) || continue
+    rd = reduceddependencies(first(parents(op)))
+    if nreduceddeps ≠ 0
+      length(rd) == nreduceddeps || return true
+    else
+      nreduceddeps = length(rd)
+    end
+    for rdo ∈ reduceddeps
+      for s ∈ rdo
+        s ∈ rd || return true
+      end
+    end
+    push!(reduceddeps, rd)
+  end
+  false
+end
 # Last in order is the inner most loop
 function choose_order_cost(ls::LoopSet)
   resize!(ls.loop_order, length(ls.loopsymbols))
@@ -1331,15 +1351,16 @@ function choose_order_cost(ls::LoopSet)
     tc = Inf
   end
   uorder, uvec, uc = choose_unroll_order(ls, tc, sld)
+  mismatched = mismatchedstorereductions(ls)
   if num_loops(ls) > 1 && tc ≤ uc
     @assert ls.loop_order.bestorder === torder
     # copyto!(ls.loop_order.bestorder, torder)
-    return torder, tunroll, ttile, tvec, tU, tT, tc, shouldinline
+    return torder, tunroll, ttile, tvec, tU, tT, Core.ifelse(mismatched, Inf, tc), shouldinline
     # return torder, tvec, 4, 4#5, 5
   else
     copyto!(ls.loop_order.bestorder, uorder)
     UF, uunroll = determine_unroll_factor(ls, uorder, uvec)
-    return uorder, uunroll, Symbol("##undefined##"), uvec, UF, -1, uc, true
+    return uorder, uunroll, Symbol("##undefined##"), uvec, UF, -1, Core.ifelse(mismatched, Inf, uc), true
   end
 end
 function choose_order(ls::LoopSet)

test/manyloopreductions.jl

@@ -1,4 +1,43 @@
 
+function mismatchedreductions_noturbo!(𝛥r392, 𝛥x923, 𝛥ℛ, ℛ, r392, x923, 𝒶𝓍k=1:2, 𝒶𝓍n=1:3, 𝒶𝓍j=1:9, 𝒶𝓍m=1:9, 𝒶𝓍i=1:3)
+  @inbounds @fastmath for k = 𝒶𝓍k
+    for i = 𝒶𝓍i
+      for m = 𝒶𝓍m
+        for j = 𝒶𝓍j
+          for n = 𝒶𝓍n
+            ℰ𝓍1 = conj(x923[m, k, n])
+            ℰ𝓍2 = 𝛥ℛ[n, j, m, i] * ℰ𝓍1
+            ℰ𝓍3 = conj(r392[i, j, k])
+            ℰ𝓍4 = 𝛥ℛ[n, j, m, i] * ℰ𝓍3
+            𝛥r392[i, j, k] = 𝛥r392[i, j, k] + ℰ𝓍2
+            𝛥x923[m, k, n] = 𝛥x923[m, k, n] + ℰ𝓍4
+          end
+        end
+      end
+    end
+  end
+  𝛥r392, 𝛥x923
+end
+function mismatchedreductions!(𝛥r392, 𝛥x923, 𝛥ℛ, ℛ, r392, x923, 𝒶𝓍k=1:2, 𝒶𝓍n=1:3, 𝒶𝓍j=1:9, 𝒶𝓍m=1:9, 𝒶𝓍i=1:3)
+  @turbo for k = 𝒶𝓍k
+    for i = 𝒶𝓍i
+      for m = 𝒶𝓍m
+        for j = 𝒶𝓍j
+          for n = 𝒶𝓍n
+            ℰ𝓍1 = conj(x923[m, k, n])
+            ℰ𝓍2 = 𝛥ℛ[n, j, m, i] * ℰ𝓍1
+            ℰ𝓍3 = conj(r392[i, j, k])
+            ℰ𝓍4 = 𝛥ℛ[n, j, m, i] * ℰ𝓍3
+            𝛥r392[i, j, k] = 𝛥r392[i, j, k] + ℰ𝓍2
+            𝛥x923[m, k, n] = 𝛥x923[m, k, n] + ℰ𝓍4
+          end
+        end
+      end
+    end
+  end
+  𝛥r392, 𝛥x923
+end
+
 @testset "Many Loop Reductions" begin
   A = rand((2:6)...);
   N = ndims(A)
@@ -53,5 +92,16 @@
     end
     @test B ≈ sum(A, dims = dims)
   end
+
+  r392 = rand(3,9,2);
+  x923 = rand(9,2,3);
+  K = rand(3,9,9,3);
+  𝛥r392_1, 𝛥x923_1, 𝛥r392_2, 𝛥x923_2, 𝛥ℛ = similar(r392), similar(x923), similar(r392), similar(x923), copy(K);
+  𝛥r392_1 .= -1; 𝛥x923_1 .= -1; 𝛥r392_2 .= -1; 𝛥x923_2 .= -1;
+
+  mismatchedreductions_noturbo!(𝛥r392_1, 𝛥x923_1, 𝛥ℛ, K, r392, x923)
+  @time mismatchedreductions!(𝛥r392_2, 𝛥x923_2, 𝛥ℛ, K, r392, x923)
+  @test 𝛥r392_1 ≈ 𝛥r392_2
+  @test 𝛥x923_1 ≈ 𝛥x923_2
 end