Major progress on overhauling LoopVectorization's lowering of indexing. Still need to finnish handling of loop remainders (this branch will currently segfault if not evenly divisble by the unrolling factors).

Author: chriselrod

src/LoopVectorization.jl

@@ -3,8 +3,8 @@ module LoopVectorization
 using VectorizationBase, SIMDPirates, SLEEFPirates, UnPack, OffsetArrays
 using VectorizationBase: REGISTER_SIZE, REGISTER_COUNT, extract_data, num_vector_load_expr,
     mask, masktable, pick_vector_width_val, valmul, valrem, valmuladd, valmulsub, valadd, valsub, _MM,
-    maybestaticlength, maybestaticsize, staticm1, staticp1, subsetview, vzero, stridedpointer_for_broadcast,
-    Static, StaticUnitRange, StaticLowerUnitRange, StaticUpperUnitRange, unwrap, maybestaticrange,
+    maybestaticlength, maybestaticsize, staticm1, staticp1, staticmul, subsetview, vzero, stridedpointer_for_broadcast,
+    Static, Zero, StaticUnitRange, StaticLowerUnitRange, StaticUpperUnitRange, unwrap, maybestaticrange,
     AbstractColumnMajorStridedPointer, AbstractRowMajorStridedPointer, AbstractSparseStridedPointer, AbstractStaticStridedPointer,
     PackedStridedPointer, SparseStridedPointer, RowMajorStridedPointer, StaticStridedPointer, StaticStridedStruct,
     maybestaticfirst, maybestaticlast, scalar_less, scalar_greater, noalias!, gesp, gepbyte
@@ -47,9 +47,10 @@ include("add_compute.jl")
 include("add_constants.jl")
 include("add_ifelse.jl")
 include("determinestrategy.jl")
+include("loopstartstopmanager.jl")
 include("lower_compute.jl")
 include("lower_constant.jl")
-include("zero.jl")
+# include("zero.jl")
 include("lower_memory_common.jl")
 include("lower_load.jl")
 include("lower_store.jl")

src/add_constants.jl

@@ -36,7 +36,7 @@ function add_constant!(ls::LoopSet, mpref::ArrayReferenceMetaPosition, elementby
     temp = gensym(:intermediateconstref)
     vloadcall = Expr(:call, lv(:vload), mpref.mref.ptr)
     if length(getindices(op)) > 0
-        push!(vloadcall.args, mem_offset(op, UnrollArgs(0, Symbol(""), Symbol(""), Symbol(""), nothing), false, false))
+        push!(vloadcall.args, mem_offset(op, UnrollArgs(0, Symbol(""), Symbol(""), Symbol(""), 0, nothing), false, false))
     end
     pushpreamble!(ls, Expr(:(=), temp, vloadcall))
     pushpreamble!(ls, op, temp)

src/determinestrategy.jl

@@ -454,7 +454,7 @@ end
 function isoptranslation(ls::LoopSet, op::Operation, unrollsyms::UnrollSymbols)
     @unpack u₁loopsym, u₂loopsym, vectorized = unrollsyms
     (vectorized == u₁loopsym || vectorized == u₂loopsym) && return false, false
-    (u₁loopsym ∈ loopdependencies(op) && u₂loopsym ∈ loopdependencies(op)) || return false, false
+    (isu₁unrolled(op) && isu₂unrolled(op)) || return false, false
     istranslation = false
     inds = getindices(op); li = op.ref.loopedindex
     translationplus = false
@@ -597,6 +597,7 @@ function evaluate_cost_tile(
     N = length(order)
     @assert N ≥ 2 "Cannot tile merely $N loops!"
     @unpack u₁loopsym, u₂loopsym, vectorized = unrollsyms
+    cacheunrolled!(ls, u₁loopsym, u₂loopsym, vectorized)
     # u₂loopsym = order[1]
     # u₁loopsym = order[2]
     ops = operations(ls)
@@ -647,8 +648,8 @@ function evaluate_cost_tile(
             rd = reduceddependencies(op)
             hasintersection(rd, @view(nested_loop_syms[1:end-length(rd)])) && return 0,0,Inf,false
             included_vars[id] = true
-            depends_on_u₁ = u₁loopsym ∈ loopdependencies(op)
-            depends_on_u₂ = u₂loopsym ∈ loopdependencies(op)
+            depends_on_u₁ = isu₁unrolled(op)
+            depends_on_u₂ = isu₂unrolled(op)
             # cost is reduced by unrolling u₁ if it is interior to u₁loop (true if either u₁reached, or if depends on u₂ [or u₁]) and doesn't depend on u₁
             reduced_by_unrolling[1,id] = (u₁reached | depends_on_u₂) & !depends_on_u₁
             reduced_by_unrolling[2,id] = (u₂reached | depends_on_u₁) & !depends_on_u₂

src/graphs.jl

@@ -9,15 +9,16 @@ struct UnrollArgs{T <: Union{Nothing,Int}}
     u₁loopsym::Symbol
     u₂loopsym::Symbol
     vectorized::Symbol
+    u₂max::Int
     suffix::T
 end
-function UnrollArgs(U::Int, unrollsyms::UnrollSymbols, suffix)
+function UnrollArgs(u₁::Int, unrollsyms::UnrollSymbols, u₂max::Int, suffix)
     @unpack u₁loopsym, u₂loopsym, vectorized = unrollsyms
-    UnrollArgs(U, u₁loopsym, u₂loopsym, vectorized, suffix)
+    UnrollArgs(u₁, u₁loopsym, u₂loopsym, vectorized, u₂max, suffix)
 end
 function UnrollArgs(ua::UnrollArgs, u::Int)
-    @unpack u₁loopsym, u₂loopsym, vectorized, suffix = ua
-    UnrollArgs(u, u₁loopsym, u₂loopsym, vectorized, suffix)
+    @unpack u₁loopsym, u₂loopsym, vectorized, u₂max, suffix = ua
+    UnrollArgs(u, u₁loopsym, u₂loopsym, vectorized, u₂max, suffix)
 end
 # UnrollSymbols(ua::UnrollArgs) = UnrollSymbols(ua.u₁loopsym, ua.u₂loopsym, ua.vectorized)
 
@@ -69,14 +70,8 @@ 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, itersymbol)
+function startloop(loop::Loop, itersymbol)
     startexact = loop.startexact
-    # if isvectorized
-    #     if startexact
-    #         Expr(:(=), itersymbol, Expr(:call, lv(:_MM), VECTORWIDTHSYMBOL, loop.starthint))
-    #     else
-    #         Expr(:(=), itersymbol, Expr(:call, lv(:_MM), VECTORWIDTHSYMBOL, loop.startsym))
-    #     end
     if startexact
         Expr(:(=), itersymbol, loop.starthint - 1)
     else
@@ -97,31 +92,44 @@ addexpr(ex::Number, incr::Number) = ex + incr
 subexpr(ex, incr) = Expr(:call, lv(:vsub), ex, incr)
 subexpr(ex::Number, incr::Number) = ex - incr
 subexpr(ex, incr::Number) = addexpr(ex,  -incr)
-function vec_looprange(loop::Loop, UF::Int, mangledname::Symbol)
+
+staticmulincr(ptr, incr) = Expr(:call, lv(:staticmul), Expr(:call, :eltype, ptr), incr)
+callpointer(sym) = Expr(:call, :pointer, sym)
+function vec_looprange(loopmax, UF::Int, mangledname::Symbol, ptrcomp::Bool)
     incr = if isone(UF)
         Expr(:call, lv(:valsub), VECTORWIDTHSYMBOL, 1)
     else
         Expr(:call, lv(:valmulsub), VECTORWIDTHSYMBOL, UF, 1)
     end
-    if loop.stopexact # split for type stability
-        Expr(:call, :<, mangledname, subexpr(loop.stophint, incr))
+    incr = ptrcomp ? staticmulincr(mangledname, incr) : incr
+    compexpr = subexpr(loopmax, incr)
+    if ptrcomp
+        Expr(:call, :<, callpointer(mangledname), compexpr)
     else
-        Expr(:call, :<, mangledname, subexpr(loop.stopsym, incr))
+        Expr(:call, :<, mangledname, compexpr)
     end
 end
 
-function looprange(stopcon, incr::Int, mangledname::Symbol)
+function looprange(stopcon, incr::Int, mangledname::Symbol, ptrcomp::Bool)
     incr = 1 - incr
     if iszero(incr)
-        Expr(:call, :<, mangledname, stopcon)
-    elseif isone(incr)
-        Expr(:call, :≤, mangledname, stopcon)
+        if ptrcomp
+            Expr(:call, :<, callpointer(mangledname), stopcon)
+        else
+            Expr(:call, :<, mangledname, stopcon)
+        end
+    elseif ptrcomp
+        Expr(:call, :<, callpointer(mangledname), addexpr(stopcon, staticmulincr(mangledname, incr)))
     else
-        Expr(:call, :<, mangledname, addexpr(stopcon, incr))
+        if isone(incr)
+            Expr(:call, :≤, mangledname, stopcon)
+        else
+            Expr(:call, :<, mangledname, addexpr(stopcon, incr))
+        end
     end
 end
 function looprange(loop::Loop, incr::Int, mangledname::Symbol)
-    loop.stopexact ? looprange(loop.stophint, incr, mangledname) : looprange(loop.stopsym, incr, mangledname)
+    loop.stopexact ? looprange(loop.stophint, incr, mangledname, false) : looprange(loop.stopsym, incr, mangledname, false)
 end
 function terminatecondition(
     loop::Loop, us::UnrollSpecification, n::Int, mangledname::Symbol, inclmask::Bool, UF::Int = unrollfactor(us, n)
@@ -130,22 +138,36 @@ function terminatecondition(
         looprange(loop, UF, mangledname)
     elseif inclmask
         looprange(loop, 1, mangledname)
+    elseif loop.stopexact
+        vec_looprange(loop.stophint, UF, mangledname, false) # may not be u₂loop
     else
-        vec_looprange(loop, UF, mangledname) # may not be u₂loop
+        vec_looprange(loop.stopsym, UF, mangledname, false) # may not be u₂loop
     end
 end
 function incrementloopcounter(us::UnrollSpecification, n::Int, mangledname::Symbol, UF::Int = unrollfactor(us, n))
     if isvectorized(us, n)
         if UF == 1
             Expr(:(=), mangledname, Expr(:call, lv(:valadd), VECTORWIDTHSYMBOL, mangledname))
-            # Expr(:(=), mangledname, Expr(:macrocall, Symbol("@show"), LineNumberNode(@__LINE__,Symbol(@__FILE__)), Expr(:call, lv(:valadd), VECTORWIDTHSYMBOL, mangledname)))
         else
             Expr(:(=), mangledname, Expr(:call, lv(:valmuladd), VECTORWIDTHSYMBOL, UF, mangledname))
         end
     else
         Expr(:(=), mangledname, Expr(:call, lv(:vadd), mangledname, UF))
     end
 end
+function incrementloopcounter!(q, us::UnrollSpecification, n::Int, UF::Int = unrollfactor(us, n))
+    if isvectorized(us, n)
+        if UF == 1
+            push!(q.args, Expr(:call, lv(:unwrap), VECTORWIDTHSYMBOL))
+        else
+            push!(q.args, Expr(:call, lv(:valmul), VECTORWIDTHSYMBOL, UF))
+        end
+    elseif isone(UF)
+        push!(q.args, Expr(:call, Expr(:curly, lv(:Static), UF)))
+    else
+        push!(q.args, UF)
+    end
+end
 
 # load/compute/store × isunrolled × istiled × pre/post loop × Loop number
 struct LoopOrder <: AbstractArray{Vector{Operation},5}
@@ -294,19 +316,22 @@ function LoopSet(mod::Symbol)
     )
 end
 
+cacheunrolled!(ls::LoopSet, u₁loop, u₂loop, vectorized) = foreach(op -> setunrolled!(op, u₁loop, u₂loop, vectorized), operations(ls))
+
 num_loops(ls::LoopSet) = length(ls.loops)
 function oporder(ls::LoopSet)
     N = length(ls.loop_order.loopnames)
     reshape(ls.loop_order.oporder, (2,2,2,N))
 end
 names(ls::LoopSet) = ls.loop_order.loopnames
+reversenames(ls::LoopSet) = ls.loop_order.bestorder
 function getloopid(ls::LoopSet, s::Symbol)::Int
     for (loopnum,sym) ∈ enumerate(ls.loopsymbols)
         s === sym && return loopnum
     end
 end
 getloop(ls::LoopSet, s::Symbol) = ls.loops[getloopid(ls, s)]
-getloop(ls::LoopSet, i::Integer) = ls.loops[i]
+# getloop(ls::LoopSet, i::Integer) = ls.loops[i]
 getloopsym(ls::LoopSet, i::Integer) = ls.loopsymbols[i]
 Base.length(ls::LoopSet, s::Symbol) = length(getloop(ls, s))
 

src/lower_load.jl

@@ -1,20 +1,19 @@
 function lower_load_scalar!(
-    q::Expr, op::Operation, ua::UnrollArgs, umin::Int = 0
+    q::Expr, op::Operation, ua::UnrollArgs, umin::Int, inds_calc_by_ptr_offset::Vector{Bool}
 )
     loopdeps = loopdependencies(op)
     @unpack u₁, u₁loopsym, u₂loopsym, vectorized, suffix = ua
     @assert vectorized ∉ loopdeps
     # mvar, opu₁, opu₂ = variable_name_and_unrolled(op, u₁loop, u₂loop, suffix)
     mvar = variable_name(op, suffix)
-    opu₁ = u₁loopsym ∈ loopdeps
-    unrolled = (opu₁ || u₂loopsym ∈ loopdeps)
-    ptr = unrolled ? offset_refname(op, ua) : refname(op)
+    opu₁ = isu₁unrolled(op)
+    ptr = refname(op)
     U = opu₁ ? u₁ : 1
     if instruction(op).instr !== :conditionalload
         for u ∈ umin:U-1
             varname = varassignname(mvar, u, opu₁)
             td = UnrollArgs(ua, u)
-            push!(q.args, Expr(:(=), varname, Expr(:call, lv(:vload), ptr, mem_offset_u(op, td, unrolled))))
+            push!(q.args, Expr(:(=), varname, Expr(:call, lv(:vload), ptr, mem_offset_u(op, td, inds_calc_by_ptr_offset))))
         end
     else
         opu₂ = !isnothing(suffix) && u₂loopsym ∈ loopdeps
@@ -24,20 +23,20 @@ function lower_load_scalar!(
             condsym = varassignname(condvar, u, condu₁)
             varname = varassignname(mvar, u, u₁loopsym ∈ loopdependencies(op))
             td = UnrollArgs(ua, u)
-            load = Expr(:call, lv(:vload), ptr, mem_offset_u(op, td, unrolled))
+            load = Expr(:call, lv(:vload), ptr, mem_offset_u(op, td, inds_calc_by_ptr_offset))
             cload = Expr(:if, condsym, load, Expr(:call, :zero, Expr(:call, :eltype, ptr)))
             push!(q.args, Expr(:(=), varname, cload))
         end
     end
     nothing
 end
 function pushvectorload!(
-    q::Expr, op::Operation, var::Symbol, td::UnrollArgs, U::Int, vectorized::Symbol, mask, u₁unrolled::Bool, unrolled::Bool
+    q::Expr, op::Operation, var::Symbol, td::UnrollArgs, U::Int, vectorized::Symbol, mask, u₁unrolled::Bool, inds_calc_by_ptr_offset::Vector{Bool}
 )
     @unpack u₁, u₁loopsym, u₂loopsym, suffix = td
-    ptr = unrolled ? offset_refname(op, td) : refname(op)
+    ptr = refname(op)
     vecnotunrolled = vectorized !== u₁loopsym
-    name, mo = name_memoffset(var, op, td, u₁unrolled, unrolled)
+    name, mo = name_memoffset(var, op, td, u₁unrolled, inds_calc_by_ptr_offset)
     instrcall = Expr(:call, lv(:vload), ptr, mo)
 
     iscondstore = instruction(op).instr === :conditionalload
@@ -67,20 +66,27 @@ function pushvectorload!(
 end
 function prefetchisagoodidea(ls::LoopSet, op::Operation, td::UnrollArgs)
     # return false
-    @unpack u₁, u₁loopsym, u₂loopsym, vectorized, suffix = td
+    @unpack u₁, u₁loopsym, u₂loopsym, vectorized, u₂max, suffix = td
     vectorized ∈ loopdependencies(op) || return 0
     u₂loopsym === Symbol("##undefined##") && return 0
     dontskip = (64 ÷ VectorizationBase.REGISTER_SIZE) - 1
-    (!isnothing(suffix) && (vectorized === u₂loopsym) && !iszero(suffix & dontskip)) && return 0
-    innermostloopsym = last(ls.loop_order.bestorder)
+    # u₂loopsym is vectorized
+    # u₁vectorized = vectorized === u₁loopsym
+    u₂vectorized = vectorized === u₂loopsym
+    (!isnothing(suffix) && u₂vectorized && !iszero(suffix & dontskip)) && return 0
+    innermostloopsym = first(names(ls))
     loopedindex = op.ref.loopedindex
     if length(loopedindex) > 1 && first(loopedindex)
         indices = getindices(op)
         if first(indices) === vectorized && last(indices) === innermostloopsym
+            # We want at least 4 reuses per load
+            uses = ifelse(isu₁unrolled(op), 1, u₁)
+            uses = ifelse(isu₂unrolled(op), uses, uses * u₂max)
+            uses < 4 && return 0
             innermostloopindv = findall(map(isequal(innermostloopsym), getindices(op)))
             isone(length(innermostloopindv)) || return 0
             innermostloopind = first(innermostloopindv)
-            if prod(s -> length(getloop(ls, s)), @view(indices[1:innermostloopind-1])) ≥ 120 && length(getloop(ls, innermostloopind)) ≥ 120
+            if prod(s -> length(getloop(ls, s)), @view(indices[1:innermostloopind-1])) ≥ 120 && length(getloop(ls, innermostloopsym)) ≥ 120
                 if op.ref.ref.offsets[innermostloopind] < 120
                     for opp ∈ operations(ls)
                         iscompute(opp) && load_constrained(opp, u₁loopsym, u₂loopsym) && return 0
@@ -97,9 +103,9 @@ function lower_load_vectorized!(
 )
     @unpack u₁, u₁loopsym, u₂loopsym, vectorized, suffix = td
     loopdeps = loopdependencies(op)
-    @assert vectorized ∈ loopdeps
-    opu₁ = u₁loopsym ∈ loopdeps
-    unrolled = opu₁ || u₂loopsym ∈ loopdeps
+    @assert isvectorized(op)
+    opu₁ = isu₁unrolled(op)
+    inds_calc_by_ptr_offset = indices_calculated_by_pointer_offsets(ls, op.ref)
     if opu₁
         umin = umin
         U = u₁
@@ -111,13 +117,13 @@ function lower_load_vectorized!(
     var = variable_name(op, suffix)
     for u ∈ umin:U-1
         td = UnrollArgs(td, u)
-        pushvectorload!(q, op, var, td, U, vectorized, mask, opu₁, unrolled)
+        pushvectorload!(q, op, var, td, U, vectorized, mask, opu₁, inds_calc_by_ptr_offset)
     end
     prefetchind = prefetchisagoodidea(ls, op, td)
     if !iszero(prefetchind)
         dontskip = (64 ÷ VectorizationBase.REGISTER_SIZE) - 1
-        ptr = offset_refname(op, td)
-        innermostloopsym = last(ls.loop_order.bestorder)
+        ptr = refname(op)
+        innermostloopsym = first(names(ls))
         us = ls.unrollspecification[]
         prefetch_multiplier = 4
         prefetch_distance = u₁loopsym === innermostloopsym ? us.u₁ : ( u₂loopsym === innermostloopsym ? us.u₂ : 1 )
@@ -128,7 +134,7 @@ function lower_load_vectorized!(
         # for u ∈ umin:min(umin,U-1)
             (u₁loopsym === vectorized && !iszero(u & dontskip)) && continue
             offsets[prefetchind] = inner_offset + prefetch_distance
-            mo = mem_offset_u(op, UnrollArgs(td, u), true)
+            mo = mem_offset_u(op, UnrollArgs(td, u), inds_calc_by_ptr_offset)
             instrcall = Expr(:call, lv(:prefetch0), ptr, mo)
             push!(q.args, instrcall)
         end
@@ -158,7 +164,7 @@ function lower_load!(
                 varnew = variable_name(op, suffix)
                 varold = variable_name(operations(ls)[id], suffix + mno)
                 opold = operations(ls)[id]
-                if u₁loopsym ∈ loopdependencies(op)
+                if isu₁unrolled(op)
                     for u ∈ 0:u₁-1
                         push!(q.args, Expr(:(=), Symbol(varnew, u), Symbol(varold, u)))
                     end
@@ -173,12 +179,11 @@ function lower_load!(
             umin = 0
         end
     else
-        maybegesp_call!(q, op, td)
         umin = 0
     end
-    if vectorized ∈ loopdependencies(op)
+    if isvectorized(op)
         lower_load_vectorized!(q, ls, op, td, mask, umin)
     else
-        lower_load_scalar!(q, op, td, umin)
+        lower_load_scalar!(q, op, td, umin, indices_calculated_by_pointer_offsets(ls, op.ref))
     end
 end