When _avx_! isn't inlined, cut out a few assembly instructions from setting up the call

Author: chriselrod

src/add_compute.jl

@@ -184,7 +184,7 @@ function add_reduction_update_parent!(
         if reduct_zero === :zero
             push!(ls.preamble_zeros, (identifier(reductinit), IntOrFloat))
         else
-            push!(ls.preamble_funcofeltypes, (identifier(reductinit), reduct_zero))
+            push!(ls.preamble_funcofeltypes, (identifier(reductinit), instrclass))
         end
     else
         reductinit = parent
@@ -384,7 +384,7 @@ function add_pow!(
     end
     if pint == 0
         op = Operation(length(operations(ls)), var, elementbytes, LOOPCONSTANT, constant, NODEPENDENCY, Symbol[], NOPARENTS)
-        push!(ls.preamble_funcofeltypes, (identifier(op),:one))
+        push!(ls.preamble_funcofeltypes, (identifier(op),MULTIPLICATIVE_IN_REDUCTIONS))
         return pushop!(ls, op)
     elseif pint == 1
         return add_compute!(ls, var, :identity, [xop], elementbytes)

src/condense_loopset.jl

@@ -145,16 +145,25 @@ function loop_boundaries(ls::LoopSet)
     lbd
 end
 
+tuple_expr(v) = tuple_expr(identity, v)
+function tuple_expr(f, v)
+    t = Expr(:tuple)
+    for vᵢ ∈ v
+        push!(t.args, f(vᵢ))
+    end
+    t
+end
+
 function argmeta_and_consts_description(ls::LoopSet, arraysymbolinds)
     Expr(
-        :curly, :Tuple,
+        :tuple,
         length(arraysymbolinds),
-        Expr(:curly, :Tuple, ls.outer_reductions...),
-        Expr(:curly, :Tuple, first.(ls.preamble_symsym)...),
-        Expr(:curly, :Tuple, ls.preamble_symint...),
-        Expr(:curly, :Tuple, ls.preamble_symfloat...),
-        Expr(:curly, :Tuple, ls.preamble_zeros...),
-        Expr(:curly, :Tuple, ls.preamble_funcofeltypes...)
+        tuple_expr(ls.outer_reductions),
+        tuple_expr(first, ls.preamble_symsym),
+        tuple_expr(ls.preamble_symint),
+        tuple_expr(ls.preamble_symfloat),
+        tuple_expr(ls.preamble_zeros),
+        tuple_expr(ls.preamble_funcofeltypes)
     )
 end
 
@@ -210,9 +219,10 @@ function check_if_empty(ls::LoopSet, q::Expr)
     Expr(:if, Expr(:call, :!, Expr(:call, :any, :isempty, lb)), q)
 end
 
+val(x) = Expr(:call, Expr(:curly, :Val, x))
 # Try to condense in type stable manner
 function generate_call(ls::LoopSet, inline_unroll::NTuple{3,Int8}, debug::Bool = false)
-    operation_descriptions = Expr(:curly, :Tuple)
+    operation_descriptions = Expr(:tuple)
     varnames = Symbol[]; ids = Vector{Int}(undef, length(operations(ls)))
     for op ∈ operations(ls)
         instr = instruction(op)
@@ -221,27 +231,29 @@ function generate_call(ls::LoopSet, inline_unroll::NTuple{3,Int8}, debug::Bool =
         push!(operation_descriptions.args, OperationStruct!(varnames, ids, ls, op))
     end
     arraysymbolinds = Symbol[]
-    arrayref_descriptions = Expr(:curly, :Tuple)
+    arrayref_descriptions = Expr(:tuple)
     foreach(ref -> push!(arrayref_descriptions.args, ArrayRefStruct(ls, ref, arraysymbolinds, ids)), ls.refs_aliasing_syms)
     argmeta = argmeta_and_consts_description(ls, arraysymbolinds)
     loop_bounds = loop_boundaries(ls)
-    loop_syms = Expr(:curly, :Tuple, map(QuoteNode, ls.loopsymbols)...)
+    loop_syms = tuple_expr(QuoteNode, ls.loopsymbols)
     inline, u₁, u₂ = inline_unroll
-
     func = debug ? lv(:_avx_loopset_debug) : lv(:_avx_!)
     lbarg = debug ? Expr(:call, :typeof, loop_bounds) : loop_bounds
     q = Expr(
-        :call, func, Expr(:call, Expr(:curly, :Val, Expr(:tuple, inline, u₁, u₂, Expr(:call, lv(:unwrap), VECTORWIDTHSYMBOL)))),
-        operation_descriptions, arrayref_descriptions, argmeta, loop_syms, lbarg
+        :call, func, val(Expr(:tuple, inline, u₁, u₂, Expr(:call, lv(:unwrap), VECTORWIDTHSYMBOL))),
+        val(operation_descriptions), val(arrayref_descriptions), val(argmeta), val(loop_syms), lbarg
     )
-    debug && deleteat!(q.args, 2)
-    foreach(ref -> push!(q.args, vptr(ref)), ls.refs_aliasing_syms)
+    # debug && deleteat!(q.args, 2)
+    vargs_as_tuple = !debug
+    extra_args = vargs_as_tuple ? Expr(:tuple) : q
+    foreach(ref -> push!(extra_args.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)
-    add_external_functions!(q, ls)
-    debug && return q
+    foreach(is -> push!(extra_args.args, last(is)), ls.preamble_symsym)
+    append!(extra_args.args, arraysymbolinds)
+    add_reassigned_syms!(extra_args, ls)
+    add_external_functions!(extra_args, ls)
+    # debug && return q
+    vargs_as_tuple && push!(q.args, extra_args)
     vecwidthdefq = Expr(:block)
     define_eltype_vec_width!(vecwidthdefq, ls, nothing)
     Expr(:block, vecwidthdefq, q)
@@ -318,7 +330,6 @@ function setup_call_inline(ls::LoopSet, inline::Int8 = zero(Int8), U::Int8 = zer
 end
 function setup_call_debug(ls::LoopSet)
     # avx_loopset(instr, ops, arf, AM, LB, vargs)
-
     pushpreamble!(ls, generate_call(ls, (zero(Int8),zero(Int8),zero(Int8)), true))
     Expr(:block, ls.prepreamble, ls.preamble)
 end

src/graphs.jl

@@ -253,7 +253,7 @@ struct LoopSet
     preamble_symint::Vector{Tuple{Int,Int}}
     preamble_symfloat::Vector{Tuple{Int,Float64}}
     preamble_zeros::Vector{Tuple{Int,NumberType}}
-    preamble_funcofeltypes::Vector{Tuple{Int,Symbol}}
+    preamble_funcofeltypes::Vector{Tuple{Int,Float64}}
     includedarrays::Vector{Symbol}
     includedactualarrays::Vector{Symbol}
     syms_aliasing_refs::Vector{Symbol}
@@ -308,7 +308,7 @@ function pushpreamble!(ls::LoopSet, op::Operation, v::Number)
     if iszero(v)
         push!(ls.preamble_zeros, (id, typ))
     elseif isone(v)
-        push!(ls.preamble_funcofeltypes, (id, :one))
+        push!(ls.preamble_funcofeltypes, (id, MULTIPLICATIVE_IN_REDUCTIONS))
     elseif v isa Integer
         push!(ls.preamble_symint, (id, convert(Int,v)))
     else
@@ -321,7 +321,7 @@ pushpreamble!(ls::LoopSet, ex::Expr) = push!(ls.preamble.args, ex)
 #     if RHS.head === :call && first(RHS.args) === :zero
 #         push!(ls.preamble_zeros, (identifier(op), IntOrFloat))
 #     elseif RHS.head === :call && first(RHS.args) === :one
-#         push!(ls.preamble_funcofeltypes, (identifier(op), :one))
+#         push!(ls.preamble_funcofeltypes, (identifier(op), MULTIPLICATIVE_IN_REDUCTIONS))
 #     else
 #         pushpreamble!(ls, Expr(:(=), c, RHS))
 #         pushpreamble!(ls, op, c)
@@ -670,7 +670,7 @@ function add_operation!(
             if f === :zero
                 push!(ls.preamble_zeros, (identifier(op), IntOrFloat))
             else
-                push!(ls.preamble_funcofeltypes, (identifier(op), :one))
+                push!(ls.preamble_funcofeltypes, (identifier(op), MULTIPLICATIVE_IN_REDUCTIONS))
             end
             op
         else
@@ -706,7 +706,7 @@ function add_operation!(
             if f === :zero
                 push!(ls.preamble_zeros, (identifier(op), IntOrFloat))
             else
-                push!(ls.preamble_funcofeltypes, (identifier(op), :one))
+                push!(ls.preamble_funcofeltypes, (identifier(op), MULTIPLICATIVE_IN_REDUCTIONS))
             end
             op
         else

src/lower_constant.jl

@@ -165,7 +165,7 @@ function lower_licm_constants!(ls::LoopSet)
         end
     end
     for (id,f) ∈ ls.preamble_funcofeltypes
-        setop!(ls, ops[id], Expr(:call, f, ELTYPESYMBOL))
+        setop!(ls, ops[id], Expr(:call, reduction_zero(f), ELTYPESYMBOL))
     end
 end
 

src/reconstruct_loopset.jl

@@ -198,24 +198,24 @@ end
 function num_parameters(AM)
     num_param::Int = AM[1]
     # num_param += length(AM[2].parameters)
-    num_param + length(AM[3].parameters)
+    num_param + length(AM[3])
 end
 function gen_array_syminds(AM)
     Symbol[Symbol("##arraysymbolind##"*i*'#') for i ∈ 1:(AM[1])::Int]
 end
 function process_metadata!(ls::LoopSet, AM, num_arrays::Int)
     opoffsets = ls.operation_offsets
-    expandbyoffset!(ls.outer_reductions, AM[2].parameters, opoffsets)
-    for (i,si) ∈ enumerate(AM[3].parameters)
+    expandbyoffset!(ls.outer_reductions, AM[2], opoffsets)
+    for (i,si) ∈ enumerate(AM[3])
         sii = si::Int
         s = gensym(:symlicm)
         push!(ls.preamble_symsym, (opoffsets[sii] + 1, s))
         pushpreamble!(ls, Expr(:(=), s, Expr(:macrocall, Symbol("@inbounds"), LineNumberNode(@__LINE__,Symbol(@__FILE__)), Expr(:ref, :vargs, num_arrays + i))))
     end
-    expandbyoffset!(ls.preamble_symint, AM[4].parameters, opoffsets)
-    expandbyoffset!(ls.preamble_symfloat, AM[5].parameters, opoffsets)
-    expandbyoffset!(ls.preamble_zeros, AM[6].parameters, opoffsets)
-    expandbyoffset!(ls.preamble_funcofeltypes, AM[7].parameters, opoffsets)
+    expandbyoffset!(ls.preamble_symint, AM[4], opoffsets)
+    expandbyoffset!(ls.preamble_symfloat, AM[5], opoffsets)
+    expandbyoffset!(ls.preamble_zeros, AM[6], opoffsets)
+    expandbyoffset!(ls.preamble_funcofeltypes, AM[7], opoffsets)
     nothing
 end
 function expandbyoffset!(indexpand::Vector{T}, inds, offsets::Vector{Int}, expand::Bool = true) where {T <: Union{Int,Tuple{Int,<:Any}}}
@@ -437,7 +437,7 @@ function sizeofeltypes(v, num_arrays)::Int
 end
 
 function avx_loopset(instr::Vector{Instruction}, ops::Vector{OperationStruct}, arf::Vector{ArrayRefStruct},
-                     AM::Core.SimpleVector, LPSYM::Core.SimpleVector, LB::Core.SimpleVector, @nospecialize(vargs))
+                     AM::Vector{Any}, LPSYM::Vector{Any}, LB::Core.SimpleVector, @nospecialize(vargs))
     ls = LoopSet(:LoopVectorization)
     num_arrays = length(arf)
     elementbytes = sizeofeltypes(vargs, num_arrays)
@@ -464,18 +464,33 @@ function avx_body(ls::LoopSet, UNROLL::Tuple{Int8,Int8,Int8,Int})
     q
 end
 
-function _avx_loopset_debug(::Type{OPS}, ::Type{ARF}, ::Type{AM}, ::Type{LPSYM}, ::Type{LB}, vargs...) where {OPS, ARF, AM, LPSYM, LB}
+function _avx_loopset_debug(::Val{UNROLL}, ::Val{OPS}, ::Val{ARF}, ::Val{AM}, ::Val{LPSYM}, ::Type{LB}, vargs...) where {UNROLL, OPS, ARF, AM, LPSYM, LB}
     @show OPS ARF AM LPSYM LB vargs
-    _avx_loopset(OPS.parameters, ARF.parameters, AM.parameters, LPSYM.parameters, LB.parameters, typeof.(vargs))
+    inline, u₁, u₂, W = UNROLL
+    ls = _avx_loopset(OPS, ARF, AM, LPSYM, LB.parameters, typeof.(vargs))
+    ls.vector_width[] = W
+    ls
+end
+function tovector(@nospecialize(t))
+    v = Vector{Any}(undef, length(t))
+    for i ∈ eachindex(v)
+        tᵢ = t[i]
+        if tᵢ isa Tuple # reduce specialization?
+            v[i] = tovector(tᵢ)
+        else
+            v[i] = tᵢ
+        end
+    end
+    v
 end
-function _avx_loopset(OPSsv::Core.SimpleVector, ARFsv::Core.SimpleVector, AMsv::Core.SimpleVector, LPSYMsv::Core.SimpleVector, LBsv::Core.SimpleVector, @nospecialize(vargs))
+function _avx_loopset(@nospecialize(OPSsv), @nospecialize(ARFsv), @nospecialize(AMsv), @nospecialize(LPSYMsv), LBsv::Core.SimpleVector, @nospecialize(vargs))
     nops = length(OPSsv) ÷ 3
     instr = Instruction[Instruction(OPSsv[3i+1], OPSsv[3i+2]) for i ∈ 0:nops-1]
     ops = OperationStruct[ OPSsv[3i] for i ∈ 1:nops ]
     avx_loopset(
         instr, ops,
         ArrayRefStruct[ARFsv...],
-        AMsv, LPSYMsv, LBsv, vargs
+        tovector(AMsv), tovector(LPSYMsv), LBsv, vargs
     )
 end
 """
@@ -497,9 +512,9 @@ Execute an `@avx` block. The block's code is represented via the arguments:
   `StaticLowerUnitRange(1)` because the lower bound of the iterator can be determined to be 1.
 - `vargs...` holds the encoded pointers of all the arrays (see `VectorizationBase`'s various pointer types).
 """
-@generated function _avx_!(::Val{UNROLL}, ::Type{OPS}, ::Type{ARF}, ::Type{AM}, ::Type{LPSYM}, lb::LB, vargs...) where {UNROLL, OPS, ARF, AM, LPSYM, LB}
+@generated function _avx_!(::Val{UNROLL}, ::Val{OPS}, ::Val{ARF}, ::Val{AM}, ::Val{LPSYM}, lb::LB, vargs::Tuple{Vararg{Any,K}}) where {UNROLL, OPS, ARF, AM, LPSYM, LB, K}
     # 1 + 1 # Irrelevant line you can comment out/in to force recompilation...
-    ls = _avx_loopset(OPS.parameters, ARF.parameters, AM.parameters, LPSYM.parameters, LB.parameters, vargs)
+    ls = _avx_loopset(OPS, ARF, AM, LPSYM, LB.parameters, vargs.parameters)
     # return @show avx_body(ls, UNROLL)
     # @show UNROLL, OPS, ARF, AM, LPSYM, LB
     avx_body(ls, UNROLL)