Add if/else support.

Author: chriselrod

Project.toml

@@ -1,7 +1,7 @@
 name = "LoopVectorization"
 uuid = "bdcacae8-1622-11e9-2a5c-532679323890"
 authors = ["Chris Elrod <[email protected]>"]
-version = "0.4.2"
+version = "0.4.3"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -14,8 +14,8 @@ VectorizationBase = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f"
 [compat]
 MacroTools = "0"
 Parameters = "0"
-SIMDPirates = "0.2"
-SLEEFPirates = "0.2"
+SIMDPirates = "0.3"
+SLEEFPirates = "0.3"
 VectorizationBase = "0.2"
 julia = "1.3"
 

src/LoopVectorization.jl

@@ -39,7 +39,7 @@ include("condense_loopset.jl")
 include("reconstruct_loopset.jl")
 include("constructors.jl")
 
-# include("precompile.jl")
-# _precompile_()
+include("precompile.jl")
+_precompile_()
 
 end # module

src/_avx.jl

@@ -59,7 +59,6 @@ function add_reduction!(
     get!(ls.opdict, var) do
         add_constant!(ls, var, elementbytes)
     end
-    # pushparent!(parents, deps, reduceddeps, parent)
 end
 function search_tree(opv::Vector{Operation}, var::Symbol) # relies on cycles being forbidden
     for opp ∈ opv

src/add_compute.jl

@@ -19,13 +19,7 @@ function add_if!(ls::LoopSet, LHS::Symbol, RHS::Expr, elementbytes::Int = 8, mpr
     add_compute!(ls, LHS, :vifelse, [condop, trueop, falseop], elementbytes)
 end
 
-function add_andblock!(ls::LoopSet, ex::Expr)
-    condop = add_compute!(ls, gensym(:mask), first(ex.args), elementbytes)
-    condeval = last(ex.args)::Expr
-    @assert condeval.head === :(=)
-    @assert length(condeval.args) == 2
-    LHS = condeval.args[1]
-    RHS = condeval.args[2]
+function add_andblock!(ls::LoopSet, condop::Operation, LHS, RHS, elementbytes::Int)
     rhsop = add_compute!(ls, gensym(:iftruerhs), RHS, elementbytes)
     if LHS isa Symbol
         altop = getop(ls, LHS)
@@ -36,22 +30,39 @@ function add_andblock!(ls::LoopSet, ex::Expr)
         throw("Don't know how to assign onto $LHS.")
     end        
 end
-function add_orblock!(ls, ex::Expr)
-    condop = add_compute!(ls, gensym(:mask), first(ex.args), elementbytes)
-    condeval = last(ex.args)::Expr
+function add_andblock!(ls::LoopSet, condexpr::Expr, condeval::Expr, elementbytes::Int)
+    condop = add_compute!(ls, gensym(:mask), condexpr, elementbytes)
     @assert condeval.head === :(=)
     @assert length(condeval.args) == 2
     LHS = condeval.args[1]
     RHS = condeval.args[2]
-    rhsop = add_compute!(ls, gensym(:iftruerhs), RHS, elementbytes)
+    add_andblock!(ls, condop, LHS, RHS, elementbytes)
+end
+function add_andblock!(ls::LoopSet, ex::Expr, elementbytes::Int)
+    add_andblock!(ls, first(ex.args)::Expr, last(ex.args)::Expr, elementbytes)
+end
+
+function add_orblock!(ls::LoopSet, condop::Operation, LHS, RHS, elementbytes::Int)
+    rhsop = add_compute!(ls, gensym(:iffalserhs), RHS, elementbytes)
     if LHS isa Symbol
         altop = getop(ls, LHS)
         return add_compute!(ls, LHS, :vifelse, [condop, altop, rhsop], elementbytes)
     elseif LHS isa Expr && LHS.head === :ref
-        negatedcondop = add_compute!(ls, gensym(:negated_mask), :vnot, [condop], elementbytes)
+        negatedcondop = add_compute!(ls, gensym(:negated_mask), :~, [condop], elementbytes)
         return add_conditional_store!(ls, LHS, negatedcondop, rhsop, elementbytes)
     else
         throw("Don't know how to assign onto $LHS.")
-    end        
+    end
+end
+function add_orblock!(ls::LoopSet, condexpr::Expr, condeval::Expr, elementbytes::Int)
+    condop = add_compute!(ls, gensym(:mask), condexpr, elementbytes)
+    @assert condeval.head === :(=)
+    @assert length(condeval.args) == 2
+    LHS = condeval.args[1]
+    RHS = condeval.args[2]
+    add_orblock!(ls, condop, LHS, RHS, elementbytes)
+end
+function add_orblock!(ls::LoopSet, ex::Expr, elementbytes::Int)
+    add_orblock!(ls, first(ex.args)::Expr, last(ex.args)::Expr, elementbytes)
 end
 

src/add_ifelse.jl

@@ -20,25 +20,21 @@ function add_store!(
     ldref = mpref.loopdependencies
     reduceddeps = mpref.reduceddeps
     parent = getop(ls, var, ldref, elementbytes)
-    # pushfirst!(parents, parent)
     pvar = parent.variable
-    nops = length(ls.operations)
-    id = nops
+    id = length(ls.operations)
     if pvar ∉ ls.syms_aliasing_refs
         push!(ls.syms_aliasing_refs, pvar)
         push!(ls.refs_aliasing_syms, mpref.mref)
-        # add_unique_store!(ls, mref, parents, ldref, reduceddeps, elementbytes)
     else
-        # try to cse store
-        # different from cse load, because the other op here must be a store
-        ref = mpref.mref.ref
+        # try to cse store, by replacing the previous one
+         ref = mpref.mref.ref
         for opp ∈ operations(ls)
             isstore(opp) || continue
-            if ref == opp.ref.ref# && return cse_store!(ls, identifier(opp), mref, parents, ldref, reduceddeps, elementbytes)
+            if ref == opp.ref.ref
                 id = opp.identifier
+                break
             end
         end
-        # add_unique_store!(ls, mref, parents, ldref, reduceddeps, elementbytes)        
     end
     pushparent!(parents, ldref, reduceddeps, parent)
     op = Operation( id, name(mpref), elementbytes, :setindex!, memstore, mpref )#loopdependencies, reduceddeps, parents, mpref.mref )
@@ -71,14 +67,49 @@ end
 # For now, it is illegal to load from a conditional store.
 # if you want that sort of behavior, do a conditional reassignment, and store that result unconditionally.
 function add_conditional_store!(ls::LoopSet, LHS, condop::Operation, storeop::Operation, elementbytes::Int)
-    array, raw_indices = ref_from_ref(ex)
+    array, raw_indices = ref_from_ref(LHS)
     ref = ArrayReference(array, raw_indices)
     mref = ArrayReferenceMeta(
         ref, fill(true, length(getindices(ref)))
     )
-    parents = [storeop, condop]
     ldref = convert(Vector{Symbol}, getindices(ref))
-    op = Operation( ls, name(mref), elementbytes, :condtionalstore!, memstore, ldref, NODEPENDENCY, parents, mref )
-    add_unique_store!(ls, op)    
+
+    pvar = storeop.variable
+    id = length(ls.operations)
+    if pvar ∉ ls.syms_aliasing_refs
+        push!(ls.syms_aliasing_refs, pvar)
+        push!(ls.refs_aliasing_syms, mref)
+        storeparents = [storeop, condop]
+    else
+        # for now, we don't try to cse the store
+        # later, as an optimization, we could:
+        # 1. cse the store
+        # 2. use the mask to combine the vector we're trying to store here with the vector that would have been stored in the now cse-ed 1.
+        # 3. use a regular (non-masked) store on that vector.
+        ref = mpref.mref.ref
+        for opp ∈ operations(ls)
+            isstore(opp) || continue
+            if ref == opp.ref.ref# && return cse_store!(ls, identifier(opp), mref, parents, ldref, reduceddeps, elementbytes)
+                id = opp.identifier
+                break
+            end
+        end
+        if id != length(ls.operations) # then there was a previous store
+            prevstore = getop(ls, id + 1)
+            storeop = add_compute!(ls, gensym(:combinedstoreop), Instruction(:vifelse), [condop, storeop, first(parents(prevstore))], elementbytes)
+            storeparents = [storeop]
+            storeinstr = if prevstore.instruction.instr === :conditionalstore!
+                push!(storeparents, add_compute!(ls, gensym(:combinedmask), Instruction(:|), [condop, last(parents(prevstore))], elementbytes))
+                :conditionalstore!
+            else
+                :setindex!
+            end
+            op = Operation( id, name(mref), elementbytes, storeinstr, memstore, ldref, NODEPENDENCY, storeparents, mref )
+            cse_store!(ls, op)
+        end
+    end
+    
+    op = Operation( id, name(mref), elementbytes, :conditionalstore!, memstore, ldref, NODEPENDENCY, storeparents, mref )
+    add_unique_store!(ls, op)
 end
 

src/add_stores.jl

@@ -113,6 +113,7 @@ const COST = Dict{Symbol,InstructionCost}(
     :vabs2 => InstructionCost(4,0.5),
     :(==) => InstructionCost(1, 0.5),
     :isequal => InstructionCost(1, 0.5),
+    :(~) => InstructionCost(1, 0.5),
     :(&) => InstructionCost(1, 0.5),
     :(|) => InstructionCost(1, 0.5),
     :(>) => InstructionCost(1, 0.5),
@@ -139,6 +140,7 @@ const COST = Dict{Symbol,InstructionCost}(
     :sqrt_fast => InstructionCost(15,4.0,-2.0),
     :log => InstructionCost(20,20.0,40.0,20),
     :exp => InstructionCost(20,20.0,20.0,18),
+    :^ => InstructionCost(40,40.0,40.0,26), # FIXME
     :sin => InstructionCost(18,15.0,68.0,23),
     :cos => InstructionCost(18,15.0,68.0,26),
     :sincos => InstructionCost(25,22.0,70.0,26),

src/costs.jl

@@ -374,13 +374,14 @@ function register_single_loop!(ls::LoopSet, looprange::Expr)
         N = gensym(Symbol(:loop, itersym))
         pushpreamble!(ls, Expr(:(=), N, Expr(:call, lv(:maybestaticlength), r.args[2])))
         Loop(itersym, 0, N)
-    elseif f === :OneTo || f === Expr(:(.), :Base, :OneTo)
+    elseif f === :OneTo || f == Expr(:(.), :Base, QuoteNode(:OneTo))
         otN = r.args[2]
         if otN isa Integer
             Loop(itersym, 0, otN)
         else
+            otN isa Expr && maybestatic!(otN)
             N = gensym(Symbol(:loop, itersym))
-            pushpreamble!(ls, Expr(:(=), N, maybestatic!(otN)))
+            pushpreamble!(ls, Expr(:(=), N, otN))
             Loop(itersym, 0, N)
         end
     else
@@ -463,7 +464,7 @@ function add_operation!(
             add_compute!(ls, LHS_sym, RHS, elementbytes, LHS_ref)
         end
     elseif RHS.head === :if
-        add_if!(ls, LHS, RHS, elementbytes, LHS_ref)
+        add_if!(ls, LHS_sym, RHS, elementbytes, LHS_ref)
     else
         throw("Expression not recognized:\n$x")
     end
@@ -509,9 +510,9 @@ function Base.push!(ls::LoopSet, ex::Expr, elementbytes::Int = 8)
     elseif ex.head === :for
         add_loop!(ls, ex)
     elseif ex.head === :&&
-        add_andblock!(ls, ex)
+        add_andblock!(ls, ex, elementbytes)
     elseif ex.head === :||
-        add_orblock!(ls, ex)
+        add_orblock!(ls, ex, elementbytes)
     else
         throw("Don't know how to handle expression:\n$ex")
     end

src/graphs.jl

@@ -49,9 +49,9 @@ function lower_conditionalstore_scalar!(
     var = pvariable_name(op, suffix)
     cond = last(parents(op))
     condvar = if suffix === nothing || tiled ∉ loopdependencies(cond)
-        pvariable_name(op, nothing)
+        variable_name(cond, nothing)
     else
-        pvariable_name(op, suffix)
+        variable_name(cond, suffix)
     end
     condunrolled = unrolled ∈ loopdependencies(cond)
     ptr = refname(op)
@@ -81,10 +81,11 @@ function lower_conditionalstore_vectorized!(
     vecnotunrolled = vectorized !== unrolled
     cond = last(parents(op))
     condvar = if suffix === nothing || tiled ∉ loopdependencies(cond)
-        pvariable_name(op, nothing)
+        variable_name(cond, nothing)
     else
-        pvariable_name(op, suffix)
+        variable_name(cond, suffix)
     end
+    # @show parents(op) cond condvar
     condunrolled = unrolled ∈ loopdependencies(cond)
     for u ∈ zero(Int32):Base.unsafe_trunc(Int32,U-1)
         td = UnrollArgs(u, unrolled, tiled, suffix)
@@ -143,7 +144,6 @@ function lower_store!(
     q::Expr, op::Operation, vectorized::Symbol, W::Symbol, unrolled::Symbol, tiled::Symbol, U::Int,
     suffix::Union{Nothing,Int}, mask::Union{Nothing,Symbol,Unsigned} = nothing
 )
-    # @show unrolled, tiled, U
     isunrolled = unrolled ∈ loopdependencies(op)
     U = isunrolled ? U : 1
     if instruction(op).instr !== :conditionalstore!