Progress in supporting VectorizationBase 0.18

Author: chriselrod

Project.toml

@@ -1,7 +1,7 @@
 name = "LoopVectorization"
 uuid = "bdcacae8-1622-11e9-2a5c-532679323890"
 authors = ["Chris Elrod <[email protected]>"]
-version = "0.11.1"
+version = "0.11.2"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
@@ -24,8 +24,8 @@ Requires = "1"
 SLEEFPirates = "0.6.7"
 ThreadingUtilities = "0.2.3"
 UnPack = "1"
-VectorizationBase = "0.17"
-julia = "1.5"
+VectorizationBase = "0.18"
+julia = "1.6"
 
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"

src/LoopVectorization.jl

@@ -5,10 +5,10 @@ module LoopVectorization
 # end
 
 using VectorizationBase, SLEEFPirates, UnPack, OffsetArrays
-using VectorizationBase: register_size, register_count, dynamic_register_size, dynamic_register_count, data,
-    mask, pick_vector_width_val, MM,
+using VectorizationBase: register_size, register_count, cache_linesize, has_opmask_registers,
+    mask, pick_vector_width, MM, data,
     maybestaticlength, maybestaticsize, staticm1, staticp1, staticmul, vzero,
-    Zero, maybestaticrange, offsetprecalc, lazymul,
+    maybestaticrange, offsetprecalc, lazymul,
     maybestaticfirst, maybestaticlast, scalar_less, scalar_greaterequal, gep, gesp, pointerforcomparison, NativeTypes,
     vfmadd, vfmsub, vfnmadd, vfnmsub, vfmadd_fast, vfmsub_fast, vfnmadd_fast, vfnmsub_fast, vfmadd231, vfmsub231, vfnmadd231, vfnmsub231,
     vfma_fast, vmuladd_fast, vdiv_fast, vadd_fast, vsub_fast, vmul_fast,
@@ -30,7 +30,7 @@ using Base.FastMath: add_fast, sub_fast, mul_fast, div_fast, inv_fast, abs2_fast
 
 
 using ArrayInterface
-using ArrayInterface: OptionallyStaticUnitRange, Zero, One#, static_length
+using ArrayInterface: OptionallyStaticUnitRange, Zero, One, StaticBool, True, False#, static_length
 const Static = ArrayInterface.StaticInt
 
 using Requires

src/broadcast.jl

@@ -355,13 +355,15 @@ end
 # size of dest determines loops
 # function vmaterialize!(
 @generated function vmaterialize!(
-    dest::AbstractArray{T,N}, bc::BC, ::Val{Mod}
-) where {T <: NativeTypes, N, BC <: Union{Broadcasted,Product}, Mod}
+    dest::AbstractArray{T,N}, bc::BC,
+    ::Val{Mod}, ::StaticInt{RS}, ::StaticInt{RC}, ::StaticInt{CLS}
+) where {T <: NativeTypes, N, BC <: Union{Broadcasted,Product}, Mod, RS, RC, CLS}
     # 2+1
     # we have an N dimensional loop.
     # need to construct the LoopSet
     # @show typeof(dest)
     ls = LoopSet(Mod)
+    set_hw!(ls, RS, RC, CLS)
     loopsyms = [gensym!(ls, "n") for n ∈ 1:N]
     ls.isbroadcast[] = true
     add_broadcast_loops!(ls, loopsyms, :dest)
@@ -385,11 +387,13 @@ end
      # ls
 end
 @generated function vmaterialize!(
-    dest′::Union{Adjoint{T,A},Transpose{T,A}}, bc::BC, ::Val{Mod}
-) where {T <: NativeTypes, N, A <: AbstractArray{T,N}, BC <: Union{Broadcasted,Product}, Mod}
+    dest′::Union{Adjoint{T,A},Transpose{T,A}}, bc::BC,
+    ::Val{Mod}, ::StaticInt{RS}, ::StaticInt{RC}, ::StaticInt{CLS}
+) where {T <: NativeTypes, N, A <: AbstractArray{T,N}, BC <: Union{Broadcasted,Product}, Mod, RS, RC, CLS}
     # we have an N dimensional loop.
     # need to construct the LoopSet
     ls = LoopSet(Mod)
+    set_hw!(ls, RS, RC, CLS)
     loopsyms = [gensym!(ls, "n") for n ∈ 1:N]
     ls.isbroadcast[] = true
     pushprepreamble!(ls, Expr(:(=), :dest, Expr(:call, :parent, :dest′)))
@@ -411,7 +415,8 @@ end
 end
 # these are marked `@inline` so the `@avx` itself can choose whether or not to inline.
 @inline function vmaterialize!(
-    dest::AbstractArray{T,N}, bc::Broadcasted{Base.Broadcast.DefaultArrayStyle{0},Nothing,typeof(identity),Tuple{T2}}, ::Val{Mod}
+    dest::AbstractArray{T,N}, bc::Broadcasted{Base.Broadcast.DefaultArrayStyle{0},Nothing,typeof(identity),Tuple{T2}},
+    ::Val{Mod}, RS::Static, RC::Static, CLS::Static
 ) where {T <: NativeTypes, N, T2 <: Number, Mod}
     arg = T(first(bc.args))
     @avx for i ∈ eachindex(dest)
@@ -420,7 +425,8 @@ end
     dest
 end
 @inline function vmaterialize!(
-    dest′::Union{Adjoint{T,A},Transpose{T,A}}, bc::Broadcasted{Base.Broadcast.DefaultArrayStyle{0},Nothing,typeof(identity),Tuple{T2}}, ::Val{Mod}
+    dest′::Union{Adjoint{T,A},Transpose{T,A}}, bc::Broadcasted{Base.Broadcast.DefaultArrayStyle{0},Nothing,typeof(identity),Tuple{T2}},
+    ::Val{Mod}, RS::Static, RC::Static, CLS::Static
 ) where {T <: NativeTypes, N, A <: AbstractArray{T,N}, T2 <: Number, Mod}
     arg = T(first(bc.args))
     dest = parent(dest′)
@@ -430,10 +436,10 @@ end
     dest′
 end
 
-@inline function vmaterialize(bc::Broadcasted, ::Val{Mod}) where {Mod}
+@inline function vmaterialize(bc::Broadcasted, ::Val{Mod}, RS::Static, RC::Static, CLS::Static) where {Mod}
     ElType = Base.Broadcast.combine_eltypes(bc.f, bc.args)
-    vmaterialize!(similar(bc, ElType), bc, Val{Mod}())
+    vmaterialize!(similar(bc, ElType), bc, Val{Mod}(), RS, RC, CLS)
 end
 
-vmaterialize!(dest, bc, ::Val{mod}) where {mod} = Base.Broadcast.materialize!(dest, bc)
+vmaterialize!(dest, bc, ::Val{mod}, ::StaticInt, ::StaticInt, ::StaticInt) where {mod} = Base.Broadcast.materialize!(dest, bc)
 

src/condense_loopset.jl

@@ -239,8 +239,21 @@ function generate_call(ls::LoopSet, inline_unroll::NTuple{3,Int8}, debug::Bool =
     inline, u₁, u₂ = inline_unroll
     func = debug ? lv(:_avx_loopset_debug) : lv(:_avx_!)
     lbarg = debug ? Expr(:call, :typeof, loop_bounds) : loop_bounds
+    unroll_param_tup = Expr(
+        :tuple, inline, u₁, u₂,
+        Expr(:call, lv(:unwrap), VECTORWIDTHSYMBOL),
+        Expr(:call, lv(:unwrap), Expr(:call, lv(:register_size))),
+        Expr(:call, lv(:unwrap),
+            Expr(:call, lv(:ifelse),
+                Expr(:call, lv(:unwrap), Expr(:call, lv(:has_opmask_registers))),
+                Expr(:call, lv(:unwrap), Expr(:call, lv(:register_count))),
+                Expr(:call, lv(:unwrap), Expr(:call, :(-), Expr(:call, lv(:register_count)), Expr(:call, lv(:One))))
+            )
+        ),
+        Expr(:call, lv(:unwrap), Expr(:call, lv(:cache_linesize)))
+    )
     q = Expr(
-        :call, func, val(Expr(:tuple, inline, u₁, u₂, Expr(:call, lv(:unwrap), VECTORWIDTHSYMBOL))),
+        :call, func, val(unroll_param_tup),
         val(operation_descriptions), val(arrayref_descriptions), val(argmeta), val(loop_syms)
     )
     # debug && deleteat!(q.args, 2)

src/constructors.jl

@@ -29,7 +29,12 @@ function add_ci_call!(q::Expr, @nospecialize(f), args, syms, i, mod = nothing)
             push!(call.args, arg)
         end
     end
-    mod === nothing || push!(call.args, Expr(:call, Expr(:curly, :Val, QuoteNode(mod))))
+    if mod !== nothing # indicates it's `vmaterialize(!)`
+        reg_size = Expr(:call, lv(:register_size))
+        reg_count = Expr(:call, lv(:register_size))
+        cache_lnsze = Expr(:call, lv(:cache_linesize))
+        push!(call.args, Expr(:call, Expr(:curly, :Val, QuoteNode(mod))), reg_size, reg_count, cache_lnsze)
+    end
     push!(q.args, Expr(:(=), syms[i], call))
 end
 
@@ -65,6 +70,12 @@ function LoopSet(q::Expr, mod::Symbol = :Main)
 end
 LoopSet(q::Expr, m::Module) = LoopSet(macroexpand(m, q)::Expr, Symbol(m))
 
+function loopset(q::Expr) # for interactive use only
+    ls = LoopSet(q)
+    set_hw!(ls)
+    ls
+end
+
 """
     @avx
 
@@ -235,13 +246,15 @@ This macro accepts the `inline` and `unroll` keyword arguments like `@avx`, but
 macro _avx(q)
     q = macroexpand(__module__, q)
     ls = LoopSet(q, __module__)
+    set_hw!(ls)
     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, check_empty, u₁, u₂ = check_macro_kwarg(arg)
     ls = LoopSet(q, __module__)
+    set_hw!(ls)
     esc(Expr(:block, ls.prepreamble, lower(ls, u₁ % Int, u₂ % Int, -1)))
 end