Compute register pressure with Float64, and hopefully improve performance for @avx and user-defined functions.

Author: chriselrod

benchmark/looptests.jl

@@ -278,26 +278,22 @@ end
 
 
 function filter2d!(out::AbstractMatrix, A::AbstractMatrix, kern)
-    rng1k, rng2k = axes(kern)
-    rng1,  rng2  = axes(out)
-    @inbounds @fastmath for j in rng2, i in rng1
+    @inbounds @fastmath for J in CartesianIndices(out)
         tmp = zero(eltype(out))
-        for jk in rng2k, ik in rng1k
-            tmp += A[i+ik,j+jk]*kern[ik,jk]
+        for I ∈ CartesianIndices(kern)
+            tmp += A[I + J] * kern[I]
         end
-        out[i,j] = tmp
+        out[J] = tmp
     end
     out
 end
 function filter2davx!(out::AbstractMatrix, A::AbstractMatrix, kern)
-    rng1k, rng2k = axes(kern)
-    rng1,  rng2  = axes(out)
-    @avx for j in rng2, i in rng1
+    @avx for J in CartesianIndices(out)
         tmp = zero(eltype(out))
-        for jk in rng2k, ik in rng1k
-            tmp += A[i+ik,j+jk]*kern[ik,jk]
+        for I ∈ CartesianIndices(kern)
+            tmp += A[I + J] * kern[I]
         end
-        out[i,j] = tmp
+        out[J] = tmp
     end
     out
 end

src/condense_loopset.jl

@@ -181,7 +181,9 @@ function add_external_functions!(q::Expr, ls::LoopSet)
     for op ∈ operations(ls)
         if iscompute(op)
             instr = instruction(op)
-            if instr.mod !== :LoopVectorization
+            if instr.mod === :Main
+                push!(q.args, instr.instr)
+            elseif instr.mod !== :LoopVectorization
                 push!(q.args, Expr(:(.), instr.mod, QuoteNode(instr.instr)))
             end
         end

src/costs.jl

@@ -18,6 +18,8 @@ Base.convert(::Type{Expr}, instr::Instruction) = Expr(:(.), instr.mod, QuoteNode
 function Base.Expr(instr::Instruction, args...)
     if instr.mod === :LoopVectorization
         Expr(:call, lv(instr.instr), args...)::Expr
+    elseif instr.mod === :Main
+        Expr(:call, instr.instr, args...)::Expr
     else
         Expr(:call, convert(Expr, instr), args...)::Expr
     end

src/determinestrategy.jl

@@ -349,10 +349,10 @@ end
 function solve_tilesize(
     ls::LoopSet, unrolled::Symbol, tiled::Symbol,
     cost_vec::AbstractVector{Float64},
-    reg_pressure::AbstractVector{Int},
+    reg_pressure::AbstractVector{Float64},
     W::Int, vectorized::Symbol
 )
-    maxTbase = maxUbase = 4
+    maxTbase = maxUbase = 4#8
     maxT = maxTbase#8
     maxU = maxUbase#8
     tiledloop = getloop(ls, tiled)
@@ -405,8 +405,8 @@ function stride_penalty(ls::LoopSet, order::Vector{Symbol})
     end
     stridepenalty# * 1e-9
 end
-function convolution_cost_factor(ls::LoopSet, op::Operation, u1::Symbol, u2::Symbol, v::Symbol)
-    (u1 ∈ loopdependencies(op) && u2 ∈ loopdependencies(op)) || return 1.0
+function convolution_cost_factor(ls::LoopSet, op::Operation, u1::Symbol, u2::Symbol)
+    (u1 ∈ loopdependencies(op) && u2 ∈ loopdependencies(op)) || return 1.0, 1.0
     istranslation = false
     inds = getindices(op); li = op.ref.loopedindex
     for i ∈ eachindex(li)
@@ -417,7 +417,7 @@ function convolution_cost_factor(ls::LoopSet, op::Operation, u1::Symbol, u2::Sym
             end
         end
     end
-    istranslation ? 0.25 : 1.0
+    istranslation ? (0.25, 1.0) : (1.0, 1.0)
 end
 # Just tile outer two loops?
 # But optimal order within tile must still be determined
@@ -484,8 +484,8 @@ function evaluate_cost_tile(
         istiled = unrolledtiled[2,id]
         rt, lat, rp = cost(ls, op, vectorized, Wshift, size_T)
         if isload(op)
-            factor = convolution_cost_factor(ls, op, unrolled, tiled, vectorized)
-            rt *= factor#; rp *= factor;
+            factor1, factor2 = convolution_cost_factor(ls, op, unrolled, tiled)
+            rt *= factor1; rp *= factor2;
         end
         # @show op rt, lat, rp
         rp = opisininnerloop ? rp : 0 # we only care about register pressure within the inner most loop

src/graphs.jl

@@ -192,7 +192,7 @@ struct LoopSet
     syms_aliasing_refs::Vector{Symbol}
     refs_aliasing_syms::Vector{ArrayReferenceMeta}
     cost_vec::Matrix{Float64}
-    reg_pres::Matrix{Int}
+    reg_pres::Matrix{Float64}
     included_vars::Vector{Bool}
     place_after_loop::Vector{Bool}
     W::Symbol
@@ -296,7 +296,7 @@ function LoopSet(mod::Symbol)# = :LoopVectorization)
         Symbol[], Symbol[], Symbol[],
         ArrayReferenceMeta[],
         Matrix{Float64}(undef, 4, 2),
-        Matrix{Int}(undef, 4, 2),
+        Matrix{Float64}(undef, 4, 2),
         Bool[], Bool[],
         gensym(:W), gensym(:T), mod
     )
@@ -427,12 +427,6 @@ function register_single_loop!(ls::LoopSet, looprange::Expr)
                 U = add_loop_bound!(ls, itersym, upper, true)
                 Loop(itersym, L, U)
             end
-        elseif f === :eachindex
-            N = gensym(Symbol(:loopeachindex, itersym))
-            pushpreamble!(ls, Expr(:(=), N, Expr(:call, lv(:maybestaticrange), r)))
-            L = add_loop_bound!(ls, itersym, Expr(:call, lv(:maybestaticfirst), N), false)
-            U = add_loop_bound!(ls, itersym, Expr(:call, lv(:maybestaticlast), N), true)
-            Loop(itersym, L, U)
         elseif f === :OneTo || isscopedname(f, :Base, :OneTo)
             otN = r.args[2]
             if otN isa Integer
@@ -444,7 +438,11 @@ function register_single_loop!(ls::LoopSet, looprange::Expr)
                 Loop(itersym, 1, N)
             end
         else
-            throw("Unrecognized loop range type: $r.")
+            N = gensym(Symbol(:loop, itersym))
+            pushpreamble!(ls, Expr(:(=), N, Expr(:call, lv(:maybestaticrange), r)))
+            L = add_loop_bound!(ls, itersym, Expr(:call, lv(:maybestaticfirst), N), false)
+            U = add_loop_bound!(ls, itersym, Expr(:call, lv(:maybestaticlast), N), true)
+            Loop(itersym, L, U)
         end
     elseif isa(r, Symbol)
         # Treat similar to `eachindex`

src/reconstruct_loopset.jl

@@ -158,7 +158,7 @@ end
 
 function num_parameters(AM)
     num_param::Int = AM[1]
-    num_param += length(AM[2].parameters)
+    # num_param += length(AM[2].parameters)
     num_param + length(AM[3].parameters)
 end
 function gen_array_syminds(AM)
@@ -335,11 +335,10 @@ function add_ops!(
         opsymbol = opsymbols[os.symid]
         add_op!(ls, instr[i], ops, nopsv, expandedv, i, mrefs, opsymbol, elementbytes)
     end
-    num_params = add_parents_to_ops!(ls, ops, constoffset)
+    add_parents_to_ops!(ls, ops, constoffset)
     # for op in operations(ls)
         # @show op
     # end
-    num_params
 end
 
 # elbytes(::VectorizationBase.AbstractPointer{T}) where {T} = sizeof(T)::Int
@@ -385,7 +384,7 @@ function avx_loopset(instr, ops, arf, AM, LPSYM, LB, vargs)
     pushpreamble!(ls, Expr(:(=), ls.T, Expr(:call, :promote_type, [Expr(:call, :eltype, vptr(mref)) for mref ∈ mrefs]...)))
     pushpreamble!(ls, Expr(:(=), ls.W, Expr(:call, lv(:pick_vector_width_val), [Expr(:call, :eltype, vptr(mref)) for mref ∈ mrefs]...)))
     num_params = num_arrays + num_parameters(AM)
-    num_params = add_ops!(ls, instr, ops, mrefs, opsymbols, num_params, nopsv, expandedv, elementbytes)
+    add_ops!(ls, instr, ops, mrefs, opsymbols, num_params, nopsv, expandedv, elementbytes)
     process_metadata!(ls, AM, length(arf))
     add_array_symbols!(ls, arraysymbolinds, num_arrays + length(ls.preamble_symsym))
     num_params = extract_external_functions!(ls, num_params)

test/miscellaneous.jl

@@ -193,7 +193,17 @@
             ret[j] = clenshaw(x[j], coeff)
         end
     end
-
+# ret = y2; coeff = c;
+#     LoopVectorization.@avx_debug for j in 1:length(ret)
+#             ret[j] = clenshaw(x[j], coeff)
+#     end
+#     t = β₁ = β₂ = ρ = s = 0.0; weights = rand(1); nodes = rand(1); lomnibus(args...) = +(args...)
+# LoopVectorization.@avx_debug for i ∈ eachindex(weights, nodes)
+#         s += weights[i] * lomnibus(nodes[i], t, β₁, β₂, ρ)
+#     end    
+# @macroexpand @avx for i ∈ eachindex(weights, nodes)
+#         s += weights[i] * lomnibus(nodes[i], t, β₁, β₂, ρ)
+#     end    
     function softmax3_core!(lse, qq, xx, tmpmax, maxk, nk)
         for k in Base.OneTo(maxk)
             @inbounds for i in eachindex(lse)