Fixed bug that made it possible for an op to point to the wrong ref after condense-> reconstucting, added CSE for constants, and improved cost evaluation of tiling. Fixes #45.

Author: chriselrod

src/add_compute.jl

@@ -118,9 +118,9 @@ function add_reduction_update_parent!(
         reductsym = gensym(:reduction)
         reductinit = add_constant!(ls, gensym(:reductzero), loopdependencies(parent), reductsym, elementbytes, :numericconstant)
         if reduct_zero === :zero
-            push!(ls.preamble_zeros, identifier(reductinit))
+            push!(ls.preamble_zeros, (identifier(reductinit), IntOrFloat))
         elseif reduct_zero === :one
-            push!(ls.preamble_ones, identifier(reductinit))
+            push!(ls.preamble_ones, (identifier(reductinit), IntOrFloat))
         else
             if reductzero === :true || reductzero === :false
                 pushpreamble!(ls, Expr(:(=), name(reductinit), reductzero))
@@ -185,11 +185,7 @@ function add_compute!(
                 add_parent!(parents, deps, reduceddeps, ls, arg, elementbytes, position)
             end
         elseif arg ∈ ls.loopsymbols
-            loopsym = gensym(arg)
-            pushpreamble!(ls, Expr(:(=), loopsym, LoopValue()))
-            loopsymop = add_simple_load!(ls, gensym(loopsym), ArrayReference(loopsym, [arg]), elementbytes, false)
-            push!(ls.syms_aliasing_refs, name(loopsymop))
-            push!(ls.refs_aliasing_syms, loopsymop.ref)
+            loopsymop = add_loopvalue!(ls, arg, elementbytes)
             pushparent!(parents, deps, reduceddeps, loopsymop)
         else
             add_parent!(parents, deps, reduceddeps, ls, arg, elementbytes, position)

src/add_constants.jl

@@ -8,6 +8,33 @@ function add_constant!(ls::LoopSet, var, elementbytes::Int = 8)
     pushpreamble!(ls, Expr(:(=), sym, var))
     add_constant!(ls, sym, elementbytes)
 end
+function add_constant!(ls::LoopSet, var::Number, elementbytes::Int = 8)
+    op = Operation(length(operations(ls)), gensym(:loopconstnumber), elementbytes, LOOPCONSTANT, constant, NODEPENDENCY, Symbol[], NOPARENTS)
+    ops = operations(ls)
+    typ = var isa Integer ? HardInt : HardFloat
+    if iszero(var)
+        for (id,typ_) ∈ ls.preamble_zeros
+            (instruction(ops[id]) === LOOPCONSTANT && typ == typ_) && return ops[id]
+        end
+        push!(ls.preamble_zeros, (identifier(op),typ))
+    elseif isone(var)
+        for (id,typ_) ∈ ls.preamble_ones
+            (instruction(ops[id]) === LOOPCONSTANT && typ == typ_) && return ops[id]
+        end
+        push!(ls.preamble_ones, (identifier(op),typ))
+    elseif var isa Integer
+        for (id,ivar) ∈ ls.preamble_symint
+            (instruction(ops[id]) === LOOPCONSTANT && ivar == var) && return ops[id]
+        end
+        push!(ls.preamble_symint, (identifier(op), var))
+    else#if var isa FloatX
+        for (id,fvar) ∈ ls.preamble_symfloat
+            (instruction(ops[id]) === LOOPCONSTANT && fvar == var) && return ops[id]
+        end
+        push!(ls.preamble_symfloat, (identifier(op), var))
+    end
+    pushop!(ls, op)
+end
 function add_constant!(ls::LoopSet, var::Symbol, mpref::ArrayReferenceMetaPosition, elementbytes::Int)
     op = Operation(length(operations(ls)), var, elementbytes, LOOPCONSTANT, constant, NODEPENDENCY, Symbol[], NOPARENTS, mpref.mref)
     add_vptr!(ls, op)
@@ -36,12 +63,6 @@ function add_constant!(
     ls::LoopSet, value::Number, deps::Vector{Symbol}, assignedsym::Symbol, elementbytes::Int
 )
     op = add_constant!(ls, gensym(Symbol(value)), deps, assignedsym, elementbytes, :numericconstant)
-    if iszero(value)
-        push!(ls.preamble_zeros, identifier(op))
-    elseif isone(value)
-        push!(ls.preamble_ones, identifier(op))
-    else
-        pushpreamble!(ls, op, value)
-    end
+    pushpreamble!(ls, op, value)
     op
 end

src/add_loads.jl

@@ -54,6 +54,22 @@ function add_load_getindex!(ls::LoopSet, var::Symbol, ex::Expr, elementbytes::In
     add_load!(ls, var, array, rawindices, elementbytes)
 end
 
+function add_loopvalue!(ls::LoopSet, arg::Symbol, elementbytes::Int)
+    # check for CSE opportunity
+    loopsym = Symbol("##LOOPSYMBOL##", arg)
+    ar = ArrayReference(loopsym, [arg])
+    for op ∈ operations(ls)
+        if isload(op) && op.ref.ref == ar
+            return op
+        end
+    end
+    pushpreamble!(ls, Expr(:(=), loopsym, LoopValue()))
+    loopsymop = add_simple_load!(ls, gensym(loopsym), ar, elementbytes, false)
+    push!(ls.syms_aliasing_refs, name(loopsymop))
+    push!(ls.refs_aliasing_syms, loopsymop.ref)
+    loopsymop
+end
+
 
 struct LoopValue end
 @inline VectorizationBase.stridedpointer(::LoopValue) = LoopValue()

src/broadcast.jl

@@ -73,7 +73,7 @@ function add_broadcast!(
     # set Cₘₙ = 0
     # setC = add_constant!(ls, zero(promote_type(recursive_eltype(A), recursive_eltype(B))), cloopsyms, mC, elementbytes)
     setC = add_constant!(ls, gensym(:zero), cloopsyms, mC, elementbytes, :numericconstant)
-    push!(ls.preamble_zeros, identifier(setC))
+    push!(ls.preamble_zeros, (identifier(setC), IntOrFloat))
     # compute Cₘₙ += Aₘₖ * Bₖₙ
     reductop = Operation(
         ls, mC, elementbytes, :vmuladd, compute, reductdeps, Symbol[k], Operation[loadA, loadB, setC]

src/condense_loopset.jl

@@ -55,10 +55,10 @@ isload(os::OperationStruct) = os.node_type == memload
 isstore(os::OperationStruct) = os.node_type == memstore
 iscompute(os::OperationStruct) = os.node_type == compute
 isconstant(os::OperationStruct) = os.node_type == constant
-function findmatchingarray(ls::LoopSet, array::Symbol)
+function findmatchingarray(ls::LoopSet, mref::ArrayReferenceMeta)
     id = 0x01
-    for as ∈ ls.refs_aliasing_syms
-        vptr(as) === array && return id
+    for r ∈ ls.refs_aliasing_syms
+        r == mref && return id
         id += 0x01
     end
     0x00
@@ -96,7 +96,7 @@ function OperationStruct!(varnames::Vector{Symbol}, ls::LoopSet, op::Operation)
     rd = reduceddeps_uint(ls, op)
     cd = childdeps_uint(ls, op)
     p = parents_uint(ls, op)
-    array = accesses_memory(op) ? findmatchingarray(ls, vptr(op.ref)) : 0x00
+    array = accesses_memory(op) ? findmatchingarray(ls, op.ref) : 0x00
     OperationStruct(
         ld, rd, cd, p, op.node_type, array, findindoradd!(varnames, name(op))
     )
@@ -228,7 +228,7 @@ end
 end
 
 # Try to condense in type stable manner
-function generate_call(ls::LoopSet, IUT)
+function generate_call(ls::LoopSet, IUT, debug::Bool = false)
     operation_descriptions = Expr(:curly, :Tuple)
     varnames = Symbol[]
     for op ∈ operations(ls)
@@ -243,12 +243,19 @@ function generate_call(ls::LoopSet, IUT)
     argmeta = argmeta_and_consts_description(ls, arraysymbolinds)
     loop_bounds = loop_boundaries(ls)
     inline, U, T = IUT
-    if inline
-        q = Expr(:call, lv(:_avx_!), Expr(:call, Expr(:curly, :Val, (U,T))), operation_descriptions, arrayref_descriptions, argmeta, loop_bounds)
+    if inline | debug
+        func = debug ? lv(:_avx_loopset) : lv(:_avx_!)
+        q = Expr(
+            :call, func, Expr(:call, Expr(:curly, :Val, (U,T))),
+            operation_descriptions, arrayref_descriptions, argmeta, loop_bounds
+        )
         foreach(ref -> push!(q.args, vptr(ref)), ls.refs_aliasing_syms)
     else
         arraydescript = Expr(:tuple)
-        q = Expr(:call, lv(:__avx__!), Expr(:call, Expr(:curly, :Val, (U,T))), operation_descriptions, arrayref_descriptions, argmeta, loop_bounds, arraydescript)
+        q = Expr(
+            :call, lv(:__avx__!), Expr(:call, Expr(:curly, :Val, (U,T))),
+            operation_descriptions, arrayref_descriptions, argmeta, loop_bounds, arraydescript
+        )
         for array ∈ ls.includedactualarrays
             push!(q.args, Expr(:call, lv(:unwrap_array), array))
             push!(arraydescript.args, Expr(:call, lv(:array_wrapper), array))
@@ -361,6 +368,11 @@ function setup_call_inline(ls::LoopSet, U = zero(Int8), T = zero(Int8))
     append!(ls.preamble.args, q.args)
     ls.preamble
 end
+function setup_call_debug(ls::LoopSet)
+    # avx_loopset(instr, ops, arf, AM, LB, vargs)
+    pushpreamble!(ls, generate_call(ls, (true,zero(Int8),zero(Int8)), true))
+    ls.preamble
+end
 function setup_call(ls::LoopSet, inline = Int8(2), U = zero(Int8), T = zero(Int8))
     # We outline/inline at the macro level by creating/not creating an anonymous function.
     # The old API instead was based on inlining or not inline the generated function, but

src/constructors.jl

@@ -151,4 +151,7 @@ macro _avx(arg, q)
 end
 
 
+macro avx_debug(q)
+    esc(LoopVectorization.setup_call_debug(LoopSet(q, __module__)))
+end
 

src/determinestrategy.jl

@@ -91,12 +91,12 @@ function evaluate_cost_unroll(
             rd = reduceddependencies(op)
             hasintersection(rd, nested_loop_syms[1:end-length(rd)]) && return Inf
             included_vars[id] = true
-            
+            # @show op first(cost(op, vectorized, Wshift, size_T)), iter
             total_cost += iter * first(cost(op, vectorized, Wshift, size_T))
             total_cost > max_cost && return total_cost # abort if more expensive; we only want to know the cheapest
         end
     end
-    total_cost
+    total_cost + stride_penalty(ls, order)
 end
 
 # only covers vectorized ops; everything else considered lifted?
@@ -198,7 +198,10 @@ function determine_unroll_factor(
 end
 
 function tile_cost(X, U, T, UL, TL)
-    X[1] + X[4] + X[2] * (num_iterations(TL, T)/TL) + X[3] * (num_iterations(UL, U)/UL)
+    Tfactor = (num_iterations(TL, T)/TL)
+    Ufactor = (num_iterations(UL, U)/UL)
+    # X[1]*Tfactor*Ufactor + X[4] + X[2] * Tfactor + X[3] * Ufactor
+    X[1] + X[4] + X[2] * Tfactor + X[3] * Ufactor
 end
 function solve_tilesize(X, R, UL, TL)
     @inbounds any(iszero, (R[1],R[2],R[3])) && return -1,-1,Inf #solve_smalltilesize(X, R, Umax, Tmax)
@@ -338,7 +341,7 @@ function stride_penalty(ls::LoopSet, op::Operation, order::Vector{Symbol})
     num_loops = length(order)
     contigsym = first(loopdependencies(op))
     contigsym == Symbol("##DISCONTIGUOUSSUBARRAY##") && return 0
-    iter = 0
+    iter = 1
     for i ∈ 0:num_loops - 1
         loopsym = order[num_loops - i]
         loopsym === contigsym && return iter
@@ -386,34 +389,29 @@ function evaluate_cost_tile(
     reg_pressure = reg_pres_buf(ls)
     # @inbounds reg_pressure[2] = 1
     # @inbounds reg_pressure[3] = 1
-    unrollediter = length(ls, unrolled)
-    tilediter = length(ls, tiled)
-    unrollediter = unrolled === vectorized ? num_iterations(unrollediter, W) : unrollediter # tiled cannot be vectorized, so do not check
-    iter::Int = tilediter * unrollediter
+    iter::Int = 1
     for n ∈ 1:N
         itersym = order[n]
         # Add to set of defined symbles
         push!(nested_loop_syms, itersym)
-        stepsize = 1
-        if n > 2
-            itersymlooplen = length(ls, itersym)
-            iter *= itersym === vectorized ? num_iterations(itersymlooplen, W) : itersymlooplen
-        end
+        looplength = length(ls, itersym)
+        iter *= itersym === vectorized ? num_iterations(looplength, W) : looplength
         # check which vars we can define at this level of loop nest
         for (id, op) ∈ enumerate(ops)
             # isconstant(op) && continue
             # @assert id == identifier(op)+1 # testing, for now
             # won't define if already defined...
             included_vars[id] && continue
             # it must also be a subset of defined symbols
-            loopdependencies(op) ⊆ nested_loop_syms || continue
+            all(ld -> ld ∈ nested_loop_syms, loopdependencies(op)) || continue
             # # @show nested_loop_syms
             # # @show reduceddependencies(op)
             rd = reduceddependencies(op)
             hasintersection(rd, nested_loop_syms[1:end-length(rd)]) && return 0,0,Inf
             included_vars[id] = true
             unrolledtiled[1,id] = unrolled ∈ loopdependencies(op)
             unrolledtiled[2,id] = tiled ∈ loopdependencies(op)
+            # @show op iter, unrolledtiled[:,id]
             iters[id] = iter
             innerloop ∈ loopdependencies(op) && set_upstream_family!(descendentsininnerloop, op, true)
         end

src/graphs.jl

@@ -140,6 +140,8 @@ Base.size(lo::LoopOrder) = (2,2,2,length(lo.loopnames))
 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...]]
 
+@enum NumberType::Int8 HardInt HardFloat IntOrFloat INVALID
+
 # Must make it easy to iterate
 # outer_reductions is a vector of indices (within operation vectors) of the reduction operation, eg the vmuladd op in a dot product
 # O(N) search is faster at small sizes
@@ -154,8 +156,8 @@ struct LoopSet
     preamble_symsym::Vector{Tuple{Int,Symbol}}
     preamble_symint::Vector{Tuple{Int,Int}}
     preamble_symfloat::Vector{Tuple{Int,Float64}}
-    preamble_zeros::Vector{Int}
-    preamble_ones::Vector{Int}
+    preamble_zeros::Vector{Tuple{Int,NumberType}}
+    preamble_ones::Vector{Tuple{Int,NumberType}}
     includedarrays::Vector{Symbol}
     includedactualarrays::Vector{Symbol}
     syms_aliasing_refs::Vector{Symbol}
@@ -199,21 +201,41 @@ function pushpreamble!(ls::LoopSet, op::Operation, v::Symbol)
     end
     nothing
 end
-pushpreamble!(ls::LoopSet, op::Operation, v::Integer) = push!(ls.preamble_symint, (identifier(op),convert(Int,v)))
-pushpreamble!(ls::LoopSet, op::Operation, v::Real) = push!(ls.preamble_symfloat, (identifier(op),convert(Float64,v)))
+function pushpreamble!(ls::LoopSet, op::Operation, v::Number)
+    typ = v isa Integer ? HardInt : HardFloat
+    id = identifier(op)
+    if iszero(v)
+        push!(ls.preamble_zeros, (id, typ))
+    elseif isone(v)
+        push!(ls.preamble_ones, (id, typ))
+    elseif v isa Integer
+        push!(ls.preamble_symint, (id, convert(Int,v)))
+    else
+        push!(ls.preamble_symfloat, (id, convert(Float64,v)))
+    end
+end
 pushpreamble!(ls::LoopSet, ex::Expr) = push!(ls.preamble.args, ex)
 function pushpreamble!(ls::LoopSet, op::Operation, RHS::Expr)
     c = gensym(:licmconst)
     if RHS.head === :call && first(RHS.args) === :zero
-        push!(ls.preamble_zeros, identifier(op))
+        push!(ls.preamble_zeros, (identifier(op), IntOrFloat))
     elseif RHS.head === :call && first(RHS.args) === :one
-        push!(ls.preamble_ones, identifier(op))
+        push!(ls.preamble_ones, (identifier(op), IntOrFloat))
     else
         pushpreamble!(ls, Expr(:(=), c, RHS))
         pushpreamble!(ls, op, c)
     end
     nothing
 end
+function zerotype(ls::LoopSet, op::Operation)
+    opid = identifier(op)
+    for (id,typ) ∈ ls.preamble_zeros
+        id == opid && return typ
+    end
+    INVALID
+end
+
+
 
 includesarray(ls::LoopSet, array::Symbol) = array ∈ ls.includedarrays
 
@@ -425,7 +447,7 @@ function add_operation!(
         elseif f === :zero || f === :one
             c = gensym(f)
             op = add_constant!(ls, c, ls.loopsymbols[1:position], LHS, elementbytes, :numericconstant)
-            push!(f === :zero ? ls.preamble_zeros : ls.preamble_ones, identifier(op))
+            push!(f === :zero ? ls.preamble_zeros : ls.preamble_ones, (identifier(op), IntOrFloat))
             op
         else
             add_compute!(ls, LHS, RHS, elementbytes, position)
@@ -452,7 +474,7 @@ function add_operation!(
         elseif f === :zero || f === :one
             c = gensym(f)
             op = add_constant!(ls, c, ls.loopsymbols[1:position], LHS_sym, elementbytes, :numericconstant)
-            push!(f === :zero ? ls.preamble_zeros : ls.preamble_ones, identifier(op))
+            push!(f === :zero ? ls.preamble_zeros : ls.preamble_ones, (identifier(op), IntOrFloat))
             op
         else
             add_compute!(ls, LHS_sym, RHS, elementbytes, position, LHS_ref)