Make determinestrategy more flexible in allowing different unroll/tile/vectorized orders when tiling.

Author: chriselrod

src/LoopVectorization.jl

@@ -9,7 +9,7 @@ using VectorizationBase: REGISTER_SIZE, REGISTER_COUNT, extract_data, num_vector
     maybestaticfirst, maybestaticlast
 using SIMDPirates: VECTOR_SYMBOLS, evadd, evmul, vrange, reduced_add, reduced_prod, reduce_to_add, reduce_to_prod,
     sizeequivalentfloat, sizeequivalentint, vadd!, vsub!, vmul!, vfdiv!, vfmadd!, vfnmadd!, vfmsub!, vfnmsub!,
-    vfmadd231, vfmsub231, vfnmadd231, vfnmsub231, #prefetch,
+    vfmadd231, vfmsub231, vfnmadd231, vfnmsub231, sizeequivalentfloat, sizeequivalentint, #prefetch,
     vmullog2, vmullog10, vdivlog2, vdivlog10, vmullog2add!, vmullog10add!, vdivlog2add!, vdivlog10add!, vfmaddaddone
 using Base.Broadcast: Broadcasted, DefaultArrayStyle
 using LinearAlgebra: Adjoint, Transpose

src/constructors.jl

@@ -88,6 +88,7 @@ true
 
 """
 macro avx(q)
+    q = macroexpand(__module__, q)
     q2 = if q.head === :for
         setup_call(LoopSet(q, __module__))
     else# assume broadcast
@@ -134,12 +135,14 @@ end
 macro avx(arg, q)
     @assert q.head === :for
     @assert arg.head === :(=)
+    q = macroexpand(__module__, q)
     inline, U, T = check_macro_kwarg(arg)
     ls = LoopSet(q, __module__)
     esc(setup_call(ls, inline, U, T))
 end
 macro avx(arg1, arg2, q)
     @assert q.head === :for
+    q = macroexpand(__module__, q)
     inline, U, T = check_macro_kwarg(arg1)
     inline, U, T = check_macro_kwarg(arg2, inline, U, T)
     esc(setup_call(LoopSet(q, __module__), inline, U, T))
@@ -154,15 +157,18 @@ While `@avx` punts to a generated function to enable type-based analysis, `_@avx
 works on just the expressions. This requires that it makes a number of default assumptions.
 """
 macro _avx(q)
+    q = macroexpand(__module__, q)
     esc(lower(LoopSet(q, __module__)))
 end
 macro _avx(arg, q)
     @assert q.head === :for
+    q = macroexpand(__module__, q)
     inline, U, T = check_macro_kwarg(arg)
     esc(lower(LoopSet(q, __module__), U, T))
 end
 
 
 macro avx_debug(q)
+    q = macroexpand(__module__, q)
     esc(LoopVectorization.setup_call_debug(LoopSet(q, __module__)))
 end

src/determinestrategy.jl

@@ -43,14 +43,15 @@ function unitstride(ls::LoopSet, op::Operation, s::Symbol)
 end
 
 function register_pressure(op::Operation)
-    if isconstant(op)
+    if isconstant(op) || isloopvalue(op)
         0
     else
         instruction_cost(instruction(op)).register_pressure
     end
 end
 function cost(ls::LoopSet, op::Operation, unrolled::Symbol, Wshift::Int, size_T::Int = op.elementbytes)
     isconstant(op) && return 0.0, 0, 1
+    isloopvalue(op) && return 0.0, 0, 1
     # Wshift == dependson(op, unrolled) ? Wshift : 0
     # c = first(cost(instruction(op), Wshift, size_T))::Int
     instr = Instruction(:LoopVectorization, instruction(op).instr)
@@ -73,8 +74,8 @@ function cost(ls::LoopSet, op::Operation, unrolled::Symbol, Wshift::Int, size_T:
             # else # vmov(a/u)pd
             end
         elseif instr === :setindex! # broadcast or reductionstore; if store we want to penalize reduction
-            srt *= 2
-            sl *= 2
+            srt *= 3
+            sl *= 3
         end
     end
     srt, sl, srp
@@ -400,16 +401,30 @@ 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
+    istranslation = false
+    inds = getindices(op); li = op.ref.loopedindex
+    for i ∈ eachindex(li)
+        if !li[i]
+            opp = findparent(ls, inds[i + (first(inds) === Symbol("##DISCONTIGUOUSSUBARRAY##"))])
+            if instruction(opp).instr ∈ (:+, :-) && u1 ∈ loopdependencies(opp) && u2 ∈ loopdependencies(opp)
+                istranslation = true
+            end
+        end
+    end
+    istranslation ? 0.25 : 1.0
+end
 # Just tile outer two loops?
 # But optimal order within tile must still be determined
 # as well as size of the tiles.
 function evaluate_cost_tile(
-    ls::LoopSet, order::Vector{Symbol}, vectorized::Symbol
+    ls::LoopSet, order::Vector{Symbol}, unrolled::Symbol, tiled::Symbol, vectorized::Symbol
 )
     N = length(order)
     @assert N ≥ 2 "Cannot tile merely $N loops!"
-    tiled = order[1]
-    unrolled = order[2]
+    # tiled = order[1]
+    # unrolled = order[2]
     ops = operations(ls)
     nops = length(ops)
     included_vars = fill!(resize!(ls.included_vars, nops), false)
@@ -464,6 +479,10 @@ function evaluate_cost_tile(
         isunrolled = unrolledtiled[1,id]
         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;
+        end
         rp = opisininnerloop ? rp : 0 # we only care about register pressure within the inner most loop
         rt *= iters[id]
         if isunrolled && istiled # no cost decrease; cost must be repeated
@@ -556,54 +575,56 @@ function choose_tile(ls::LoopSet)
     lo = LoopOrders(ls)
     # @show lo.syms ls.loop_order.bestorder
     best_order = copyto!(ls.loop_order.bestorder, lo.syms)
-    best_vec = first(best_order) # filler
+    best_unrolled = best_tiled = best_vec = first(best_order) # filler
     new_order, state = iterate(lo) # right now, new_order === best_order
     U, T, lowest_cost = 0, 0, Inf
+    nloops = length(new_order)
     while true
-        for new_vec ∈ @view(new_order[2:end]) # view to skip first
-            U_temp, T_temp, cost_temp = evaluate_cost_tile(ls, new_order, new_vec)
-            if cost_temp < lowest_cost
-                lowest_cost = cost_temp
-                U, T = U_temp, T_temp
-                best_vec = new_vec
-                copyto!(best_order, new_order)
-                save_tilecost!(ls)
+        for new_vec ∈ new_order # view to skip first
+            for nt ∈ 1:nloops-1
+                new_tiled = new_order[nt]
+                for new_unrolled ∈ @view(new_order[nt+1:end])
+                    U_temp, T_temp, cost_temp = evaluate_cost_tile(ls, new_order, new_unrolled, new_tiled, new_vec)
+                    if cost_temp < lowest_cost
+                        lowest_cost = cost_temp
+                        U, T = U_temp, T_temp
+                        best_vec = new_vec
+                        best_tiled = new_tiled
+                        best_unrolled = new_unrolled
+                        copyto!(best_order, new_order)
+                        save_tilecost!(ls)
+                    end
+                end
             end
         end
         iter = iterate(lo, state)
-        iter === nothing && return best_order, best_vec, U, T, lowest_cost
+        iter === nothing && return best_order, best_unrolled, best_tiled, best_vec, U, T, lowest_cost
         new_order, state = iter
     end
 end
 # Last in order is the inner most loop
 function choose_order(ls::LoopSet)
     if num_loops(ls) > 1
-        torder, tvec, tU, tT, tc = choose_tile(ls)
+        torder, tunroll, ttile, tvec, tU, tT, tc = choose_tile(ls)
     else
         tc = Inf
     end
     uorder, uvec, uc = choose_unroll_order(ls, tc)
     if num_loops(ls) > 1 && tc ≤ uc
-        return torder, tvec, min(tU, tT), tT
+        return torder, tunroll, ttile, tvec, min(tU, tT), tT
         # return torder, tvec, 4, 4#5, 5
     else
-        return uorder, uvec, determine_unroll_factor(ls, uorder, first(uorder), uvec), -1
+        return uorder, first(uorder), Symbol("##undefined##"), uvec, determine_unroll_factor(ls, uorder, first(uorder), uvec), -1
     end
 end
 
 function register_pressure(ls::LoopSet)
-    # uses unroll of 1 if not tiling
-    if num_loops(ls) > 1
-        torder, tvec, tU, tT, tc = choose_tile(ls)
+    order, unroll, vec, U, T = choose_order(ls)
+    if T == -1
+        sum(register_pressure, operations(ls))
     else
-        tc = Inf
-    end
-    uorder, uvec, uc = choose_unroll_order(ls, tc)
-    if num_loops(ls) > 1 && tc ≤ uc # tile
         rp = @view ls.reg_pressure[:,1]
         tU * tT * rp[1] + tU * rp[2] + rp[3] + rp[4]
-    else
-        sum(register_pressure, operations(ls))
-    end    
+    end
 end
 

src/graphs.jl

@@ -36,6 +36,30 @@
 
 # end
 
+struct UnrollSpecification
+    unrolledloopnum::Int
+    tiledloopnum::Int
+    vectorizedloopnum::Int
+    U::Int
+    T::Int
+end
+# UnrollSpecification(ls::LoopSet, unrolled::Loop, vectorized::Symbol, U, T) = UnrollSpecification(ls, unrolled.itersymbol, vectorized, U, T)
+function UnrollSpecification(us::UnrollSpecification, U, T)
+    @unpack unrolledloopnum, tiledloopnum, vectorizedloopnum = us
+    UnrollSpecification(unrolledloopnum, tiledloopnum, vectorizedloopnum, U, T)
+end
+function UnrollSpecification(us::UnrollSpecification; U = us.U, T = us.T)
+    @unpack unrolledloopnum, tiledloopnum, vectorizedloopnum = us
+    UnrollSpecification(unrolledloopnum, tiledloopnum, vectorizedloopnum, U, T)
+end
+isunrolled(us::UnrollSpecification, n::Int) = us.unrolledloopnum == n
+istiled(us::UnrollSpecification, n::Int) = !isunrolled(us, n) && us.tiledloopnum == n
+isvectorized(us::UnrollSpecification, n::Int) = us.vectorizedloopnum == n
+function unrollfactor(us::UnrollSpecification, n::Int)
+    @unpack unrolledloopnum, tiledloopnum, U, T = us
+    (unrolledloopnum == n) ? U : ((tiledloopnum == n) ? T : 1)
+end
+
 struct Loop
     itersymbol::Symbol
     starthint::Int
@@ -57,7 +81,7 @@ function Loop(itersymbol::Symbol, start::Union{Int,Symbol}, stop::Union{Int,Symb
 end
 Base.length(loop::Loop) = 1 + loop.stophint - loop.starthint
 isstaticloop(loop::Loop) = loop.startexact & loop.stopexact
-function startloop(loop::Loop, isvectorized, W, itersymbol = loop.itersymbol)
+function startloop(loop::Loop, isvectorized, W, itersymbol)
     startexact = loop.startexact
     if isvectorized
         if startexact
@@ -71,41 +95,48 @@ function startloop(loop::Loop, isvectorized, W, itersymbol = loop.itersymbol)
         Expr(:(=), itersymbol, Expr(:call, lv(:unwrap), loop.startsym))
     end
 end
-function vec_looprange(loop::Loop, isunrolled::Bool, W::Symbol, U::Int)
+function vec_looprange(loop::Loop, W::Symbol, UF::Int, mangledname::Symbol)
+    isunrolled = UF > 1
     incr = if isunrolled
-        Expr(:call, lv(:valmuladd), W, U, -2)
+        Expr(:call, lv(:valmuladd), W, UF, -2)
     else
         Expr(:call, lv(:valsub), W, 2)
     end
     if loop.stopexact # split for type stability
-        Expr(:call, :<, loop.itersymbol, Expr(:call, :-, loop.stophint, incr))
+        Expr(:call, :<, mangledname, Expr(:call, :-, loop.stophint, incr))
     else
-        Expr(:call, :<, loop.itersymbol, Expr(:call, :-, loop.stopsym, incr))
+        Expr(:call, :<, mangledname, Expr(:call, :-, loop.stopsym, incr))
     end
 end
 function looprange(loop::Loop, incr::Int, mangledname::Symbol)
-    incr -= 2
+    incr = 2 - incr
     if iszero(incr)
         Expr(:call, :<, mangledname, loop.stopexact ? loop.stophint : loop.stopsym)
     else
-        Expr(:call, :<, mangledname, loop.stopexact ? loop.stophint - incr : Expr(:call, :-, loop.stopsym, incr))
+        Expr(:call, :<, mangledname, loop.stopexact ? loop.stophint + incr : Expr(:call, :+, loop.stopsym, incr))
     end
 end
 function terminatecondition(
-    loop::Loop, W::Symbol, U::Int, T::Int, isvectorized::Bool, isunrolled::Bool, istiled::Bool,
-    mangledname::Symbol = loop.itersymbol, mask::Nothing = nothing
+    loop::Loop, us::UnrollSpecification, n::Int, W::Symbol, mangledname::Symbol, inclmask::Bool, UF::Int = unrollfactor(us, n)
 )
-    if isvectorized
-        vec_looprange(loop, isunrolled, W, U) # may not be tiled
+    if !isvectorized(us, n)
+        looprange(loop, UF, mangledname)
+    elseif inclmask
+        looprange(loop, 1, mangledname)
     else
-        looprange(loop, isunrolled ? U : (istiled ? T : 1), mangledname)
+        vec_looprange(loop, W, UF, mangledname) # may not be tiled
     end
 end
-function terminatecondition(
-    loop::Loop, W::Symbol, U::Int, T::Int, isvectorized::Bool, isunrolled::Bool, istiled::Bool,
-    mangledname::Symbol, mask::Symbol
-)
-    looprange(loop, 1, mangledname)
+function incrementloopcounter(us::UnrollSpecification, n::Int, W::Symbol, mangledname::Symbol, UF::Int = unrollfactor(us, n))
+    if isvectorized(us, n)
+        if UF == 1
+            Expr(:(=), mangledname, Expr(:call, lv(:valadd), W, mangledname))
+        else
+            Expr(:+=, mangledname, Expr(:call, lv(:valmul), W, UF))
+        end
+    else
+        Expr(:+=, mangledname, UF)
+    end
 end
 
 # load/compute/store × isunrolled × istiled × pre/post loop × Loop number
@@ -556,3 +587,11 @@ function Base.push!(ls::LoopSet, ex::Expr, elementbytes::Int, position::Int)
         throw("Don't know how to handle expression:\n$ex")
     end
 end
+
+function UnrollSpecification(ls::LoopSet, unrolled::Symbol, tiled::Symbol, vectorized::Symbol, U, T)
+    order = names(ls)
+    nu = findfirst(isequal(unrolled), order)::Int
+    nt = T == -1 ? nu : findfirst(isequal(tiled), order)::Int
+    nv = findfirst(isequal(vectorized), order)::Int
+    UnrollSpecification(nu, nt, nv, U, T)
+end

src/lower_compute.jl

@@ -58,7 +58,7 @@ function lower_compute!(
     Uiter = opunrolled ? U - 1 : 0
     isreduct = isreduction(op)
     if !isnothing(suffix) && isreduct && tiledouterreduction == -1
-        instrfid = findfirst(isequal(instr.instr), (:vfmadd_fast, :vfnmadd_fast, :vfmsub_fast, :vfnmsub_fast))
+        instrfid = findfirst(isequal(instr.instr), (:vfmadd, :vfnmadd, :vfmsub, :vfnmsub))
         if instrfid !== nothing
             instr = Instruction((:vfmadd231, :vfnmadd231, :vfmsub231, :vfnmsub231)[instrfid])
         end