Merge pull request #61 from chriselrod/loopvecsupport

Added some methods to help support LoopVectorization.jl

Author: chriselrod

.gitignore

@@ -3,3 +3,4 @@
 *.jl.mem
 deps/deps.jl
 Manifest.toml
+*~

Project.toml

@@ -1,6 +1,6 @@
 name = "ArrayInterface"
 uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-version = "2.12.1"
+version = "2.13.0"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -16,10 +16,11 @@ Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BlockBandedMatrices = "ffab5731-97b5-5995-9138-79e8c1846df0"
 LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
+OffsetArrays = "6fe1bfb0-de20-5000-8ca7-80f57d26f881"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 SuiteSparse = "4607b0f0-06f3-5cda-b6b1-a6196a1729e9"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "LabelledArrays", "StaticArrays", "BandedMatrices", "BlockBandedMatrices", "SuiteSparse", "Random", "Aqua"]
+test = ["Test", "LabelledArrays", "StaticArrays", "BandedMatrices", "BlockBandedMatrices", "SuiteSparse", "Random", "OffsetArrays", "Aqua"]

README.md

@@ -134,12 +134,80 @@ Otherwise, returns `nothing`. For example, `known_step(UnitRange{Int})` returns
 If `length` of an instance of type `T` is known at compile time, return it.
 Otherwise, return `nothing`.
 
-## Static(N::Int)
+## device(::Type{T})
+
+Indicates the most efficient way to access elements from the collection in low level code.
+For `GPUArrays`, will return `ArrayInterface.GPU()`.
+For `AbstractArray` supporting a `pointer` method, returns `ArrayInterface.CPUPointer()`.
+For other `AbstractArray`s and `Tuple`s, returns `ArrayInterface.CPUIndex()`.
+Otherwise, returns `nothing`.
+
+## contiguous_axis(::Type{T})
+
+Returns the axis of an array of type `T` containing contiguous data.
+If no axis is contiguous, it returns `Contiguous{-1}`.
+If unknown, it returns `nothing`.
+
+## contiguous_axis_indicator(::Type{T})
+
+Returns a tuple of boolean `Val`s indicating whether that axis is contiguous.
+
+## contiguous_batch_size(::Type{T})
+
+Returns the size of contiguous batches if `!isone(stride_rank(T, contiguous_axis(T)))`.
+If `isone(stride_rank(T, contiguous_axis(T)))`, then it will return `ContiguousBatch{0}()`.
+If `contiguous_axis(T) == -1`, it will return `ContiguousBatch{-1}()`.
+If unknown, it will return `nothing`.
+
+## stride_rank(::Type{T})
+
+Returns the rank of each stride.
+
+## dense_dims(::Type{T})
+Returns a tuple of indicators for whether each axis is dense.
+An axis `i` of array `A` is dense if `stride(A, i) * size(A, i) == stride(A, j)` where `j` is the axis (if it exists) such that `stride_rank(A)[i] + 1 == stride_rank(A)[j]`.
+
+## ArrayInterface.size(A)
+
+Returns the size of `A`. If the size of any axes are known at compile time,
+these should be returned as `StaticInt`s. For example:
+```julia
+julia> using StaticArrays, ArrayInterface
+
+julia> A = @SMatrix rand(3,4);
+
+julia> ArrayInterface.size(A)
+(StaticInt{3}(), StaticInt{4}())
+```
+
+## ArrayInterface.strides(A)
+
+Returns the strides of array `A`. If any strides are known at compile time,
+these should be returned as `StaticInt`s. For example:
+```julia
+julia> using ArrayInterface
+
+julia> A = rand(3,4);
+
+julia> ArrayInterface.strides(A)
+(StaticInt{1}(), 3)
+```
+## offsets(A)
+
+Returns offsets of indices with respect to 0. If values are known at compile time,
+it should return them as `StaticInt`s.
+For example, if `A isa Base.Matrix`, `offsets(A) === (StaticInt(1), StaticInt(1))`.
+
+## can_avx(f)
+
+Is the function `f` whitelisted for `LoopVectorization.@avx`?
+
+## StaticInt(N::Int)
 
 Creates a static integer with value known at compile time. It is a number,
-supporting basic arithmetic. Many operations with two `Static` integers
-will produce another `Static` integer. If one of the arguments to a
-function call isn't static (e.g., `Static(4) + 3`) then the `Static`
+supporting basic arithmetic. Many operations with two `StaticInt` integers
+will produce another `StaticInt` integer. If one of the arguments to a
+function call isn't static (e.g., `StaticInt(4) + 3`) then the `StaticInt`
 number will promote to a dynamic value.
 
 # List of things to add

src/ArrayInterface.jl

@@ -33,8 +33,9 @@ If `length` of an instance of type `T` is known at compile time, return it.
 Otherwise, return `nothing`.
 """
 known_length(x) = known_length(typeof(x))
-known_length(::Type{<:NTuple{N,<:Any}}) where {N} = N
 known_length(::Type{<:NamedTuple{L}}) where {L} = length(L)
+known_length(::Type{T}) where {T<:Base.Slice} = known_length(parent_type(T))
+known_length(::Type{<:Tuple{Vararg{Any,N}}}) where {N} = N
 known_length(::Type{<:Number}) = 1
 function known_length(::Type{T}) where {T}
     if parent_type(T) <: T
@@ -52,7 +53,7 @@ _known_length(x::Tuple{Vararg{Int}}) = prod(x)
 """
     can_change_size(::Type{T}) -> Bool
 
-Returns `true` if the size of `T` can change, in which case operations
+Returns `true` if the Base.size of `T` can change, in which case operations
 such as `pop!` and `popfirst!` are available for collections of type `T`.
 """
 can_change_size(x) = can_change_size(typeof(x))
@@ -102,7 +103,7 @@ function Base.setindex(x::AbstractVector,v,i::Int)
 end
 
 function Base.setindex(x::AbstractMatrix,v,i::Int,j::Int)
-  n,m = size(x)
+  n,m = Base.size(x)
   x .* (i .!== 1:n) .* (j  .!== i:m)' .+ v .* (i .== 1:n) .* (j  .== i:m)'
 end
 
@@ -202,7 +203,7 @@ Return: (I,J) #indexable objects
 Find sparsity pattern of special matrices, the same as the first two elements of findnz(::SparseMatrixCSC)
 """
 function findstructralnz(x::Diagonal)
-  n=size(x,1)
+  n = Base.size(x,1)
   (1:n,1:n)
 end
 
@@ -412,15 +413,15 @@ function Base.getindex(ind::BandedBlockBandedMatrixIndex,i::Int)
 end
 
 function findstructralnz(x::Bidiagonal)
-  n=size(x,1)
+  n= Base.size(x,1)
   isup= x.uplo=='U' ? true : false
   rowind=BidiagonalIndex(n+n-1,isup)
   colind=BidiagonalIndex(n+n-1,!isup)
   (rowind,colind)
 end
 
 function findstructralnz(x::Union{Tridiagonal,SymTridiagonal})
-  n=size(x,1)
+  n= Base.size(x,1)
   rowind=TridiagonalIndex(n+n-1+n-1,n,true)
   colind=TridiagonalIndex(n+n-1+n-1,n,false)
   (rowind,colind)
@@ -447,26 +448,26 @@ fast_matrix_colors(A::Type{<:Union{Diagonal,Bidiagonal,Tridiagonal,SymTridiagona
     matrix_colors(A::Union{Array,UpperTriangular,LowerTriangular})
 
 The color vector for dense matrix and triangular matrix is simply
-`[1,2,3,...,size(A,2)]`
+`[1,2,3,..., Base.size(A,2)]`
 """
 function matrix_colors(A::Union{Array,UpperTriangular,LowerTriangular})
-    eachindex(1:size(A,2)) # Vector size matches number of rows
+    eachindex(1:Base.size(A,2)) # Vector Base.size matches number of rows
 end
 
 function _cycle(repetend,len)
     repeat(repetend,div(len,length(repetend))+1)[1:len]
 end
 
 function matrix_colors(A::Diagonal)
-    fill(1,size(A,2))
+    fill(1, Base.size(A,2))
 end
 
 function matrix_colors(A::Bidiagonal)
-    _cycle(1:2,size(A,2))
+    _cycle(1:2, Base.size(A,2))
 end
 
 function matrix_colors(A::Union{Tridiagonal,SymTridiagonal})
-    _cycle(1:3,size(A,2))
+    _cycle(1:3, Base.size(A,2))
 end
 
 """
@@ -540,9 +541,64 @@ function restructure(x,y)
 end
 
 function restructure(x::Array,y)
-  reshape(convert(Array,y),size(x)...)
+  reshape(convert(Array,y), Base.size(x)...)
 end
 
+abstract type AbstractDevice end
+abstract type AbstractCPU <: AbstractDevice end
+struct CPUPointer <: AbstractCPU end
+struct CheckParent end
+struct CPUIndex <: AbstractCPU end
+struct GPU <: AbstractDevice end
+"""
+device(::Type{T})
+
+Indicates the most efficient way to access elements from the collection in low level code.
+For `GPUArrays`, will return `ArrayInterface.GPU()`.
+For `AbstractArray` supporting a `pointer` method, returns `ArrayInterface.CPUPointer()`.
+For other `AbstractArray`s and `Tuple`s, returns `ArrayInterface.CPUIndex()`.
+Otherwise, returns `nothing`.
+"""
+device(A) = device(typeof(A))
+device(::Type) = nothing
+device(::Type{<:Tuple}) = CPUIndex()
+# Relies on overloading for GPUArrays that have subtyped `StridedArray`.
+device(::Type{<:StridedArray}) = CPUPointer()
+function device(::Type{T}) where {T <: AbstractArray}
+    P = parent_type(T)
+    T === P ? CPUIndex() : device(P)
+end
+
+
+"""
+defines_strides(::Type{T}) -> Bool
+
+Is strides(::T) defined?
+"""
+defines_strides(::Type) = false
+defines_strides(x) = defines_strides(typeof(x))
+defines_strides(::Type{<:StridedArray}) = true
+defines_strides(::Type{A}) where {A <: Union{<:Transpose,<:Adjoint,<:SubArray,<:PermutedDimsArray}} = defines_strides(parent_type(A))
+
+"""
+can_avx(f)
+
+Returns `true` if the function `f` is guaranteed to be compatible with `LoopVectorization.@avx` for supported element and array types.
+While a return value of `false` does not indicate the function isn't supported, this allows a library to conservatively apply `@avx`
+only when it is known to be safe to do so.
+
+```julia
+function mymap!(f, y, args...)
+    if can_avx(f)
+        @avx @. y = f(args...)
+    else
+        @. y = f(args...)
+    end
+end
+```
+"""
+can_avx(::Any) = false
+
 """
     insert(collection, index, item)
 
@@ -654,6 +710,7 @@ function __init__()
     ismutable(::Type{<:StaticArrays.StaticArray}) = false
     can_setindex(::Type{<:StaticArrays.StaticArray}) = false
     ismutable(::Type{<:StaticArrays.MArray}) = true
+    ismutable(::Type{<:StaticArrays.SizedArray}) = true
 
     function lu_instance(_A::StaticArrays.StaticMatrix{N,N}) where {N}
       A = StaticArrays.SArray(_A)
@@ -675,6 +732,26 @@ function __init__()
     known_last(::Type{StaticArrays.SOneTo{N}}) where {N} = N
     known_length(::Type{StaticArrays.SOneTo{N}}) where {N} = N
 
+    device(::Type{<:StaticArrays.MArray}) = CPUPointer()
+    contiguous_axis(::Type{<:StaticArrays.StaticArray}) = Contiguous{1}()
+    contiguous_batch_size(::Type{<:StaticArrays.StaticArray}) = ContiguousBatch{0}()
+    stride_rank(::Type{T}) where {N, T <: StaticArrays.StaticArray{<:Any,<:Any,N}} = StrideRank{ntuple(identity, Val{N}())}()
+    dense_dims(::Type{<:StaticArrays.StaticArray{S,T,N}}) where {S,T,N} = DenseDims{ntuple(_ -> true, Val(N))}()
+    defines_strides(::Type{<:StaticArrays.MArray}) = true
+    @generated function size(A::StaticArrays.StaticArray{S}) where {S}
+        t = Expr(:tuple); Sp = S.parameters
+        for n in 1:length(Sp)
+            push!(t.args, Expr(:call, Expr(:curly, :StaticInt, Sp[n])))
+        end
+        t
+    end
+    @generated function strides(A::StaticArrays.StaticArray{S}) where {S}
+        t = Expr(:tuple, Expr(:call, Expr(:curly, :StaticInt, 1))); Sp = S.parameters; x = 1
+        for n in 1:length(Sp)-1
+            push!(t.args, Expr(:call, Expr(:curly, :StaticInt, (x *= Sp[n]))))
+        end
+        t
+    end
     @require Adapt="79e6a3ab-5dfb-504d-930d-738a2a938a0e" begin
       function Adapt.adapt_storage(::Type{<:StaticArrays.SArray{S}},xs::Array) where S
           StaticArrays.SArray{S}(xs)
@@ -694,7 +771,7 @@ function __init__()
     aos_to_soa(x::AbstractArray{<:Tracker.TrackedReal,N}) where N = Tracker.collect(x)
   end
 
-  @require CuArrays="3a865a2d-5b23-5a0f-bc46-62713ec82fae" begin
+    @require CuArrays="3a865a2d-5b23-5a0f-bc46-62713ec82fae" begin
     @require Adapt="79e6a3ab-5dfb-504d-930d-738a2a938a0e" begin
       include("cuarrays.jl")
     end
@@ -717,7 +794,7 @@ function __init__()
   @require BandedMatrices="aae01518-5342-5314-be14-df237901396f" begin
     function findstructralnz(x::BandedMatrices.BandedMatrix)
       l,u=BandedMatrices.bandwidths(x)
-      rowsize,colsize=size(x)
+      rowsize,colsize= Base.size(x)
       rowind=BandedMatrixIndex(rowsize,colsize,l,u,true)
       colind=BandedMatrixIndex(rowsize,colsize,l,u,false)
       (rowind,colind)
@@ -730,7 +807,7 @@ function __init__()
     function matrix_colors(A::BandedMatrices.BandedMatrix)
         l,u=BandedMatrices.bandwidths(A)
         width=u+l+1
-        _cycle(1:width,size(A,2))
+        _cycle(1:width, Base.size(A,2))
     end
 
   end
@@ -794,9 +871,19 @@ function __init__()
       end
     end
   end
+  @require OffsetArrays="6fe1bfb0-de20-5000-8ca7-80f57d26f881" begin
+      size(A::OffsetArrays.OffsetArray) = size(parent(A))
+      strides(A::OffsetArrays.OffsetArray) = strides(parent(A))
+      # offsets(A::OffsetArrays.OffsetArray) = map(+, A.offsets, offsets(parent(A)))
+      device(::OffsetArrays.OffsetArray) = CheckParent()
+      contiguous_axis(A::OffsetArrays.OffsetArray) = contiguous_axis(parent(A))
+      contiguous_batch_size(A::OffsetArrays.OffsetArray) = contiguous_batch_size(parent(A))
+      stride_rank(A::OffsetArrays.OffsetArray) = stride_rank(parent(A))
+  end
 end
 
 include("static.jl")
 include("ranges.jl")
+include("stridelayout.jl")
 
 end

src/cuarrays.jl

@@ -9,5 +9,8 @@ function Base.setindex(x::CuArrays.CuArray,v,i::Int)
 end
 
 function restructure(x::CuArrays.CuArray,y)
-  reshape(Adapt.adapt(parameterless_type(x),y),size(x)...)
+  reshape(Adapt.adapt(parameterless_type(x),y), Base.size(x)...)
 end
+
+Device(::Type{<:CuArrays.CuArray}) = GPU()
+

src/cuarrays2.jl

@@ -9,5 +9,8 @@ function Base.setindex(x::CUDA.CuArray,v,i::Int)
 end
 
 function restructure(x::CUDA.CuArray,y)
-  reshape(Adapt.adapt(parameterless_type(x),y),size(x)...)
+  reshape(Adapt.adapt(parameterless_type(x),y), Base.size(x)...)
 end
+
+Device(::Type{<:CUDA.CuArray}) = GPU()
+