Update to Polyester 0.4, fixes #330.

Author: chriselrod

Project.toml

@@ -1,12 +1,13 @@
 name = "LoopVectorization"
 uuid = "bdcacae8-1622-11e9-2a5c-532679323890"
 authors = ["Chris Elrod <[email protected]>"]
-version = "0.12.66"
+version = "0.12.67"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 IfElse = "615f187c-cbe4-4ef1-ba3b-2fcf58d6d173"
+LayoutPointers = "10f19ff3-798f-405d-979b-55457f8fc047"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
@@ -19,16 +20,17 @@ UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 VectorizationBase = "3d5dd08c-fd9d-11e8-17fa-ed2836048c2f"
 
 [compat]
-ArrayInterface = "3.1.9 - 3.1.23"
+ArrayInterface = "3.1.9 - 3.1.23, 3.1.25"
 DocStringExtensions = "0.8"
 IfElse = "0.1"
+LayoutPointers = "0.1.2"
 OffsetArrays = "1.4.1"
-Polyester = "0.3"
+Polyester = "0.4.0"
 Requires = "1"
 SLEEFPirates = "0.6.23"
 Static = "0.2, 0.3"
-StrideArraysCore = "0.1.12"
+StrideArraysCore = "0.2"
 ThreadingUtilities = "0.4.5"
 UnPack = "1"
-VectorizationBase = "0.20.36"
+VectorizationBase = "0.21.1"
 julia = "1.5"

src/LoopVectorization.jl

@@ -10,7 +10,7 @@ using VectorizationBase: register_size, register_count, cache_linesize, cache_si
     maybestaticfirst, maybestaticlast, gep, gesp, NativeTypes, #llvmptr,
     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,
-    relu, stridedpointer, stridedpointer_preserve, StridedPointer, StridedBitPointer, AbstractStridedPointer, _vload, _vstore!,
+    relu, stridedpointer, StridedPointer, StridedBitPointer, AbstractStridedPointer, _vload, _vstore!,
     reduced_add, reduced_prod, reduce_to_add, reduce_to_prod, reduced_max, reduced_min, reduce_to_max, reduce_to_min,
     reduced_all, reduced_any, reduce_to_all, reduce_to_any,
     vsum, vprod, vmaximum, vminimum, vany, vall, unwrap, Unroll, VecUnroll,
@@ -22,9 +22,11 @@ using VectorizationBase: register_size, register_count, cache_linesize, cache_si
     contract_and, collapse_and,
     contract_or,  collapse_or,
     num_threads, num_cores,
-    max_mask#,zero_mask
+    max_mask, maybestaticsize#,zero_mask
 
-using VectorizationBase: maybestaticsize # for compatibility
+
+
+using LayoutPointers: stridedpointer_preserve, GroupedStridedPointers
 
 using IfElse: ifelse
 
@@ -39,7 +41,7 @@ using Base.FastMath: add_fast, sub_fast, mul_fast, div_fast, inv_fast, abs2_fast
 using SLEEFPirates: log_fast, log2_fast, log10_fast, pow, sin_fast, cos_fast, sincos_fast
 
 using ArrayInterface
-using ArrayInterface: OptionallyStaticUnitRange, OptionallyStaticRange, Zero, One, StaticBool, True, False, reduce_tup, indices, UpTri, LoTri
+using ArrayInterface: OptionallyStaticUnitRange, OptionallyStaticRange, Zero, One, StaticBool, True, False, reduce_tup, indices, UpTri, LoTri, strides, offsets, size, StrideIndex
 using StrideArraysCore: CloseOpen, PtrArray
 # @static if VERSION ≥ v"1.6.0-rc1" #TODO: delete `else` when dropping 1.5 support
     # using ArrayInterface: static_step

src/broadcast.jl

@@ -7,21 +7,21 @@ end
 Base.@propagate_inbounds Base.getindex(A::LowDimArray, i::Vararg{Union{Integer,CartesianIndex},K}) where {K} = getindex(A.data, i...)
 @inline Base.size(A::LowDimArray) = Base.size(A.data)
 @inline Base.size(A::LowDimArray, i) = Base.size(A.data, i)
-@inline Base.strides(A::LowDimArray) = strides(A.data)
+
 @inline ArrayInterface.parent_type(::Type{LowDimArray{D,T,N,A}}) where {T,D,N,A} = A
-@inline ArrayInterface.strides(A::LowDimArray) = ArrayInterface.strides(A.data)
+@inline Base.strides(A::LowDimArray) = map(Int, strides(A))
 @inline ArrayInterface.device(A::LowDimArray) = ArrayInterface.CPUPointer()
 @generated function ArrayInterface.size(A::LowDimArray{D,T,N}) where {D,T,N}
-    t = Expr(:tuple)
-    gf = GlobalRef(Core,:getfield)
-    for n ∈ 1:N
-        if n > length(D) || D[n]
-            push!(t.args, Expr(:call, gf, :s, n, false))
-        else
-            push!(t.args, Expr(:call, Expr(:curly, lv(:StaticInt), 1)))
-        end
+  t = Expr(:tuple)
+  gf = GlobalRef(Core,:getfield)
+  for n ∈ 1:N
+    if n > length(D) || D[n]
+      push!(t.args, Expr(:call, gf, :s, n, false))
+    else
+      push!(t.args, Expr(:call, Expr(:curly, lv(:StaticInt), 1)))
     end
-    Expr(:block, Expr(:meta,:inline), :(s = ArrayInterface.size(parent(A))), t)
+  end
+  Expr(:block, Expr(:meta,:inline), :(s = ArrayInterface.size(parent(A))), t)
 end
 Base.parent(A::LowDimArray) = getfield(A, :data)
 Base.unsafe_convert(::Type{Ptr{T}}, A::LowDimArray{D,T}) where {D,T} = pointer(parent(A))
@@ -31,103 +31,174 @@ ArrayInterface.stride_rank(A::LowDimArray) = ArrayInterface.stride_rank(parent(A
 ArrayInterface.offsets(A::LowDimArray) = ArrayInterface.offsets(parent(A))
 
 @generated function _lowdimfilter(::Val{D}, tup::Tuple{Vararg{Any,N}}) where {D,N}
-    t = Expr(:tuple)
-    gf = GlobalRef(Core,:getfield)
-    for n ∈ 1:N
-        if n > length(D) || D[n]
-            push!(t.args, Expr(:call, gf, :tup, n, false))
-        end
+  t = Expr(:tuple)
+  gf = GlobalRef(Core,:getfield)
+  for n ∈ 1:N
+    if n > length(D) || D[n]
+      push!(t.args, Expr(:call, gf, :tup, n, false))
     end
-    Expr(:block, Expr(:meta,:inline), t)
+  end
+  Expr(:block, Expr(:meta,:inline), t)
 end
 
 struct ForBroadcast{T,N,A<:AbstractArray{T,N}} <: AbstractArray{T,N}
-    data::A
+  data::A
 end
 @inline Base.parent(fb::ForBroadcast) = getfield(fb, :data)
+@inline ArrayInterface.parent_type(::Type{ForBroadcast{T,N,A}}) where {T,N,A} = A
 Base.@propagate_inbounds Base.getindex(A::ForBroadcast, i::Vararg{Any,K}) where {K} = parent(A)[i...]
 const LowDimArrayForBroadcast{D,T,N,A} = ForBroadcast{T,N,LowDimArray{D,T,N,A}}
 @inline function VectorizationBase.contiguous_axis(fb::LowDimArrayForBroadcast{D,T,N,A}) where {D,T,N,A}
-    _contiguous_axis(Val{D}(), VectorizationBase.contiguous_axis(parent(parent(fb))))
+  _contiguous_axis(Val{D}(), VectorizationBase.contiguous_axis(parent(parent(fb))))
 end
 @inline forbroadcast(A::AbstractArray) = ForBroadcast(A)
 @inline forbroadcast(A::AbstractRange) = A
 @inline forbroadcast(A) = A
 @inline forbroadcast(A::Adjoint) = forbroadcast(parent(A))
 @inline forbroadcast(A::Transpose) = forbroadcast(parent(A))
-Base.size(fb::ForBroadcast) = size(parent(fb))
+@inline function ArrayInterface.strides(A::Union{LowDimArray,ForBroadcast})
+  B = parent(A)
+  _strides(size(A), strides(B), VectorizationBase.val_stride_rank(B), VectorizationBase.val_dense_dims(B))
+end
+
 
 # @inline function VectorizationBase.contiguous_batch_size(fb::LowDimArrayForBroadcast{D,T,N,A}) where {D,T,N,A}
 #     _contiguous_axis(Val{D}(), VectorizationBase.contiguous_batch_size(parent(parent(fb))))
 # end
 @generated function _contiguous_axis(::Val{D}, ::StaticInt{C}) where {D,C}
-    Dlen = length(D)
-    (C > 0) || return Expr(:block,Expr(:meta,:inline), staticexpr(C))
-    if C ≤ Dlen
-        D[C] || return Expr(:block,Expr(:meta,:inline), staticexpr(-1))
-    end
-    Cnew = 0
-    for n ∈ 1:C
-        Cnew += ((n > Dlen)) || D[n]
-    end
-    Expr(:block,Expr(:meta,:inline), staticexpr(Cnew))
+  Dlen = length(D)
+  (C > 0) || return Expr(:block,Expr(:meta,:inline), staticexpr(C))
+  if C ≤ Dlen
+    D[C] || return Expr(:block,Expr(:meta,:inline), staticexpr(-1))
+  end
+  Cnew = 0
+  for n ∈ 1:C
+    Cnew += ((n > Dlen)) || D[n]
+  end
+  Expr(:block,Expr(:meta,:inline), staticexpr(Cnew))
 end
-@inline function VectorizationBase.val_stride_rank(fb::LowDimArrayForBroadcast{D}) where {D}
-    VectorizationBase.asvalint(_lowdimfilter(Val(D), ArrayInterface.stride_rank(parent(parent(fb)))))
+@inline function ArrayInterface.stride_rank(::Type{LowDimArrayForBroadcast{D,T,N,A}}) where {D,T,N,A}
+  _lowdimfilter(Val(D), ArrayInterface.stride_rank(A))
 end
-@inline function VectorizationBase.val_dense_dims(fb::LowDimArrayForBroadcast{D}) where {D}
-    VectorizationBase.asvalbool(_lowdimfilter(Val(D), ArrayInterface.dense_dims(parent(parent(fb)))))
+@inline function ArrayInterface.dense_dims(::Type{LowDimArrayForBroadcast{D,T,N,A}}) where {D,T,N,A}
+  _lowdimfilter(Val(D), ArrayInterface.dense_dims(A))
 end
-@inline function VectorizationBase.bytestrides(fb::LowDimArrayForBroadcast{D}) where {D}
-    p = parent(parent(fb))
-    s = _bytestrides(ArrayInterface.size(p), ArrayInterface.strides(p), p)
-    _lowdimfilter(Val(D), s)
+@inline function ArrayInterface.strides(fb::LowDimArrayForBroadcast{D}) where {D}
+  _lowdimfilter(Val(D), strides(parent(fb)))
 end
 @inline function ArrayInterface.offsets(fb::LowDimArrayForBroadcast{D}) where {D}
-    _lowdimfilter(Val(D), ArrayInterface.offsets(parent(parent(fb))))
+  _lowdimfilter(Val(D), ArrayInterface.offsets(parent(parent(fb))))
+end
+@inline function ArrayInterface.StrideIndex(a::A) where {A<:LowDimArrayForBroadcast}
+  _stride_index(ArrayInterface.stride_rank(A), ArrayInterface.contiguous_axis(A), a)
+end
+@inline function _stride_index(r::Tuple{Vararg{StaticInt,N}}, ::StaticInt{C}, A) where {N,C}
+  StrideIndex{N,ArrayInterface.known(r),C}(A)
 end
 
+for f ∈ [ # groupedstridedpointer support
+    :(ArrayInterface.contiguous_axis),
+    :(ArrayInterface.contiguous_batch_size),
+    :(ArrayInterface.device),
+    :(ArrayInterface.stride_rank),
+]
+    @eval @inline $f(::Type{ForBroadcast{T,N,A}}) where {T,N,A} = $f(A)
+end
 for f ∈ [ # groupedstridedpointer support
     :(VectorizationBase.memory_reference),
     :(ArrayInterface.contiguous_axis),
     :(ArrayInterface.contiguous_batch_size),
     :(ArrayInterface.device),
-    :(VectorizationBase.val_stride_rank),
+    :(ArrayInterface.stride_rank),
     :(VectorizationBase.val_dense_dims),
-    :(ArrayInterface.offsets)
+    :(ArrayInterface.offsets),
+     :(Base.size)#, :(ArrayInterface.strides)
 ]
     @eval @inline $f(fb::ForBroadcast) = $f(getfield(fb, :data))
 end
-@inline _bytestrides(s,paren) = VectorizationBase.bytestrides(paren)
-@generated function _bytestrides(s::Tuple{Vararg{Integer,N}}, x::Tuple{Vararg{Integer,N}}, paren::AbstractArray{T,N}) where {T,N}
-    q = Expr(:block, Expr(:meta,:inline))
-    strd_tup = Expr(:tuple)
-    gf = GlobalRef(Core, :getfield)
-    ifel = GlobalRef(Core, :ifelse)
-    st = staticexpr(sizeof(T))
-    for n ∈ 1:N
-        s_type = s.parameters[n]
-        if s_type <: Static
-            if s_type === One
-                push!(strd_tup.args, Expr(:call, lv(:Zero)))
-            else
-                push!(strd_tup.args, :($gf(x, $n, false) * $st))
-            end
+
+function is_column_major(x)
+  for (i, j) ∈ enumerate(x)
+    i == j || return false
+  end
+  true
+end
+is_row_major(x) = is_column_major(reverse(x))
+# @inline _bytestrides(s,paren) = VectorizationBase.bytestrides(paren)
+function _strides_expr(@nospecialize(s), @nospecialize(x), R::Vector{Int}, D::Vector{Bool})
+  N = length(R)
+  q = Expr(:block, Expr(:meta,:inline))
+  strd_tup = Expr(:tuple)
+  gf = GlobalRef(Core, :getfield)
+  ifel = GlobalRef(Core, :ifelse)
+  Nrange = 1:1:N # type stability w/ respect to reverse
+  use_stride_acc = true
+  stride_acc::Int = 1
+  if is_column_major(R)
+  elseif is_row_major(R)
+    Nrange = reverse(Nrange)
+  else # not worth my time optimizing this case at the moment...
+    # will write something generic stride-rank agnostic eventually
+    use_stride_acc = false
+    stride_acc = 0
+  end
+  sₙ_value::Int = 0
+  for n ∈ Nrange
+    xₙ_type = x[n]
+    # xₙ_type = typeof(x).parameters[n]
+    xₙ_static = xₙ_type <: Static
+    xₙ_value::Int = xₙ_static ? (xₙ_type.parameters[1])::Int : 0
+    s_type = s[n]
+    # s_type = typeof(s).parameters[n]
+    sₙ_static = s_type <: Static
+    if sₙ_static
+      sₙ_value = s_type.parameters[1]
+      if s_type === One
+        push!(strd_tup.args, Expr(:call, lv(:Zero)))
+      elseif stride_acc ≠ 0
+        push!(strd_tup.args, staticexpr(stride_acc))
+      else
+        push!(strd_tup.args, :($gf(x, $n, false)))
+      end
+    else
+      if xₙ_static
+        push!(strd_tup.args, staticexpr(xₙ_value))
+      elseif stride_acc ≠ 0
+        push!(strd_tup.args, staticexpr(stride_acc))
+      else
+        push!(strd_tup.args, :($ifel(isone($gf(s, $n, false)), zero($xₙ_type), $gf(x, $n, false))))
+      end
+    end
+    if (n ≠ last(Nrange)) && use_stride_acc
+      nnext = n + step(Nrange)
+      if D[nnext]
+        if xₙ_static & sₙ_static
+          stride_acc = xₙ_value * sₙ_value
+        elseif sₙ_static
+          if stride_acc ≠ 0
+            stride_acc *= sₙ_value
+          end
         else
-            Xₙ_type = x.parameters[n]
-            if Xₙ_type <: Static # FIXME; what to do here? Dynamic dispatch? 
-                push!(strd_tup.args, :($gf(x, $n, false)*$st))
-            else
-                push!(strd_tup.args, :($ifel(isone($gf(s, $n, false)), zero($Xₙ_type), $gf(x, $n, false)*$st)))
-            end
+          stride_acc = 0
         end
+      else
+        stride_acc = 0
+      end
     end
-    push!(q.args, strd_tup)
-    q
-end
-@inline function VectorizationBase.bytestrides(fb::ForBroadcast)
-    p = getfield(fb,:data)
-    _bytestrides(ArrayInterface.size(p), ArrayInterface.strides(p), p)
+  end
+  push!(q.args, strd_tup)
+  q
+end
+@generated function _strides(
+  s::Tuple{Vararg{Integer,N}}, x::Tuple{Vararg{Integer,N}}, ::Val{R}, ::Val{D}
+) where {N,R,D}
+  Rv = Vector{Int}(undef, N)
+  Dv = Vector{Bool}(undef, N)
+  for n in 1:N
+    Rv[n] = R[n]
+    Dv[n] = D[n]
+  end
+  _strides_expr(s.parameters, x.parameters, Rv, Dv)
 end
 
 struct Product{A,B}