Fix crash from #134.

Author: chriselrod

src/lower_store.jl

@@ -3,37 +3,43 @@ using VectorizationBase: vnoaliasstore!
 
 @inline vstoreadditivereduce!(args...) = vnoaliasstore!(args...)
 @inline vstoremultiplicativevereduce!(args...) = vnoaliasstore!(args...)
-@inline function vstoreadditivereduce!(ptr::VectorizationBase.AbstractStridedPointer, v::VectorizationBase.SVec, i::NTuple{N,<:Integer}) where {N}
+@inline function vstoreadditivereduce!(ptr::VectorizationBase.AbstractStridedPointer, v::VectorizationBase.SVec, i::Tuple{Vararg{Union{Integer,Static}}})
     vnoaliasstore!(ptr, SIMDPirates.vsum(v), i)
 end
-@inline function vstoreadditivereduce!(ptr::VectorizationBase.AbstractStridedPointer, v::VectorizationBase.SVec, i::NTuple{N,<:Integer}, m::VectorizationBase.Mask) where {N}
+@inline function vstoreadditivereduce!(ptr::VectorizationBase.AbstractStridedPointer, v::VectorizationBase.SVec, i::Tuple{Vararg{Union{Integer,Static}}}, m::VectorizationBase.Mask)
     vnoaliasstore!(ptr, SIMDPirates.vsum(v), i, m)
 end
-@inline function vstoremultiplicativevereduce!(ptr::VectorizationBase.AbstractStridedPointer, v::VectorizationBase.SVec, i::NTuple{N,<:Integer}) where {N}
+@inline function vstoremultiplicativevereduce!(ptr::VectorizationBase.AbstractStridedPointer, v::VectorizationBase.SVec, i::Tuple{Vararg{Union{Integer,Static}}})
     vnoaliasstore!(ptr, SIMDPirates.vprod(v), i)
 end
-@inline function vstoremultiplicativevereduce!(ptr::VectorizationBase.AbstractStridedPointer, v::VectorizationBase.SVec, i::NTuple{N,<:Integer}, m::VectorizationBase.Mask) where {N}
+@inline function vstoremultiplicativevereduce!(ptr::VectorizationBase.AbstractStridedPointer, v::VectorizationBase.SVec, i::Tuple{Vararg{Union{Integer,Static}}}, m::VectorizationBase.Mask)
     vnoaliasstore!(ptr, SIMDPirates.vprod(v), i, m)
 end
 
-function storeinstr(op::Operation)
+function storeinstr(op::Operation, vectorized::Symbol)
+    # defaultstoreop = :vstore!
+    defaultstoreop = :vnoaliasstore!
+    vectorized ∉ reduceddependencies(op) && return lv(defaultstoreop)
+    vectorized ∈ loopdependencies(op) && return lv(defaultstoreop)
+    # vectorized is not a loopdep, but is a reduced dep
     opp = first(parents(op))
-    if instruction(opp).instr === :identity
-        opp = first(parents(opp))
+    while vectorized ∉ loopdependencies(opp)
+        oppold = opp
+        for oppp ∈ parents(opp)
+            if vectorized ∈ reduceddependencies(oppp)
+                @assert opp !== oppp "More than one parent is a reduction over the vectorized variable."
+                opp = oppp
+            end
+        end
+        @assert opp !== oppold "Failed to find any parents "
     end
-    defaultstoreop = :vnoaliasstore!
-    # defaultstoreop = :vstore!
-    instr = if iszero(length(reduceddependencies(opp)))
+    instr_class = reduction_instruction_class(instruction(opp))
+    instr = if instr_class === ADDITIVE_IN_REDUCTIONS
+        :vstoreadditivereduce!
+    elseif instr_class === MULTIPLICATIVE_IN_REDUCTIONS
+        :vstoremultiplicativevereduce!
+    else #FIXME
         defaultstoreop
-    else
-        instr_class = reduction_instruction_class(instruction(opp))
-        if instr_class === ADDITIVE_IN_REDUCTIONS
-            :vstoreadditivereduce!
-        elseif instr_class === MULTIPLICATIVE_IN_REDUCTIONS
-            :vstoremultiplicativevereduce!
-        else #FIXME
-            defaultstoreop
-        end
     end
     lv(instr)
 end
@@ -117,7 +123,7 @@ function lower_conditionalstore_scalar!(
         varname = varassignname(mvar, u, opu₁)
         condvarname = varassignname(condvar, u, condu₁)
         td = UnrollArgs(ua, u)
-        push!(q.args, Expr(:&&, condvarname, Expr(:call, storeinstr(op), ptr, varname, mem_offset_u(op, td, inds_calc_by_ptr_offset))))
+        push!(q.args, Expr(:&&, condvarname, Expr(:call, storeinstr(op, vectorized), ptr, varname, mem_offset_u(op, td, inds_calc_by_ptr_offset))))
     end
     nothing
 end
@@ -145,7 +151,7 @@ function lower_conditionalstore_vectorized!(
         td = UnrollArgs(ua, u)
         name, mo = name_memoffset(mvar, op, td, opu₁, inds_calc_by_ptr_offset)
         condvarname = varassignname(condvar, u, condu₁)
-        instrcall = Expr(:call, storeinstr(op), ptr, name, mo)
+        instrcall = Expr(:call, storeinstr(op, vectorized), ptr, name, mo)
         if mask !== nothing && (vecnotunrolled || u == U - 1)
             push!(instrcall.args, Expr(:call, :&, condvarname, mask))
         else
@@ -166,7 +172,7 @@ function lower_store_scalar!(
     for u ∈ 0:u₁-1
         varname = varassignname(mvar, u, opu₁)
         td = UnrollArgs(ua, u)
-        push!(q.args, Expr(:call, storeinstr(op), ptr, varname, mem_offset_u(op, td, inds_calc_by_ptr_offset)))
+        push!(q.args, Expr(:call, storeinstr(op, vectorized), ptr, varname, mem_offset_u(op, td, inds_calc_by_ptr_offset)))
     end
     nothing
 end
@@ -191,7 +197,7 @@ function lower_store_vectorized!(
     for u ∈ umin:U-1
         td = UnrollArgs(ua, u)
         name, mo = name_memoffset(mvar, op, td, opu₁, inds_calc_by_ptr_offset)
-        instrcall = Expr(:call, storeinstr(op), ptr, name, mo)
+        instrcall = Expr(:call, storeinstr(op, vectorized), ptr, name, mo)
         if mask !== nothing && (vecnotunrolled || u == U - 1)
             push!(instrcall.args, mask)
         end

test/miscellaneous.jl

@@ -692,6 +692,35 @@ function splitintonoloop_reference(U = randn(2,2), E1 = randn(2))
     end
     U, E1
 end
+function findreducedparentfornonvecstoreavx!(U::AbstractMatrix{T}, E1::AbstractVector{T}) where T
+    n,k = size(U)
+    _s = zero(T)
+    a = 1.0
+    @avx for j = 1:k
+        for i = 1:n
+            t = tanh(a * U[i,j])
+            U[i,j] = t
+            _s += a * (1 - t^2)
+        end
+        E1[j] = _s / n
+    end
+    U,E1
+end
+function findreducedparentfornonvecstore!(U::AbstractMatrix{T}, E1::AbstractVector{T}) where T
+    n,k = size(U)
+    _s = zero(T)
+    a = 1.0
+    for j = 1:k
+        for i = 1:n
+            t = tanh(a * U[i,j])
+            U[i,j] = t
+            _s += a * (1 - t^2)
+        end
+        E1[j] = _s / n
+    end
+    U,E1
+end
+
 
 
 
@@ -898,11 +927,15 @@ end
         R .+= randn.(T); Rc = copy(R);
         @test maxavx!(R, Q, true) == max.(vec(maximum(Q, dims=(2,3))), Rc)
 
-        U1 = randn(5,7); E1 = randn(7);
-        U2, E2 = splitintonoloop_reference(copy(U1), copy(E1));
-        splitintonoloop(U1, E1);
+        U0 = randn(5,7); E0 = randn(7);
+        U1, E1 = splitintonoloop_reference(copy(U0), copy(E0));
+        U2, E2 = splitintonoloop(copy(U0), copy(E0));
         @test U1 ≈ U2
         @test E1 ≈ E2
+        U3, E3 = findreducedparentfornonvecstoreavx!(copy(U0), copy(E0));
+        findreducedparentfornonvecstore!(U0, E0);
+        @test U0 ≈ U3
+        @test E0 ≈ E3        
     end
     for T ∈ [Int16, Int32, Int64]
         n = 8sizeof(T) - 1