Fixed a lot of bugs.

Author: chriselrod

src/LoopVectorization.jl

@@ -1,11 +1,11 @@
 module LoopVectorization
 
 using VectorizationBase, SIMDPirates, SLEEFPirates, MacroTools, Parameters
-using VectorizationBase: REGISTER_SIZE, extract_data, num_vector_load_expr
-using SIMDPirates: VECTOR_SYMBOLS, evadd, evmul
+using VectorizationBase: REGISTER_SIZE, extract_data, num_vector_load_expr, mask
+using SIMDPirates: VECTOR_SYMBOLS, evadd, evmul, vrange, reduced_add, reduced_prod
 using MacroTools: prewalk, postwalk
 
-export vectorizable, @vectorize, @vvectorize
+export vectorizable, @vectorize, @vvectorize, @avx
 
 function isdense end #
 

src/constructors.jl

@@ -1,10 +1,6 @@
 
 ### This file contains convenience functions for constructing LoopSets.
 
-function walk_body!(ls::LoopSet, body::Expr)
-    
-
-end
 function Base.copyto!(ls::LoopSet, q::Expr)
     q.head === :for || throw("Expression must be a for loop.")
     add_loop!(ls, q)
@@ -18,5 +14,7 @@ function LoopSet(q::Expr)
     ls
 end
 
-
+macro avx(q)
+    esc(lower(LoopSet(q)))
+end
 

src/costs.jl

@@ -123,7 +123,7 @@ const REDUCTION_TRANSLATION = Dict{Symbol,Symbol}(
     :(-) => :evadd,
     :vsub => :evadd,
     :(/) => :evmul,
-    :vdiv => :evmul,
+    :vfdiv => :evmul,
     :muladd => :evadd,
     :fma => :evadd,
     :vmuladd => :evadd,
@@ -141,7 +141,7 @@ const REDUCTION_ZERO = Dict{Symbol,Symbol}(
     :(-) => :zero,
     :vsub => :zero,
     :(/) => :one,
-    :vdiv => :one,
+    :vfdiv => :one,
     :muladd => :zero,
     :fma => :zero,
     :vmuladd => :zero,
@@ -151,12 +151,46 @@ const REDUCTION_ZERO = Dict{Symbol,Symbol}(
     :vfnmadd => :zero,
     :vfnmsub => :zero    
 )
-# const SIMDPIRATES_COST = Dict{Symbol,InstructionCost}()
-# const SLEEFPIRATES_COST = Dict{Symbol,InstructionCost}()
+# Fast functions, because common pattern is
+const REDUCTION_SCALAR_COMBINE = Dict{Symbol,Expr}(
+    :(+) => :(LoopVectorization.reduced_add),
+    :vadd => :(LoopVectorization.reduced_add),
+    :(*) => :(LoopVectorization.reduced_prod),
+    :vmul => :(LoopVectorization.reduced_prod),
+    :(-) => :(LoopVectorization.reduced_add),
+    :vsub => :(LoopVectorization.reduced_add),
+    :(/) => :(LoopVectorization.reduced_prod),
+    :vfdiv => :(LoopVectorization.reduced_prod),
+    :muladd => :(LoopVectorization.reduced_add),
+    :fma => :(LoopVectorization.reduced_add),
+    :vmuladd => :(LoopVectorization.reduced_add),
+    :vfma => :(LoopVectorization.reduced_add),
+    :vfmadd => :(LoopVectorization.reduced_add),
+    :vfmsub => :(LoopVectorization.reduced_add),
+    :vfnmadd => :(LoopVectorization.reduced_add),
+    :vfnmsub => :(LoopVectorization.reduced_add)
+)
+
+const FUNCTION_MODULES = Dict{Symbol,Expr}(
+    :vadd => :(LoopVectorization.vadd),
+    :vmul => :(LoopVectorization.vmul),
+    :vsub => :(LoopVectorization.vsub),
+    :vfdiv => :(LoopVectorization.vfdiv),
+    :vmuladd => :(LoopVectorization.vmuladd),
+    :vfma => :(LoopVectorization.vfma),
+    :vfmadd => :(LoopVectorization.vfmadd),
+    :vfmsub => :(LoopVectorization.vfmsub),
+    :vfnmadd => :(LoopVectorization.vfnmadd),
+    :vfnmsub => :(LoopVectorization.vfnmsub),
+    :vsqrt => :(LoopVectorization.vsqrt),
+    :log => :(LoopVectorization.SIMDPirates.vlog),
+    :exp => :(LoopVectorization.SIMDPirates.vexp),
+    :sin => :(LoopVectorization.SLEEFPirates.sin),
+    :cos => :(LoopVectorization.SLEEFPirates.cos),
+    :sincos => :(LoopVectorization.SLEEFPirates.sincos)
+)
+function callfun(f::Symbol)
+    Expr(:call, get(FUNCTION_MODULES, f, f))::Expr
+end
 
-# const MODULE_LOOKUP = Dict{Symbol,Dict{Symbol,InstructionCost}}(
-    # :Base => BASE_COST,
-    # :SIMDPirates => SIMDPIRATES_COST,
-    # :SLEEFPirates => SLEEFPIRATES_COST
-# )
 

src/determinestrategy.jl

@@ -88,22 +88,21 @@ end
 
 # only covers unrolled ops; everything else considered lifted?
 function depchain_cost!(
-    skip::Vector{Bool}, op::Operation, unrolled::Symbol, Wshift::Int, size_T::Int, sl::Int = 0, rt::Float64 = 0.0
+    skip::Vector{Bool}, op::Operation, unrolled::Symbol, Wshift::Int, size_T::Int, rt::Float64 = 0.0, sl::Int = 0
 )
     skip[identifier(op)] = true
     # depth first search
     for opp ∈ parents(op)
         skip[identifier(opp)] && continue
-        sl, rt = depchain_cost!(skip, opp, unrolled, Wshift, size_T, sl, rt)
+        rt, sl = depchain_cost!(skip, opp, unrolled, Wshift, size_T, rt, sl)
     end
     # Basically assuming memory and compute don't conflict, but everything else does
     # Ie, ignoring the fact that integer and floating point operations likely don't either
-    if accesses_memory(op)
-        return sl, rt
+    if iscompute(op)
+        rtᵢ, slᵢ = cost(op, unrolled, Wshift, size_T)
+        rt += rtᵢ; sl += slᵢ
     end
-    # @show instruction(op)
-    rtᵢ, slᵢ = cost(op, unrolled, Wshift, size_T)
-    sl + slᵢ, rt + rtᵢ
+    rt, sl
 end
 
 function determine_unroll_factor(
@@ -116,21 +115,35 @@ function determine_unroll_factor(
     # The assumption here is that unrolling provides no real benefit, unless it is needed to enable OOO execution by breaking up these dependency chains
     num_reductions = sum(isreduction, operations(ls))
     # @show num_reductions
-    iszero(num_reductions) && return 1
+    if iszero(num_reductions) # the 4 is a hack, based on the idea that there is some cost to moving through columns
+        return length(order) == 1 ? 1 : 4
+    end
     # So if num_reductions > 0, we set the unroll factor to be high enough so that the CPU can be kept busy
     # if there are, U = max(1, round(Int, max(latency) * throughput / num_reductions)) = max(1, round(Int, latency / (recip_throughput * num_reductions)))
     # We also make sure register pressure is not too high.
     latency = 0
-    recip_throughput = 0.0
+    compute_recip_throughput = 0.0
     visited_nodes = fill(false, length(operations(ls)))
+    load_recip_throughput = 0.0
+    store_recip_throughput = 0.0
     for op ∈ operations(ls)
-        if isreduction(op) && dependson(op, unrolled)
-            sl, rt = depchain_cost!(visited_nodes, op, unrolled, Wshift, size_T)
+        dependson(op, unrolled) || continue
+        if isreduction(op)
+            rt, sl = depchain_cost!(visited_nodes, op, unrolled, Wshift, size_T)
             latency = max(sl, latency)
-            recip_throughput += rt
+            compute_recip_throughput += rt
+        elseif isload(op)
+            load_recip_throughput += first(cost(op, unrolled, Wshift, size_T))
+        elseif isstore(op)
+            store_recip_throughput += first(cost(op, unrolled, Wshift, size_T))
         end
     end
-    max(1, round(Int, latency / (recip_throughput * num_reductions) ) )  
+    recip_throughput = max(
+        compute_recip_throughput,
+        load_recip_throughput,
+        store_recip_throughput
+    )
+    max(1, round(Int, latency / (recip_throughput * num_reductions) ) )
 end
 
 function tile_cost(X, U, T)

src/graphs.jl

@@ -1,4 +1,4 @@
-
+lv(x) = Expr(:(.), :LoopVectorization, QuoteNode(x))
 
 isdense(::Type{<:DenseArray}) = true
 
@@ -72,6 +72,7 @@ Base.@propagate_inbounds Base.getindex(lo::LoopOrder, i::Int) = lo.oporder[i]
 Base.@propagate_inbounds Base.getindex(lo::LoopOrder, i...) = lo.oporder[LinearIndices(size(lo))[i...]]
 
 # Must make it easy to iterate
+# outer_reductions is a vector of indixes (within operation vectors) of the reduction operation, eg the vmuladd op in a dot product
 struct LoopSet
     loops::Dict{Symbol,Loop} # sym === loops[sym].itersymbol
     opdict::Dict{Symbol,Operation}
@@ -80,7 +81,13 @@ struct LoopSet
     loop_order::LoopOrder
     # stridesets::Dict{ShortVector{Symbol},ShortVector{Symbol}}
     preamble::Expr # TODO: add preamble to lowering
-    includedarrays::Vector{Symbol}
+    includedarrays::Vector{Tuple{Symbol,Int}}
+end
+function includesarray(ls::LoopSet, array::Symbol)
+    for (a,i) ∈ ls.includedarrays
+        a === array && return i
+    end
+    -1
 end
 function LoopSet()
     LoopSet(
@@ -90,7 +97,7 @@ function LoopSet()
         Int[],
         LoopOrder(),
         Expr(:block,),
-        Symbol[]
+        Tuple{Symbol,Int}[]
     )
 end
 num_loops(ls::LoopSet) = length(ls.loops)
@@ -106,13 +113,13 @@ looprangesym(ls::LoopSet, s::Symbol) = ls.loops[s].rangesym
 getop(ls::LoopSet, s::Symbol) = ls.opdict[s]
 getop(ls::LoopSet, i::Int) = ls.operations[i + 1]
 
-function looprange(ls::LoopSet, s::Symbol, incr::Int = 1)
+function looprange(ls::LoopSet, s::Symbol, incr::Int = 1, mangledname::Symbol = s)
     loop = ls.loops[s]
     incr -= 1
     if iszero(incr)
-        Expr(:call, :<, s, loop.hintexact ? loop.rangehint : loop.rangesym)
+        Expr(:call, :<, mangledname, loop.hintexact ? loop.rangehint : loop.rangesym)
     else
-        Expr(:call, :<, s, loop.hintexact ? loop.rangehint - incr : Expr(:call, :-, loop.rangesym, incr))
+        Expr(:call, :<, mangledname, loop.hintexact ? loop.rangehint - incr : Expr(:call, :-, loop.rangesym, incr))
     end
 end
 function Base.length(ls::LoopSet, is::Symbol)
@@ -193,10 +200,10 @@ function add_loop!(ls::LoopSet, q::Expr, elementbytes::Int = 8)
         Base.push!(ls, q, elementbytes)
     end
 end
-function add_vptr!(ls::LoopSet, indexed::Symbol)
-    if indexed ∉ ls.includedarrays
-        push!(ls.includedarrays, indexed)
-        push!(ls.preamble.args, Expr(:(=), Symbol(:vptr_, indexed), Expr(:call, Expr(:(.), :VectorizationBase, QuoteNode(:stridedpointer)), indexed)))
+function add_vptr!(ls::LoopSet, indexed::Symbol, id::Int)
+    if includesarray(ls, indexed) < 0
+        push!(ls.includedarrays, (indexed, id))
+        push!(ls.preamble.args, Expr(:(=), Symbol("##vptr##_", indexed), Expr(:call, lv(:stridedpointer), indexed)))
     end
     nothing
 end
@@ -205,7 +212,7 @@ function add_load!(
     ls::LoopSet, var::Symbol, indexed::Symbol, indices::AbstractVector, elementbytes::Int = 8
 )
     op = Operation( length(operations(ls)), var, elementbytes, :getindex, memload, indices, [indexed], NOPARENTS )
-    add_vptr!(ls, indexed)
+    add_vptr!(ls, indexed, identifier(op))
     pushop!(ls, op, var)
 end
 function add_load_ref!(ls::LoopSet, var::Symbol, ex::Expr, elementbytes::Int = 8)
@@ -240,12 +247,12 @@ end
 ### if it is a literal, that literal is either var"##ZERO#Float##", var"##ONE#Float##", or has to have been assigned to var in the preamble.
 # if it is a literal, that literal has to have been assigned to var in the preamble.
 function add_constant!(ls::LoopSet, var::Symbol, elementbytes::Int = 8)
-    pushop!(ls, Operation(length(operations(ls)), var, elementbytes, Symbol("##CONSTANT##"), constant, NODEPENDENCY, NODEPENDENCY, NOPARENTS), var)
+    pushop!(ls, Operation(length(operations(ls)), var, elementbytes, Symbol("##CONSTANT##"), constant, NODEPENDENCY, Symbol[], NOPARENTS), var)
 end
 function add_constant!(ls::LoopSet, var, elementbytes::Int = 8)
     sym = gensym(:temp)
     push!(ls.preamble.args, Expr(:(=), sym, var))
-    pushop!(ls, Operation(length(operations(ls)), sym, elementbytes, Symbol("##CONSTANT##"), constant, NODEPENDENCY, NODEPENDENCY, NOPARENTS), sym)
+    pushop!(ls, Operation(length(operations(ls)), sym, elementbytes, Symbol("##CONSTANT##"), constant, NODEPENDENCY, Symbol[], NOPARENTS), sym)
 end
 # This version has loop dependencies. var gets assigned to sym when lowering.
 function add_constant!(ls::LoopSet, var::Symbol, deps::Vector{Symbol}, sym::Symbol = gensym(:constant), elementbytes::Int = 8)
@@ -260,9 +267,28 @@ end
 function pushparent!(parents::Vector{Operation}, deps::Vector{Symbol}, reduceddeps::Vector{Symbol}, parent::Operation)
     push!(parents, parent)
     mergesetv!(deps, loopdependencies(parent))
-    isload(parent) || mergesetv!(reduceddeps, reduceddependencies(parent))
+    if !(isload(parent) || isconstant(parent))
+        mergesetv!(reduceddeps, reduceddependencies(parent))
+    end
     nothing
 end
+function maybe_cse_load!(ls::LoopSet, expr::Expr, elementbytes::Int = 8)
+    if expr.head === :ref
+        array = first(expr.args)::Symbol
+        args = @view expr.args[2:end]
+    elseif expr.head === :call && first(expr.args) === :getindex
+        array = (expr.args[2])::Symbol
+        args = @view expr.args[3:end]
+    else
+        return add_operation!(ls, gensym(:temporary), expr, elementbytes)
+    end
+    id = includesarray(ls, array)
+    if id > 0
+        ls.operations[id]
+    else
+        add_load!( ls, gensym(:temporary), array, args, elementbytes )
+    end
+end
 function add_parent!(
     parents::Vector{Operation}, deps::Vector{Symbol}, reduceddeps::Vector{Symbol}, ls::LoopSet, var, elementbytes::Int = 8
 )
@@ -271,9 +297,8 @@ function add_parent!(
             # might add constant
             add_constant!(ls, var, elementbytes)
         end
-    elseif var isa Expr
-        temp = gensym(:temporary)
-        add_operation!(ls, temp, var, elementbytes)
+    elseif var isa Expr #CSE candidate
+        maybe_cse_load!(ls, var, elementbytes)
     else # assumed constant
         add_constant!(ls, var, elementbytes)
     end
@@ -283,13 +308,23 @@ function add_reduction!(
     parents::Vector{Operation}, deps::Vector{Symbol}, reduceddeps::Vector{Symbol}, ls::LoopSet, var::Symbol, elementbytes::Int = 8
 )
     get!(ls.opdict, var) do
-        p = add_constant!(ls, var, elementbytes)
-        push!(ls.outer_reductions, identifier(p))
-        p
+        add_constant!(ls, var, elementbytes)
+        # push!(ls.outer_reductions, identifier(p))
+        # p
     end
     # pushparent!(parents, deps, reduceddeps, parent)
 end
-
+function add_reduction_update_parent!(
+    parents::Vector{Operation}, deps::Vector{Symbol}, reduceddeps::Vector{Symbol}, ls::LoopSet,
+    var::Symbol, instr::Symbol, elementbytes::Int = 8
+)
+    parent = getop(ls, var)
+    setdiffv!(reduceddeps, deps, loopdependencies(parent))
+    pushparent!(parents, deps, reduceddeps, parent) # deps and reduced deps will not be disjoint
+    op = Operation(length(operations(ls)), var, elementbytes, instr, compute, deps, reduceddeps, parents)
+    parent.instruction === Symbol("##CONSTANT##") && push!(ls.outer_reductions, identifier(op))
+    pushop!(ls, op, var) # note this overwrites the entry in the operations dict, but not the vector
+end
 function add_compute!(ls::LoopSet, var::Symbol, ex::Expr, elementbytes::Int = 8)
     @assert ex.head === :call
     instr = instruction(first(ex.args))::Symbol
@@ -308,20 +343,18 @@ function add_compute!(ls::LoopSet, var::Symbol, ex::Expr, elementbytes::Int = 8)
         end
     end
     if reduction # arg[reduction] is the reduction
-        parent = getop(ls, var)
-        setdiffv!(reduceddeps, deps, loopdependencies(parent))
-        pushparent!(parents, deps, reduceddeps, parent) # deps and reduced deps will not be disjoint
-        # append!(reduceddependencies(parent), reduceddeps)
+        add_reduction_update_parent!(parents, deps, reduceddeps, ls, var, instr, elementbytes)
+    else
+        op = Operation(length(operations(ls)), var, elementbytes, instr, compute, deps, reduceddeps, parents)
+        pushop!(ls, op, var)
     end
-    op = Operation(length(operations(ls)), var, elementbytes, instr, compute, deps, reduceddeps, parents)
-    pushop!(ls, op, var) # note this overwrites the entry in the operations dict, but not the vector
 end
 function add_store!(
     ls::LoopSet, indexed::Symbol, var::Symbol, indices::AbstractVector, elementbytes::Int = 8
 )
     parent = getop(ls, var)
     op = Operation( length(operations(ls)), indexed, elementbytes, :setindex!, memstore, indices, reduceddependencies(parent), [parent] )
-    add_vptr!(ls, indexed)
+    add_vptr!(ls, indexed, identifier(op))
     pushop!(ls, op, var)
 end
 function add_store_ref!(ls::LoopSet, var::Symbol, ex::Expr, elementbytes::Int = 8)
@@ -414,11 +447,11 @@ function fillorder!(ls::LoopSet, order::Vector{Symbol}, loopistiled::Bool)
             isunrolled = (unrolled ∈ loopdependencies(op)) + 1
             istiled = (loopistiled ? (tiled ∈ loopdependencies(op)) : false) + 1
             optype = Int(op.node_type) + 1
-            after_loop = (length(reduceddependencies(op)) > 0) + 1
+            after_loop = isload(op) ? 1 : (length(reduceddependencies(op)) > 0) + 1
             push!(lo[optype,isunrolled,istiled,after_loop,_n], op)
         end
     end    
-    @show 3, ro, order
+    # 3, ro, order
 end