Merge pull request #90 from SciML/loopvecbase

More changes to support VectorizationBase/LoopVectorization

Author: chriselrod

.github/workflows/ci.yml

@@ -0,0 +1,69 @@
+name: CI
+on:
+  pull_request:
+    branches:
+      - master
+  push:
+    branches:
+      - master
+    tags: '*'
+jobs:
+  test:
+    name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        version:
+          - '1.2'
+          - '1'
+          - 'nightly'
+        os:
+          - ubuntu-latest
+        arch:
+          - x64
+    steps:
+      - uses: actions/checkout@v2
+      - uses: julia-actions/setup-julia@v1
+        with:
+          version: ${{ matrix.version }}
+          arch: ${{ matrix.arch }}
+      - uses: actions/cache@v1
+        env:
+          cache-name: cache-artifacts
+        with:
+          path: ~/.julia/artifacts
+          key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }}
+          restore-keys: |
+            ${{ runner.os }}-test-${{ env.cache-name }}-
+            ${{ runner.os }}-test-
+            ${{ runner.os }}-
+      - uses: julia-actions/julia-buildpkg@v1
+      - uses: julia-actions/julia-runtest@v1
+      - uses: julia-actions/julia-processcoverage@v1
+      - uses: codecov/codecov-action@v1
+        with:
+          file: lcov.info
+  docs:
+    name: Documentation
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v2
+      - uses: julia-actions/setup-julia@v1
+        with:
+          version: '1'
+      - run: |
+          julia --project=docs -e '
+            using Pkg
+            Pkg.develop(PackageSpec(path=pwd()))
+            Pkg.instantiate()'
+      # - run: |
+      #     julia --project=docs -e '
+      #       using Documenter: doctest
+      #       using ArrayInterface
+      #       doctest(ArrayInterface)' 
+      # - run: julia --project=docs docs/make.jl
+        env:
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          DOCUMENTER_KEY: ${{ secrets.DOCUMENTER_KEY }}
+

Project.toml

@@ -1,6 +1,6 @@
 name = "ArrayInterface"
 uuid = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
-version = "2.13.8"
+version = "2.14.0"
 
 [deps]
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

src/ArrayInterface.jl

@@ -25,6 +25,7 @@ parent_type(::Type{<:LinearAlgebra.AbstractTriangular{T,S}}) where {T,S} = S
 parent_type(::Type{<:PermutedDimsArray{T,N,I1,I2,A}}) where {T,N,I1,I2,A} = A
 parent_type(::Type{Slice{T}}) where {T} = T
 parent_type(::Type{T}) where {T} = T
+parent_type(::Type{R}) where {S, T, A <: AbstractArray{S}, N, R <: Base.ReinterpretArray{T, N, S, A}} = A
 
 """
     known_length(::Type{T})
@@ -880,10 +881,11 @@ function __init__()
       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))
+      parent_type(::Type{O}) where {T,N,A<:AbstractArray{T,N},O<:OffsetArrays.OffsetArray{T,N,A}} = A
+      device(::Type{<:OffsetArrays.OffsetArray}) = CheckParent()
+      contiguous_axis(::Type{A}) where {A <: OffsetArrays.OffsetArray} = contiguous_axis(parent_type(A))
+      contiguous_batch_size(::Type{A}) where {A <: OffsetArrays.OffsetArray} = contiguous_batch_size(parent_type(A))
+      stride_rank(::Type{A}) where {A <: OffsetArrays.OffsetArray} = stride_rank(parent_type(A))
   end
 end
 

src/dimensions.jl

@@ -30,10 +30,20 @@ Return the names of the dimensions for `x`.
     end
 end
 @inline function dimnames(::Type{T}) where {T<:Union{Transpose,Adjoint}}
-    return _transpose_dimnames(dimnames(parent_type(T)))
+    return _transpose_dimnames(Val(dimnames(parent_type(T))))
 end
-_transpose_dimnames(x::Tuple{Symbol,Symbol}) = (last(x), first(x))
-_transpose_dimnames(x::Tuple{Symbol}) = (:_, first(x))
+# inserting the Val here seems to help inferability; I got a test failure without it.
+function _transpose_dimnames(::Val{S}) where {S}
+    if length(S) == 1
+        (:_, first(S))
+    elseif length(S) == 2
+        (last(S), first(S))
+    else
+        throw("Can't transpose $S of dim $(length(S)).")
+    end
+end
+@inline _transpose_dimnames(x::Tuple{Symbol,Symbol}) = (last(x), first(x))
+@inline _transpose_dimnames(x::Tuple{Symbol}) = (:_, first(x))
 
 @inline function dimnames(::Type{T}) where {I,T<:PermutedDimsArray{<:Any,<:Any,I}}
     return map(i -> dimnames(parent_type(T), i), I)
@@ -143,6 +153,8 @@ julia> ArrayInterface.size(A)
 """
 size(A) = Base.size(A)
 size(A, d) = Base.size(A, to_dims(A, d))
+size(x::LinearAlgebra.Adjoint{T,V}) where {T, V <: AbstractVector{T}} = (One(), static_length(x))
+size(x::LinearAlgebra.Transpose{T,V}) where {T, V <: AbstractVector{T}} = (One(), static_length(x))
 
 """
     axes(A, d)

src/static.jl

@@ -22,21 +22,23 @@ Base.Integer(x::StaticInt{N}) where {N} = x
 (::Type{T})(x::StaticInt{N}) where {T<:Integer,N} = T(N)
 (::Type{T})(x::Int) where {T<:StaticInt} = StaticInt(x)
 Base.convert(::Type{StaticInt{N}}, ::StaticInt{N}) where {N} = StaticInt{N}()
+Base.float(::StaticInt{N}) where {N} = Float64(N)
 
-Base.promote_rule(::Type{<:StaticInt}, ::Type{T}) where {T <: AbstractIrrational} = promote_rule(Int, T)
-Base.promote_rule(::Type{T}, ::Type{<:StaticInt}) where {T <: AbstractIrrational} = promote_rule(T, Int)
+Base.promote_rule(::Type{<:StaticInt}, ::Type{T}) where {T <: Number} = promote_type(Int, T)
+Base.promote_rule(::Type{<:StaticInt}, ::Type{T}) where {T <: AbstractIrrational} = promote_type(Int, T)
+# Base.promote_rule(::Type{T}, ::Type{<:StaticInt}) where {T <: AbstractIrrational} = promote_rule(T, Int)
 for (S,T) ∈ [(:Complex,:Real), (:Rational, :Integer), (:(Base.TwicePrecision),:Any)]
-    @eval Base.promote_rule(::Type{$S{T}}, ::Type{<:StaticInt}) where {T <: $T} = promote_rule($S{T}, Int)
+    @eval Base.promote_rule(::Type{$S{T}}, ::Type{<:StaticInt}) where {T <: $T} = promote_type($S{T}, Int)
 end
 Base.promote_rule(::Type{Union{Nothing,Missing}}, ::Type{<:StaticInt}) = Union{Nothing, Missing, Int}
-Base.promote_rule(::Type{T}, ::Type{<:StaticInt}) where {T >: Union{Missing,Nothing}} = promote_rule(T, Int)
-Base.promote_rule(::Type{T}, ::Type{<:StaticInt}) where {T >: Nothing} = promote_rule(T, Int)
-Base.promote_rule(::Type{T}, ::Type{<:StaticInt}) where {T >: Missing} = promote_rule(T, Int)
+Base.promote_rule(::Type{T}, ::Type{<:StaticInt}) where {T >: Union{Missing,Nothing}} = promote_type(T, Int)
+Base.promote_rule(::Type{T}, ::Type{<:StaticInt}) where {T >: Nothing} = promote_type(T, Int)
+Base.promote_rule(::Type{T}, ::Type{<:StaticInt}) where {T >: Missing} = promote_type(T, Int)
 for T ∈ [:Bool, :Missing, :BigFloat, :BigInt, :Nothing, :Any]
 # let S = :Any    
     @eval begin
-        Base.promote_rule(::Type{S}, ::Type{$T}) where {S <: StaticInt} = promote_rule(Int, $T)
-        Base.promote_rule(::Type{$T}, ::Type{S}) where {S <: StaticInt} = promote_rule($T, Int)
+        Base.promote_rule(::Type{S}, ::Type{$T}) where {S <: StaticInt} = promote_type(Int, $T)
+        Base.promote_rule(::Type{$T}, ::Type{S}) where {S <: StaticInt} = promote_type($T, Int)
     end
 end
 Base.promote_rule(::Type{<:StaticInt}, ::Type{<:StaticInt}) = Int
@@ -85,6 +87,12 @@ end
 for f ∈ [:(+), :(-), :(*), :(/), :(÷), :(%), :(<<), :(>>), :(>>>), :(&), :(|), :(⊻)]
     @eval @generated Base.$f(::StaticInt{M}, ::StaticInt{N}) where {M,N} = Expr(:call, Expr(:curly, :StaticInt, $f(M, N)))
 end
+for f ∈ [:(<<), :(>>), :(>>>)]
+    @eval begin
+        @inline Base.$f(::StaticInt{M}, x::UInt) where {M} = $f(M, x)
+        @inline Base.$f(x::Integer, ::StaticInt{M}) where {M} = $f(x, M)
+    end
+end
 for f ∈ [:(==), :(!=), :(<), :(≤), :(>), :(≥)]
     @eval begin
         @inline Base.$f(::StaticInt{M}, ::StaticInt{N}) where {M,N} = $f(M, N)

src/stridelayout.jl

@@ -65,6 +65,14 @@ _contiguous_axis(::Any, ::Nothing) = nothing
     Expr(:call, Expr(:curly, :Contiguous, new_contig))
 end
 
+# contiguous_if_one(::Contiguous{1}) = Contiguous{1}()
+# contiguous_if_one(::Any) = Contiguous{-1}()
+function contiguous_axis(::Type{R}) where {T, N, S, A <: Array{S}, R <: Base.ReinterpretArray{T, N, S, A}}
+    isbitstype(S) ? Contiguous{1}() : nothing
+    # contiguous_if_one(contiguous_axis(parent_type(R)))
+end
+
+
 """
 contiguous_axis_indicator(::Type{T}) -> Tuple{Vararg{<:Val}}
 
@@ -108,6 +116,8 @@ _contiguous_batch_size(::Any, ::Any, ::Any) = nothing
     end
 end
 
+contiguous_batch_size(::Type{R}) where {T, N, S, A <: Array{S}, R <: Base.ReinterpretArray{T, N, S, A}} = ContiguousBatch{0}()
+
 struct StrideRank{R} end
 Base.@pure StrideRank(R::NTuple{<:Any,Int}) = StrideRank{R}()
 _get(::StrideRank{R}) where {R} = R
@@ -164,6 +174,7 @@ _stride_rank(::Any, ::Any) = nothing
     Expr(:call, Expr(:curly, :StrideRank, ranktup))
 end
 stride_rank(x, i) = stride_rank(x)[i]
+stride_rank(::Type{R}) where {T, N, S, A <: Array{S}, R <: Base.ReinterpretArray{T, N, S, A}} = StrideRank{ntuple(identity, Val{N}())}()
 
 """
 is_column_major(A) -> Val{true/false}()
@@ -248,6 +259,16 @@ julia> A = rand(3,4);
 
 julia> ArrayInterface.strides(A)
 (StaticInt{1}(), 3)
+
+Additionally, the behavior differs from `Base.strides` for adjoint vectors:
+
+julia> x = rand(5);
+
+julia> ArrayInterface.strides(x')
+(StaticInt{1}(), StaticInt{1}())
+
+This is to support the pattern of using just the first stride for linear indexing, `x[i]`,
+while still producing correct behavior when using valid cartesian indices, such as `x[1,i]`.
 ```
 """
 strides(A) = Base.strides(A)
@@ -264,6 +285,16 @@ offsets(::Any) = (StaticInt{1}(),) # Assume arbitrary Julia data structures use
 @inline strides(A::Vector{<:Any}) = (StaticInt(1),)
 @inline strides(A::Array{<:Any,N}) where {N} = (StaticInt(1), Base.tail(Base.strides(A))...)
 @inline strides(A::AbstractArray) = _strides(A, Base.strides(A), contiguous_axis(A))
+
+@inline function strides(x::LinearAlgebra.Adjoint{T,V}) where {T, V <: AbstractVector{T}}
+    strd = stride(parent(x), One())
+    (strd, strd)
+end
+@inline function strides(x::LinearAlgebra.Transpose{T,V}) where {T, V <: AbstractVector{T}}
+    strd = stride(parent(x), One())
+    (strd, strd)
+end
+                
 @generated function _strides(A::AbstractArray{T,N}, s::NTuple{N}, ::Contiguous{C}) where {T,N,C}
     if C ≤ 0 || C > N
         return Expr(:block, Expr(:meta,:inline), :s)
@@ -282,6 +313,22 @@ offsets(::Any) = (StaticInt{1}(),) # Assume arbitrary Julia data structures use
     end
 end
 
+if VERSION ≥ v"1.6.0-DEV.1581"
+    @generated function _strides(_::Base.ReinterpretArray{T, N, S, A, true}, s::NTuple{N}, ::Contiguous{1}) where {T, N, S, D, A <: Array{S,D}}
+        stup = Expr(:tuple, :(One()))
+        if D < N
+            push!(stup.args, Expr(:call, Expr(:curly, :StaticInt, sizeof(S) ÷ sizeof(T))))
+        end
+        for n ∈ 2+(D < N):N
+            push!(stup.args, Expr(:ref, :s, n))
+        end
+        quote
+            $(Expr(:meta,:inline))
+            @inbounds $stup
+        end        
+    end
+end
+
 @inline function offsets(x, i)
     inds = indices(x, i)
     start = known_first(inds)
@@ -304,7 +351,7 @@ end
 @inline strides(B::PermutedDimsArray{T,N,I1,I2,A}) where {T,N,I1,I2,A<:AbstractArray{T,N}} = permute(strides(parent(B)), Val{I1}())
 @inline stride(A::AbstractArray, ::StaticInt{N}) where {N} = strides(A)[N]
 @inline stride(A::AbstractArray, ::Val{N}) where {N} = strides(A)[N]
-stride(A, i) = Base.stride(A, i)
+stride(A, i) = Base.stride(A, i) # for type stability
 
 size(B::S) where {N,NP,T,A<:AbstractArray{T,NP},I,S <: SubArray{T,N,A,I}} = _size(size(parent(B)), B.indices, map(static_length, B.indices))
 strides(B::S) where {N,NP,T,A<:AbstractArray{T,NP},I,S <: SubArray{T,N,A,I}} = _strides(strides(parent(B)), B.indices)
@@ -324,8 +371,10 @@ end
 @generated function _strides(A::Tuple{Vararg{Any,N}}, inds::I) where {N, I<:Tuple}
     t = Expr(:tuple)
     for n in 1:N
-        if I.parameters[n] <: AbstractRange
+        if I.parameters[n] <: AbstractUnitRange
             push!(t.args, Expr(:ref, :A, n))
+        elseif I.parameters[n] <: AbstractRange
+            push!(t.args, Expr(:call, :(*), Expr(:ref, :A, n), Expr(:call, :static_step, Expr(:ref, :inds, n))))
         elseif !(I.parameters[n] <: Integer)
             return nothing
         end