When unrolling, calculate initial offset pointer, and then calculate remaining loads with respect to that.

Author: chriselrod

Project.toml

@@ -17,7 +17,7 @@ DocStringExtensions = "0.8"
 OffsetArrays = "1"
 SIMDPirates = "0.7.24"
 SLEEFPirates = "0.4.8"
-UnPack = "0"
+UnPack = "0,1"
 VectorizationBase = "0.11.3"
 julia = "1.1"
 

src/LoopVectorization.jl

@@ -3,11 +3,11 @@ 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, valadd, valsub, _MM,
-    maybestaticlength, maybestaticsize, staticm1, subsetview, vzero, stridedpointer_for_broadcast,
+    maybestaticlength, maybestaticsize, staticm1, staticp1, subsetview, vzero, stridedpointer_for_broadcast,
     Static, StaticUnitRange, StaticLowerUnitRange, StaticUpperUnitRange, unwrap, maybestaticrange,
     AbstractColumnMajorStridedPointer, AbstractRowMajorStridedPointer, AbstractSparseStridedPointer, AbstractStaticStridedPointer,
     PackedStridedPointer, SparseStridedPointer, RowMajorStridedPointer, StaticStridedPointer, StaticStridedStruct,
-    maybestaticfirst, maybestaticlast, scalar_less, scalar_greater, noalias!
+    maybestaticfirst, maybestaticlast, scalar_less, scalar_greater, noalias!, gesp
 using SIMDPirates: VECTOR_SYMBOLS, evadd, evsub, evmul, evfdiv, vrange, 
     reduced_add, reduced_prod, reduce_to_add, reduced_max, reduced_min, vsum, vprod, vmaximum, vminimum,
     sizeequivalentfloat, sizeequivalentint, vadd!, vsub!, vmul!, vfdiv!, vfmadd!, vfnmadd!, vfmsub!, vfnmsub!,

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)))
+        push!(vloadcall.args, mem_offset(op, UnrollArgs(0, Symbol(""), Symbol(""), Symbol(""), nothing), false, false))
     end
     pushpreamble!(ls, Expr(:(=), temp, vloadcall))
     pushpreamble!(ls, op, temp)

src/condense_loopset.jl

@@ -195,6 +195,11 @@ function add_external_functions!(q::Expr, ls::LoopSet)
     end
 end
 
+function check_if_empty(ls::LoopSet, q::Expr)
+    lb = loop_boundaries(ls)
+    Expr(:if, Expr(:call, :!, Expr(:call, :any, :isempty, lb)), q)
+end
+
 # 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)
@@ -247,6 +252,8 @@ Additionally, define `pointer` and `stride` methods.
 @inline check_args(A::PermutedDimsArray) = check_args(parent(A))
 @inline check_args(A::StridedArray) = check_type(eltype(A))
 @inline check_args(A::AbstractRange) = check_type(eltype(A))
+@inline check_args(A::BitVector) = true
+@inline check_args(A::BitMatrix) = true
 @inline function check_args(A::AbstractArray)
     M = parentmodule(typeof(A))
     if parent(A) === A # SparseMatrix, StaticArray, etc
@@ -305,15 +312,16 @@ 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))
-    ls.preamble
+    Expr(:block, ls.prepreamble, ls.preamble)
 end
-function setup_call(ls::LoopSet, q = nothing, inline::Int8 = zero(Int8), u₁::Int8 = zero(Int8), u₂::Int8 = zero(Int8))
+function setup_call(ls::LoopSet, q = nothing, inline::Int8 = zero(Int8), check_empty::Bool = false, u₁::Int8 = zero(Int8), u₂::Int8 = zero(Int8))
     # We outline/inline at the macro level by creating/not creating an anonymous function.
     # The old API instead was based on inlining or not inline the generated function, but
     # the generated function must be inlined into the initial loop preamble for performance reasons.
     # Creating an anonymous function and calling it also achieves the outlining, while still
     # inlining the generated function into the loop preamble.
     call = setup_call_inline(ls, inline, u₁, u₂)
+    call = check_empty ? check_if_empty(ls, call) : call
     isnothing(q) && return Expr(:block, ls.prepreamble, call)
     Expr(:block, ls.prepreamble, Expr(:if, check_args_call(ls), call, make_fast_and_crashy(q)))
 end

src/constructors.jl

@@ -112,6 +112,10 @@ where `inline=false` is faster, so the best setting may require experimentation.
 tries to guess. Currently the algorithm is simple: roughly, if there are more than two dynamically sized loops
 or and no convolutions, it will probably not force inlining. Otherwise, it probably will.
 
+`check_empty` (default is `false`) determines whether or not it will check if any of the iterators are empty.
+If false, you must ensure yourself that they are not empty, else the behavior of the loop is undefined and
+(like with `@inbounds`) segmentation faults are likely.
+
 `unroll` is an integer that specifies the loop unrolling factor, or a
 tuple `(u₁, u₂) = (4, 2)` signaling that the generated code should unroll more than
 one loop. `u₁` is the unrolling factor for the first unrolled loop and `u₂` for the next (if present),
@@ -165,30 +169,46 @@ function check_unroll(arg)
     end
     u₁, u₂
 end
-function check_macro_kwarg(arg, inline::Int8 = zero(Int8), u₁::Int8 = zero(Int8), u₂::Int8 = zero(Int8))
+function check_checkempty(arg)
+    arg.args[1] === :check_empty ? (arg.args[2])::Bool : nothing
+end
+function check_macro_kwarg(arg, inline::Int8 = zero(Int8), check_empty::Bool = false, u₁::Int8 = zero(Int8), u₂::Int8 = zero(Int8))
     @assert arg.head === :(=)
     i = check_inline(arg)
     if iszero(i)
-        u₁, u₂ = check_unroll(arg)
+        ce = check_checkempty(arg)
+        if isnothing(ce)
+            u₁, u₂ = check_unroll(arg)
+        else
+            check_empty = ce
+        end
     else
         inline = i
     end
-    inline, u₁, u₂
+    inline, check_empty, u₁, u₂
 end
 macro avx(arg, q)
     @assert q.head === :for
     @assert arg.head === :(=)
     q = macroexpand(__module__, q)
-    inline, u₁, u₂ = check_macro_kwarg(arg)
+    inline, check_empty, u₁, u₂ = check_macro_kwarg(arg)
     ls = LoopSet(q, __module__)
-    esc(setup_call(ls, q, inline, u₁, u₂))
+    esc(setup_call(ls, q, inline, check_empty, u₁, u₂))
 end
 macro avx(arg1, arg2, q)
     @assert q.head === :for
     q = macroexpand(__module__, q)
-    inline, u₁, u₂ = check_macro_kwarg(arg1)
-    inline, u₁, u₂ = check_macro_kwarg(arg2, inline, u₁, u₂)
-    esc(setup_call(LoopSet(q, __module__), q, inline, u₁, u₂))
+    inline, check_empty, u₁, u₂ = check_macro_kwarg(arg1)
+    inline, check_empty, u₁, u₂ = check_macro_kwarg(arg2, inline, check_empty, u₁, u₂)
+    esc(setup_call(LoopSet(q, __module__), q, inline, check_empty, u₁, u₂))
+end
+macro avx(arg1, arg2, arg3, q)
+    @assert q.head === :for
+    q = macroexpand(__module__, q)
+    inline, check_empty, u₁, u₂ = check_macro_kwarg(arg1)
+    inline, check_empty, u₁, u₂ = check_macro_kwarg(arg2, inline, check_empty, u₁, u₂)
+    inline, check_empty, u₁, u₂ = check_macro_kwarg(arg3, inline, check_empty, u₁, u₂)
+    esc(setup_call(LoopSet(q, __module__), q, inline, check_empty, u₁, u₂))
 end
 
 
@@ -197,20 +217,25 @@ end
 
 This macro transforms loops similarly to [`@avx`](@ref).
 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.
+works on just the expressions. This requires that it makes a number of default assumptions. Use of `@avx` is preferred.
+
+This macro accepts the `inline` and `unroll` keyword arguments like `@avx`, but ignores the `check_empty` argument.
 """
 macro _avx(q)
     q = macroexpand(__module__, q)
-    esc(lower_and_split_loops(LoopSet(q, __module__), -1))
+    ls = LoopSet(q, __module__)
+    esc(Expr(:block, ls.prepreamble, lower_and_split_loops(ls, -1)))
 end
 macro _avx(arg, q)
     @assert q.head === :for
     q = macroexpand(__module__, q)
-    inline, u₁, u₂ = check_macro_kwarg(arg)
-    esc(lower(LoopSet(q, __module__), u₁ % Int, u₂ % Int, -1))
+    inline, check_empty, u₁, u₂ = check_macro_kwarg(arg)
+    ls = LoopSet(q, __module__)
+    esc(Expr(:block, ls.prepreamble, lower(ls, u₁ % Int, u₂ % Int, -1)))
 end
 
 macro avx_debug(q)
     q = macroexpand(__module__, q)
-    esc(LoopVectorization.setup_call_debug(LoopSet(q, __module__)))
+    ls = LoopSet(q, __module__)
+    esc(LoopVectorization.setup_call_debug(ls))
 end

src/determinestrategy.jl

@@ -767,7 +767,7 @@ function choose_unroll_order(ls::LoopSet, lowest_cost::Float64 = Inf)
             cost_temp = evaluate_cost_unroll(ls, new_order, new_vec, lowest_cost)
             if cost_temp < lowest_cost
                 lowest_cost = cost_temp
-                best_order = new_order
+                copyto!(best_order, new_order)
                 best_vec = new_vec
             end
         end

src/graphs.jl

@@ -71,13 +71,13 @@ Base.length(loop::Loop) = 1 + loop.stophint - loop.starthint
 isstaticloop(loop::Loop) = loop.startexact & loop.stopexact
 function startloop(loop::Loop, isvectorized, 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
-    elseif 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)
     else
         Expr(:(=), itersymbol, Expr(:call, lv(:unwrap), loop.startsym))
@@ -91,17 +91,17 @@ function vec_looprange(loop::Loop, UF::Int, mangledname::Symbol)
         Expr(:call, lv(:valsub), VECTORWIDTHSYMBOL, 2)
     end
     if loop.stopexact # split for type stability
-        Expr(:call, lv(:scalar_less), mangledname, Expr(:call, :-, loop.stophint, incr))
+        Expr(:call, :<, mangledname, Expr(:call, lv(:vsub), loop.stophint, incr))
     else
-        Expr(:call, lv(:scalar_less), mangledname, Expr(:call, :-, loop.stopsym, incr))
+        Expr(:call, :<, mangledname, Expr(:call, lv(:vsub), loop.stopsym, incr))
     end
 end
 function looprange(loop::Loop, incr::Int, mangledname::Symbol)
     incr = 2 - incr
     if iszero(incr)
-        Expr(:call, lv(:scalar_less), mangledname, loop.stopexact ? loop.stophint : loop.stopsym)
+        Expr(:call, :<, mangledname, loop.stopexact ? loop.stophint : loop.stopsym)
     else
-        Expr(:call, lv(:scalar_less), mangledname, loop.stopexact ? loop.stophint + incr : Expr(:call, :+, loop.stopsym, incr))
+        Expr(:call, :<, mangledname, loop.stopexact ? loop.stophint + incr : Expr(:call, lv(:vadd), loop.stopsym, incr))
     end
 end
 function terminatecondition(
@@ -119,11 +119,12 @@ function incrementloopcounter(us::UnrollSpecification, n::Int, mangledname::Symb
     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(:valmul), VECTORWIDTHSYMBOL, UF))
+            Expr(:(=), mangledname, Expr(:call, lv(:valmuladd), VECTORWIDTHSYMBOL, UF, mangledname))
         end
     else
-        Expr(:+=, mangledname, UF)
+        Expr(:(=), mangledname, Expr(:call, lv(:vadd), mangledname, UF))
     end
 end
 
@@ -367,7 +368,7 @@ function add_loop_bound!(ls::LoopSet, itersym::Symbol, bound, upper::Bool = true
     (bound isa Symbol && upper) && return bound
     bound isa Expr && maybestatic!(bound)
     N = gensym(string(itersym) * (upper ? "_loop_upper_bound" : "_loop_lower_bound"))
-    pushpreamble!(ls, Expr(:(=), N, bound))
+    pushprepreamble!(ls, Expr(:(=), N, bound))
     N
 end
 
@@ -410,20 +411,20 @@ function register_single_loop!(ls::LoopSet, looprange::Expr)
             else
                 otN isa Expr && maybestatic!(otN)
                 N = gensym("loop" * string(itersym))
-                pushpreamble!(ls, Expr(:(=), N, otN))
+                pushprepreamble!(ls, Expr(:(=), N, otN))
                 Loop(itersym, 1, N)
             end
         else
             N = gensym("loop" * string(itersym))
-            pushpreamble!(ls, Expr(:(=), N, Expr(:call, lv(:maybestaticrange), r)))
+            pushprepreamble!(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`
         N = gensym("loop" * string(itersym))
-        pushpreamble!(ls, Expr(:(=), N, Expr(:call, lv(:maybestaticrange), r)))
+        pushprepreamble!(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 = Loop(itersym, L, U)

src/lower_compute.jl

@@ -71,9 +71,9 @@ end
 function add_loopvalue!(instrcall::Expr, loopval::Symbol, vectorized::Symbol, u::Int)
     if loopval === vectorized
         if isone(u)
-            push!(instrcall.args, Expr(:call, lv(:valadd), VECTORWIDTHSYMBOL, loopval))
+            push!(instrcall.args, Expr(:call, lv(:valadd), VECTORWIDTHSYMBOL, _MMind(loopval)))
         else
-            push!(instrcall.args, Expr(:call, lv(:valmuladd), VECTORWIDTHSYMBOL, u, loopval))
+            push!(instrcall.args, Expr(:call, lv(:valmuladd), VECTORWIDTHSYMBOL, u, _MMind(loopval)))
         end
     else
         push!(instrcall.args, Expr(:call, :+, loopval, u))
@@ -85,6 +85,8 @@ function add_loopvalue!(instrcall::Expr, loopval, ua::UnrollArgs, u::Int)
         add_loopvalue!(instrcall, loopval, vectorized, u)
     elseif !isnothing(suffix) && suffix > 0 && loopval === u₂loopsym
         add_loopvalue!(instrcall, loopval, vectorized, suffix)
+    elseif loopval === vectorized
+        push!(instrcall.args, _MMind(loopval))
     else
         push!(instrcall.args, loopval)
     end

src/lower_load.jl

@@ -6,14 +6,15 @@ function lower_load_scalar!(
     @assert vectorized ∉ loopdeps
     # mvar, opu₁, opu₂ = variable_name_and_unrolled(op, u₁loop, u₂loop, suffix)
     mvar = variable_name(op, suffix)
-    ptr = refname(op)
     opu₁ = u₁loopsym ∈ loopdeps
+    unrolled = (opu₁ || u₂loopsym ∈ loopdeps)
+    ptr = unrolled ? offset_refname(op) : 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))))
+            push!(q.args, Expr(:(=), varname, Expr(:call, lv(:vload), ptr, mem_offset_u(op, td, unrolled))))
         end
     else
         opu₂ = !isnothing(suffix) && u₂loopsym ∈ loopdeps
@@ -23,20 +24,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))
+            load = Expr(:call, lv(:vload), ptr, mem_offset_u(op, td, unrolled))
             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
+    q::Expr, op::Operation, var::Symbol, td::UnrollArgs, U::Int, vectorized::Symbol, mask, u₁unrolled::Bool, unrolled::Bool
 )
     @unpack u₁, u₁loopsym, u₂loopsym, suffix = td
-    ptr = refname(op)
+    ptr = unrolled ? offset_refname(op) : refname(op)
     vecnotunrolled = vectorized !== u₁loopsym
-    name, mo = name_memoffset(var, op, td)
+    name, mo = name_memoffset(var, op, td, u₁unrolled, unrolled)
     instrcall = Expr(:call, lv(:vload), ptr, mo)
 
     iscondstore = instruction(op).instr === :conditionalload
@@ -96,7 +97,9 @@ function lower_load_vectorized!(
     @unpack u₁, u₁loopsym, u₂loopsym, vectorized, suffix = td
     loopdeps = loopdependencies(op)
     @assert vectorized ∈ loopdeps
-    if u₁loopsym ∈ loopdeps
+    opu₁ = u₁loopsym ∈ loopdeps
+    unrolled = opu₁ || u₂loopsym ∈ loopdeps
+    if opu₁
         umin = umin
         U = u₁
     else
@@ -107,25 +110,24 @@ 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)
+        pushvectorload!(q, op, var, td, U, vectorized, mask, opu₁, unrolled)
     end
     prefetchind = prefetchisagoodidea(ls, op, td)
     if !iszero(prefetchind)
         dontskip = (64 ÷ VectorizationBase.REGISTER_SIZE) - 1
-        ptr = refname(op)
+        ptr = offset_refname(op)
         innermostloopsym = last(ls.loop_order.bestorder)
         us = ls.unrollspecification[]
         prefetch_multiplier = 4
         prefetch_distance = u₁loopsym === innermostloopsym ? us.u₁ : ( u₂loopsym === innermostloopsym ? us.u₂ : 1 )
         prefetch_distance *= prefetch_multiplier
         offsets = op.ref.ref.offsets
         inner_offset = offsets[prefetchind]
-        
         for u ∈ umin:U-1
         # for u ∈ umin:min(umin,U-1)
             (u₁loopsym === vectorized && !iszero(u & dontskip)) && continue
             offsets[prefetchind] = inner_offset + prefetch_distance
-            mo = last(name_memoffset(var, op, UnrollArgs(td, u)))
+            mo = mem_offset_u(op, UnrollArgs(td, u), true)
             instrcall = Expr(:call, lv(:prefetch0), ptr, mo)
             push!(q.args, instrcall)
         end
@@ -170,6 +172,7 @@ function lower_load!(
             umin = 0
         end
     else
+        maybegesp_call!(q, op, td)
         umin = 0
     end
     if vectorized ∈ loopdependencies(op)