Merge pull request #122 from SciML/bitarray

Add stridelayout support for `BitArray`

Author: chriselrod

src/ArrayInterface.jl

@@ -644,6 +644,7 @@ defines_strides(::Type{<:StridedArray}) = true
 function defines_strides(::Type{<:SubArray{T,N,P,I}}) where {T,N,P,I}
     return stride_preserving_index(I) === True()
 end
+defines_strides(::Type{<:BitArray}) = true
 
 """
 can_avx(f)

src/ranges.jl

@@ -539,7 +539,7 @@ Base.Slice(Static(1):100)
 ```
 """
 @generated function reduce_tup(f::F, inds::Tuple{Vararg{Any,N}}) where {F,N}
-    q = Expr(:block, Expr(:meta, :inline))
+    q = Expr(:block, Expr(:meta, :inline, :propagate_inbounds))
     if N == 1
         push!(q.args, :(inds[1]))
         return q
@@ -566,7 +566,7 @@ Base.Slice(Static(1):100)
     q
 end
 
-@inline function _pick_range(x, y)
+@propagate_inbounds function _pick_range(x, y)
     fst = _try_static(static_first(x), static_first(y))
     lst = _try_static(static_last(x), static_last(y))
     return Base.Slice(OptionallyStaticUnitRange(fst, lst))
@@ -591,19 +591,19 @@ specified, then the indices for visiting each index of `x` are returned.
 end
 @inline indices(x::AbstractUnitRange{<:Integer}) = Base.Slice(OptionallyStaticUnitRange(x))
 
-function indices(x::Tuple)
+@propagate_inbounds function indices(x::Tuple)
     inds = map(eachindex, x)
     return reduce_tup(_pick_range, inds)
 end
 
 @inline indices(x, d) = indices(axes(x, d))
 
-@inline function indices(x::Tuple{Vararg{Any,N}}, dim) where {N}
+@propagate_inbounds function indices(x::Tuple{Vararg{Any,N}}, dim) where {N}
     inds = map(x_i -> indices(x_i, dim), x)
     return reduce_tup(_pick_range, inds)
 end
 
-@inline function indices(x::Tuple{Vararg{Any,N}}, dim::Tuple{Vararg{Any,N}}) where {N}
+@propagate_inbounds function indices(x::Tuple{Vararg{Any,N}}, dim::Tuple{Vararg{Any,N}}) where {N}
     inds = map(indices, x, dim)
     return reduce_tup(_pick_range, inds)
 end

src/stridelayout.jl

@@ -46,8 +46,10 @@ function contiguous_axis(::Type{T}) where {T}
         return contiguous_axis(parent_type(T))
     end
 end
-contiguous_axis(::Type{<:Array}) = StaticInt{1}()
-contiguous_axis(::Type{<:Tuple}) = StaticInt{1}()
+contiguous_axis(::Type{<:Array}) = One()
+contiguous_axis(::Type{<:BitArray}) = One()
+contiguous_axis(::Type{<:AbstractRange}) = One()
+contiguous_axis(::Type{<:Tuple}) = One()
 function contiguous_axis(::Type{T}) where {T<:VecAdjTrans}
     c = contiguous_axis(parent_type(T))
     if c === nothing
@@ -138,6 +140,8 @@ function stride_rank(::Type{T}) where {T}
     end
 end
 stride_rank(::Type{Array{T,N}}) where {T,N} = nstatic(Val(N))
+stride_rank(::Type{BitArray{N}}) where {N} = nstatic(Val(N))
+stride_rank(::Type{<:AbstractRange}) = (One(),)
 stride_rank(::Type{<:Tuple}) = (One(),)
 
 stride_rank(::Type{T}) where {T<:VecAdjTrans} = (StaticInt(2), StaticInt(1))
@@ -187,8 +191,10 @@ function _contiguous_batch_size(::StaticInt{D}, ::R) where {D,R<:Tuple}
     end
 end
 
-contiguous_batch_size(::Type{Array{T,N}}) where {T,N} = StaticInt{0}()
-contiguous_batch_size(::Type{<:Tuple}) = StaticInt{0}()
+contiguous_batch_size(::Type{Array{T,N}}) where {T,N} = Zero()
+contiguous_batch_size(::Type{BitArray{N}}) where {N} = Zero()
+contiguous_batch_size(::Type{<:AbstractRange}) = Zero()
+contiguous_batch_size(::Type{<:Tuple}) = Zero()
 function contiguous_batch_size(::Type{T}) where {T<:Union{Transpose,Adjoint}}
     return contiguous_batch_size(parent_type(T))
 end
@@ -216,6 +222,7 @@ Returns `Val{true}` if elements of `A` are stored in column major order. Otherwi
 is_column_major(A) = is_column_major(stride_rank(A), contiguous_batch_size(A))
 is_column_major(sr::Nothing, cbs) = False()
 is_column_major(sr::R, cbs) where {R} = _is_column_major(sr, cbs)
+is_column_major(::AbstractRange) = False()
 
 # cbs > 0
 _is_column_major(sr::R, cbs::StaticInt) where {R} = False()
@@ -239,6 +246,8 @@ end
 _all_dense(::Val{N}) where {N} = ntuple(_ -> True(), Val{N}())
 
 dense_dims(::Type{Array{T,N}}) where {T,N} = _all_dense(Val{N}())
+dense_dims(::Type{BitArray{N}}) where {N} = _all_dense(Val{N}())
+dense_dims(::Type{<:AbstractRange}) = (True(),)
 dense_dims(::Type{<:Tuple}) = (True(),)
 function dense_dims(::Type{T}) where {T<:VecAdjTrans}
     dense = dense_dims(parent_type(T))
@@ -413,6 +422,7 @@ function strides(x)
 end
 #@inline strides(A) = _strides(A, Base.strides(A), contiguous_axis(A))
 
+strides(::AbstractRange) = (One(),)
 function strides(x::VecAdjTrans)
     st = first(strides(parent(x)))
     return (st, st)
@@ -457,6 +467,8 @@ function strides(a::A, dim::Integer) where {A}
     end
 end
 
+
+
 @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) # for type stability

test/runtests.jl

@@ -303,6 +303,7 @@ using OffsetArrays
     # R = reshape(view(x, 1:100), (10,10));
     # A = zeros(3,4,5);
     A = Wrapper(reshape(view(x, 1:60), (3,4,5)))
+    B = A .== 0;
     D1 = view(A, 1:2:3, :, :)  # first dimension is discontiguous
     D2 = view(A, :, 2:2:4, :)  # first dimension is contiguous
 
@@ -312,6 +313,8 @@ using OffsetArrays
     @test !@inferred(ArrayInterface.defines_strides(view(A, :, [1,2],1)))
 
     @test @inferred(device(A)) === ArrayInterface.CPUPointer()
+    @test @inferred(device(B)) === ArrayInterface.CPUIndex()
+    @test @inferred(device(-1:19)) === ArrayInterface.CPUIndex()
     @test @inferred(device((1,2,3))) === ArrayInterface.CPUIndex()
     @test @inferred(device(PermutedDimsArray(A,(3,1,2)))) === ArrayInterface.CPUPointer()
     @test @inferred(device(view(A, 1, :, 2:4))) === ArrayInterface.CPUPointer()
@@ -330,6 +333,8 @@ using OffsetArrays
     =#
     @test @inferred(contiguous_axis(@SArray(zeros(2,2,2)))) === ArrayInterface.StaticInt(1)
     @test @inferred(contiguous_axis(A)) === ArrayInterface.StaticInt(1)
+    @test @inferred(contiguous_axis(B)) === ArrayInterface.StaticInt(1)
+    @test @inferred(contiguous_axis(-1:19)) === ArrayInterface.StaticInt(1)
     @test @inferred(contiguous_axis(D1)) === ArrayInterface.StaticInt(-1)
     @test @inferred(contiguous_axis(D2)) === ArrayInterface.StaticInt(1)
     @test @inferred(contiguous_axis(PermutedDimsArray(A,(3,1,2)))) === ArrayInterface.StaticInt(2)
@@ -350,6 +355,8 @@ using OffsetArrays
 
     @test @inferred(ArrayInterface.contiguous_axis_indicator(@SArray(zeros(2,2,2)))) == (true,false,false)
     @test @inferred(ArrayInterface.contiguous_axis_indicator(A)) == (true,false,false)
+    @test @inferred(ArrayInterface.contiguous_axis_indicator(B)) == (true,false,false)
+    @test @inferred(ArrayInterface.contiguous_axis_indicator(-1:10)) == (true,)
     @test @inferred(ArrayInterface.contiguous_axis_indicator(PermutedDimsArray(A,(3,1,2)))) == (false,true,false)
     @test @inferred(ArrayInterface.contiguous_axis_indicator(@view(PermutedDimsArray(A,(3,1,2))[2,1:2,:]))) == (true,false)
     @test @inferred(ArrayInterface.contiguous_axis_indicator(@view(PermutedDimsArray(A,(3,1,2))[2,1:2,:])')) == (false,true)
@@ -362,6 +369,8 @@ using OffsetArrays
 
     @test @inferred(contiguous_batch_size(@SArray(zeros(2,2,2)))) === ArrayInterface.StaticInt(0)
     @test @inferred(contiguous_batch_size(A)) === ArrayInterface.StaticInt(0)
+    @test @inferred(contiguous_batch_size(B)) === ArrayInterface.StaticInt(0)
+    @test @inferred(contiguous_batch_size(-1:18)) === ArrayInterface.StaticInt(0)
     @test @inferred(contiguous_batch_size(PermutedDimsArray(A,(3,1,2)))) === ArrayInterface.StaticInt(0)
     @test @inferred(contiguous_batch_size(@view(PermutedDimsArray(A,(3,1,2))[2,1:2,:]))) === ArrayInterface.StaticInt(0)
     @test @inferred(contiguous_batch_size(@view(PermutedDimsArray(A,(3,1,2))[2,1:2,:])')) === ArrayInterface.StaticInt(0)
@@ -372,6 +381,8 @@ using OffsetArrays
 
     @test @inferred(stride_rank(@SArray(zeros(2,2,2)))) == (1, 2, 3)
     @test @inferred(stride_rank(A)) == (1,2,3)
+    @test @inferred(stride_rank(B)) == (1,2,3)
+    @test @inferred(stride_rank(-4:4)) == (1,)
     @test @inferred(stride_rank(view(A,:,:,1))) === (static(1), static(2))
     @test @inferred(stride_rank(view(A,:,:,1))) === ((ArrayInterface.StaticInt(1),ArrayInterface.StaticInt(2)))
     @test @inferred(stride_rank(PermutedDimsArray(A,(3,1,2)))) == (3, 1, 2)
@@ -400,6 +411,8 @@ using OffsetArrays
 
     @test @inferred(ArrayInterface.is_column_major(@SArray(zeros(2,2,2)))) === True()
     @test @inferred(ArrayInterface.is_column_major(A)) === True()
+    @test @inferred(ArrayInterface.is_column_major(B)) === True()
+    @test @inferred(ArrayInterface.is_column_major(-4:7)) === False()
     @test @inferred(ArrayInterface.is_column_major(PermutedDimsArray(A,(3,1,2)))) === False()
     @test @inferred(ArrayInterface.is_column_major(@view(PermutedDimsArray(A,(3,1,2))[2,1:2,:]))) === True()
     @test @inferred(ArrayInterface.is_column_major(@view(PermutedDimsArray(A,(3,1,2))[2,1:2,:])')) === False()
@@ -408,11 +421,12 @@ using OffsetArrays
     @test @inferred(ArrayInterface.is_column_major(@view(PermutedDimsArray(A,(3,1,2))[2:3,2,:])')) === True()
     @test @inferred(ArrayInterface.is_column_major(@view(PermutedDimsArray(A,(3,1,2))[:,1:2,1])')) === True()
     @test @inferred(ArrayInterface.is_column_major(@view(PermutedDimsArray(A,(3,1,2))[:,2,1])')) === False()
-    @test @inferred(ArrayInterface.is_column_major(1:10)) === False()
     @test @inferred(ArrayInterface.is_column_major(2.3)) === False()
 
     @test @inferred(dense_dims(@SArray(zeros(2,2,2)))) == (true,true,true)
     @test @inferred(dense_dims(A)) == (true,true,true)
+    @test @inferred(dense_dims(B)) == (true,true,true)
+    @test @inferred(dense_dims(-3:9)) == (true,)
     @test @inferred(dense_dims(PermutedDimsArray(A,(3,1,2)))) == (true,true,true)
     @test @inferred(dense_dims(@view(PermutedDimsArray(A,(3,1,2))[2,1:2,:]))) == (true,false)
     @test @inferred(dense_dims(@view(PermutedDimsArray(A,(3,1,2))[2,1:2,:])')) == (false,true)
@@ -434,9 +448,9 @@ using OffsetArrays
     @test @inferred(dense_dims(view(DummyZeros(3,4), :, 1))) === nothing
     @test @inferred(dense_dims(view(DummyZeros(3,4), :, 1)')) === nothing
 
-    B = Array{Int8}(undef, 2,2,2,2);
-    doubleperm = PermutedDimsArray(PermutedDimsArray(B,(4,2,3,1)), (4,2,1,3));
-    @test collect(strides(B))[collect(stride_rank(doubleperm))] == collect(strides(doubleperm))
+    C = Array{Int8}(undef, 2,2,2,2);
+    doubleperm = PermutedDimsArray(PermutedDimsArray(C,(4,2,3,1)), (4,2,1,3));
+    @test collect(strides(C))[collect(stride_rank(doubleperm))] == collect(strides(doubleperm))
 
     @test @inferred(ArrayInterface.indices(OffsetArray(view(PermutedDimsArray(A, (3,1,2)), 1, :, 2:4)', 3, -173),1)) === Base.Slice(ArrayInterface.OptionallyStaticUnitRange(4,6))
     @test @inferred(ArrayInterface.indices(OffsetArray(view(PermutedDimsArray(A, (3,1,2)), 1, :, 2:4)', 3, -173),2)) === Base.Slice(ArrayInterface.OptionallyStaticUnitRange(-172,-170))