Handle offsets alongside +/- expressions in indices, and still unroll when we have load-into-stores; stridepenalty function needs work based on actual size of the arrays.

Author: chriselrod

Project.toml

@@ -1,7 +1,7 @@
 name = "LoopVectorization"
 uuid = "bdcacae8-1622-11e9-2a5c-532679323890"
 authors = ["Chris Elrod <[email protected]>"]
-version = "0.8.9"
+version = "0.8.10"
 
 [deps]
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"

src/determinestrategy.jl

@@ -26,10 +26,10 @@ function unitstride(ls::LoopSet, op::Operation, s::Symbol)
     fi = first(inds)
     if fi === Symbol("##DISCONTIGUOUSSUBARRAY##")
         return false
-    elseif !first(li)
-        # We must check if this
-        parent = findparent(ls, fi)
-        indexappearences(parent, s) > 1 && return false
+    # elseif !first(li)
+    #     # We must check if this
+    #     parent = findparent(ls, fi)
+    #     indexappearences(parent, s) > 1 && return false
     end
     for i ∈ 2:length(inds)
         if li[i]
@@ -217,8 +217,7 @@ function unroll_no_reductions(ls, order, vectorized)
         elseif isstore(op)
             store_rt += first(cost(ls, op, vectorized, Wshift, size_T))
         end
-    end
-    # @show compute_rt, load_rt, store_rt
+    end 
     # heuristic guess
     # roundpow2(min(4, round(Int, (compute_rt + load_rt + 1) / compute_rt)))
     memory_rt = load_rt + store_rt
@@ -374,7 +373,7 @@ function solve_unroll(X, R, u₁L, u₂L, u₁step, u₂step)
     u₂low = max(u₂step, floor(Int, u₂float)) # must be at least 1
     u₁high = solve_unroll_constT(R, u₂low) + u₁step
     u₂high = solve_unroll_constU(R, u₁low) + u₂step
-    maxunroll = REGISTER_COUNT == 32 ? 10 : 6
+    maxunroll = REGISTER_COUNT == 32 ? (((X₂ > 0) & (X₃ > 0)) ? 10 : 8) : 6
     u₁low = min(u₁low, maxunroll)
     u₂low = min(u₂low, maxunroll)
     u₁high = min(u₁high, maxunroll)
@@ -534,18 +533,19 @@ function stride_penalty(ls::LoopSet, op::Operation, order::Vector{Symbol}, loopf
     penalty
 end
 function stride_penalty(ls::LoopSet, order::Vector{Symbol})
-    stridepenalty = 0.0
+    stridepenaltydict = Dict{Symbol,Vector{Float64}}()
     loopfreqs = Vector{Int}(undef, length(order))
     loopfreqs[1] = 1
     for i ∈ 2:length(order)
         loopfreqs[i] = loopfreqs[i-1] * length(getloop(ls, order[i]))
     end
     for op ∈ operations(ls)
         if accesses_memory(op)
-            stridepenalty += stride_penalty(ls, op, order, loopfreqs)
+            v = get!(() -> Float64[], stridepenaltydict, op.ref.ref.array)
+            push!(v, stride_penalty(ls, op, order, loopfreqs))
         end
     end
-    stridepenalty# * 1e-9
+    sum(maximum, values(stridepenaltydict)) #* prod(length, ls.loops) / 1024^length(order)
 end
 function isoptranslation(ls::LoopSet, op::Operation, unrollsyms::UnrollSymbols)
     @unpack u₁loopsym, u₂loopsym, vectorized = unrollsyms
@@ -627,6 +627,7 @@ function load_elimination_cost_factor!(
     @unpack u₁loopsym, u₂loopsym, vectorized = unrollsyms
     if !iszero(first(isoptranslation(ls, op, unrollsyms)))
         rt, lat, rp = cost(ls, op, vectorized, Wshift, size_T)
+        rto = rt
         rt *= iters
             # rt *= factor1; rp *= factor2;
         choose_to_inline[] = true
@@ -645,8 +646,8 @@ function load_elimination_cost_factor!(
         # end
         # # (0.25, REGISTER_COUNT == 32 ? 1.2 : 1.0)
         # (0.25, 1.0)
-        cost_vec[1] += 0.1rt
-        reg_pressure[1] += 0.51rp
+        cost_vec[1] -= 0.1prod(length, ls.loops)
+        reg_pressure[1] += 0.25rp
         cost_vec[2] += rt
         reg_pressure[2] += rp
         cost_vec[3] += rt
@@ -658,13 +659,13 @@ function load_elimination_cost_factor!(
         false
     end
 end
-# function loadintostore(ls::LoopSet, op::Operation)
-#     isload(op) || return false # leads to bad behavior more than it helps
-#     for opp ∈ operations(ls)
-#         isstore(opp) && opp.ref == op.ref && return true
-#     end
-#     false
-# end
+function loadintostore(ls::LoopSet, op::Operation)
+    isload(op) || return false # leads to bad behavior more than it helps
+    for opp ∈ operations(ls)
+        isstore(opp) && opp.ref == op.ref && return true
+    end
+    false
+end
 function store_load_deps!(deps::Vector{Symbol}, op::Operation, compref = op.ref)
     for opp ∈ parents(op)
         foreach(ld -> ((ld ∈ deps) || push!(deps, ld)), loopdependencies(opp))
@@ -799,7 +800,6 @@ function evaluate_cost_tile(
             all(ld -> ld ∈ nested_loop_syms, loopdependencies(op)) || continue
             rd = reduceddependencies(op)
             if hasintersection(rd, @view(nested_loop_syms[1:end-length(rd)]))
-                # @show rd, op itersym, nested_loop_syms @view(nested_loop_syms[1:end-length(rd)])
                 return 0,0,Inf,false
             end
             if isstore(op)
@@ -860,7 +860,7 @@ function evaluate_cost_tile(
     costpenalty = (sum(reg_pressure) > REGISTER_COUNT) ? 2 : 1
     u₁v = vectorized === u₁loopsym; u₂v = vectorized === u₂loopsym
     round_uᵢ = prefetch_good_idea ? (u₁v ? 1 : (u₂v ? 2 : 0)) : 0
-    if irreducible_storecosts / sum(cost_vec) ≥ 0.25
+    if (irreducible_storecosts / sum(cost_vec) ≥ 0.25) && !any(op -> loadintostore(ls, op), operations(ls))
         u₁, u₂ = (1, 1)
         ucost = unroll_cost(cost_vec, 1, 1, length(getloop(ls, u₁loopsym)), length(getloop(ls, u₂loopsym)))
     else

src/lowering.jl

@@ -379,17 +379,17 @@ function lower_unrolled_dynamic(ls::LoopSet, us::UnrollSpecification, n::Int, in
         end
     elseif iszero(UFt)
         Expr( :block, q )
-    # elseif !nisvectorized && !loopisstatic && UF ≥ 8
-    #     rem_uf = UF - 1
-    #     UF = rem_uf >> 1
-    #     UFt = rem_uf - UF
-    #     ust = nisunrolled ? UnrollSpecification(us, UFt, u₂) : UnrollSpecification(us, u₁, UFt)
-    #     newblock = lower_block(ls, ust, n, remmask, UFt)
-    #     # comparison = unrollremcomparison(ls, loop, UFt, n, nisvectorized, remfirst)
-    #     comparison = terminatecondition(ls, us, n, inclmask, UFt)
-    #     UFt = 1
-    #     UF += 1 - iseven(rem_uf)
-    #     Expr( :block, q, Expr(iseven(rem_uf) ? :while : :if, comparison, newblock), remblock )
+    elseif !nisvectorized && !loopisstatic && UF ≥ 10
+        rem_uf = UF - 1
+        UF = rem_uf >> 1
+        UFt = rem_uf - UF
+        ust = nisunrolled ? UnrollSpecification(us, UFt, u₂) : UnrollSpecification(us, u₁, UFt)
+        newblock = lower_block(ls, ust, n, remmask, UFt)
+        # comparison = unrollremcomparison(ls, loop, UFt, n, nisvectorized, remfirst)
+        comparison = terminatecondition(ls, us, n, inclmask, UFt)
+        UFt = 1
+        UF += 1 - iseven(rem_uf)
+        Expr( :block, q, Expr(iseven(rem_uf) ? :while : :if, comparison, newblock), remblock )
     else
         # if (usorig.u₁ == us.u₁) && (usorig.u₂ == us.u₂) && !isstaticloop(loop) && !inclmask# && !ls.loadelimination[]
         #     # Expr(:block, sl, assumeloopiteratesatleastonce(loop), Expr(:while, tc, body))

src/memory_ops_common.jl

@@ -85,35 +85,46 @@ function subset_vptr!(ls::LoopSet, vptr::Symbol, indnum::Int, ind, previndices,
 end
 
 function addoffset!(ls, indices, offsets, loopedindex, loopdependencies, ind, offset)
-    if typemin(Int8) ≤ offset ≤ typemax(Int8)
-        push!(indices, ind);
-        push!(offsets, offset % Int8)
-        push!(loopedindex, true)
-        push!(loopdependencies, ind)
-        true
-    else
-        false
-    end
+    (typemin(Int8) ≤ offset ≤ typemax(Int8)) || return false
+    push!(indices, ind);
+    push!(offsets, offset % Int8)
+    push!(loopedindex, true)
+    push!(loopdependencies, ind)
+    true
+end
+function addoffsetexpr!(ls, parents, indices, offsets, loopedindex, loopdependencies, reduceddeps, ind, offset, elementbytes)
+    (typemin(Int8) ≤ offset ≤ typemax(Int8)) || return false
+    parent = add_operation!(ls, gensym(:indexpr), ind, elementbytes, length(ls.loopsymbols))
+    pushparent!(parents, loopdependencies, reduceddeps, parent)
+    push!(indices, name(parent)); 
+    push!(offsets, offset % Int8)
+    push!(loopedindex, false)
+    true
 end
 
 function checkforoffset!(
-    ls::LoopSet, indices::Vector{Symbol}, offsets::Vector{Int8}, loopedindex::Vector{Bool}, loopdependencies::Vector{Symbol}, ind::Expr
+    ls::LoopSet, parents::Vector{Operation}, indices::Vector{Symbol}, offsets::Vector{Int8}, loopedindex::Vector{Bool},
+    loopdependencies::Vector{Symbol}, reduceddeps::Vector{Symbol}, ind::Expr, elementbytes::Int
 )
     ind.head === :call || return false
     f = first(ind.args)
     (((f === :+) || (f === :-)) && (length(ind.args) == 3)) || return false
     factor = f === :+ ? 1 : -1
     arg1 = ind.args[2]
     arg2 = ind.args[3]
-    if arg1 isa Integer && isone(factor)
+    if arg1 isa Integer# && isone(factor)
         if arg2 isa Symbol && arg2 ∈ ls.loopsymbols
             addoffset!(ls, indices, offsets, loopedindex, loopdependencies, arg2, arg1 * factor)
+        elseif arg2 isa Expr
+            addoffsetexpr!(ls, parents, indices, offsets, loopedindex, loopdependencies, reduceddeps, arg2, arg1 * factor, elementbytes)
         else
             false
         end
     elseif arg2 isa Integer
         if arg1 isa Symbol && arg1 ∈ ls.loopsymbols
             addoffset!(ls, indices, offsets, loopedindex, loopdependencies, arg1, arg2 * factor)
+        elseif arg1 isa Expr
+            addoffsetexpr!(ls, parents, indices, offsets, loopedindex, loopdependencies, reduceddeps, arg1, arg2 * factor, elementbytes)
         else
             false
         end
@@ -158,7 +169,7 @@ function array_reference_meta!(ls::LoopSet, array::Symbol, rawindices, elementby
             length(indices) == 0 && push!(indices, DISCONTIGUOUS)
         elseif ind isa Expr
             #FIXME: position (in loopnest) wont be length(ls.loopsymbols) in general
-            if !checkforoffset!(ls, indices, offsets, loopedindex, loopdependencies, ind)
+            if !checkforoffset!(ls, parents, indices, offsets, loopedindex, loopdependencies, reduceddeps, ind, elementbytes)
                 parent = add_operation!(ls, gensym(:indexpr), ind, elementbytes, length(ls.loopsymbols))
                 pushparent!(parents, loopdependencies, reduceddeps, parent)
                 push!(indices, name(parent));