Reworked handling of loop constants, reducing to two primary strategies (one for loop constants assinged outside of the loops, another for assignments nested within the loops.

chriselrod · chriselrod · commit ce3147e55a83 · 2020-01-23T14:33:08.000-05:00
diff --git a/benchmark/looptests.jl b/benchmark/looptests.jl
@@ -11,7 +11,7 @@ function jgemm!(C, A, B)
     end
 end
 @inline function gemmavx!(C, A, B)
-    @avx inline=true for i ∈ 1:size(A,1), j ∈ 1:size(B,2)
+    @avx for i ∈ 1:size(A,1), j ∈ 1:size(B,2)
         Cᵢⱼ = zero(eltype(C))
         for k ∈ 1:size(A,2)
             Cᵢⱼ += A[i,k] * B[k,j]
diff --git a/src/LoopVectorization.jl b/src/LoopVectorization.jl
@@ -36,7 +36,7 @@ include("constructors.jl")
 
 export @_avx, _avx, @_avx_, avx_!
 
-include("precompile.jl")
-_precompile_()
+# include("precompile.jl")
+# _precompile_()
 
 end # module
diff --git a/src/add_compute.jl b/src/add_compute.jl
@@ -94,7 +94,15 @@ function add_reduction_update_parent!(
         reduct_zero = reduction_zero(instrclass)
         reductcombine = reduction_scalar_combine(instrclass)
         reductsym = gensym(:reduction)
-        reductinit = add_constant!(ls, Expr(:call, reduct_zero, ls.T), loopdependencies(parent), reductsym, reduct_zero, elementbytes)
+        reductinit = add_constant!(ls, gensym(:reductzero), loopdependencies(parent), reductsym, elementbytes, :numericconstant)
+        if reduct_zero === :zero
+            push!(ls.preamble_zeros, identifier(reductinit))
+        elseif reduct_zero === :one
+            push!(ls.preamble_ones, identifier(reductinit))
+        else
+            pushpreamble!(ls, Expr(:(=), name(reductinit), reductzero))
+            pushpreamble!(ls, op, name, reductinit)
+        end
         if isconstant(parent) && reduct_zero === parent.instruction.mod #we can use parent op as initialization.
             reductcombine = reduction_combine_to(instrclass)
         end
diff --git a/src/add_constants.jl b/src/add_constants.jl
@@ -1,38 +1,45 @@
 function add_constant!(ls::LoopSet, var::Symbol, elementbytes::Int = 8)
     op = Operation(length(operations(ls)), var, elementbytes, LOOPCONSTANT, constant, NODEPENDENCY, Symbol[], NOPARENTS)
-    pushpreamble!(ls, op, mangledvar(op))
+    pushpreamble!(ls, op, var)
     pushop!(ls, op, var)
 end
 function add_constant!(ls::LoopSet, var, elementbytes::Int = 8)
-    sym = gensym(:temp)
-    op = Operation(length(operations(ls)), sym, elementbytes, LOOPCONSTANT, constant, NODEPENDENCY, Symbol[], NOPARENTS)
-    temp = gensym(:intermediateconst)
-    pushpreamble!(ls, Expr(:(=), temp, var))
-    pushpreamble!(ls, op, temp)
-    pushop!(ls, op, sym)
+    sym = gensym(:loopconstant)
+    pushpreamble!(ls, Expr(:(=), sym, var))
+    add_constant!(ls, sym, elementbytes)
 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)
     temp = gensym(:intermediateconstref)
     pushpreamble!(ls, Expr(:(=), temp, Expr(:call, lv(:load), mpref.mref.ptr, mem_offset(op, UnrollArgs(zero(Int32), Symbol(""), Symbol(""), nothing)))))
     pushpreamble!(ls, op, temp)
-    pushop!(ls, op, var)
+    pushop!(ls, op, temp)
 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), f::Symbol = Symbol(""), elementbytes::Int = 8)
-    # length(deps) == 0 && push!(ls.preamble.args, Expr(:(=), sym, var))
-    pushop!(ls, Operation(length(operations(ls)), sym, elementbytes, Instruction(f,var), constant, deps, NODEPENDENCY, NOPARENTS), sym)
+function add_constant!(
+    ls::LoopSet, value::Symbol, deps::Vector{Symbol}, assignedsym::Symbol = gensym(:constant), elementbytes::Int = 8, f::Symbol = Symbol("")
+)
+    op = Operation(length(operations(ls)), assignedsym, elementbytes, Instruction(f, value), constant, deps, NODEPENDENCY, NOPARENTS)
+    pushop!(ls, op, assignedsym)
+end
+function add_constant!(
+    ls::LoopSet, value, deps::Vector{Symbol}, assignedsym::Symbol = gensym(:constant), elementbytes::Int = 8, f::Symbol = Symbol("")
+)
+    intermediary = gensym(:intermediate) # hack, passing meta info here
+    pushpreamble!(ls, Expr(:(=), intermediary, value))
+    add_constant!(ls, intermediary, deps, assignedsym, f, elementbytes)
 end
-
 function add_constant!(
-    ls::LoopSet, var, deps::Vector{Symbol}, sym::Symbol = gensym(:constant), f::Symbol = Symbol(""), elementbytes::Int = 8
+    ls::LoopSet, value::Number, deps::Vector{Symbol}, assignedsym::Symbol = gensym(:constant), elementbytes::Int = 8
 )
-    sym2 = gensym(:temp) # hack, passing meta info here
-    op = Operation(length(operations(ls)), sym, elementbytes, Instruction(f, sym2), constant, deps, NODEPENDENCY, NOPARENTS)
-    # @show f, sym, name(op), mangledvar(op)
-    # temp = gensym(:temp2)
-    # pushpreamble!(ls, Expr(:(=), temp, var))
-    pushpreamble!(ls, op, var)#temp)
-    pushop!(ls, op, sym)
+    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
+    op
 end
diff --git a/src/broadcast.jl b/src/broadcast.jl
@@ -72,7 +72,7 @@ function add_broadcast!(
     loadB = add_broadcast!(ls, gensym(:B), mB, bloopsyms, B, elementbytes)
     # 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, :zero, elementbytes)
+    setC = add_constant!(ls, gensym(:zero), cloopsyms, mC, elementbytes, :numericconstant)
     push!(ls.preamble_zeros, identifier(setC))
     # compute Cₘₙ += Aₘₖ * Bₖₙ
     reductop = Operation(
@@ -136,11 +136,9 @@ function add_broadcast!(
     add_simple_load!(ls, destname, ArrayReference(bcname, @view(loopsyms[1:N])), elementbytes)
 end
 function add_broadcast!(
-    ls::LoopSet, destname::Symbol, bcname::Symbol, loopsyms::Vector{Symbol}, ::Type{T}, elementbytes::Int = 8
-) where {T<:Union{Integer,Float32,Float64}}
-    op = add_constant!(ls, destname, elementbytes) # or replace elementbytes with sizeof(T) ? u
-    pushpreamble!(ls, Expr(:(=), mangledvar(op), bcname))
-    op
+    ls::LoopSet, ::Symbol, bcname::Symbol, loopsyms::Vector{Symbol}, ::Type{T}, elementbytes::Int = 8
+) where {T<:Number}
+    add_constant!(ls, bcname, elementbytes) # or replace elementbytes with sizeof(T) ? u
 end
 function add_broadcast!(
     ls::LoopSet, destname::Symbol, bcname::Symbol, loopsyms::Vector{Symbol},
@@ -168,7 +166,7 @@ function add_broadcast!(
     reduceddeps = Symbol[]
     for (i,arg) ∈ enumerate(args)
         argname = gensym(:arg)
-        pushpreamble!(ls, Expr(:macrocall, Symbol("@inbounds"), LineNumberNode(@__LINE__,@__FILE__), Expr(:(=), argname, Expr(:ref, bcargs, i))))
+        pushpreamble!(ls, Expr(:(=), argname, Expr(:macrocall, Symbol("@inbounds"), LineNumberNode(@__LINE__,@__FILE__), Expr(:ref, bcargs, i))))
         # dynamic dispatch
         parent = add_broadcast!(ls, gensym(:temp), argname, loopsyms, arg, elementbytes)::Operation
         pushparent!(parents, deps, reduceddeps, parent)
diff --git a/src/condense_loopset.jl b/src/condense_loopset.jl
@@ -116,7 +116,7 @@ function loop_boundaries(ls::LoopSet)
         elseif stopexact
             Expr(:call, Expr(:curly, lv(:StaticUpperUnitRange), loop.stophint), loop.startsym)
         else
-            Expr(:call, Expr(:call, :(:), loop.startsym, loop.stopsym))
+            Expr(:call, :(:), loop.startsym, loop.stopsym)
         end
         push!(lbd.args, lexpr)
     end
@@ -159,6 +159,14 @@ function loopset_return_value(ls::LoopSet, ::Val{extract}) where {extract}
     end
 end
 
+function add_reassigned_syms!(q::Expr, ls::LoopSet)
+    for op ∈ operations(ls)
+        if isconstant(op)
+            instr = instruction(op)
+            (instr == LOOPCONSTANT || instr.mod === :numericconstant) || push!(q.args, instr.instr)
+        end
+    end
+end
 
 # Try to condense in type stable manner
 function generate_call(ls::LoopSet, IUT)
@@ -177,14 +185,14 @@ function generate_call(ls::LoopSet, IUT)
     loop_bounds = loop_boundaries(ls)
 
     q = Expr(:call, lv(:_avx_!), Expr(:call, Expr(:curly, :Val, IUT)), operation_descriptions, arrayref_descriptions, argmeta, loop_bounds)
-
     foreach(ref -> push!(q.args, vptr(ref)), ls.refs_aliasing_syms)
     foreach(is -> push!(q.args, last(is)), ls.preamble_symsym)
     append!(q.args, arraysymbolinds)
+    add_reassigned_syms!(q, ls)
     q
 end
 
-function setup_call(ls::LoopSet, inline = one(Int8), U = zero(Int8), T = zero(Int8))
+function setup_call(ls::LoopSet, inline = Int8(2), U = zero(Int8), T = zero(Int8))
     call = generate_call(ls, (inline,U,T))
     hasouterreductions = length(ls.outer_reductions) > 0
     if hasouterreductions
diff --git a/src/costs.jl b/src/costs.jl
@@ -23,6 +23,7 @@ function Base.isless(instr1::Instruction, instr2::Instruction)
         isless(instr1.mod, instr2.mod)
     end
 end
+Base.isequal(ins1::Instruction, ins2::Instruction) = (ins1.instr === ins2.instr) && (ins1.mod === ins2.mod)
 
 const LOOPCONSTANT = Instruction(gensym())
 
diff --git a/src/determinestrategy.jl b/src/determinestrategy.jl
@@ -405,13 +405,12 @@ function evaluate_cost_tile(
     end
     for (id, op) ∈ enumerate(ops)
         iters[id] == -99.9 && continue
-        descendentsininnerloop[id] || continue
+        opisininnerloop = descendentsininnerloop[id]
         isunrolled = unrolledtiled[1,id]
         istiled = unrolledtiled[2,id]
         rt, lat, rp = cost(op, vectorized, Wshift, size_T)
-            # @show instruction(op), rt, lat, rp, iter
+        rp = opisininnerloop ? rp : 0 # we only care about register pressure within the inner most loop
         rt *= iters[id]
-            # @show isunrolled, istiled
         if isunrolled && istiled # no cost decrease; cost must be repeated
             cost_vec[1] += rt
             reg_pressure[1] += rp
diff --git a/src/graphs.jl b/src/graphs.jl
@@ -261,18 +261,12 @@ looprangehint(ls::LoopSet, s::Symbol) = length(getloop(ls, s))
 looprangesym(ls::LoopSet, s::Symbol) = getloop(ls, s).rangesym
 function getop(ls::LoopSet, var::Symbol, elementbytes::Int = 8)
     get!(ls.opdict, var) do
-        # might add constant
-        op = add_constant!(ls, var, elementbytes)
-        pushpreamble!(ls, op, var)
-        op
+        add_constant!(ls, var, elementbytes)
     end
 end
 function getop(ls::LoopSet, var::Symbol, deps, elementbytes::Int = 8)
     get!(ls.opdict, var) do
-        # might add constant
-        op = add_constant!(ls, var, deps, gensym(:constant), Symbol(""), elementbytes)
-        pushpreamble!(ls, op, var)
-        op
+        add_constant!(ls, var, deps, gensym(:constant), elementbytes)
     end
 end
 getop(ls::LoopSet, i::Int) = ls.operations[i + 1]
@@ -370,7 +364,7 @@ function register_single_loop!(ls::LoopSet, looprange::Expr)
         else # neither are integers
             L = add_loop_bound!(ls, itersym, lower, false)
             U = add_loop_bound!(ls, itersym, upper, true)
-            Loop(itersym, L, N)
+            Loop(itersym, L, U)
         end
     elseif f === :eachindex
         N = gensym(Symbol(:loop, itersym))
@@ -434,8 +428,7 @@ function add_operation!(
             add_load_getindex!(ls, LHS, RHS, elementbytes)
         elseif f === :zero || f === :one
             c = gensym(f)
-            # pushpreamble!(ls, Expr(:(=), c, RHS))
-            op = add_constant!(ls, c, copy(ls.loopsymbols), LHS, f, elementbytes)
+            op = add_constant!(ls, c, copy(ls.loopsymbols), LHS, elementbytes, :numericconstant)
             push!(f === :zero ? ls.preamble_zeros : ls.preamble_ones, identifier(op))
             op
         else
@@ -459,8 +452,7 @@ function add_operation!(
             add_load!(ls, LHS_sym, LHS_ref, elementbytes)
         elseif f === :zero || f === :one
             c = gensym(f)
-            # pushpreamble!(ls, Expr(:(=), c, RHS))
-            op = add_constant!(ls, c, copy(ls.loopsymbols), LHS_sym, f, elementbytes)
+            op = add_constant!(ls, c, copy(ls.loopsymbols), LHS_sym, elementbytes, :numericconstant)
             push!(f === :zero ? ls.preamble_zeros : ls.preamble_ones, identifier(op))
             op
         else
@@ -488,33 +480,7 @@ function Base.push!(ls::LoopSet, ex::Expr, elementbytes::Int = 8)
             if RHS isa Expr
                 add_operation!(ls, LHS, RHS, elementbytes)
             else
-                deps = copy(ls.loopsymbols)
-                if RHS isa Number
-                    fisone = false
-                    fiszero = false
-                    instr = if iszero(RHS)
-                        fiszero = true
-                        :zero
-                    elseif isone(RHS)
-                        fisone = true
-                        :one
-                    else
-                        :numericalconstant
-                    end
-                    op = add_constant!(ls, gensym(instr), deps, LHS, instr, elementbytes)
-                    if fiszero
-                        push!(ls.preamble_zeros, identifier(op))
-                    elseif fisone
-                        push!(ls.preamble_ones, identifier(op))
-                    else
-                        pushpreamble!(ls, op, RHS)
-                    end
-                    op
-                elseif RHS isa Symbol
-                    add_constant!(ls, RHS, deps, LHS, :constsym, elementbytes)
-                else
-                    add_constant!(ls, RHS, deps, LHS, :constmisc, elementbytes)
-                end
+                add_constant!(ls, RHS, copy(ls.loopsymbols), LHS, elementbytes)
             end
         elseif LHS isa Expr
             @assert LHS.head === :ref
diff --git a/src/lowering.jl b/src/lowering.jl
@@ -877,36 +877,37 @@ end
 @inline sizeequivalentint(::Type{Float16}, x::Int64) = Int16(x)
 @inline sizeequivalentint(::Type{Float16}, x::Int32) = Int16(x)
 
+function setop!(ls, op, val)
+    if instruction(op) === LOOPCONSTANT# && mangledvar(op) !== val
+        pushpreamble!(ls, Expr(:(=), mangledvar(op), val))
+    else
+        pushpreamble!(ls, Expr(:(=), instruction(op).instr, val))
+    end
+    nothing
+end
+function setconstantop!(ls, op, val)
+    if instruction(op) === LOOPCONSTANT# && mangledvar(op) !== val
+        pushpreamble!(ls, Expr(:(=), mangledvar(op), val))
+    end
+    nothing
+end
+
 function lower_licm_constants!(ls::LoopSet)
     ops = operations(ls)
-    for (id,sym) ∈ ls.preamble_symsym
-        op = ops[id]
-        mv = mangledvar(op)
-        # mv === sym || pushpreamble!(ls, Expr(:(=), instruction(op).instr, sym))
-        if mv !== sym
-            if length(ls.includedarrays) == 0 && instruction(op).mod === Symbol("")
-                pushpreamble!(ls, Expr(:(=), instruction(op).instr, sym))
-            else
-                pushpreamble!(ls, Expr(:(=), mv, sym))
-            end
-            # pushpreamble!(ls, instruction(op) === LOOPCONSTANT ? Expr(:(=), mv, sym) : Expr(:(=), instruction(op).instr, sym))
-        end
+    for (id, sym) ∈ ls.preamble_symsym
+        setconstantop!(ls, ops[id], sym)
     end
     for (id,intval) ∈ ls.preamble_symint
-        op = ops[id]
-        pushpreamble!(ls, Expr(:(=), mangledvar(op), Expr(:call, lv(:sizeequivalentint), ls.T, intval)))
+        setop!(ls, ops[id], Expr(:call, lv(:sizeequivalentint), ls.T, intval))
     end
     for (id,floatval) ∈ ls.preamble_symfloat
-        op = ops[id]
-        pushpreamble!(ls, Expr(:(=), mangledvar(op), Expr(:call, lv(:sizeequivalentfloat), ls.T, floatval)))
+        setop!(ls, ops[id], Expr(:call, lv(:sizeequivalentfloat), ls.T, intval))
     end
     for id ∈ ls.preamble_zeros
-        op = ops[id]
-        pushpreamble!(ls, Expr(:(=), instruction(op).instr, Expr(:call, :zero, ls.T)))
+        setop!(ls, ops[id], Expr(:call, :zero, ls.T))
     end
     for id ∈ ls.preamble_ones
-        op = ops[id]
-        pushpreamble!(ls, Expr(:(=), instruction(op).instr, Expr(:call, :one, ls.T)))
+        setop!(ls, ops[id], Expr(:call, :one, ls.T))
     end
 end
 function setup_preamble!(ls::LoopSet, W::Symbol, typeT::Symbol, vectorized::Symbol, unrolled::Symbol, tiled::Symbol, U::Int)
diff --git a/src/reconstruct_loopset.jl b/src/reconstruct_loopset.jl
diff --git a/test/runtests.jl b/test/runtests.jl
