Use combinemasks function; dispatches to & given shared types, and ? when mixing bools and others.

Author: chriselrod

src/lower_store.jl

@@ -70,6 +70,10 @@ function lower_conditionalstore_scalar!(
     end
     nothing
 end
+@inline combinemasks(a::Unsigned, b::Unsigned) = a & b
+@inline combinemasks(a::Unsigned, b::Bool) = b ? a : zero(a)
+@inline combinemasks(a::Bool, b::Unsigned) = a ? b : zero(b)
+@inline combinemasks(a::Bool, b::Bool) = a & b
 function lower_conditionalstore_vectorized!(
     q::Expr, op::Operation, vectorized::Symbol, W::Symbol, unrolled::Symbol, tiled::Symbol, U::Int,
     suffix::Union{Nothing,Int}, mask::Union{Nothing,Symbol,Unsigned}, isunrolled::Bool
@@ -101,7 +105,7 @@ function lower_conditionalstore_vectorized!(
         condvarname = varassignname(condvar, u, condunrolled)
         instrcall = Expr(:call, lv(:vstore!), ptr, name, mo)
         if mask !== nothing && (vecnotunrolled || u == U - 1)
-            push!(instrcall.args, Expr(:call, :&, condvarname, mask))
+            push!(instrcall.args, Expr(:call, lv(:combinemasks), condvarname, mask))
         else
             push!(instrcall.args, condvarname)
         end

src/operation_evaluation_order.jl

@@ -1,3 +1,4 @@
+
 function set_upstream_family!(adal::Vector{T}, op::Operation, val::T, ld::Vector{Symbol}, id::Int) where {T}
     adal[identifier(op)] == val && return # must already have been set
     # ld != loopdependencies(op) &&

test/runtests.jl

@@ -1340,7 +1340,6 @@ end
         end
         qq[:,Base.OneTo(maxk)] ./= vec(lse)
     end
-
     add_1_dim(x::AbstractArray) = reshape(x, size(x)..., 1)
     check_finite(x::AbstractArray) = all(isfinite.(x)) || throw(error("x not finite!"))
     function softmax3_setup!(q::AA, lse::A, tmpmax::A, x::AA, maxk=size(q, ndims(q)) ) where {T<:Real, A<:Array{T}, AA<:AbstractArray{T}}
@@ -1492,7 +1491,7 @@ end
         @test y1 ≈ y2
 
 
-        ni, nj, nk = (100, 100, 10)
+        ni, nj, nk = (127, 113, 13)
         x = rand(T, ni, nj, nk);
         q1 = similar(x);
         q2 = similar(x);