Extend strides for ReshapedArray with strided parent. (#44507)

* Extend `strides(::ReshapedArray)` with non-contiguous strided parent
* Make sure `strides(::StridedReinterpretArray)` calls `size_to_strides`

Co-authored-by: Tim Holy <[email protected]>

base/reinterpretarray.jl

@@ -152,23 +152,15 @@ strides(a::Union{DenseArray,StridedReshapedArray,StridedReinterpretArray}) = siz
 stride(A::Union{DenseArray,StridedReshapedArray,StridedReinterpretArray}, k::Integer) =
  k ≤ ndims(A) ? strides(A)[k] : length(A)
 
-function strides(a::ReshapedReinterpretArray)
- ap = parent(a)
- els, elp = elsize(a), elsize(ap)
- stp = strides(ap)
- els == elp && return stp
- els < elp && return (1, _checked_strides(stp, els, elp)...)
+function strides(a::ReinterpretArray{T,<:Any,S,<:AbstractArray{S},IsReshaped}) where {T,S,IsReshaped}
+ _checkcontiguous(Bool, a) && return size_to_strides(1, size(a))
+ stp = strides(parent(a))
+ els, elp = sizeof(T), sizeof(S)
+ els == elp && return stp # 0dim parent is also handled here.
+ IsReshaped && els < elp && return (1, _checked_strides(stp, els, elp)...)
  stp[1] == 1 || throw(ArgumentError("Parent must be contiguous in the 1st dimension!"))
- return _checked_strides(tail(stp), els, elp)
-end
-
-function strides(a::NonReshapedReinterpretArray)
- ap = parent(a)
- els, elp = elsize(a), elsize(ap)
- stp = strides(ap)
- els == elp && return stp
- stp[1] == 1 || throw(ArgumentError("Parent must be contiguous in the 1st dimension!"))
- return (1, _checked_strides(tail(stp), els, elp)...)
+ st′ = _checked_strides(tail(stp), els, elp)
+ return IsReshaped ? st′ : (1, st′...)
 end
 
 @inline function _checked_strides(stp::Tuple, els::Integer, elp::Integer)

base/reshapedarray.jl

@@ -295,14 +295,51 @@ unsafe_convert(::Type{Ptr{T}}, V::SubArray{T,N,P,<:Tuple{Vararg{Union{RangeIndex
  unsafe_convert(Ptr{T}, V.parent) + (first_index(V)-1)*sizeof(T)
 
 
-_checkcontiguous(::Type{Bool}, A::AbstractArray) = size_to_strides(1, size(A)...) == strides(A)
-_checkcontiguous(::Type{Bool}, A::Array) = true
+_checkcontiguous(::Type{Bool}, A::AbstractArray) = false
+# `strides(A::DenseArray)` calls `size_to_strides` by default.
+# Thus it's OK to assume all `DenseArray`s are contiguously stored.
+_checkcontiguous(::Type{Bool}, A::DenseArray) = true
 _checkcontiguous(::Type{Bool}, A::ReshapedArray) = _checkcontiguous(Bool, parent(A))
 _checkcontiguous(::Type{Bool}, A::FastContiguousSubArray) = _checkcontiguous(Bool, parent(A))
 
 function strides(a::ReshapedArray)
- # We can handle non-contiguous parent if it's a StridedVector
- ndims(parent(a)) == 1 && return size_to_strides(only(strides(parent(a))), size(a)...)
- _checkcontiguous(Bool, a) || throw(ArgumentError("Parent must be contiguous."))
- size_to_strides(1, size(a)...)
+ _checkcontiguous(Bool, a) && return size_to_strides(1, size(a)...)
+ apsz::Dims = size(a.parent)
+ apst::Dims = strides(a.parent)
+ msz, mst, n = merge_adjacent_dim(apsz, apst) # Try to perform "lazy" reshape
+ n == ndims(a.parent) && return size_to_strides(mst, size(a)...) # Parent is stridevector like
+ return _reshaped_strides(size(a), 1, msz, mst, n, apsz, apst)
+end
+
+function _reshaped_strides(::Dims{0}, reshaped::Int, msz::Int, ::Int, ::Int, ::Dims, ::Dims)
+ reshaped == msz && return ()
+ throw(ArgumentError("Input is not strided."))
+end
+function _reshaped_strides(sz::Dims, reshaped::Int, msz::Int, mst::Int, n::Int, apsz::Dims, apst::Dims)
+ st = reshaped * mst
+ reshaped = reshaped * sz[1]
+ if length(sz) > 1 && reshaped == msz && sz[2] != 1
+ msz, mst, n = merge_adjacent_dim(apsz, apst, n + 1)
+ reshaped = 1
+ end
+ sts = _reshaped_strides(tail(sz), reshaped, msz, mst, n, apsz, apst)
+ return (st, sts...)
+end
+
+merge_adjacent_dim(::Dims{0}, ::Dims{0}) = 1, 1, 0
+merge_adjacent_dim(apsz::Dims{1}, apst::Dims{1}) = apsz[1], apst[1], 1
+function merge_adjacent_dim(apsz::Dims{N}, apst::Dims{N}, n::Int = 1) where {N}
+ sz, st = apsz[n], apst[n]
+ while n < N
+ szₙ, stₙ = apsz[n+1], apst[n+1]
+ if sz == 1
+ sz, st = szₙ, stₙ
+ elseif stₙ == st * sz || szₙ == 1
+ sz *= szₙ
+ else
+ break
+ end
+ n += 1
+ end
+ return sz, st, n
 end

stdlib/LinearAlgebra/src/blas.jl

@@ -147,18 +147,19 @@ end
 # Level 1
 # A help function to pick the pointer and inc for 1d like inputs.
 @inline function vec_pointer_stride(x::AbstractArray, stride0check = nothing)
- isdense(x) && return pointer(x), 1 # simpify runtime check when possibe
- ndims(x) == 1 || strides(x) == Base.size_to_strides(stride(x, 1), size(x)...) ||
- throw(ArgumentError("only support vector like inputs"))
- st = stride(x, 1)
+ Base._checkcontiguous(Bool, x) && return pointer(x), 1 # simpify runtime check when possibe
+ st, ptr = checkedstride(x), pointer(x)
  isnothing(stride0check) || (st == 0 && throw(stride0check))
- ptr = st > 0 ? pointer(x) : pointer(x, lastindex(x))
+ ptr += min(st, 0) * sizeof(eltype(x)) * (length(x) - 1)
  ptr, st
 end
-isdense(x) = x isa DenseArray
-isdense(x::Base.FastContiguousSubArray) = isdense(parent(x))
-isdense(x::Base.ReshapedArray) = isdense(parent(x))
-isdense(x::Base.ReinterpretArray) = isdense(parent(x))
+function checkedstride(x::AbstractArray)
+ szs::Dims = size(x)
+ sts::Dims = strides(x)
+ _, st, n = Base.merge_adjacent_dim(szs, sts)
+ n === ndims(x) && return st
+ throw(ArgumentError("only support vector like inputs"))
+end
 ## copy
 
 """

stdlib/LinearAlgebra/test/blas.jl

@@ -18,9 +18,14 @@ function pack(A, uplo)
 end
 
 @testset "vec_pointer_stride" begin
- a = zeros(4,4,4)
- @test BLAS.asum(view(a,1:2:4,:,:)) == 0 # vector like
+ a = float(rand(1:20,4,4,4))
+ @test BLAS.asum(a) == sum(a) # dense case
+ @test BLAS.asum(view(a,1:2:4,:,:)) == sum(view(a,1:2:4,:,:)) # vector like
+ @test BLAS.asum(view(a,1:3,2:2,3:3)) == sum(view(a,1:3,2:2,3:3))
+ @test BLAS.asum(view(a,1:1,1:3,1:1)) == sum(view(a,1:1,1:3,1:1))
+ @test BLAS.asum(view(a,1:1,1:1,1:3)) == sum(view(a,1:1,1:1,1:3))
  @test_throws ArgumentError BLAS.asum(view(a,1:3:4,:,:)) # non-vector like
+ @test_throws ArgumentError BLAS.asum(view(a,1:2,1:1,1:3))
 end
 Random.seed!(100)
 ## BLAS tests - testing the interface code to BLAS routines

test/abstractarray.jl

@@ -1575,22 +1575,54 @@ end
  @test length(rr) == length(r)
 end
 
+struct FakeZeroDimArray <: AbstractArray{Int, 0} end
+Base.strides(::FakeZeroDimArray) = ()
+Base.size(::FakeZeroDimArray) = ()
 @testset "strides for ReshapedArray" begin
  # Type-based contiguous check is tested in test/compiler/inline.jl
+ function check_strides(A::AbstractArray)
+ # Make sure stride(A, i) is equivalent with strides(A)[i] (if 1 <= i <= ndims(A))
+ dims = ntuple(identity, ndims(A))
+ map(i -> stride(A, i), dims) == @inferred(strides(A)) || return false
+ # Test strides via value check.
+ for i in eachindex(IndexLinear(), A)
+ A[i] === Base.unsafe_load(pointer(A, i)) || return false
+ end
+ return true
+ end
  # General contiguous check
  a = view(rand(10,10), 1:10, 1:10)
- @test strides(vec(a)) == (1,)
+ @test check_strides(vec(a))
  b = view(parent(a), 1:9, 1:10)
- @test_throws "Parent must be contiguous." strides(vec(b))
+ @test_throws "Input is not strided." strides(vec(b))
  # StridedVector parent
  for n in 1:3
  a = view(collect(1:60n), 1:n:60n)
- @test strides(reshape(a, 3, 4, 5)) == (n, 3n, 12n)
- @test strides(reshape(a, 5, 6, 2)) == (n, 5n, 30n)
+ @test check_strides(reshape(a, 3, 4, 5))
+ @test check_strides(reshape(a, 5, 6, 2))
  b = view(parent(a), 60n:-n:1)
- @test strides(reshape(b, 3, 4, 5)) == (-n, -3n, -12n)
- @test strides(reshape(b, 5, 6, 2)) == (-n, -5n, -30n)
+ @test check_strides(reshape(b, 3, 4, 5))
+ @test check_strides(reshape(b, 5, 6, 2))
  end
+ # StridedVector like parent
+ a = randn(10, 10, 10)
+ b = view(a, 1:10, 1:1, 5:5)
+ @test check_strides(reshape(b, 2, 5))
+ # Other StridedArray parent
+ a = view(randn(10,10), 1:9, 1:10)
+ @test check_strides(reshape(a,3,3,2,5))
+ @test check_strides(reshape(a,3,3,5,2))
+ @test check_strides(reshape(a,9,5,2))
+ @test check_strides(reshape(a,3,3,10))
+ @test check_strides(reshape(a,1,3,1,3,1,5,1,2))
+ @test check_strides(reshape(a,3,3,5,1,1,2,1,1))
+ @test_throws "Input is not strided." strides(reshape(a,3,6,5))
+ @test_throws "Input is not strided." strides(reshape(a,3,2,3,5))
+ @test_throws "Input is not strided." strides(reshape(a,3,5,3,2))
+ @test_throws "Input is not strided." strides(reshape(a,5,3,3,2))
+ # Zero dimensional parent
+ a = reshape(FakeZeroDimArray(),1,1,1)
+ @test @inferred(strides(a)) == (1, 1, 1)
 end
 
 @testset "stride for 0 dims array #44087" begin