Interpret lonely loads in if/else statements as being conditional loads.

Author: chriselrod

src/add_ifelse.jl

@@ -2,7 +2,7 @@
 ## Currently, if/else will create its own local scope
 ## Assignments will not register in the loop's main scope
 ## although stores and return values will.
-
+negateop!(ls::LoopSet, condop::Operation, elementbytes::Int) = add_compute!(ls, gensym(:negated_mask), :~, [condop], elementbytes)
 
 function add_if!(ls::LoopSet, LHS::Symbol, RHS::Expr, elementbytes::Int, position::Int, mpref::Union{Nothing,ArrayReferenceMetaPosition} = nothing)
     # for now, just simple 1-liners
@@ -14,18 +14,24 @@ function add_if!(ls::LoopSet, LHS::Symbol, RHS::Expr, elementbytes::Int, positio
         add_operation!(ls, gensym(:mask), condition, mpref, elementbytes, position)
     end
     iftrue = RHS.args[2]
-    trueop = if iftrue isa Expr
-        (iftrue isa Expr && iftrue.head !== :call) && throw("Only calls or constant expressions are currently supported in if/else blocks.")
-        add_operation!(ls, Symbol(:iftrue), iftrue, elementbytes, position)
+    if iftrue isa Expr
+        trueop = add_operation!(ls, Symbol(:iftrue), iftrue, elementbytes, position)
+        if iftrue.head === :ref && all(ld -> ld ∈ loopdependencies(trueop), loopdependencies(condop))
+            trueop.instruction = Instruction(:conditionalload)
+            push!(parents(trueop), condop)
+        end
     else
-        getop(ls, iftrue, elementbytes)
+        trueop = getop(ls, iftrue, elementbytes)
     end
     iffalse = RHS.args[3]
-    falseop = if iffalse isa Expr
-        (iffalse isa Expr && iffalse.head !== :call) && throw("Only calls or constant expressions are currently supported in if/else blocks.")
-        add_operation!(ls, Symbol(:iffalse), iffalse, elementbytes, position)
+    if iffalse isa Expr
+        falseop = add_operation!(ls, Symbol(:iffalse), iffalse, elementbytes, position)
+        if iffalse.head === :ref && all(ld -> ld ∈ loopdependencies(falseop), loopdependencies(condop))
+            falseop.instruction = Instruction(:conditionalload)
+            push!(parents(falseop), negateop!(ls, condop, elementbytes))
+        end
     else
-        getop(ls, iffalse, elementbytes)
+        falseop = getop(ls, iffalse, elementbytes)
     end
     add_compute!(ls, LHS, :vifelse, [condop, trueop, falseop], elementbytes)
 end
@@ -67,7 +73,7 @@ function add_andblock!(ls::LoopSet, ex::Expr, elementbytes::Int, position::Int)
 end
 
 function add_orblock!(ls::LoopSet, condop::Operation, LHS, rhsop::Operation, elementbytes::Int, position::Int)
-    negatedcondop = add_compute!(ls, gensym(:negated_mask), :~, [condop], elementbytes)
+    negatedcondop = negateop!(ls, condop, elementbytes)
     if LHS isa Symbol
         altop = getop(ls, LHS, elementbytes)
         # return add_compute!(ls, LHS, :vifelse, [condop, altop, rhsop], elementbytes)

src/lower_load.jl

@@ -1,26 +1,63 @@
 function lower_load_scalar!(
-    q::Expr, op::Operation, vectorized::Symbol, W::Symbol, unrolled::Symbol, tiled::Symbol, U::Int,
-    suffix::Union{Nothing,Int}, mask::Union{Nothing,Symbol,Unsigned} = nothing, umin::Int = 0
+    q::Expr, op::Operation, vectorized::Symbol, W::Symbol, unrolled::Symbol,
+    tiled::Symbol, U::Int, suffix::Union{Nothing,Int}, umin::Int = 0
 )
     loopdeps = loopdependencies(op)
     @assert vectorized ∉ loopdeps
     var = variable_name(op, suffix)
     ptr = refname(op)
     isunrolled = unrolled ∈ loopdeps
     U = isunrolled ? U : 1
-    for u ∈ umin:U-1
-        varname = varassignname(var, u, isunrolled)
-        td = UnrollArgs(u, unrolled, tiled, suffix)
-        push!(q.args, Expr(:(=), varname, Expr(:call, lv(:vload), ptr, mem_offset_u(op, td))))
+    if instruction(op).instr !== :conditionalload
+        for u ∈ umin:U-1
+            varname = varassignname(var, u, isunrolled)
+            td = UnrollArgs(u, unrolled, tiled, suffix)
+            push!(q.args, Expr(:(=), varname, Expr(:call, lv(:vload), ptr, mem_offset_u(op, td))))
+        end
+    else
+        condop = last(parents(op))
+        condvar = variable_name(condop, suffix)
+        condunrolled = any(isequal(unrolled), loopdependencies(condop))
+        for u ∈ umin:U-1
+            condsym = condunrolled ? Symbol(condvar, u) : condvar
+            varname = varassignname(var, u, isunrolled)
+            td = UnrollArgs(u, unrolled, tiled, suffix)
+            load = Expr(:call, lv(:vload), ptr, mem_offset_u(op, td))
+            cload = Expr(:if, condsym, load, Expr(:call, :zero, Expr(:call, :eltype, ptr)))
+            push!(q.args, Expr(:(=), varname, cload))
+        end
     end
     nothing
 end
-function pushvectorload!(q::Expr, op::Operation, var::Symbol, td::UnrollArgs, U::Int, W::Symbol, mask, vecnotunrolled::Bool)
-    @unpack u, unrolled = td
+function pushvectorload!(
+    q::Expr, op::Operation, var::Symbol, td::UnrollArgs, U::Int, W::Symbol, vectorized::Symbol, mask
+)
+    @unpack u, unrolled, suffix = td
     ptr = refname(op)
+    vecnotunrolled = vectorized !== unrolled
     name, mo = name_memoffset(var, op, td, W, vecnotunrolled)
     instrcall = Expr(:call, lv(:vload), ptr, mo)
-    if mask !== nothing && (vecnotunrolled || u == U - 1)
+
+    iscondstore = instruction(op).instr === :conditionalload
+    maskend = mask !== nothing && (vecnotunrolled || u == U - 1)
+    if iscondstore
+        condop = last(parents(op))
+        # @show condop
+        condsym = variable_name(condop, suffix)
+        condsym = any(isequal(unrolled), loopdependencies(condop)) ? Symbol(condsym, u) : condsym        
+        if vectorized ∈ loopdependencies(condop)
+            if maskend
+                push!(instrcall.args, Expr(:call, :&, condsym, mask))
+            else
+                push!(instrcall.args, condsym)
+            end
+        else
+            if maskend
+                push!(instrcall.args, mask)
+            end
+            instrcall = Expr(:if, condsym, instrcall, Expr(:call, lv(:vzero), W, Expr(:call, :eltype, ptr)))
+        end
+    elseif maskend
         push!(instrcall.args, mask)
     end
     push!(q.args, Expr(:(=), name, instrcall))
@@ -40,10 +77,9 @@ function lower_load_vectorized!(
     end
     # Urange = unrolled ∈ loopdeps ? 0:U-1 : 0
     var = variable_name(op, suffix)
-    vecnotunrolled = vectorized !== unrolled
     for u ∈ umin:U-1
         td = UnrollArgs(u, unrolled, tiled, suffix)
-        pushvectorload!(q, op, var, td, U, W, mask, vecnotunrolled)
+        pushvectorload!(q, op, var, td, U, W, vectorized, mask)
     end
     nothing
 end
@@ -73,6 +109,6 @@ function lower_load!(
     if vectorized ∈ loopdependencies(op)
         lower_load_vectorized!(q, op, vectorized, W, unrolled, tiled, U, suffix, mask, umin)
     else
-        lower_load_scalar!(q, op, vectorized, W, unrolled, tiled, U, suffix, mask, umin)
+        lower_load_scalar!(q, op, vectorized, W, unrolled, tiled, U, suffix, umin)
     end
 end

src/lower_store.jl

@@ -41,6 +41,7 @@ end
 # const STOREOP = :vstore!
 variable_name(op::Operation, ::Nothing) = mangledvar(op)
 variable_name(op::Operation, suffix) = Symbol(mangledvar(op), suffix, :_)
+# variable_name(op::Operation, suffix, u::Int) = (n = variable_name(op, suffix); u < 0 ? n : Symbol(n, u))
 function reduce_range!(q::Expr, toreduct::Symbol, instr::Instruction, Uh::Int, Uh2::Int)
     for u ∈ 0:Uh-1
         tru = Symbol(toreduct, u)

test/miscellaneous.jl

@@ -681,7 +681,6 @@ using Test
             ifirst, ilast = first(irng), last(irng)
             ifirst > ilast && return s
             @avx tile=(1,1) for Ipost in Rpost
-                # Handle all other entries
                 for Ipre in Rpre
                     s[Ipre, ifirst, Ipost] = x[Ipre, ifirst, Ipost]
                     for i = ifirst+1:ilast
@@ -694,15 +693,10 @@ using Test
         function smoothdim_ifelse_avx!(s, x, α, Rpre, irng::AbstractUnitRange, Rpost)
             ifirst, ilast = first(irng), last(irng)
             ifirst > ilast && return s
-            @avx tile=(1,1) for Ipost in Rpost
-                # Handle all other entries
-                for i = ifirst:ilast
-                    for Ipre in Rpre
-                        xi = x[Ipre, i, Ipost]
-                        xim = ifelse(i == ifirst, xi, x[Ipre, i-1, Ipost])
-                        s[Ipre, i, Ipost] = α*xi + (1-α)*xim
-                    end
-                end
+            @avx tile=(1,1) for Ipost in Rpost, i = ifirst:ilast, Ipre in Rpre
+                xi = x[Ipre, i, Ipost]
+                xim = i > ifirst ? x[Ipre, i-1, Ipost] : xi
+                s[Ipre, i, Ipost] = α*xi + (1-α)*xim
             end
             s
         end
@@ -714,10 +708,10 @@ using Test
             # @show d
             Rpre  = CartesianIndices(axes(x)[1:d-1]);
             Rpost = CartesianIndices(axes(x)[d+1:end]);
-            smoothdim!(dest1, x, α, Rpre, axes(x, d), Rpost)
-            smoothdim_avx!(dest2, x, α, Rpre, axes(x, d), Rpost)
+            smoothdim!(dest1, x, α, Rpre, axes(x, d), Rpost);
+            smoothdim_avx!(dest2, x, α, Rpre, axes(x, d), Rpost);
             @test dest1 ≈ dest2
-            fill!(dest2, NaN); smoothdim_ifelse_avx!(dest2, x, α, Rpre, axes(x, d), Rpost)
+            fill!(dest2, NaN); smoothdim_ifelse_avx!(dest2, x, α, Rpre, axes(x, d), Rpost);
             @test dest1 ≈ dest2
         end
     end