diff --git a/.github/workflows/ci.yml b/.github/workflows/CI.yml similarity index 99% rename from .github/workflows/ci.yml rename to .github/workflows/CI.yml index 084929f..c3b8389 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/CI.yml @@ -26,7 +26,7 @@ jobs: fail-fast: false matrix: version: - - '1.6' + - '1.10' - '1' - 'pre' os: diff --git a/.gitignore b/.gitignore index d0b53d9..a4e2d19 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,5 @@ .ipynb_checkpoints .vscode Manifest.toml +Manifest-v*.toml +LocalPreferences.toml diff --git a/Project.toml b/Project.toml index 848cd29..0b358a1 100644 --- a/Project.toml +++ b/Project.toml @@ -1,26 +1,36 @@ name = "ArraysOfArrays" uuid = "65a8f2f4-9b39-5baf-92e2-a9cc46fdf018" -version = "0.6.6" +version = "0.7.0" [deps] -Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" -ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" -StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" [weakdeps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" +FixedSizeArrays = "3821ddf9-e5b5-40d5-8e25-6813ab96b5e2" +GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" +KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" [extensions] ArraysOfArraysAdaptExt = "Adapt" ArraysOfArraysChainRulesCoreExt = "ChainRulesCore" +ArraysOfArraysFixedSizeArraysExt = "FixedSizeArrays" +ArraysOfArraysGPUArraysCoreExt = "GPUArraysCore" +ArraysOfArraysGPUKernelsExt = ["GPUArraysCore", "KernelAbstractions"] +ArraysOfArraysKernelAbstractionsExt = "KernelAbstractions" +ArraysOfArraysMooncakeExt = "Mooncake" ArraysOfArraysStaticArraysCoreExt = "StaticArraysCore" [compat] Adapt = "1, 2, 3, 4" ChainRulesCore = "1" +FixedSizeArrays = "1" +GPUArraysCore = "0.2" +KernelAbstractions = "0.9" +Mooncake = "0.5" StaticArraysCore = "1" Statistics = "1" -julia = "1.6" +julia = "1.10" diff --git a/README.md b/README.md index efd24e9..e0b70b5 100644 --- a/README.md +++ b/README.md @@ -10,7 +10,7 @@ A Julia package for efficient storage and handling of nested arrays. ArraysOfArrays provides two different types of nested arrays: `ArrayOfSimilarArrays` and `VectorOfArrays`. -An `ArrayOfSimilarArrays` offers a duality of view between representing the same data as both a flat multi-dimensional array and as an array of equally-sized arrays. A `VectorOfArrays` represents a vector of arrays of equal dimensionality but different size. Internally, both types store their data in flat arrays that are accessible to the user `flatview()`. +An `ArrayOfSimilarArrays` offers a duality of view between representing the same data as both a flat multi-dimensional array and as an array of equally-sized arrays. A `VectorOfArrays` represents a vector of arrays of equal dimensionality but different size. Internally, both types store their data in flat arrays that are accessible to the user via `flatview()`. ## Documentation diff --git a/docs/src/index.md b/docs/src/index.md index 114dfce..70a66c9 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -4,13 +4,39 @@ A Julia package for efficient storage and handling of nested arrays. ArraysOfArr This package also defines and exports the following new functions applicable to nested arrays in general: -* [`nestedview`](@ref) and [`flatview`](@ref) switch between a flat and a nested view of the same data. +* [`sliced`](@ref) and [`flatview`](@ref) switch between a flat and a nested (sliced) view of the same data. +* [`partitioned`](@ref) creates a view of a vector as a vector of arrays that may differ in size. +* [`getsplitmode`](@ref), [`splitup`](@ref) and [`fused`](@ref) provide a general API to split arrays into nested (sliced or partitioned) form and to fuse them back into flat form. +* [`stacked`](@ref) and [`unstackmode`](@ref) are similar to `Base.stack`, but can return the original underlying data of sliced arrays without copying it. +* [`vecflattened`](@ref) concatenates the elements of nested arrays into a single vector. * [`innersize`](@ref) returns the size of the elements of an array, provided they all have equal size. -* [`deepgetindex`](@ref), [`deepsetindex!`](@ref) and [`deepview`](@ref) provide index-based access across multiple layers of nested arrays * [`innermap`](@ref) and [`deepmap`](@ref) apply a function to the elements of the inner (resp. innermost) arrays. -* [`abstract_nestedarray_type`](@ref) returns the type of nested `AbstractArray`s for a given innermost element type with multiple layers of nesting. * [`consgroupedview`](@ref) computes a grouping of equal consecutive elements on a vector and applies it to another vector or (named or unnamed) tuple of vectors. +## Operations at a given nesting depth + +Similar to axis-targeted operations in Python's AwkwardArrays, but with array-of-arrays nesting semantics: + +* [`mapat(f, Val(d), As...)`](@ref) maps `f` over the objects at nesting depth `d` (`d = 1` ≡ `map`, `d = 2` ≡ `innermap`). +* [`bcastat(f, Val(d), args...)`](@ref) broadcasts `f` at depth `d`: nested arguments align, shallower arrays contribute one value per element of their level, scalars broadcast over everything. +* [`innermapreduce`](@ref), [`innerreduce`](@ref) and [`innersum`](@ref) reduce over the contents of each element array. +* [`innersizes`](@ref) and [`innerlengths`](@ref) return per-element sizes/lengths (elements need not be of equal size). + +For split arrays these operate on the underlying flat data — a single (GPU-compatible) operation per nesting level. Outer-level broadcasts like `(x -> 2 .* x).(A)` keep their usual Julia semantics (`f` receives whole element arrays), but return a `VectorOfArrays` when the results are arrays. + +## Which flattening function do I want? + +* [`fused(A)`](@ref): the original underlying array, `splitup(fused(A), getsplitmode(A)) == A`. +* [`flatview(A)`](@ref): the underlying storage, requires a memory-ordered layout. +* [`stacked(A)`](@ref): elements joined along new trailing dimensions, like `Base.stack`. +* [`vecflattened(A)`](@ref): elements concatenated into a single vector, like `reduce(vcat, A)`. + +All four are zero-copy where possible and so may return arrays that share memory with `A`. In contrast, `Base.stack(A)` and `reduce(vcat, A)` always return independent arrays. + +## GPU support + +Both array types work with GPU-resident data. An `ArrayOfSimilarArrays` backed by a GPU array is fully functional. For a `VectorOfArrays`, the shape information (`elem_ptr` and `kernel_size`) can either stay on the host — element access then returns device views — or live on the device as well (e.g. via `Adapt.adapt`), which is the layout to use inside GPU kernels. `KernelAbstractions.get_backend` returns the backend of the underlying data. + ## [ArrayOfSimilarArrays](@id section_ArrayOfSimilarArrays) @@ -18,7 +44,7 @@ An `ArrayOfSimilarArrays` offers a duality of view between representing the same ```julia A_flat = rand(2,3,4,5,6) -A_nested = nestedview(A_flat, 2) +A_nested = sliced(A_flat, 2) ``` creates a view of `A_flat` as an array of arrays: @@ -44,9 +70,9 @@ all(x -> x == 4.2, A_flat[:, :, 2, 4, 3]) The following type aliases are defined: -* `VectorOfSimilarArrays{T,M} = AbstractArrayOfSimilarArrays{T,M,1}` -* `ArrayOfSimilarVectors{T,N} = AbstractArrayOfSimilarArrays{T,1,N}` -* `VectorOfSimilarVectors{T} = AbstractArrayOfSimilarArrays{T,1,1}` +* `VectorOfSimilarArrays{T,M} = ArrayOfSimilarArrays{T,M,1}` +* `ArrayOfSimilarVectors{T,N} = ArrayOfSimilarArrays{T,1,N}` +* `VectorOfSimilarVectors{T} = ArrayOfSimilarArrays{T,1,1}` For each of the types there is also an abstract type (`AbstractArrayOfSimilarArrays`, etc.). @@ -57,7 +83,7 @@ If a `VectorOfSimilarArrays` is backed by an `ElasticArrays.ElasticArray`, addit ```julia using ElasticArrays -A_nested = nestedview(ElasticArray{Float64}(undef, 2, 3, 0), 2) +A_nested = sliced(ElasticArray{Float64}(undef, 2, 3, 0), 2) for i in 1:4 push!(A_nested, rand(2, 3)) @@ -92,7 +118,7 @@ VA_flat = flatview(VA) VA_flat isa Vector{Float64} ``` -Calling `getindex` on `A_nested` returns a view into `A_flat`: +Calling `getindex` on `VA` returns a view into `VA_flat`: ```julia VA_flat = flatview(VA) @@ -108,7 +134,7 @@ all(x -> x == 4.2, VA[2]) ### Type aliases The following type aliases are defined: -* `VectorOfVectors{T,VT,VI,VD} = VectorOfArrays{T,1,VT,VI,VD}` +* `PartsView{T,VT,VI,VD,ET} = VectorOfArrays{T,1,0,VT,VI,VD,ET}` ### Appending data and resizing diff --git a/ext/ArraysOfArraysAdaptExt.jl b/ext/ArraysOfArraysAdaptExt.jl index 6bd2e65..69f37c0 100644 --- a/ext/ArraysOfArraysAdaptExt.jl +++ b/ext/ArraysOfArraysAdaptExt.jl @@ -5,7 +5,7 @@ module ArraysOfArraysAdaptExt import Adapt using Adapt: adapt -using ArraysOfArrays: ArrayOfSimilarArrays, VectorOfArrays +using ArraysOfArrays: ArrayOfSimilarArrays, VectorOfArrays, SplitParts using ArraysOfArrays: no_consistency_checks @@ -25,4 +25,12 @@ function Adapt.adapt_structure(to, A::VectorOfArrays) end +function Adapt.adapt_structure(to, smode::SplitParts) + SplitParts( + adapt(to, smode.elem_ptr), + adapt(to, smode.kernel_size) + ) +end + + end # module ArraysOfArraysAdaptExt diff --git a/ext/ArraysOfArraysChainRulesCoreExt.jl b/ext/ArraysOfArraysChainRulesCoreExt.jl index 6d2e143..83a6863 100644 --- a/ext/ArraysOfArraysChainRulesCoreExt.jl +++ b/ext/ArraysOfArraysChainRulesCoreExt.jl @@ -2,18 +2,125 @@ module ArraysOfArraysChainRulesCoreExt -import ChainRulesCore -using ChainRulesCore: NoTangent, unthunk +using ChainRulesCore: ChainRulesCore, NoTangent, AbstractThunk, Thunk, unthunk, @thunk, @non_differentiable -using ArraysOfArrays: ArrayOfSimilarArrays -using ArraysOfArrays: flatview +using ArraysOfArrays: getsplitmode, is_memordered_splitmode, splitup, fused, stacked, unstackmode, + flatview, innersize, vecflattened, partitioned +using ArraysOfArrays: NonSplitMode, AbstractSplitMode +using ArraysOfArrays: AbstractArrayOfSimilarArrays, ArrayOfSimilarArrays, VectorOfArrays + + +struct _MappedMaybeThunk{F, T} <: AbstractThunk + f::F + x::T +end +ChainRulesCore.unthunk(mt::_MappedMaybeThunk) = mt.f(unthunk(mt.x)) +mapthunk(f::F, x::T) where {F,T} = _MappedMaybeThunk{F,T}(f, x) +mapthunk(::Type{F}, x::T) where {F,T} = _MappedMaybeThunk{Type{F},T}(F, x) + + +@non_differentiable getsplitmode(::Any) +@non_differentiable unstackmode(::Any) +@non_differentiable innersize(::Any) +@non_differentiable is_memordered_splitmode(::Any) + + +function ChainRulesCore.rrule(::typeof(splitup), A::AbstractArray, smode::NonSplitMode) + return splitup(A, smode), _nosplit_pullback +end +_nosplit_pullback(ΔΩ) = NoTangent(), ΔΩ, NoTangent() + +function ChainRulesCore.rrule(::typeof(splitup), A::AbstractArray, smode::AbstractSplitMode) + return splitup(A, smode), _splitview_pullback +end +_splitview_pullback(ΔΩ) = NoTangent(), mapthunk(fused, ΔΩ), NoTangent() + + + +ChainRulesCore.rrule(::typeof(fused), A::AbstractArray) = fused(A), _nofuse_pullback +_nofuse_pullback(ΔΩ) = NoTangent(), ΔΩ + +function ChainRulesCore.rrule(::typeof(fused), A::AbstractArray{<:AbstractArray}) + return fused(A), Base.Fix2(_fused_pullback, getsplitmode(A)) +end +_fused_pullback(ΔΩ, smode) = NoTangent(), mapthunk(Base.Fix2(splitup, smode), ΔΩ) + + +ChainRulesCore.rrule(::typeof(stacked), A::AbstractArray) = stacked(A), _nostack_pullback +_nostack_pullback(ΔΩ) = NoTangent(), ΔΩ + +function ChainRulesCore.rrule(::typeof(stacked), A::AbstractArray{<:AbstractArray}) + return stacked(A), Base.Fix2(_stacked_pullback, unstackmode(A)) +end +function ChainRulesCore.rrule(::typeof(stack), A::AbstractArrayOfSimilarArrays) + return stack(A), Base.Fix2(_stacked_pullback, unstackmode(A)) +end +_stacked_pullback(ΔΩ, smode) = NoTangent(), mapthunk(Base.Fix2(splitup, smode), ΔΩ) + + + +function ChainRulesCore.rrule(::typeof(vecflattened), A::AbstractArrayOfSimilarArrays) + smode = getsplitmode(A) + sz = size(fused(A)) + function aosa_vecflattened_pullback(ΔΩ) + Δdata = reshape(unthunk(ΔΩ), sz) + return NoTangent(), splitup(Δdata, smode) + end + return vecflattened(A), aosa_vecflattened_pullback +end + +function ChainRulesCore.rrule(::typeof(vecflattened), A::VectorOfArrays) + smode = getsplitmode(A) # makes a defensive copy, safe to close over + datalen = length(A.data) + covered_from = first(A.elem_ptr) + covered_len = last(A.elem_ptr) - covered_from + function voa_vecflattened_pullback(ΔΩ) + Δ = unthunk(ΔΩ) + Δdata = if covered_len == datalen + Δ + else + padded = fill!(similar(Δ, datalen), zero(eltype(Δ))) + copyto!(padded, covered_from, Δ, firstindex(Δ), covered_len) + padded + end + return NoTangent(), splitup(Δdata, smode) + end + return vecflattened(A), voa_vecflattened_pullback +end + + +function ChainRulesCore.rrule(::typeof(partitioned), A::AbstractVector, lengths::AbstractVector{<:Integer}) + return partitioned(A, lengths), _partitioned_pullback_for(A) +end + +function ChainRulesCore.rrule(::typeof(partitioned), A::AbstractVector, shapes::AbstractVector{<:Dims}) + return partitioned(A, shapes), _partitioned_pullback_for(A) +end + +function _partitioned_pullback_for(A::AbstractVector) + datalen = length(A) + function partitioned_pullback(ΔΩ) + Δflat = vecflattened(unthunk(ΔΩ)) + ΔA = if length(Δflat) == datalen + Δflat + else + # Partial partitions leave uncovered data, which gets no tangent + # contribution: + padded = fill!(similar(Δflat, datalen), zero(eltype(Δflat))) + copyto!(padded, firstindex(padded), Δflat, firstindex(Δflat), length(Δflat)) + padded + end + return NoTangent(), ΔA, NoTangent() + end + return partitioned_pullback +end function _aosa_ctor_fromflat_pullback(ΔΩ) NoTangent(), flatview(convert(ArrayOfSimilarArrays, unthunk(ΔΩ))) end -function ChainRulesCore.rrule(::Type{ArrayOfSimilarArrays{T,M,N}}, flat_data::AbstractArray{U,L}) where {T,M,N,L,U} +function ChainRulesCore.rrule(::Type{ArrayOfSimilarArrays{T,M,N}}, flat_data::AbstractArray{U}) where {T,M,N,U} return ArrayOfSimilarArrays{T,M,N}(flat_data), _aosa_ctor_fromflat_pullback end diff --git a/ext/ArraysOfArraysFixedSizeArraysExt.jl b/ext/ArraysOfArraysFixedSizeArraysExt.jl new file mode 100644 index 0000000..a740024 --- /dev/null +++ b/ext/ArraysOfArraysFixedSizeArraysExt.jl @@ -0,0 +1,14 @@ +# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). + +module ArraysOfArraysFixedSizeArraysExt + +using FixedSizeArrays: FixedSizeVector + +import ArraysOfArrays + +# FixedSizeVectors cannot be resized, so sharing them between a +# VectorOfArrays and a split mode is safe and no defensive copy is +# required: +ArraysOfArrays._shapeinfo_copy(x::FixedSizeVector) = x + +end # module ArraysOfArraysFixedSizeArraysExt diff --git a/ext/ArraysOfArraysGPUArraysCoreExt.jl b/ext/ArraysOfArraysGPUArraysCoreExt.jl new file mode 100644 index 0000000..49fbf7d --- /dev/null +++ b/ext/ArraysOfArraysGPUArraysCoreExt.jl @@ -0,0 +1,31 @@ +# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). + +module ArraysOfArraysGPUArraysCoreExt + +using GPUArraysCore: AbstractGPUArray, @allowscalar + +import ArraysOfArrays + +# Two O(1) scalar reads during construction are acceptable: +ArraysOfArrays._scalar_first_last(x::AbstractGPUArray) = @allowscalar (first(x), last(x)) + +# Vectorized formulation of the O(n) partition checks, to avoid scalar +# indexing of GPU arrays: +function ArraysOfArrays._partition_sizes_valid(elem_ptr::AbstractGPUArray{<:Integer}, kernel_size::AbstractVector) + ep_lo = view(elem_ptr, firstindex(elem_ptr):(lastindex(elem_ptr) - 1)) + ep_hi = view(elem_ptr, (firstindex(elem_ptr) + 1):lastindex(elem_ptr)) + len = ep_hi .- ep_lo + + klen = prod.(kernel_size) + # kernel_size may still live on the host: + klen_dev = klen isa AbstractGPUArray ? klen : copyto!(similar(elem_ptr, eltype(klen), size(klen)), klen) + + valid = (len .>= 0) .& ( + (klen_dev .== 1) .| ( + (klen_dev .!= 0) .& (mod.(len, max.(klen_dev, one(eltype(klen_dev)))) .== 0) + ) + ) + return all(valid) +end + +end # module ArraysOfArraysGPUArraysCoreExt diff --git a/ext/ArraysOfArraysGPUKernelsExt.jl b/ext/ArraysOfArraysGPUKernelsExt.jl new file mode 100644 index 0000000..05c8862 --- /dev/null +++ b/ext/ArraysOfArraysGPUKernelsExt.jl @@ -0,0 +1,65 @@ +# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). + +module ArraysOfArraysGPUKernelsExt + +using GPUArraysCore: AbstractGPUArray + +import KernelAbstractions as KA +using KernelAbstractions: @kernel, @index, @Const + +using ArraysOfArrays: ArraysOfArrays, VectorOfArrays, innerlengths + + +@kernel function _segmented_mapreduce_kernel!(out, f, op, @Const(data), @Const(elem_ptr), init) + i = @index(Global, Linear) + j0 = elem_ptr[i] + j1 = elem_ptr[i + 1] - 1 + if init isa ArraysOfArrays._NoInit + # Empty element arrays have been excluded beforehand: + acc = f(data[j0]) + for j in (j0 + 1):j1 + acc = op(acc, f(data[j])) + end + else + acc = init + for j in j0:j1 + acc = op(acc, f(data[j])) + end + end + out[i] = acc +end + + +function _on_backend(backend, x::AbstractArray) + KA.get_backend(x) == backend && return x + y = KA.allocate(backend, eltype(x), size(x)) + copyto!(y, x) + return y +end + + +# Single-pass segmented reduction for GPU-resident data: +function ArraysOfArrays._innermapreduce_impl(f, op, init, A::VectorOfArrays{T,N,M,<:AbstractGPUArray}) where {T,N,M} + data = A.data + backend = KA.get_backend(data) + elem_ptr = _on_backend(backend, A.elem_ptr) + n = length(A) + + if init isa ArraysOfArrays._NoInit && any(iszero, innerlengths(A)) + throw(ArgumentError("Reducing over empty element arrays requires an init value")) + end + + T_f = Base.promote_op(f, T) + T_init = init isa ArraysOfArrays._NoInit ? T_f : typeof(init) + T_out = promote_type(T_f, Base.promote_op(op, T_init, T_f)) + out = KA.allocate(backend, T_out, n) + + if n > 0 + # Stream-ordered, like GPU array operations no synchronization here: + kernel! = _segmented_mapreduce_kernel!(backend) + kernel!(out, f, op, data, elem_ptr, init; ndrange = n) + end + return out +end + +end # module ArraysOfArraysGPUKernelsExt diff --git a/ext/ArraysOfArraysKernelAbstractionsExt.jl b/ext/ArraysOfArraysKernelAbstractionsExt.jl new file mode 100644 index 0000000..7bedb5f --- /dev/null +++ b/ext/ArraysOfArraysKernelAbstractionsExt.jl @@ -0,0 +1,16 @@ +# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). + +module ArraysOfArraysKernelAbstractionsExt + +import KernelAbstractions as KA + +using ArraysOfArrays: ArrayOfSimilarArrays, VectorOfArrays + +# Kernels over nested arrays run on the backend of the underlying data. The +# shape information of a VectorOfArrays may deliberately live on the host +# while its data lives on a device, so no cross-field consistency is +# required here: +KA.get_backend(A::ArrayOfSimilarArrays) = KA.get_backend(A.data) +KA.get_backend(A::VectorOfArrays) = KA.get_backend(A.data) + +end # module ArraysOfArraysKernelAbstractionsExt diff --git a/ext/ArraysOfArraysMooncakeExt.jl b/ext/ArraysOfArraysMooncakeExt.jl new file mode 100644 index 0000000..a97ddb2 --- /dev/null +++ b/ext/ArraysOfArraysMooncakeExt.jl @@ -0,0 +1,26 @@ +# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). + +module ArraysOfArraysMooncakeExt + +using Mooncake: DefaultCtx, @zero_derivative + +using ArraysOfArrays: ArraysOfArrays, AbstractSplitMode, + getsplitmode, unstackmode, is_memordered_splitmode, innersize, getslicemap, + consgrouped_ptrs + +# Mooncake can differentiate through the ArraysOfArrays types and operations +# without custom rules, so only the functions that query shape/structure +# information (and so are not differentiable with respect to array contents) +# get explicit zero-derivative rules: + +@zero_derivative DefaultCtx Tuple{typeof(getsplitmode), AbstractArray} +@zero_derivative DefaultCtx Tuple{typeof(unstackmode), AbstractArray} +@zero_derivative DefaultCtx Tuple{typeof(is_memordered_splitmode), AbstractSplitMode} +@zero_derivative DefaultCtx Tuple{typeof(innersize), AbstractArray} +@zero_derivative DefaultCtx Tuple{typeof(innersize), AbstractArray, Integer} +@zero_derivative DefaultCtx Tuple{typeof(getslicemap), AbstractArray} +@zero_derivative DefaultCtx Tuple{typeof(ArraysOfArrays.getinnerdims), Tuple{Vararg{Integer}}, AbstractSplitMode} +@zero_derivative DefaultCtx Tuple{typeof(ArraysOfArrays.getouterdims), Tuple{Vararg{Integer}}, AbstractSplitMode} +@zero_derivative DefaultCtx Tuple{typeof(consgrouped_ptrs), AbstractVector} + +end # module ArraysOfArraysMooncakeExt diff --git a/ext/ArraysOfArraysStaticArraysCoreExt.jl b/ext/ArraysOfArraysStaticArraysCoreExt.jl index c370949..62094a3 100644 --- a/ext/ArraysOfArraysStaticArraysCoreExt.jl +++ b/ext/ArraysOfArraysStaticArraysCoreExt.jl @@ -6,21 +6,27 @@ import StaticArraysCore using StaticArraysCore: StaticArray, SVector import ArraysOfArrays -using ArraysOfArrays: nestedview +using ArraysOfArrays: sliced @inline ArraysOfArrays.flatview(A::AbstractArray{SA,N}) where {S,T,M,N,SA<:StaticArray{S,T,M}} = reshape(reinterpret(T, A), size(SA)..., size(A)...) -@inline function ArraysOfArrays.nestedview(A::AbstractArray{T}, SA::Type{SVector{S,T}}) where {T,S} +@inline function ArraysOfArrays.sliced(A::AbstractArray{T}, SA::Type{SVector{S,T}}) where {T,S} size_A = size(A) size_A[1] == S || throw(DimensionMismatch("Length $S of static vector type does not match first dimension of array of size $size_A")) reshape(reinterpret(SA, A), ArraysOfArrays._tail(size_A)...) end -@inline ArraysOfArrays.nestedview(A::AbstractArray{T}, ::Type{SVector{S}}) where {T,S} = - nestedview(A, SVector{S,T}) +@inline ArraysOfArrays.sliced(A::AbstractArray{T}, ::Type{SVector{S}}) where {T,S} = + sliced(A, SVector{S,T}) + + +# Deprecated: + +Base.@deprecate ArraysOfArrays.nestedview(A::AbstractArray{T}, SA::Type{SVector{S,T}}) where {T,S} ArraysOfArrays.sliced(A, SA) false +Base.@deprecate ArraysOfArrays.nestedview(A::AbstractArray{T}, SA::Type{SVector{S}}) where {T,S} ArraysOfArrays.sliced(A, SA) false end # module ArraysOfArraysStaticArraysCoreExt diff --git a/src/ArraysOfArrays.jl b/src/ArraysOfArrays.jl index 1828f35..2a0788d 100644 --- a/src/ArraysOfArrays.jl +++ b/src/ArraysOfArrays.jl @@ -17,14 +17,9 @@ using Statistics include("util.jl") include("functions.jl") +include("base_slices.jl") include("array_of_similar_arrays.jl") include("vector_of_arrays.jl") include("broadcasting.jl") -@static if !isdefined(Base, :get_extension) - include("../ext/ArraysOfArraysAdaptExt.jl") - include("../ext/ArraysOfArraysChainRulesCoreExt.jl") - include("../ext/ArraysOfArraysStaticArraysCoreExt.jl") -end - end # module diff --git a/src/array_of_similar_arrays.jl b/src/array_of_similar_arrays.jl index 7d6a0b2..87bb345 100644 --- a/src/array_of_similar_arrays.jl +++ b/src/array_of_similar_arrays.jl @@ -1,139 +1,195 @@ # This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). """ - AbstractArrayOfSimilarArrays{T,M,N} <: AbstractArray{<:AbstractArray{T,M},N} + AbstractArrayOfSimilarArrays{T,M,N,ET<:AbstractArray{T,M}} <: AbstractSlices{ET,N} An array that contains arrays that have the same size/axes. The array is internally stored in flattened form as some kind of array of dimension -`M + N`. The flattened form can be accessed via `flatview(A)`. +`M + N`, in memory order. The flattened form can be accessed via +`flatview(A)`. -Subtypes must implement (in addition to typical array operations): +Subtypes must implement (in addition to typical array operations) - flatview(A::SomeArrayOfSimilarArrays)::AbstractArray{T,M+N} + ArraysOfArrays.fused(A::SomeArrayOfSimilarArrays)::AbstractArray{T,M+N} + +which must return the underlying flat array. All split-mode operations +([`getsplitmode`](@ref), [`stacked`](@ref), [`flatview`](@ref), +[`innersize`](@ref), [`getslicemap`](@ref), etc.) are then provided +automatically. The following type aliases are defined: -* `AbstractVectorOfSimilarArrays{T,M} = AbstractArrayOfSimilarArrays{T,M,1}` -* `AbstractArrayOfSimilarVectors{T,N} = AbstractArrayOfSimilarArrays{T,1,N}` -* `AbstractVectorOfSimilarVectors{T} = AbstractArrayOfSimilarArrays{T,1,1}` +* `AbstractVectorOfSimilarArrays{T,M,ET} = AbstractArrayOfSimilarArrays{T,M,1,ET}` +* `AbstractArrayOfSimilarVectors{T,N,ET} = AbstractArrayOfSimilarArrays{T,1,N,ET}` +* `AbstractVectorOfSimilarVectors{T,ET} = AbstractArrayOfSimilarArrays{T,1,1,ET}` """ -abstract type AbstractArrayOfSimilarArrays{T,M,N} <: AbstractArray{Array{T,M},N} end +abstract type AbstractArrayOfSimilarArrays{T,M,N,ET<:AbstractArray{T,M}} <: AbstractSlices{ET,N} end export AbstractArrayOfSimilarArrays -const AbstractVectorOfSimilarArrays{T,M} = AbstractArrayOfSimilarArrays{T,M,1} +const AbstractVectorOfSimilarArrays{T,M,ET<:AbstractArray{T,M}} = AbstractArrayOfSimilarArrays{T,M,1,ET} export AbstractVectorOfSimilarArrays -const AbstractArrayOfSimilarVectors{T,N} = AbstractArrayOfSimilarArrays{T,1,N} +const AbstractArrayOfSimilarVectors{T,N,ET<:AbstractVector{T}} = AbstractArrayOfSimilarArrays{T,1,N,ET} export AbstractArrayOfSimilarVectors -const AbstractVectorOfSimilarVectors{T} = AbstractArrayOfSimilarArrays{T,1,1} +const AbstractVectorOfSimilarVectors{T,ET<:AbstractVector{T}} = AbstractArrayOfSimilarArrays{T,1,1,ET} export AbstractVectorOfSimilarVectors """ - ArrayOfSimilarArrays{T,M,N,L,P} <: AbstractArrayOfSimilarArrays{T,M,N} + struct SplitSlices{M,N} <: AbstractSlicingMode{M,N} + +The split mode of [`ArrayOfSimilarArrays`](@ref): memory-ordered slicing +with `M` inner and `N` outer dimensions. + +Constructor: + +``` +SplitSlices{M,N}() +``` + +See also [`AbstractSlicingMode`](@ref). +""" +struct SplitSlices{M,N} <: AbstractSlicingMode{M,N} end +export SplitSlices + +is_memordered_splitmode(::SplitSlices) = true + +getinnerdims(obj::Tuple, ::SplitSlices{M,N}) where {M,N} = front_tuple(obj, Val(M)) +getouterdims(obj::Tuple, ::SplitSlices{M,N}) where {M,N} = back_tuple(obj, Val(N)) + +@inline splitup(A::AbstractArray{T}, ::SplitSlices{M,N}) where {T,M,N} = ArrayOfSimilarArrays{T,M,N}(A) + + +# Subtypes of AbstractArrayOfSimilarArrays only need to implement `fused` +# (plus the typical array operations), all split-mode operations are derived +# from it: + +function fused(A::AbstractArrayOfSimilarArrays) + throw(ArgumentError("Subtypes of AbstractArrayOfSimilarArrays like $(nameof(typeof(A))) must implement ArraysOfArrays.fused")) +end + +@inline getsplitmode(::AbstractArrayOfSimilarArrays{T,M,N}) where {T,M,N} = SplitSlices{M,N}() +@inline unstackmode(A::AbstractArrayOfSimilarArrays) = getsplitmode(A) + +@inline stacked(A::AbstractArrayOfSimilarArrays) = fused(A) + +@inline Base.parent(A::AbstractArrayOfSimilarArrays) = fused(A) + +@inline vecflattened(A::AbstractArrayOfSimilarArrays) = vec(fused(A)) +# Disambiguation: +@inline vecflattened(A::AbstractVectorOfSimilarArrays) = vec(fused(A)) +@inline vecflattened(A::AbstractArrayOfSimilarVectors) = vec(fused(A)) +@inline vecflattened(A::AbstractVectorOfSimilarVectors) = vec(fused(A)) + +function getslicemap(::AbstractArrayOfSimilarArrays{T,M,N}) where {T,M,N} + return (_ncolons(Val{M}())..., _oneto_tpl(Val{N}())...) +end + +# `stack` must return an independent array, unlike `stacked`. Julia does not +# dispatch on keyword arguments, so dispatch on the value of `dims` instead: +Base.stack(A::AbstractArrayOfSimilarArrays; dims::Union{Integer,Colon} = :) = _stack_impl(A, dims) + +# Fast path for the default layout: a single bulk copy instead of the +# element-by-element copy that generic `stack` would do: +_stack_impl(A::AbstractArrayOfSimilarArrays, ::Colon) = copy(fused(A)) +_stack_impl(A::AbstractArrayOfSimilarArrays, dims::Integer) = stack(collect(A); dims) + +# Per-element reductions reduce over the inner dimensions of the flat data: +function _innermapreduce_impl(f, op, init, A::AbstractArrayOfSimilarArrays{T,M,N}) where {T,M,N} + data = fused(A) + dims = ntuple(identity, Val(M)) + r = if init isa _NoInit + mapreduce(f, op, data; dims = dims) + else + mapreduce(f, op, data; dims = dims, init = init) + end + return reshape(r, size(A)) +end + + +@inline Base.:(==)(A::AbstractArrayOfSimilarArrays{<:Any,M,N}, B::AbstractArrayOfSimilarArrays{<:Any,M,N}) where {M,N} = (stacked(A) == stacked(B)) +@inline Base.isequal(A::AbstractArrayOfSimilarArrays{<:Any,M,N}, B::AbstractArrayOfSimilarArrays{<:Any,M,N}) where {M,N} = isequal(stacked(A), stacked(B)) +@inline Base.isapprox(A::AbstractArrayOfSimilarArrays{<:Any,M,N}, B::AbstractArrayOfSimilarArrays{<:Any,M,N}; kwargs...) where {M,N} = isapprox(stacked(A), stacked(B); kwargs...) + + +""" + ArrayOfSimilarArrays{T,M,N,P,ET} <: AbstractArrayOfSimilarArrays{T,M,N,ET} -Represents a view of an array of dimension `L = M + N` as an array of -dimension M with elements that are arrays with dimension N. All element arrays +Represents a view of an array of dimension `M + N` as an `N`-dimensional +array with elements that are `M`-dimensional arrays. All element arrays implicitly have equal size/axes. Constructors: - ArrayOfSimilarArrays{T,M,N}(flat_data::AbstractArray) - ArrayOfSimilarArrays{T,M}(flat_data::AbstractArray) + ArrayOfSimilarArrays{T,M,N}(data::AbstractArray) + ArrayOfSimilarArrays{T,M}(data::AbstractArray) The following type aliases are defined: -* `VectorOfSimilarArrays{T,M} = AbstractArrayOfSimilarArrays{T,M,1}` -* `ArrayOfSimilarVectors{T,N} = AbstractArrayOfSimilarArrays{T,1,N}` -* `VectorOfSimilarVectors{T} = AbstractArrayOfSimilarArrays{T,1,1}` +* `VectorOfSimilarArrays{T,M} = ArrayOfSimilarArrays{T,M,1}` +* `ArrayOfSimilarVectors{T,N} = ArrayOfSimilarArrays{T,1,N}` +* `VectorOfSimilarVectors{T} = ArrayOfSimilarArrays{T,1,1}` `VectorOfSimilarArrays` supports `push!()`, etc., provided the underlying -array supports resizing of it's last dimension (e.g. an `ElasticArray`). +array supports resizing of its last dimension (e.g. an `ElasticArray`). -The nested array can also be created using the function [`nestedview`](@ref) +The nested array can also be created using the function [`sliced`](@ref) and the wrapped flat array can be accessed using [`flatview`](@ref) afterwards: ```julia A_flat = rand(2,3,4,5,6) -A_nested = nestedview(A_flat, 2) +A_nested = sliced(A_flat, Val(2)) A_nested isa AbstractArray{<:AbstractArray{T,2},3} where T flatview(A_nested) === A_flat ``` """ struct ArrayOfSimilarArrays{ - T, M, N, L, - P<:AbstractArray{T,L} -} <: AbstractArrayOfSimilarArrays{T,M,N} + T, M, N, + P<:AbstractArray{T}, + ET<:AbstractArray{T,M} +} <: AbstractArrayOfSimilarArrays{T,M,N,ET} data::P - function ArrayOfSimilarArrays{T,M,N}(flat_data::AbstractArray{U,L}) where {T,M,N,L,U} - require_ndims(flat_data, _add_vals(Val{M}(), Val{N}())) - conv_parent = _convert_elype(T, flat_data) - P = typeof(conv_parent) - new{T,M,N,L,P}(conv_parent) + function ArrayOfSimilarArrays{T,M,N}(data::AbstractArray{T,L}) where {T,M,N,L} + _require_ndims(Val(L), _add_vals(Val(M), Val(N))) + P = typeof(data) + ET = Base.promote_op(view, P, _nColons(Val(M))..., _nInts(Val(N))...) + new{T,M,N,P,ET}(data) end end -function ArrayOfSimilarArrays{T,M}(flat_data::AbstractArray{U,L}) where {T,M,L,U} - _, size_outer = split_tuple(size(flat_data), Val{M}()) - N = length(size_outer) - ArrayOfSimilarArrays{T,M,N}(flat_data) -end - export ArrayOfSimilarArrays -function ArrayOfSimilarArrays{T,M,N}(A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} - B = ArrayOfSimilarArrays{T,M,N}(Array{T}(undef, innersize(A)..., size(A)...)) - copyto!(B, A) -end - -ArrayOfSimilarArrays{T}(A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = - ArrayOfSimilarArrays{T,M,N}(A) - -ArrayOfSimilarArrays(A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} = - ArrayOfSimilarArrays{T,M,N}(A) - -Base.convert(R::Type{ArrayOfSimilarArrays{T,M,N}}, flat_data::AbstractArray{U,L}) where {T,M,N,L,U} = R(flat_data) -Base.convert(R::Type{ArrayOfSimilarArrays{T,M}}, flat_data::AbstractArray{U,L}) where {T,M,L,U} = R(flat_data) - -Base.convert(R::Type{ArrayOfSimilarArrays{T,M,N}}, A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = R(A) -Base.convert(R::Type{ArrayOfSimilarArrays{T}}, A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = R(A) -Base.convert(R::Type{ArrayOfSimilarArrays}, A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} = R(A) - - -@inline function innersize(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} - front_tuple(size(A.data), Val{M}()) +function ArrayOfSimilarArrays{T,M,N}(orig_data::AbstractArray{U,L}) where {T,M,N,U,L} + conv_data = _convert_eltype(T, orig_data)::AbstractArray{T,L} + return ArrayOfSimilarArrays{T,M,N}(conv_data) end - -@inline function _innerlength(A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} - prod(innersize(A)) +function ArrayOfSimilarArrays{T,M}(data::AbstractArray{U,L}) where {T,M,U,L} + N = _val_value(_subtract_vals(Val(L), Val(M))) + ArrayOfSimilarArrays{T,M,N}(data) end - -import Base.== -(==)(A::ArrayOfSimilarArrays{T,M,N}, B::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} = - (A.data == B.data) +Base.convert(::Type{ArrayOfSimilarArrays{T,M,N}}, A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = ArrayOfSimilarArrays{T,M,N}(stacked(A)) +Base.convert(::Type{ArrayOfSimilarArrays{T}}, A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} = ArrayOfSimilarArrays{T,M,N}(stacked(A)) +Base.convert(::Type{ArrayOfSimilarArrays}, A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} = ArrayOfSimilarArrays{T,M,N}(stacked(A)) -""" - flatview(A::ArrayOfSimilarArrays{T,M,N,L,P})::P +@inline fused(A::ArrayOfSimilarArrays) = A.data -Returns the array of dimensionality `L = M + N` wrapped by `A`. The shape of -the result may be freely changed without breaking the inner consistency of -`A`. -""" -flatview(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} = A.data +function Base.Array(A::ArrayOfSimilarArrays{T,M,N,P,ET}) where {T,M,N,P,ET} + new_ET = Base.promote_op(similar, ET) + return Array{new_ET,N}(A) +end Base.size(A::ArrayOfSimilarArrays{T,M,N}) where {T,M,N} = split_tuple(size(A.data), Val{M}())[2] - Base.@propagate_inbounds Base.getindex(A::ArrayOfSimilarArrays{T,M,N}, idxs::Vararg{Integer,N}) where {T,M,N} = view(A.data, _ncolons(Val{M}())..., idxs...) @@ -141,6 +197,12 @@ Base.@propagate_inbounds Base.getindex(A::ArrayOfSimilarArrays{T,M,N}, idxs::Var Base.@propagate_inbounds Base.setindex!(A::ArrayOfSimilarArrays{T,M,N}, x::AbstractArray{U,M}, idxs::Vararg{Integer,N}) where {T,M,N,U} = setindex!(A.data, x, _ncolons(Val{M}())..., idxs...) +function Base.fill!(A::ArrayOfSimilarArrays{T,M,N}, x::AbstractArray{U,M}) where {T,M,N,U} + size(x) == innersize(A) || throw(DimensionMismatch("Can't fill ArrayOfSimilarArrays with an array of different size than its elements")) + A.data .= reshape(x, size(x)..., ntuple(_ -> 1, Val(N))...) + return A +end + Base.@propagate_inbounds function Base.unsafe_view(A::ArrayOfSimilarArrays{T,M,N}, idxs::Vararg{Union{Real, AbstractArray},N}) where {T,M,N} dataview = view(A.data, _ncolons(Val{M}())..., idxs...) L = length(size(dataview)) @@ -157,9 +219,9 @@ end function Base.similar(A::ArrayOfSimilarArrays{T,M,N}, ::Type{<:AbstractArray{U}}, dims::Dims) where {T,M,N,U} data = A.data - size_inner, size_outer = split_tuple(size(data), Val{M}()) + size_inner = front_tuple(size(data), Val{M}()) # ToDo: Don't use similar if data is an ElasticArray? - ArrayOfSimilarArrays{T,M,N}(similar(data, U, size_inner..., dims...)) + ArrayOfSimilarArrays{U,M,length(dims)}(similar(data, U, size_inner..., dims...)) end @@ -211,82 +273,19 @@ Base.map(::typeof(identity), A::ArrayOfSimilarArrays) = A Base.Broadcast.broadcasted(::typeof(identity), A::ArrayOfSimilarArrays) = A -Base.@pure _result_is_nested(idxs_outer::Tuple, idxs_inner::Tuple) = - Val{!(Base.index_dimsum(idxs_outer...) isa Tuple{}) && !(Base.index_dimsum(idxs_inner...) isa Tuple{})}() - -Base.@pure ndims_after_getindex(idxs::Tuple) = Val{length(Base.index_dimsum(idxs...))}() - - -Base.@propagate_inbounds function deepgetindex(A::ArrayOfSimilarArrays{T,M,N,L}, idxs::Vararg{Any,L}) where {T,M,N,L} - idxs_outer, idxs_inner = split_tuple(idxs, Val{N}()) - nested = _result_is_nested(idxs_outer, idxs_inner) - _deepgetindex_impl_aosa(A, idxs_outer, idxs_inner, nested) -end - -Base.@propagate_inbounds _deepgetindex_impl_aosa(A::ArrayOfSimilarArrays, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{false}) = - getindex(A.data, idxs_inner..., idxs_outer...) - -Base.@propagate_inbounds function _deepgetindex_impl_aosa(A::ArrayOfSimilarArrays, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{true}) - new_data = getindex(A.data, idxs_inner..., idxs_outer...) - nestedview(new_data, ndims_after_getindex(idxs_inner)) -end - - -Base.@propagate_inbounds function deepsetindex!(A::ArrayOfSimilarArrays{T,M,N,L}, x, idxs::Vararg{Any,L}) where {T,M,N,L} - idxs_outer, idxs_inner = split_tuple(idxs, Val{N}()) - _deepsetindex_impl_aosa!(A, x, idxs_outer, idxs_inner, _result_is_nested(idxs_outer, idxs_inner)) - A -end - -Base.@propagate_inbounds _deepsetindex_impl_aosa!(A::ArrayOfSimilarArrays, x, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{false}) = - setindex!(A.data, x, idxs_inner..., idxs_outer...) - -Base.@propagate_inbounds _deepsetindex_impl_aosa!(A::ArrayOfSimilarArrays, x, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{true}) = - _deepsetindex_impl!(A, x, idxs_outer, idxs_inner) - -# TODO: Specialized method _deepsetindex_impl_aosa!(A::ArrayOfSimilarArrays, x::ArrayOfSimilarArrays, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{true}) = - - -Base.@propagate_inbounds function deepview(A::ArrayOfSimilarArrays{T,M,N,L}, idxs::Vararg{Any,L}) where {T,M,N,L} - idxs_outer, idxs_inner = split_tuple(idxs, Val{N}()) - nested = _result_is_nested(idxs_outer, idxs_inner) - _deepview_impl_aosa(A, idxs_outer, idxs_inner, nested) -end - -Base.@propagate_inbounds _deepview_impl_aosa(A::ArrayOfSimilarArrays, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{false}) = - view(A.data, idxs_inner..., idxs_outer...) - -Base.@propagate_inbounds function _deepview_impl_aosa(A::ArrayOfSimilarArrays, idxs_outer::Tuple, idxs_inner::Tuple, nested::Val{true}) - new_data = view(A.data, idxs_inner..., idxs_outer...) - nestedview(new_data, ndims_after_getindex(idxs_inner)) -end - - - const VectorOfSimilarArrays{ - T, M, L, - P<:AbstractArray{T,L} -} = ArrayOfSimilarArrays{T,M,1,L,P} + T, M, + P<:AbstractArray{T}, + ET<:AbstractArray{T,M} +} = ArrayOfSimilarArrays{T,M,1,P,ET} export VectorOfSimilarArrays -VectorOfSimilarArrays{T}(flat_data::AbstractArray{U,L}) where {T,U,L} = - ArrayOfSimilarArrays{T,length(Base.front(size(flat_data))),1}(flat_data) - -VectorOfSimilarArrays(flat_data::AbstractArray{T,L}) where {T,L} = - ArrayOfSimilarArrays{T,length(Base.front(size(flat_data))),1}(flat_data) - -VectorOfSimilarArrays{T}(A::AbstractVector{<:AbstractArray{U,M}}) where {T,M,U} = - VectorOfSimilarArrays{T,M}(A) - -VectorOfSimilarArrays(A::AbstractVector{<:AbstractArray{T,M}}) where {T,M} = - VectorOfSimilarArrays{T,M}(A) +VectorOfSimilarArrays{T}(data::AbstractArray{U}) where {T,U} = ArrayOfSimilarArrays{T,length(Base.front(size(data))),1}(data) +VectorOfSimilarArrays(data::AbstractArray{T}) where {T} = ArrayOfSimilarArrays{T,length(Base.front(size(data))),1}(data) - -Base.convert(R::Type{VectorOfSimilarArrays{T}}, flat_data::AbstractArray{U,L}) where {T,U,L} = R(flat_data) -Base.convert(R::Type{VectorOfSimilarArrays}, flat_data::AbstractArray{T,L}) where {T,L} = R(flat_data) -Base.convert(R::Type{VectorOfSimilarArrays{T}}, A::AbstractVector{<:AbstractArray{U,M}}) where {T,M,U} = R(A) -Base.convert(R::Type{VectorOfSimilarArrays}, A::AbstractVector{<:AbstractArray{T,M}}) where {T,M} = R(A) +Base.convert(::Type{VectorOfSimilarArrays{T}}, A::AbstractVector{<:AbstractArray{U,M}}) where {T,M,U} = ArrayOfSimilarArrays{T,M,1}(stacked(A)) +Base.convert(::Type{VectorOfSimilarArrays}, A::AbstractVector{<:AbstractArray{T,M}}) where {T,M} = ArrayOfSimilarArrays{T,M,1}(stacked(A)) @inline Base.IndexStyle(::Type{<:VectorOfSimilarArrays}) = IndexLinear() @@ -315,6 +314,20 @@ end # popfirst!(V::ArrayOfSimilarArrays) = ... +# Concatenating vectors of similar arrays concatenates their underlying +# data along its last dimension, with a single allocation: + +function _vcat_vosas(Vs) + isempty(Vs) && throw(ArgumentError("reducing over an empty collection is not allowed")) + new_data = _cat_lastdim(map(fused, Vs)) + return VectorOfSimilarArrays(new_data) +end + +Base.vcat(V1::AbstractVectorOfSimilarArrays, Vs::AbstractVectorOfSimilarArrays...) = _vcat_vosas((V1, Vs...)) + +Base.reduce(::typeof(vcat), Vs::AbstractVector{<:AbstractVectorOfSimilarArrays}) = _vcat_vosas(Vs) + + function _empty_data_size(A::VectorOfSimilarArrays{T,M}) where {T,M} size_inner, size_outer = split_tuple(size(A.data), Val{M}()) empty_size_outer = map(x -> zero(x), size_outer) @@ -324,7 +337,7 @@ end function Base.empty(A::VectorOfSimilarArrays{T,M}, ::Type{<:AbstractArray{U}}) where {T,M,U} new_data_size = _empty_data_size(A) # ToDo: Don't use similar if data is an ElasticArray? - VectorOfSimilarArrays{T,M}(similar(A.data, U, new_data_size...)) + ArrayOfSimilarArrays{U,M,1}(similar(A.data, U, new_data_size...)) end function Base.empty!(A::VectorOfSimilarArrays{T,M}) where {T,M} @@ -335,70 +348,117 @@ end const ArrayOfSimilarVectors{ - T, N, L, - P<:AbstractArray{T,L} -} = ArrayOfSimilarArrays{T,1,N,L,P} + T, N, + P<:AbstractArray{T}, + ET<:AbstractVector{T} +} = ArrayOfSimilarArrays{T,1,N,P,ET} export ArrayOfSimilarVectors -ArrayOfSimilarVectors{T}(flat_data::AbstractArray{U,L}) where {T,U,L} = - ArrayOfSimilarArrays{T,1,length(Base.front(size(flat_data)))}(flat_data) - -ArrayOfSimilarVectors(flat_data::AbstractArray{T,L}) where {T,L} = - ArrayOfSimilarArrays{T,1,length(Base.front(size(flat_data)))}(flat_data) - -ArrayOfSimilarVectors{T}(A::AbstractArray{<:AbstractVector{U},N}) where {T,N,U} = - ArrayOfSimilarVectors{T,N}(A) +ArrayOfSimilarVectors{T}(data::AbstractArray{U}) where {T,U} = ArrayOfSimilarArrays{T,1,length(Base.front(size(data)))}(data) +ArrayOfSimilarVectors(data::AbstractArray{T}) where {T} = ArrayOfSimilarArrays{T,1,length(Base.front(size(data)))}(data) -ArrayOfSimilarVectors(A::AbstractArray{<:AbstractVector{T},N}) where {T,N} = - ArrayOfSimilarVectors{T,N}(A) - - -Base.convert(R::Type{ArrayOfSimilarVectors{T}}, flat_data::AbstractArray{U,L}) where {T,U,L} = R(flat_data) -Base.convert(R::Type{ArrayOfSimilarVectors}, flat_data::AbstractArray{T,L}) where {T,L} = R(flat_data) -Base.convert(R::Type{ArrayOfSimilarVectors{T}}, A::AbstractArray{<:AbstractVector{U},N}) where {T,N,U} = R(A) -Base.convert(R::Type{ArrayOfSimilarVectors}, A::AbstractArray{<:AbstractVector{T},N}) where {T,N} = R(A) +Base.convert(R::Type{ArrayOfSimilarVectors{T}}, A::AbstractArray{<:AbstractVector{U},N}) where {T,N,U} = R(stacked(A)) +Base.convert(R::Type{ArrayOfSimilarVectors}, A::AbstractArray{<:AbstractVector{T},N}) where {T,N} = R(stacked(A)) const VectorOfSimilarVectors{ T, - P<:AbstractArray{T,2} -} = ArrayOfSimilarArrays{T,1,1,2,P} + P<:AbstractArray{T,2}, + ET<:AbstractVector{T} +} = ArrayOfSimilarArrays{T,1,1,P,ET} export VectorOfSimilarVectors -VectorOfSimilarVectors{T}(flat_data::AbstractArray{U,2}) where {T,U} = - ArrayOfSimilarArrays{T,1,1}(flat_data) +VectorOfSimilarVectors{T}(data::AbstractArray{U,2}) where {T,U} = ArrayOfSimilarArrays{T,1,1}(data) +VectorOfSimilarVectors(data::AbstractArray{T,2}) where {T} = VectorOfSimilarVectors{T}(data) -VectorOfSimilarVectors(flat_data::AbstractArray{T,2}) where {T} = - VectorOfSimilarVectors{T}(flat_data) +Base.convert(R::Type{VectorOfSimilarVectors{T}}, A::AbstractVector{<:AbstractVector{U}}) where {T,U} = R(stacked(A)) +Base.convert(R::Type{VectorOfSimilarVectors}, A::AbstractVector{<:AbstractVector{T}}) where {T} = R(stacked(A)) -VectorOfSimilarVectors{T}(A::AbstractVector{<:AbstractVector{U}}) where {T,U} = - ArrayOfSimilarArrays{T,1}(A) -VectorOfSimilarVectors(A::AbstractVector{<:AbstractVector{T}}) where {T} = - VectorOfSimilarVectors{T}(A) +# Fast paths over the flat data for the full (dims = :) reductions. Julia +# does not dispatch on keyword arguments, so dispatch on the value of `dims` +# instead and forward non-Colon dims to the generic implementations: -Base.convert(R::Type{VectorOfSimilarVectors{T}}, flat_data::AbstractArray{U,2}) where {T,U} = R(flat_data) -Base.convert(R::Type{VectorOfSimilarVectors}, flat_data::AbstractArray{T,2}) where {T} = R(flat_data) -Base.convert(R::Type{VectorOfSimilarVectors{T}}, A::AbstractVector{<:AbstractVector{U}}) where {T,U} = R(A) -Base.convert(R::Type{VectorOfSimilarVectors}, A::AbstractVector{<:AbstractVector{T}}) where {T} = R(A) +Base.sum(X::AbstractVectorOfSimilarArrays; dims::Union{Colon,Integer} = :) = _aosa_sum(X, dims) - -Base.sum(X::AbstractVectorOfSimilarArrays{T,M}) where {T,M} = +_aosa_sum(X::AbstractVectorOfSimilarArrays{T,M}, ::Colon) where {T,M} = sum(flatview(X); dims = M + 1)[_ncolons(Val{M}())...] +_aosa_sum(X::AbstractVectorOfSimilarArrays, dims::Integer) = + Base.@invoke sum(X::AbstractArray; dims = dims) + +Statistics.mean(X::AbstractVectorOfSimilarArrays; dims::Union{Colon,Integer} = :) = _aosa_mean(X, dims) -Statistics.mean(X::AbstractVectorOfSimilarArrays{T,M}) where {T,M} = +_aosa_mean(X::AbstractVectorOfSimilarArrays{T,M}, ::Colon) where {T,M} = mean(flatview(X); dims = M + 1)[_ncolons(Val{M}())...] +_aosa_mean(X::AbstractVectorOfSimilarArrays, dims::Integer) = + Base.@invoke mean(X::AbstractArray; dims = dims) -Statistics.var(X::AbstractVectorOfSimilarArrays{T,M}; corrected::Bool = true, mean = nothing) where {T,M} = +Statistics.var(X::AbstractVectorOfSimilarArrays; corrected::Bool = true, mean = nothing, dims::Union{Colon,Integer} = :) = + _aosa_var(X, corrected, mean, dims) + +_aosa_var(X::AbstractVectorOfSimilarArrays{T,M}, corrected::Bool, mean, ::Colon) where {T,M} = var(flatview(X); dims = M + 1, corrected = corrected, mean = mean)[_ncolons(Val{M}())...] +_aosa_var(X::AbstractVectorOfSimilarArrays, corrected::Bool, mean, dims::Integer) = + Base.@invoke var(X::AbstractArray; corrected = corrected, mean = mean, dims = dims) + +Statistics.std(X::AbstractVectorOfSimilarArrays; corrected::Bool = true, mean = nothing, dims::Union{Colon,Integer} = :) = + _aosa_std(X, corrected, mean, dims) -Statistics.std(X::AbstractVectorOfSimilarArrays{T,M}; corrected::Bool = true, mean = nothing) where {T,M} = +_aosa_std(X::AbstractVectorOfSimilarArrays{T,M}, corrected::Bool, mean, ::Colon) where {T,M} = std(flatview(X); dims = M + 1, corrected = corrected, mean = mean)[_ncolons(Val{M}())...] +_aosa_std(X::AbstractVectorOfSimilarArrays, corrected::Bool, mean, dims::Integer) = + Base.@invoke std(X::AbstractArray; corrected = corrected, mean = mean, dims = dims) Statistics.cov(X::AbstractVectorOfSimilarVectors; corrected::Bool = true) = cov(flatview(X); dims = 2, corrected = corrected) Statistics.cor(X::AbstractVectorOfSimilarVectors) = cor(flatview(X); dims = 2) + + + +# Deprecations: + +@deprecate ArrayOfSimilarArrays{T,M,N}(A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} ArrayOfSimilarArrays{T,M,N}(stacked(A)) false +@deprecate ArrayOfSimilarArrays{T}(A::AbstractArray{<:AbstractArray{U,M},N}) where {T,M,N,U} ArrayOfSimilarArrays{T,M,N}(stacked(A)) false +@deprecate ArrayOfSimilarArrays(A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} ArrayOfSimilarArrays{T,M,N}(stacked(A)) false + +@deprecate Base.convert(::Type{ArrayOfSimilarArrays{T,M,N}}, data::AbstractArray{U,L}) where {T,M,N,U<:Number,L} ArrayOfSimilarArrays{T,M,N}(data) false +@deprecate Base.convert(::Type{ArrayOfSimilarArrays{T,M}}, data::AbstractArray{U,L}) where {T,M,U<:Number,L} ArrayOfSimilarArrays{T,M}(data) false + + +@deprecate VectorOfSimilarArrays{T}(A::AbstractVector{<:AbstractArray{U,M}}) where {T,M,U} ArrayOfSimilarArrays{T,M,1}(stacked(A)) false +@deprecate VectorOfSimilarArrays(A::AbstractVector{<:AbstractArray{T,M}}) where {T,M} ArrayOfSimilarArrays{T,M,1}(stacked(A)) false + +@deprecate Base.convert(::Type{VectorOfSimilarArrays{T}}, data::AbstractArray{U}) where {T,U<:Number} VectorOfSimilarArrays{T}(data) false +@deprecate Base.convert(::Type{VectorOfSimilarArrays}, data::AbstractArray{T}) where {T<:Number} VectorOfSimilarArrays(data) false + + +@deprecate ArrayOfSimilarVectors{T}(A::AbstractArray{<:AbstractVector{U},N}) where {T,N,U} ArrayOfSimilarArrays{T,1,N}(stacked(A)) false +@deprecate ArrayOfSimilarVectors(A::AbstractArray{<:AbstractVector{T},N}) where {T,N} ArrayOfSimilarArrays{T,1,N}(stacked(A)) false + +@deprecate Base.convert(::Type{ArrayOfSimilarVectors{T}}, data::AbstractArray{U}) where {T,U<:Number} ArrayOfSimilarVectors{T}(data) false +@deprecate Base.convert(::Type{ArrayOfSimilarVectors}, data::AbstractArray{T}) where {T<:Number} ArrayOfSimilarVectors(data) false + + +@deprecate VectorOfSimilarVectors{T}(A::AbstractVector{<:AbstractVector{U}}) where {T,U} ArrayOfSimilarArrays{T,1,1}(stacked(A)) false +@deprecate VectorOfSimilarVectors(A::AbstractVector{<:AbstractVector{T}}) where {T} ArrayOfSimilarArrays{T,1,1}(stacked(A)) false + +@deprecate Base.convert(::Type{VectorOfSimilarVectors{T}}, data::AbstractArray{U,2}) where {T,U<:Number} VectorOfSimilarVectors{T}(data) false +@deprecate Base.convert(::Type{VectorOfSimilarVectors}, data::AbstractArray{T,2}) where {T<:Number} VectorOfSimilarVectors(data) false + + +""" + nestedview(A::AbstractArray{T,M+N}, M::Integer) + nestedview(A::AbstractArray{T,2}) + +Deprecated, use [`sliced`](@ref) instead. +""" +function nestedview end +export nestedview + +@deprecate nestedview(A::AbstractArray, M::Integer) sliced(A, Val(M)) +@deprecate nestedview(A::AbstractArray, ::Val{M}) where {M} sliced(A, Val(M)) +@deprecate nestedview(A::AbstractArray{T,2}) where {T} sliced(A) diff --git a/src/base_slices.jl b/src/base_slices.jl new file mode 100644 index 0000000..34e201a --- /dev/null +++ b/src/base_slices.jl @@ -0,0 +1,95 @@ +# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). + +""" + struct BaseSlicing{M,N,TPL<:Tuple{Vararg{Union{Colon,Int}}}} <: AbstractSlicingMode{M,N} + +The split mode of `Base.Slices` (as returned by `eachslice`, `eachcol` and +`eachrow`). + +Constructor: + +``` +BaseSlicing{M,N,TPL}(slicemap::TPL) +``` + +`slicemap` equals the `slicemap` property of `Base.Slices` objects. + +See also [`AbstractSlicingMode`](@ref). +""" +struct BaseSlicing{M,N,TPL<:Tuple{Vararg{Union{Colon,Int}}}} <: AbstractSlicingMode{M,N} + slicemap::TPL +end +export BaseSlicing + +is_memordered_splitmode(::BaseSlicing{1,1,Tuple{Colon,Int}}) = true + +function is_memordered_splitmode(smode::BaseSlicing{M,N}) where {M,N} + dims = _oneto_tpl(Val(M+N)) + issorted((getinnerdims(dims, smode)..., getouterdims(dims, smode)...)) +end + + +@inline @generated function getinnerdims(obj::Tuple, ::BaseSlicing{M,N,SliceMapT}) where {M,N,SliceMapT} + expr = Expr(:tuple) + slicepars = SliceMapT.parameters + for i in 1:length(slicepars) + if slicepars[i] <: Colon + push!(expr.args, :(obj[$i])) + end + end + return expr +end + + +@inline @generated function getouterdims(obj::Tuple, smode::BaseSlicing{M,N,SliceMapT}) where {M,N,SliceMapT} + slicepars = SliceMapT.parameters + outdimidxs = Int[] + for i in 1:length(slicepars) + if slicepars[i] <: Integer + push!(outdimidxs, i) + end + end + + # With less than two outer dimensions the result is fully determined by + # SliceMapT already: + if length(outdimidxs) < 2 + return Expr(:tuple, [:(obj[$i]) for i in outdimidxs]...) + end + + # vals[j] is the entry of obj and pos[j] the outer dimension number for + # the j-th outer dimension of the parent array. The result must be + # ordered by outer dimension number, so it is vals in inverse-perm order + # of pos. pos is only known at runtime, its values are not encoded in + # SliceMapT. + vals_expr = Expr(:tuple, [:(obj[$i]) for i in outdimidxs]...) + pos_expr = Expr(:tuple, [:(slicemap[$i]) for i in outdimidxs]...) + + quote + slicemap = smode.slicemap + vals = $vals_expr + pos = $pos_expr + return _invpermuted(vals, pos) + end +end + +@inline function _invpermuted(vals::NTuple{N,Any}, pos::NTuple{N,Int}) where N + ntuple(k -> vals[findfirst(==(k), pos)::Int], Val(N)) +end + + +@inline getslicemap(A::Slices) = A.slicemap + +@inline function getsplitmode(A::Slices) + M = ndims(eltype(A)) + N = ndims(A) + slicemap = getslicemap(A) + BaseSlicing{M,N,typeof(slicemap)}(slicemap) +end + +function splitup(A::AbstractArray, smode::BaseSlicing) + slicemap = smode.slicemap + axs = getouterdims(axes(A), smode) + return Slices(A, slicemap, axs) +end + +@inline fused(A::Slices) = parent(A) diff --git a/src/broadcasting.jl b/src/broadcasting.jl index f253fb5..95b83a6 100644 --- a/src/broadcasting.jl +++ b/src/broadcasting.jl @@ -4,6 +4,47 @@ const _RefLike{T} = Union{Tuple{T}, Ref{T}} +""" + ArraysOfArrays.NestedArrayStyle{N}() + +Broadcast style of nested array types like [`VectorOfArrays`](@ref) and +[`ArrayOfSimilarArrays`](@ref). + +Broadcasts over a single outer dimension that produce array-valued results +(`f` receives whole element arrays and returns arrays, as in +`(x -> 2 .* x).(A)`) return a [`VectorOfArrays`](@ref) instead of a `Vector` +of arrays, where possible. All other broadcasts behave like the default +broadcast machinery. Use [`bcastat`](@ref) to broadcast over the *contents* +of the element arrays instead. +""" +struct NestedArrayStyle{N} <: Broadcast.AbstractArrayStyle{N} end + +NestedArrayStyle{M}(::Val{N}) where {M,N} = NestedArrayStyle{N}() + +Base.Broadcast.BroadcastStyle(::Type{<:AbstractArrayOfSimilarArrays{<:Any,<:Any,N}}) where {N} = NestedArrayStyle{N}() +Base.Broadcast.BroadcastStyle(::Type{<:VectorOfArrays}) = NestedArrayStyle{1}() + +function Base.copy(bc::Broadcast.Broadcasted{NestedArrayStyle{N}}) where {N} + ElType = Broadcast.combine_eltypes(bc.f, bc.args) + if N == 1 && ElType <: Array && isconcretetype(ElType) && axes(bc, 1) isa Base.OneTo + return _collect_nested(bc, ElType) + else + # Everything else behaves like the default broadcast machinery: + return copy(convert(Broadcast.Broadcasted{Broadcast.DefaultArrayStyle{N}}, bc)) + end +end + +function _collect_nested(bc::Broadcast.Broadcasted, ::Type{Array{T,M}}) where {T,M} + dest = VectorOfArrays{T,M}() + sizehint!(dest.elem_ptr, length(axes(bc, 1)) + 1) + sizehint!(dest.kernel_size, length(axes(bc, 1))) + for i in eachindex(bc) + push!(dest, bc[i]) + end + return dest +end + + _idx_type(::VectorOfSimilarArrays) = Int _idx_type(A::VectorOfArrays) = eltype(A.elem_ptr) @@ -35,7 +76,7 @@ _noreindex_view(A::AbstractArray{T,N}, ::Vararg{Colon,N}) where {T,N} = A _generic_size(A) = size(A) _generic_size(tpl::Tuple) = (length(tpl),) -Base.Base.@propagate_inbounds function _to_indices(A, idxs) +Base.@propagate_inbounds function _to_indices(A, idxs) new_idxs = Base.to_indices(A, idxs) @boundscheck Base.checkbounds(A, new_idxs...) return new_idxs @@ -45,7 +86,7 @@ end # ToDo: Extend to vectors of arrays. function Base.Broadcast.broadcasted( ::typeof(getindex), - A::Union{VectorOfSimilarVectors, VectorOfVectors}, + A::Union{VectorOfSimilarVectors, PartsView}, Idxs::Union{AbstractVector{<:AbstractVector{<:Integer}},AbstractVector{Colon},_RefLike{<:Union{AbstractVector{<:Integer},Colon}}}... ) # Checks size compatibility: @@ -123,7 +164,7 @@ end function Base.Broadcast.broadcasted( ::typeof(findall), - A::Union{VectorOfSimilarVectors{Bool}, VectorOfVectors{Bool}} + A::Union{VectorOfSimilarVectors{Bool}, PartsView{Bool}} ) new_lengths = _similar_idx_vector(A, _idx_type(A), length(A)) new_lengths .= sum.(A) diff --git a/src/functions.jl b/src/functions.jl index 7509e9d..7fccf6b 100644 --- a/src/functions.jl +++ b/src/functions.jl @@ -2,83 +2,328 @@ """ - innermap(f::Base.Callable, A::AbstractArray{<:AbstractArray}) + abstract type AbstractSplitMode <: Function -Nested `map` at depth 2. Equivalent to `map(X -> map(f, X) A)`. +Abstract supertype for array split modes. + +Use [`getsplitmode`](@ref) to get the split mode of an array. + +Use [`splitup`](@ref) or call `smode::AbstractSplitMode` as a function to +split an array: + +```julia +splitup(A, smode) === smode(A) +``` + +See also [`splitup`](@ref) and [`fused`](@ref). + +# Implementation + +Subtypes of `AbstractSplitMode` should specialize + +* `splitup(A, smode::SomeSplitMode)` for arrays `A` +* `is_memordered_splitmode(smode::SomeSplitMode)` + +`(smode::SomeSplitMode)(A)` calls `splitup(A, smode)` by default and should +not be specialized. """ -function innermap end -export innermap +abstract type AbstractSplitMode <: Function end +export AbstractSplitMode + +(smode::AbstractSplitMode)(A::AbstractArray) = splitup(A, smode) + -innermap(f::Base.Callable, A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} = - map(X -> map(f, X), A) +""" + is_memordered_splitmode(smode::AbstractSplitMode)::Bool +Check whether `smode` splits in memory order. +If true, inner arrays are stored contiguously in memory in Julia-native +dimension order, and the same is true for the outer dimensions (no dimension +reordering). +If true, `flatview` and `fused` are equivalent. """ - deepmap(f::Base.Callable, x::Any) - deepmap(f::Base.Callable, A::AbstractArray{<:AbstractArray{<:...}}) +function is_memordered_splitmode end +export is_memordered_splitmode + -Applies `map` at the deepest possible layer of nested arrays. If `A` is not -a nested array, `deepmap` behaves identical to `Base.map`. """ -function deepmap end -export deepmap + struct NonSplitMode{N} <: AbstractSplitMode + +The split mode of unsplit collections that have `N` dimensions. + +Constructor: `NonSplitMode{N}()` +""" +struct NonSplitMode{N} <: AbstractSplitMode end +export NonSplitMode + +is_memordered_splitmode(::NonSplitMode) = true + + +""" + struct UnknownSplitMode{AT} <: AbstractSplitMode + +Split mode of generic split objects of type `AT` that have been split in an +unknown way, e.g. nested arrays of type `Array{<:Array}`. + +Since the split parts may be (typically are) non-contiguous in memory, this +split mode does not allow for `fused` or `flatview`. It is also not +inferable whether the split object should be interpreted as a sliced array, +a ragged array, or something else. + +Constructor: `UnknownSplitMode{T}()` +""" +struct UnknownSplitMode{AT} <: AbstractSplitMode end +export UnknownSplitMode + +is_memordered_splitmode(::UnknownSplitMode) = false + + +""" + getsplitmode(A::AbstractArray)::NonSplitMode + getsplitmode(A::AbstractArray{<:AbstractArray})::AbstractSplitMode + +Get the split mode of `A`. + +`splitup(fused(A), getsplitmode(A))` must equal `A`, and should have +the same type as `A` if at all possible, except if `getsplitmode(A)` is an +`UnknownSplitMode`. + +`getsplitmode` should be a zero-copy O(1) operation, if at all possible. +""" +function getsplitmode end +export getsplitmode + +@inline getsplitmode(::AbstractArray{<:Any,N}) where N = NonSplitMode{N}() + +@inline getsplitmode(A::AbstractArray{<:AbstractArray}) = UnknownSplitMode{typeof(A)}() + + +""" + splitup(A::AbstractArray, smode::AbstractSplitMode) + +View array `A` in split form, as an array of arrays. + +`splitup` should be a zero-copy O(1) operation, if at all possible. + +See also [`fused`](@ref) and [`getsplitmode`](@ref). +""" +function splitup end +export splitup + +@inline splitup(A::AbstractArray, ::NonSplitMode) = A + +function splitup(::AbstractArray, ::UnknownSplitMode) + throw(ArgumentError("splitup cannot be used with UnknownSplitMode")) +end + -deepmap(f::Base.Callable, x::Any) = map(f, x) +""" + fused(A::AbstractArray) + fused(A::AbstractArray{<:AbstractArray}) + +View array `A` in unsplit form. + +`splitup(fused(A), getsplitmode(A))` must equal `A`, and should have +the same type as `A` if at all possible, except if `getsplitmode(A)` is an +`UnknownSplitMode`. + +If `A` is not a nested array, return `A` itself. If `A` is a split array, +return the original unsplit array. + +If `is_memordered_splitmode(getsplitmode(A))` is true, `fused(A)` is +equivalent to [`flatview(A)`](@ref). + +`fused` should be a zero-copy O(1) operation, if at all possible. +""" +function fused end +export fused + +@inline fused(A::AbstractArray) = A + +@inline fused(A::AbstractArray{<:AbstractArray}) = _fused_impl(A, getsplitmode(A)) -deepmap(f::Base.Callable, A::AbstractArray{<:AbstractArray}) = - map(X -> deepmap(f, X), A) +@inline _fused_impl(obj, ::NonSplitMode) = obj + +function _fused_impl(@nospecialize(obj), ::UnknownSplitMode) + throw(ArgumentError("fused not implemented for objects of type $(nameof(typeof(obj))) with unknown split mode")) +end """ flatview(A::AbstractArray) - flatview(A::AbstractArray{<:AbstractArray{<:...}}) + flatview(A::AbstractArray{<:AbstractArray}) + +View array `A` in a flattened form, with inner dimensions first. The shape of +the flattened form will depend on the type of `A`. If `A` is not a +nested array, the return value is `A` itself. Only specific types of nested +arrays are supported. -View array `A` in a suitable flattened form. The shape of the flattened form -will depend on the type of `A`. If the `A` is not a nested array, the return -value is `A` itself. When no type-specific method is available, `flatview` -will fall back to `Base.Iterators.flatten`. +`flatview` is a zero-copy O(1) operation. + +If `is_memordered_splitmode(getsplitmode(A))` is true, `flatview(A)` is +equivalent to [`fused(A)`](@ref). + +For sliced arrays the result of `flatview(A)` will equal [`stacked(A)`](@ref). +For partitioned vectors it will equal [`vecflattened(A)`](@ref), provided +that the parts cover the underlying data completely. """ function flatview end export flatview @inline flatview(A::AbstractArray) = A +function flatview(A::AbstractArray{<:AbstractArray}) + throw(ArgumentError("flatview not implemented for nested arrays of type $(nameof(typeof(A)))")) +end + +function flatview(A::AbstractSlices{<:AbstractArray}) + smode = getsplitmode(A) + if is_memordered_splitmode(smode) + return fused(A) + else + throw(ArgumentError("flatview requires memory-ordered split/slicing, but array has split mode $smode")) + end +end + + + +""" + abstract type AbstractSlicingMode{M,N} <: AbstractSplitMode + +Abstract supertype for array slicing modes with `M` inner dimensions and `N` +outer dimensions. + +Use `getsplitmode` to get the split mode of a split array. +""" +abstract type AbstractSlicingMode{M,N} <: AbstractSplitMode end +export AbstractSlicingMode + + +""" + getslicemap(A::AbstractSlices) + +Return the slicemap of `A` with respect to `B = fused(A)`: a tuple with one +entry per dimension of `B`, `Colon()` for sliced (inner) dimensions and `k` +for dimensions indexed by dimension `k` of `A`, so that + +```julia +A[i...] == view(B, map(s -> s isa Colon ? (:) : i[s], getslicemap(A))...) +``` + +E.g. `A = eachslice(B, dims = (3,1,5))` of a five-dimensional `B` has the +slicemap `(2, :, 1, :, 3)`, since `A[i1, i2, i3] == view(B, i2, :, i1, :, i3)`. + +Equals the `slicemap` field of `Base.Slices` objects. +""" +function getslicemap end +export getslicemap + + +""" + ArraysOfArrays.getinnerdims(tpl::Tuple, smode::AbstractSplitMode) + +Get the entries of `tpl` corresponding to the inner dimensions of split +mode `smode`, in the order specified by `smode`. +""" +function getinnerdims end + +@inline getinnerdims(::Tuple, ::NonSplitMode) = () + +function getinnerdims(::Tuple, ::UnknownSplitMode) + throw(ArgumentError("getinnerdims cannot be used with UnknownSplitMode")) +end + + +""" + ArraysOfArrays.getouterdims(tpl::Tuple, smode::AbstractSplitMode) + +Get the entries of `tpl` corresponding to the outer dimensions of split +mode `smode`, in the order specified by `smode`. +""" +function getouterdims end + +@inline getouterdims(x::Tuple, ::NonSplitMode) = x + +function getouterdims(::Tuple, ::UnknownSplitMode) + throw(ArgumentError("getouterdims cannot be used with UnknownSplitMode")) +end + + +""" + unstackmode(A::AbstractArray) + unstackmode(A::AbstractArray{<:AbstractArray}) + +Get the split mode required to restore `stacked(A)` so that +`splitup(stacked(A), unstackmode(A)) == A`. + +The result of `splitup(stacked(A), unstackmode(A))` may have a different +type and underlying memory layout than `A`. +""" +function unstackmode end +export unstackmode + +@inline unstackmode(::AbstractArray{<:Any,N}) where N = NonSplitMode{N}() + +@inline unstackmode(A::AbstractArray{<:AbstractArray}) = UnknownSplitMode{typeof(A)}() + +function unstackmode(A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} + innersize(A) # Ensure element arrays have equal size + return SplitSlices{M,N}() +end -# TODO: Implement flatview on generic nested arrays via new `FlatView`, using -# deepgetindex to implement getindex, etc. -@inline flatview(A::AbstractArray{<:AbstractArray}) = Base.Iterators.flatten(A) +""" + stacked(A::AbstractArray{T,N})::AbstractArray{T,N} + stacked(A::AbstractArray{<:AbstractArray{T,M},N})::AbstractArray{T,M+N} + +Join stacked arrays of a nested array into a single array along one or more +new dimensions, return non-nested arrays unchanged. +Similar to `Base.stack`, but can return the original underlying array of +sliced arrays in more cases. """ - nestedview(A::AbstractArray{T,M+N}, M::Integer) - nestedview(A::AbstractArray{T,2}) +function stacked end +export stacked + +@inline stacked(A::AbstractArray) = A +@inline stacked(A::AbstractArray{<:AbstractArray}) = _stacked_impl(A, getsplitmode(A)) -AbstractArray{<:AbstractArray{T,M},N} +_stacked_impl(A::AbstractArray{<:AbstractArray}, ::AbstractSplitMode) = stack(A) -View array `A` in as an `N`-dimensional array of `M`-dimensional arrays by -wrapping it into an [`ArrayOfSimilarArrays`](@ref). +function _stacked_impl(A::AbstractSlices{<:AbstractArray}, smode::AbstractSlicingMode) + A_joined = fused(A) + is_memordered_splitmode(smode) ? A_joined : _stacked_permutedims(A_joined, smode) +end + +function _stacked_permutedims(A_joined::AbstractArray{T,N}, smode::AbstractSlicingMode) where {T,N} + dimnumbers = _oneto_tpl(Val(N)) + dimorder = (getinnerdims(dimnumbers, smode)..., getouterdims(dimnumbers, smode)...) + return permutedims(A_joined, dimorder) +end -It's also possible to use a `StaticVector` of length `S` as the type of the -inner arrays via - nestedview(A::AbstractArray{T}, ::Type{StaticArrays.SVector{S}}) - nestedview(A::AbstractArray{T}, ::Type{StaticArrays.SVector{S,T}}) """ -function nestedview end -export nestedview + sliced(A::AbstractArray{T,2}) + sliced(A::AbstractArray{T,M+N}, Val(M)) + sliced(A::AbstractArray{T,M+N}, M::Integer) -@inline nestedview(A::AbstractArray{T,L}, M::Integer) where {T,L} = - ArrayOfSimilarArrays{T,M}(A) +Return a sliced view of `A`, using the columns or the first `M` +dimensions as inner dimensions. -@inline nestedview(A::AbstractArray{T,L}, ::Val{M}) where {T,L,M} = - ArrayOfSimilarArrays{T,M}(A) +With StaticArrays loaded, `sliced(A, SVector{S})` returns a reinterpreted +array with `SVector{S}` elements instead of an array of views. +""" +function sliced end +export sliced -@inline nestedview(A::AbstractArray{T,2}) where {T} = - VectorOfSimilarVectors{T}(A) +@inline sliced(A::AbstractArray, M::Integer) = sliced(A, Val(M)) +@inline sliced(A::AbstractArray{T,L}, ::Val{M}) where {T,L,M} = splitup(A, SplitSlices{M,L-M}()) +@inline sliced(A::AbstractArray{T,2}) where {T} = sliced(A, Val(1)) """ - innersize(A:AbstractArray{<:AbstractArray}, [dim]) + innersize(A::AbstractArray{<:AbstractArray}, [dim]) Returns the size of the element arrays of `A`. Fails if the element arrays are not of equal size. @@ -86,6 +331,10 @@ are not of equal size. function innersize end export innersize +@inline innersize(A::AbstractArray{<:AbstractArray}, dim::Integer) = innersize(A)[dim] + +innersize(::AbstractArray{<:Number}) = () + function innersize(A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} s = if !isempty(A) let sz_A = size(first(A)) @@ -101,141 +350,340 @@ function innersize(A::AbstractArray{<:AbstractArray{T,M},N}) where {T,M,N} end end - s + return s end -@inline innersize(A::AbstractArray{<:AbstractArray}, dim::Integer) = - innersize(A)[dim] +@inline innersize(A::AbstractSlices{<:AbstractArray{T,M},N}) where {T,M,N} = getinnerdims(size(fused(A)), getsplitmode(A)) + + +""" + innersizes(A::AbstractArray{<:AbstractArray{T,M}}) + +Returns the sizes of the element arrays of `A`, as an array of `Dims{M}` +shaped like `A`. In contrast to [`innersize`](@ref), the element arrays do +not need to be of equal size. +""" +function innersizes end +export innersizes + +innersizes(A::AbstractArray{<:AbstractArray}) = map(size, A) -@inline innersize(tpl::Tuple{T}) where T = size(only(tpl)) -@inline innersize(ref::Ref) = size(only(ref)) +innersizes(A::AbstractSlices{<:AbstractArray{T,M},N}) where {T,M,N} = fill(innersize(A), size(A)) """ - deepgetindex(A::AbstractArray, idxs...) - deepgetindex(A::AbstractArray{<:AbstractArray, N}, idxs...) where {N} + innerlengths(A::AbstractArray{<:AbstractArray}) + +Returns the lengths of the element arrays of `A`, as an array of `Int` +shaped like `A`. +""" +function innerlengths end +export innerlengths + +innerlengths(A::AbstractArray{<:AbstractArray}) = map(length, A) -Recursive `getindex` on flat or nested arrays. If A is an array of arrays, -uses the first `N` entries in `idxs` on `A`, then the rest on the inner -array(s). If A is not a nested array, `deepgetindex` is equivalent to -`getindex`. +innerlengths(A::AbstractSlices{<:AbstractArray{T,M},N}) where {T,M,N} = fill(prod(innersize(A)), size(A)) -See also [`deepsetindex!`](@ref) and [`deepview`](@ref). + +""" + innermap(f, A::AbstractArray) + innermap(f, A::AbstractArray{<:AbstractArray}) + +Nested `map` at depth 2. Equivalent to `map(X -> map(f, X), A)` for arrays +of arrays, otherwise equivalent to `Base.map`. """ -function deepgetindex end -export deepgetindex +function innermap end +export innermap -Base.@propagate_inbounds deepgetindex(A::AbstractArray{T,N}, idxs::Vararg{Any,N}) where {T,N} = getindex(A, idxs...) -Base.@propagate_inbounds deepgetindex(A::AbstractArray{<:AbstractArray,N}, idxs::Vararg{Any,N}) where {N} = getindex(A, idxs...) +innermap(f, A::AbstractArray) = map(f, A) +innermap(f, A::AbstractArray{<:AbstractArray}) = map(Base.Fix1(map, f), A) +innermap(f, A::AbstractSlices{<:AbstractArray}) = _generic_innermap_impl(f, A) -Base.@propagate_inbounds function deepgetindex(A::AbstractArray{<:AbstractArray,N}, idxs...) where {N} - idxs_outer, idxs_inner = split_tuple(idxs, Val{N}()) - _deepgetindex_impl(A, idxs_outer, idxs_inner) +function _generic_innermap_impl(f, A::AbstractArray) + joined_A = fused(A) + mapped_joined_A = map(f, joined_A) + mapped_A = splitup(mapped_joined_A, getsplitmode(A)) + return mapped_A end -Base.@propagate_inbounds _deepgetindex_impl(A::AbstractArray{<:AbstractArray}, idxs_outer::NTuple{N,Real}, idxs_inner::Tuple) where {N} = - deepgetindex(getindex(A, idxs_outer...), idxs_inner...) -Base.@propagate_inbounds _deepgetindex_impl(A::AbstractArray{<:AbstractArray}, idxs_outer::NTuple{N,Any}, idxs_inner::Tuple) where {N} = - _deepgetindex_tupled.(view(A, idxs_outer...), (idxs_inner,)) +""" + abstract type AbstractPartMode{M,N} <: AbstractSplitMode + +Abstract supertype for array partition modes with `M` inner dimensions and `N` +outer dimensions. + +The mode need not represent a true partition, a partition that discards part +of the original array is allowed. The elements of the partition may also +be reshaped, depending on the mode. -Base.@propagate_inbounds _deepgetindex_tupled(A::AbstractArray, idxs::Tuple) = deepgetindex(A, idxs...) +Use `getsplitmode` to get the split mode of a split array. +""" +abstract type AbstractPartMode{M,N} <: AbstractSplitMode end +export AbstractPartMode """ - deepsetindex!(A::AbstractArray, x, idxs...) - deepsetindex!(A::AbstractArray{<:AbstractArray,N}, x, idxs...) where {N} + vecflattened(A::AbstractArray{T})::AbstractVector{T} + vecflattened(A::AbstractArray{<:AbstractArray})::AbstractVector{T} + +Concatenate nested arrays into a single vector, return non-nested vectors +unchanged. + +If `A` is a nested view of a vector, `vecflattened(A)` should return the +underlying vector in a zero-copy O(1) fashion. So in contrast to +`reduce(vcat, A)` and `mapreduce(vec, vcat, A)`, the result may share +memory with `A`. + +# Implementation -Recursive `setindex!` on flat or nested arrays. If A is an array of arrays, -uses the first `N` entries in `idxs` on `A`, then the rest on the inner -array(s). If A is not a nested array, `deepsetindex!` is equivalent to -`setindex!`. +The default implementations are -See also [`deepgetindex`](@ref) and [`deepview`](@ref). +```julia +vecflattened(A::AbstractVector) = A +vecflattened(A::AbstractArray) = vec(A) +vecflattened(A::AbstractVector{<:AbstractVector}) = reduce(vcat, A) +vecflattened(A::AbstractArray{<:AbstractArray}) = mapreduce(vec, vcat, A) +``` + +Specialize `vecflattened` for custom nested array types that can provide a +zero-copy implementation. """ -function deepsetindex! end -export deepsetindex! +function vecflattened end +export vecflattened -Base.@propagate_inbounds deepsetindex!(A::AbstractArray{T,N}, x, idxs::Vararg{Any,N}) where {T,N} = setindex!(A, x, idxs...) -Base.@propagate_inbounds deepsetindex!(A::AbstractArray{<:AbstractArray,N}, x, idxs::Vararg{Any,N}) where {N} = setindex!(A, x, idxs...) +@inline vecflattened(A::AbstractVector) = A +@inline vecflattened(A::AbstractArray) = vec(A) +@inline vecflattened(A::AbstractVector{<:AbstractVector}) = reduce(vcat, A) +@inline vecflattened(A::AbstractVector{<:AbstractArray}) = mapreduce(vec, vcat, A) +@inline vecflattened(A::AbstractArray{<:AbstractVector}) = mapreduce(vec, vcat, A) +@inline vecflattened(A::AbstractArray{<:AbstractArray}) = mapreduce(vec, vcat, A) -Base.@propagate_inbounds function deepsetindex!(A::AbstractArray{<:AbstractArray,N}, x, idxs...) where {N} - idxs_outer, idxs_inner = split_tuple(idxs, Val{N}()) - _deepsetindex_impl!(A, x, idxs_outer, idxs_inner) - A -end -Base.@propagate_inbounds function _deepsetindex_impl!(A::AbstractArray{<:AbstractArray}, x, idxs_outer::NTuple{N,Real}, idxs_inner::Tuple) where {N} - B = getindex(A, idxs_outer...) - deepsetindex!(B, x, idxs_inner...) +""" + deepmap(f, A::AbstractArray) + deepmap(f, A::AbstractArray{<:AbstractArray{<:...}}) + +Applies `map` at the deepest layer of nested arrays. If `A` is not +a nested array, `deepmap` behaves identically to `Base.map`. +""" +function deepmap end +export deepmap + +deepmap(f, A::AbstractArray) = map(f, A) +deepmap(f, A::AbstractArray{<:AbstractArray}) = map(Base.Fix1(deepmap, f), A) +deepmap(f, A::AbstractSlices{<:AbstractArray}) = _generic_deepmap_impl(f, A) + +function _generic_deepmap_impl(f, A::AbstractArray) + joined_A = fused(A) + mapped_joined_A = deepmap(f, joined_A) + mapped_A = splitup(mapped_joined_A, getsplitmode(A)) + return mapped_A end -Base.@propagate_inbounds function _deepsetindex_impl!(A::AbstractArray{<:AbstractArray}, x, idxs_outer::NTuple{N,Any}, idxs_inner::Tuple) where {N} - B = view(A, idxs_outer...) - for i in eachindex(B, x) - deepsetindex!(B[i], x[i], idxs_inner...) + +""" + mapat(f, ::Val{depth}, As::AbstractArray...) + mapat(f, depth::Integer, As::AbstractArray...) + +Nested `map` at nesting depth `depth`: apply `f` elementwise to the objects +at depth `depth` of the nested arrays `As`. Depth 1 refers to the elements +of the arrays themselves, so `mapat(f, Val(1), As...)` is equivalent to +`map(f, As...)` and `mapat(f, Val(2), A)` is equivalent to +[`innermap`](@ref)`(f, A)`. If `depth` exceeds the nesting depth of the +arrays, `f` is applied to the innermost elements, like [`deepmap`](@ref). + +The `Integer` form relies on constant propagation for type stability, use +the `Val` form when passing a non-constant depth. + +All of `As` must have the same nesting structure down to `depth`, their +split modes (see [`getsplitmode`](@ref)) must be equal on each nesting +level. + +For split arrays (like [`ArrayOfSimilarArrays`](@ref) and +[`VectorOfArrays`](@ref)) `mapat` operates on the underlying flat data +directly, in a single (GPU-compatible) operation per nesting level. + +See also [`bcastat`](@ref) for broadcast semantics with arguments of +different nesting depth. +""" +function mapat end +export mapat + +@inline mapat(f, depth::Integer, As::Vararg{AbstractArray,NA}) where {NA} = mapat(f, Val(depth), As...) + +@inline mapat(f, ::Val{1}, As::Vararg{AbstractArray,NA}) where {NA} = map(f, As...) + +function mapat(f, ::Val{depth}, As::Vararg{AbstractArray,NA}) where {depth,NA} + depth isa Integer && depth >= 1 || throw(ArgumentError("mapat depth must be a positive integer")) + smodes = map(getsplitmode, As) + smode = first(smodes) + if smode isa UnknownSplitMode + all(m -> m isa UnknownSplitMode, smodes) || throw(DimensionMismatch("mapat requires arrays with equal nesting structure, but split modes differ")) + return map((xs...) -> mapat(f, Val(depth - 1), xs...), As...) + else + all(isequal(smode), smodes) || throw(DimensionMismatch("mapat requires arrays with equal nesting structure, but split modes differ")) + return splitup(mapat(f, Val(depth - 1), map(fused, As)...), smode) end end - """ - deepview(A::AbstractArray, idxs...) - deepview(A::AbstractArray{<:AbstractArray, N}, idxs...) where {N} + innermapreduce(f, op, A::AbstractArray{<:AbstractArray}; [init]) -Recursive `view` on flat or nested arrays. If A is an array of arrays, -uses the first `N` entries in `idxs` on `A`, then the rest on the inner -array(s). If A is not a nested array, `deepview` is equivalent to `view`. +Per-element `mapreduce` over the contents of the element arrays of `A`: +returns an array shaped like `A` that contains +`mapreduce(f, op, A[i]; [init])` for each element `A[i]`. -See also [`deepgetindex`](@ref) and [`deepsetindex!`](@ref). +For split arrays (like [`ArrayOfSimilarArrays`](@ref) and +[`VectorOfArrays`](@ref)) this uses efficient (GPU-compatible) segmented +reductions over the underlying flat data where possible. """ -function deepview end -export deepview +function innermapreduce end +export innermapreduce -Base.@propagate_inbounds deepview(A::AbstractArray{T,N}, idxs::Vararg{Any,N}) where {T,N} = view(A, idxs...) -Base.@propagate_inbounds deepview(A::AbstractArray{<:AbstractArray,N}, idxs::Vararg{Any,N}) where {N} = view(A, idxs...) +innermapreduce(f, op, A::AbstractArray{<:AbstractArray}; init = _NoInit()) = _innermapreduce_impl(f, op, init, A) -Base.@propagate_inbounds function deepview(A::AbstractArray{<:AbstractArray,N}, idxs...) where {N} - idxs_outer, idxs_inner = split_tuple(idxs, Val{N}()) - _deepview_impl(A, idxs_outer, idxs_inner) +function _innermapreduce_impl(f, op, init, A::AbstractArray{<:AbstractArray}) + if init isa _NoInit + map(x -> mapreduce(f, op, x), A) + else + map(x -> mapreduce(f, op, x; init = init), A) + end end -Base.@propagate_inbounds _deepview_impl(A::AbstractArray{<:AbstractArray}, idxs_outer::NTuple{N,Real}, idxs_inner::NTuple{M,Real}) where {N,M} = - deepview(getindex(A, idxs_outer...), idxs_inner...) -Base.@propagate_inbounds _deepview_impl(A::AbstractArray{<:AbstractArray}, idxs_outer::NTuple{N,Real}, idxs_inner::NTuple{M,Any}) where {N,M} = - deepview(getindex(A, idxs_outer...), idxs_inner...) +""" + innerreduce(op, A::AbstractArray{<:AbstractArray}; [init]) -Base.@propagate_inbounds _deepview_impl(A::AbstractArray{<:AbstractArray}, idxs_outer::NTuple{N,Any}, idxs_inner::NTuple{M,Real}) where {N,M} = - throw(ArgumentError("deepview not supported yes with outer indices $idxs_outer and inner indices $idxs_inner")) +Per-element `reduce` over the contents of the element arrays of `A`, +equivalent to `innermapreduce(identity, op, A; [init])`. +""" +function innerreduce end +export innerreduce -Base.@propagate_inbounds _deepview_impl(A::AbstractArray{<:AbstractArray}, idxs_outer::NTuple{N,Any}, idxs_inner::NTuple{M,Any}) where {N,M} = - _deepview_tupled.(view(A, idxs_outer...), (idxs_inner,)) +innerreduce(op, A::AbstractArray{<:AbstractArray}; init = _NoInit()) = _innermapreduce_impl(identity, op, init, A) -Base.@propagate_inbounds _deepview_tupled(A::AbstractArray, idxs::Tuple) = deepview(A, idxs...) +""" + innersum(A::AbstractArray{<:AbstractArray}) +Per-element sum over the contents of the element arrays of `A`, empty +element arrays sum to zero. """ - abstract_nestedarray_type(T_inner::Type, ::Val{ndims_tuple}) +function innersum end +export innersum -Return the type of nested `AbstractArray`s. `T_inner` specifies the element -type of the innermost layer of arrays, `ndims_tuple` specifies the -dimensionality of each nesting layer (outer arrays first). +innersum(A::AbstractArray{<:AbstractArray}) = innermapreduce(identity, +, A) +innersum(A::AbstractArray{<:AbstractArray{T}}) where {T<:Number} = innermapreduce(identity, +, A, init = zero(T)) -If `ndims_tuple` is empty, the returns is the (typically scalar) type -`T_inner` itself. + +""" + bcastat(f, ::Val{depth}, args...) + bcastat(f, depth::Integer, args...) + +Broadcast `f` over the contents of nested arrays at nesting depth `depth`, +with AwkwardArrays-like alignment, but with array-of-arrays nesting +semantics: + +* Nested (split) arrays with equal split modes are aligned at depth `depth`. +* Arrays that match the outer structure of a shallower nesting level + contribute one value per element of that level, broadcast over everything + below it. +* Scalars and `Ref`s broadcast over everything. + +`bcastat(f, Val(1), args...)` is equivalent to `broadcast(f, args...)`, and +like [`mapat`](@ref), a `depth` that exceeds the nesting depth of the +arguments applies `f` at the innermost level. The `Integer` form relies on +constant propagation for type stability, use the `Val` form when passing a +non-constant depth. + +Nested array arguments must be split arrays (like +[`ArrayOfSimilarArrays`](@ref) and [`VectorOfArrays`](@ref)): for these, +`bcastat` compiles to a single flat (GPU-compatible) broadcast per nesting +level. """ -function abstract_nestedarray_type end -export abstract_nestedarray_type +function bcastat end +export bcastat + +@inline bcastat(f, depth::Integer, args...) = bcastat(f, Val(depth), args...) +@inline bcastat(f, ::Val{1}, args...) = broadcast(f, args...) -Base.@pure function abstract_nestedarray_type(::Type{T_inner}, outer::Val{ndims_tuple}) where {T_inner,ndims_tuple} - _abstract_nestedarray_type_impl(T_inner, ndims_tuple...) +function bcastat(f, ::Val{depth}, args...) where {depth} + depth isa Integer && depth >= 1 || throw(ArgumentError("bcastat depth must be a positive integer")) + r = _bcast_ref(args...) + if r === nothing + # No nested arguments left, apply at the innermost level: + foreach(_require_known_mode, args) + return broadcast(f, args...) + else + smode, ref_flat = r + descended = map(arg -> _bcast_descend(arg, smode, ref_flat), args) + return splitup(bcastat(f, Val(depth - 1), descended...), smode) + end +end + +function _bcast_ref(args...) + for a in args + if a isa AbstractArray + m = getsplitmode(a) + (m isa AbstractSlicingMode || m isa AbstractPartMode) && return (m, fused(a)) + end + end + return nothing +end + +_require_known_mode(@nospecialize(x)) = nothing + +function _require_known_mode(x::AbstractArray) + getsplitmode(x) isa UnknownSplitMode && throw(ArgumentError("bcastat requires nested array arguments to be split arrays with a known split mode, like VectorOfArrays or ArrayOfSimilarArrays")) + nothing +end + +function _bcast_descend(x, smode::AbstractSplitMode, ref_flat::AbstractArray) + x isa AbstractArray || return x + ndims(x) == 0 && return x + m = getsplitmode(x) + if m isa UnknownSplitMode + throw(ArgumentError("bcastat requires nested array arguments to be split arrays with a known split mode, like VectorOfArrays or ArrayOfSimilarArrays")) + elseif m isa NonSplitMode + return _bcast_expand(x, smode, ref_flat) + else + isequal(m, smode) || throw(DimensionMismatch("bcastat requires nested array arguments with equal split modes")) + return fused(x) + end +end + +function _bcast_expand(x::AbstractArray, smode::AbstractSlicingMode{M,N}, ref_flat::AbstractArray) where {M,N} + if axes(x) == getouterdims(axes(ref_flat), smode) + is_memordered_splitmode(smode) || throw(ArgumentError("bcastat does not support outer-value broadcasting for slicings that are not in memory order")) + return reshape(x, ntuple(_ -> 1, Val(M))..., size(x)...) + elseif axes(x) == axes(ref_flat) + return x + else + throw(DimensionMismatch("bcastat argument shape matches neither the outer structure nor the flat data of the nested arguments")) + end +end + +function _bcast_expand(x::AbstractArray, smode::AbstractPartMode, ref_flat::AbstractArray) + throw(ArgumentError("bcastat does not support outer-value broadcasting for split mode $(nameof(typeof(smode)))")) end -Base.@pure _abstract_nestedarray_type_impl(::Type{T_inner}) where {T_inner} = T_inner -Base.@pure _abstract_nestedarray_type_impl(::Type{T_inner}, N) where {T_inner} = AbstractArray{T_inner, N} +""" + partitioned(A::AbstractVector, lengths::AbstractVector{<:Integer}) + partitioned(A::AbstractVector, shapes::AbstractVector{<:Dims}) + +Return a partitioned view of `A`, as a vector of arrays. + +The parts are consecutive, non-overlapping views of `A`, with sizes given by +`lengths` (resulting in a vector of vectors) or `shapes` (resulting in a +vector of arrays). +""" +function partitioned end +export partitioned + +# partitioned methods are defined in vector_of_arrays.jl. -Base.@pure _abstract_nestedarray_type_impl(::Type{T_inner}, N, M, ndims_tuple...) where {T_inner} = - AbstractArray{<:_abstract_nestedarray_type_impl(T_inner, M, ndims_tuple...), N} +# ToDo: Add partitioned(A::AbstractVector, n::Integer) and partitioned(A::AbstractVector, shape::Dims) ? diff --git a/src/util.jl b/src/util.jl index 10ae248..8245511 100644 --- a/src/util.jl +++ b/src/util.jl @@ -1,20 +1,15 @@ # This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). -#= - -function _split_dims(dims::NTuple{N,Integer}) where {N} - int_dims = Int.(dims) - Base.front(int_dims), int_dims[end] -end - -=# - @inline _tail_impl(x, ys...) = (ys...,) @inline _tail(x) = _tail_impl(x...) Base.@pure _ncolons(::Val{N}) where N = ntuple(_ -> Colon(), Val{N}()) +Base.@pure _nColons(::Val{N}) where N = ntuple(_ -> Colon, Val{N}()) + +@inline _oneto_tpl(::Val{N}) where N = ntuple(identity, Val{N}()) +Base.@pure _nInts(::Val{N}) where N = ntuple(_ -> Int, Val{N}()) Base.@propagate_inbounds front_tuple(x::NTuple{N,Any}, ::Val{M}) where {N,M} = @@ -30,17 +25,49 @@ Base.@propagate_inbounds swap_front_back_tuple(x::NTuple{N,Any}, ::Val{M}) where (back_tuple(x, Val{N - M}())..., front_tuple(x, Val{M}())...) -_convert_elype(::Type{T}, A::AbstractArray{T}) where {T} = A +_convert_eltype(::Type{T}, A::AbstractArray{T}) where {T} = A + +_convert_eltype(::Type{T}, A::AbstractArray{U}) where {T,U} = broadcast(Base.Fix1(convert, T), A) -_convert_elype(::Type{T}, A::AbstractArray{U}) where {T,U} = broadcast(x -> convert(T, x), A) +Base.@pure _add_vals(::Val{A}, ::Val{B}) where {A,B} = Val{A + B}() -Base.@pure _add_vals(Val_M::Val{M}, Val_N::Val{N}) where {M,N} = - Val{length((ntuple(identity, Val_M)..., ntuple(identity, Val_N)...))}() +Base.@pure _subtract_vals(::Val{A}, ::Val{B}) where {A,B} = Val{A - B}() +@inline _require_ndims(::Val{N}, ::Val{N}) where {N} = nothing + +function _require_ndims(::Val{N1}, ::Val{N2}) where {N1,N2} + throw(ArgumentError("Require an array with $N2 dimensions, but got an array with $N1 dimensions")) +end -Base.@pure require_ndims(A::AbstractArray{T,N}, Val_N::Val{N}) where {T,N} = - nothing +Base.@pure _val_value(::Val{x}) where x = x -Base.@pure require_ndims(A::AbstractArray{T,M}, Val_N::Val{N}) where {T,M,N} = - throw(ArgumentError("Require an array with $N dimensions")) + +# Internal sentinel for "no init value given": +struct _NoInit end + + +# Concatenate arrays of equal ndims and equal size except in their last +# dimension along their last dimension, with a single allocation: +_cat_lastdim(datas) = _cat_lastdim_impl(datas, Val(ndims(first(datas)))) + +function _cat_lastdim_impl(datas, ::Val{N}) where {N} + data1 = first(datas) + inner_sz = Base.front(size(data1)) + foreach(datas) do d + ndims(d) == N && Base.front(size(d)) == inner_sz || throw(DimensionMismatch("Can't concatenate arrays with different sizes in their non-last dimensions along their last dimension")) + end + + T = mapreduce(eltype, promote_type, datas) + n_lastdim = sum(d -> size(d, N), datas) + result = similar(data1, T, (inner_sz..., n_lastdim)) + + colons = ntuple(_ -> :, Val(N - 1)) + offset = firstindex(result, N) + for d in datas + n = size(d, N) + result[colons..., offset:(offset + n - 1)] = d + offset += n + end + return result +end diff --git a/src/vector_of_arrays.jl b/src/vector_of_arrays.jl index 645a9e2..4ff8d6c 100644 --- a/src/vector_of_arrays.jl +++ b/src/vector_of_arrays.jl @@ -1,14 +1,26 @@ # This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). +Base.@propagate_inbounds _voa_elem_range(elem_ptr::AbstractVector{<:Integer}, i::Integer) = + elem_ptr[i]:(elem_ptr[i+1] - 1) + +# Element type of a `VectorOfArrays{T,N}` with data of type `VT` and element +# pointers of type `VI`, must match the return type of `getindex` exactly: +@inline function _voa_eltype(::Type{VT}, ::Type{VI}, ::Val{N}) where {VT<:AbstractVector,VI<:AbstractVector{<:Integer},N} + R = Base.promote_op(_voa_elem_range, VI, Int) + SV = Base.promote_op(view, VT, R) + return Base.promote_op(_reshape_dataview, SV, NTuple{N,Int}) +end + + """ - VectorOfArrays{T,N,M} <: AbstractVector{<:AbstractArray{T,N}} + VectorOfArrays{T,N,M,VT,VI,VD,ET<:AbstractArray{T,N}} <: AbstractVector{ET} -An `VectorOfArrays` represents a vector of `N`-dimensional arrays (that may +A `VectorOfArrays` represents a vector of `N`-dimensional arrays (that may differ in size). Internally, `VectorOfArrays` stores all elements of all -arrays in a single flat vector. `M` must equal `N - 1` +arrays in a single flat vector. `M` must equal `N - 1`. -The `VectorOfArrays` itself supports `push!`, `unshift!`, etc., but the size +The `VectorOfArrays` itself supports `push!`, `append!`, etc., but the size of each individual array in the vector is fixed. `resize!` can be used to shrink, but not to grow, as the size of the additional element arrays in the vector would be unknown. However, memory space for up to `n` arrays with a @@ -35,18 +47,19 @@ VectorOfArrays( Other suitable values for `checks` are `ArraysOfArrays.simple_consistency_checks` and `ArraysOfArrays.no_consistency_checks`. -`VectorOfVectors` is defined as an type alias: +`PartsView` is defined as an type alias: ```julia -`VectorOfVectors{T,VT,VI,VD} = VectorOfArrays{T,1,VT,VI,VD}` +`PartsView{T,VT,VI,VD,ET} = VectorOfArrays{T,1,0,VT,VI,VD,ET}` ``` """ struct VectorOfArrays{ T, N, M, VT<:AbstractVector{T}, VI<:AbstractVector{<:Integer}, - VD<:AbstractVector{Dims{M}} -} <: AbstractVector{Array{T,N}} + VD<:AbstractVector{Dims{M}}, + ET<:AbstractArray{T,N} +} <: AbstractVector{ET} data::VT elem_ptr::VI kernel_size::VD @@ -57,11 +70,13 @@ struct VectorOfArrays{ elem_ptr = [firstindex(data)] kernel_size = Vector{Dims{M}}() + ET = _voa_eltype(typeof(data), typeof(elem_ptr), Val(N)) new{ T, N, M, typeof(data), typeof(elem_ptr), - typeof(kernel_size) + typeof(kernel_size), + ET }(data, elem_ptr, kernel_size) end @@ -77,7 +92,8 @@ struct VectorOfArrays{ VD<:AbstractVector{Dims{M}} } N = length((ntuple(_ -> 0, Val{M}())..., 0)) - A = new{T,N,M,VT,VI,VD}(data, elem_ptr, kernel_size) + ET = _voa_eltype(VT, VI, Val(N)) + A = new{T,N,M,VT,VI,VD,ET}(data, elem_ptr, kernel_size) checks(A) A end @@ -104,7 +120,7 @@ Base.convert(VA::Type{VectorOfArrays}, A::AbstractVector{AbstractArray{T,N}}) wh Returns the internal element pointer vector of `A`. -Do *not* change modify the returned vector in any way, as this would break the +Do *not* modify the returned vector in any way, as this would break the inner consistency of `A`. Use with care, see [`element_ptr`](@ref) for a safe version of this function. @@ -123,51 +139,89 @@ element_ptr(A::VectorOfArrays) = deepcopy(internal_element_ptr(A)) function full_consistency_checks(A::VectorOfArrays) simple_consistency_checks(A) - all(eachindex(A.kernel_size)) do i - len = A.elem_ptr[i+1] - A.elem_ptr[i] - klen = prod(A.kernel_size[i]) - len >= 0 && (klen == 1 || mod(len, klen) == 0) - end || throw(ArgumentError("VectorOfArrays inconsistent: Content of elem_ptr and kernel_size is inconsistent")) + _partition_sizes_valid(A.elem_ptr, A.kernel_size) || throw(ArgumentError("VectorOfArrays inconsistent: Content of elem_ptr and kernel_size is inconsistent")) nothing end +# Package extensions specialize this for GPU arrays, to avoid scalar indexing: +function _partition_sizes_valid(elem_ptr::AbstractVector{<:Integer}, kernel_size::AbstractVector) + all(eachindex(kernel_size)) do i + len = elem_ptr[i+1] - elem_ptr[i] + klen = prod(kernel_size[i]) + len >= 0 && (klen == 1 || mod(len, klen) == 0) + end +end + function simple_consistency_checks(A::VectorOfArrays{T,N,M}) where {T,N,M} M == N - 1 || throw(ArgumentError("VectorOfArrays{T,N,M} inconsistent: M must equal N - 1")) firstindex(A.elem_ptr) == firstindex(A.kernel_size) || throw(ArgumentError("VectorOfArrays inconsistent: elem_ptr and kernel_size have incompatible indexing")) size(A.elem_ptr, 1) == size(A.kernel_size, 1) + 1 || throw(ArgumentError("VectorOfArrays inconsistent: elem_ptr and kernel_size have incompatible size")) - first(A.elem_ptr) >= firstindex(A.data) || throw(ArgumentError("VectorOfArrays inconsistent: First elem_ptr inconsistent with data indices")) - last(A.elem_ptr) - 1 <= lastindex(A.data) || throw(ArgumentError("VectorOfArrays inconsistent: Last elem_ptr inconsistent with data indices")) + ep_first, ep_last = _scalar_first_last(A.elem_ptr) + ep_first >= firstindex(A.data) || throw(ArgumentError("VectorOfArrays inconsistent: First elem_ptr inconsistent with data indices")) + ep_last - 1 <= lastindex(A.data) || throw(ArgumentError("VectorOfArrays inconsistent: Last elem_ptr inconsistent with data indices")) nothing end +# Package extensions specialize this for GPU arrays, to allow the two scalar +# reads explicitly: +_scalar_first_last(x::AbstractVector) = (first(x), last(x)) + function no_consistency_checks(A::VectorOfArrays) nothing end -Base.@propagate_inbounds function _elem_range_size(A::VectorOfArrays, i::Integer) - elem_ptr = A.elem_ptr - - from = elem_ptr[i] - until = elem_ptr[i+1] - to = until - 1 - len = until - from - - ksize = A.kernel_size[i] +@inline function _elem_size(ksize::Dims, len::Integer) klen = prod(ksize) len_p, klen_p = promote(len, klen) - sz_lastdim = len == 0 ? len_p : div(len_p, klen_p) - sz = (ksize..., sz_lastdim) + sz_lastdim = len == 0 ? zero(len_p) : div(len_p, klen_p) + return (ksize..., Int(sz_lastdim)) +end + +Base.@propagate_inbounds function _elem_range_size(A::VectorOfArrays, i::Integer) + r = _voa_elem_range(A.elem_ptr, i) + sz = _elem_size(A.kernel_size[i], length(r)) + (r, sz) +end + - (from:to, sz) +function innerlengths(A::VectorOfArrays) + ep = A.elem_ptr + return view(ep, firstindex(ep)+1:lastindex(ep)) .- view(ep, firstindex(ep):lastindex(ep)-1) end +innersizes(A::VectorOfArrays) = _elem_size.(A.kernel_size, innerlengths(A)) + + +# Equality must be equivalent to elementwise comparison, but can be checked +# much more efficiently via the flat data. The underlying data of A and B may +# be offset relative to their element pointers, so only element sizes and the +# data regions covered by the elements may be compared: + +function _elem_lengths_equal(elem_ptr_A::AbstractVector{<:Integer}, elem_ptr_B::AbstractVector{<:Integer}) + length(elem_ptr_A) == length(elem_ptr_B) || return false + iA, iB = firstindex(elem_ptr_A), firstindex(elem_ptr_B) + @inbounds for k in 0:(length(elem_ptr_A) - 2) + elem_ptr_A[iA+k+1] - elem_ptr_A[iA+k] == elem_ptr_B[iB+k+1] - elem_ptr_B[iB+k] || return false + end + return true +end import Base.== -(==)(A::VectorOfArrays, B::VectorOfArrays) = - A.data == B.data && A.elem_ptr == B.elem_ptr && A.kernel_size == B.kernel_size +function ==(A::VectorOfArrays{<:Any,N}, B::VectorOfArrays{<:Any,N}) where {N} + axes(A) == axes(B) || return false + A.kernel_size == B.kernel_size || return false + _elem_lengths_equal(A.elem_ptr, B.elem_ptr) || return false + return vecflattened(A) == vecflattened(B) +end + +function Base.isequal(A::VectorOfArrays{<:Any,N}, B::VectorOfArrays{<:Any,N}) where {N} + axes(A) == axes(B) && isequal(A.kernel_size, B.kernel_size) && + _elem_lengths_equal(A.elem_ptr, B.elem_ptr) && + isequal(vecflattened(A), vecflattened(B)) +end """ flatview(A::VectorOfArrays{T})::Vector{T} @@ -182,6 +236,133 @@ function flatview(A::VectorOfArrays{<:Any,N,M,<:Any,<:SubArray}) where {N,M} view(A.data, A.elem_ptr[begin]:A.elem_ptr[end]-1) end + +""" + struct SplitParts{M,VI,VD} <: AbstractPartMode{M,1} + +The split mode of [`VectorOfArrays`](@ref): a partition of a vector into +consecutive parts of possibly different size, viewed as a vector of +`M`-dimensional arrays. + +Constructor: + +``` +SplitParts( + elem_ptr::AbstractVector{<:Integer}, + kernel_size::AbstractVector{Dims{M-1}} +) +``` + +`elem_ptr` and `kernel_size` equal the equivalent properties of +`VectorOfArrays`. + +[`getsplitmode(A::VectorOfArrays)`](@ref) copies the shape information of +`A` (an O(length) operation), so the resulting mode is not affected if `A` +is resized afterwards. A `VectorOfArrays` created via +[`splitup`](@ref), on the other hand, shares the vectors of the mode it was +created from, like the `VectorOfArrays` inner constructor. + +See also [`AbstractPartMode`](@ref). +""" +struct SplitParts{ + M, + VI<:AbstractVector{<:Integer}, + VD<:AbstractVector{<:Dims} +} <: AbstractPartMode{M,1} + elem_ptr::VI + kernel_size::VD + + function SplitParts( + elem_ptr::VI, + kernel_size::VD + ) where { + Mk, + VI<:AbstractVector{<:Integer}, + VD<:AbstractVector{Dims{Mk}} + } + M = _val_value(_add_vals(Val(Mk), Val(1))) + new{M,VI,VD}(elem_ptr, kernel_size) + end +end +export SplitParts + +is_memordered_splitmode(::SplitParts) = true + +Base.:(==)(a::SplitParts, b::SplitParts) = + a.elem_ptr == b.elem_ptr && a.kernel_size == b.kernel_size + +Base.hash(x::SplitParts, h::UInt) = + hash(x.kernel_size, hash(x.elem_ptr, hash(:SplitParts, h))) + +# Defensive copy: the split mode must not be affected if the vector of +# arrays is resized later on. Package extensions may specialize this for +# vector types that cannot be resized: +_shapeinfo_copy(x::AbstractVector) = copy(x) + +getsplitmode(A::VectorOfArrays) = SplitParts(_shapeinfo_copy(A.elem_ptr), _shapeinfo_copy(A.kernel_size)) + +@inline fused(A::VectorOfArrays) = A.data + +function splitup(A::AbstractVector, smode::SplitParts) + VectorOfArrays(A, smode.elem_ptr, smode.kernel_size, simple_consistency_checks) +end + +function _bcast_expand(x::AbstractArray, smode::SplitParts, ref_flat::AbstractArray) + ep = smode.elem_ptr + n_parts = length(smode.kernel_size) + if x isa AbstractVector && length(x) == n_parts + # One value per part, broadcast over the contents of that part. The + # value of x used for data not covered by any part is arbitrary, + # such data does not contribute to the result: + idx = similar(ref_flat, Int) + idx .= firstindex(ref_flat):lastindex(ref_flat) + seg = clamp.(searchsortedlast.(Ref(ep), idx), firstindex(x), lastindex(x)) + return x[seg] + elseif x isa AbstractVector && axes(x) == axes(ref_flat) + return x + else + throw(DimensionMismatch("bcastat argument shape matches neither the outer structure nor the flat data of the nested arguments")) + end +end + + +function partitioned(A::AbstractVector, lengths::AbstractVector{<:Integer}) + elem_ptr = _elem_ptr_from_lengths(A, lengths) + kernel_size = similar(lengths, Dims{0}) + fill!(kernel_size, ()) + VectorOfArrays(A, elem_ptr, kernel_size, no_consistency_checks) +end + +function partitioned(A::AbstractVector, shapes::AbstractVector{Dims{N}}) where {N} + elem_ptr = _elem_ptr_from_lengths(A, prod.(shapes)) + kernel_size = Base.front.(shapes) + VectorOfArrays(A, elem_ptr, kernel_size, no_consistency_checks) +end + +function _elem_ptr_from_lengths(A::AbstractVector, lengths::AbstractVector{<:Integer}) + elem_ptr = similar(lengths, Int, length(lengths) + 1) + i = firstindex(elem_ptr) + elem_ptr[i] = firstindex(A) + for l in lengths + l >= 0 || throw(ArgumentError("Part lengths must not be negative")) + elem_ptr[i + 1] = elem_ptr[i] + l + i += 1 + end + last(elem_ptr) - 1 <= lastindex(A) || throw(ArgumentError("Sum of part lengths exceeds length of data vector")) + return elem_ptr +end + + +function vecflattened(A::VectorOfArrays) + ep = A.elem_ptr + view(A.data, first(ep):(last(ep) - 1)) +end + +# Fast paths, must return independent arrays, unlike `vecflattened`: +Base.mapreduce(::typeof(vec), ::typeof(vcat), A::VectorOfArrays) = copy(vecflattened(A)) +Base.reduce(::typeof(vcat), A::VectorOfArrays{T,1}) where {T} = copy(vecflattened(A)) + + Base.size(A::VectorOfArrays) = size(A.kernel_size) Base.IndexStyle(::Type{<:VectorOfArrays}) = IndexLinear() @@ -307,7 +488,11 @@ end function Base.append!(A::VectorOfArrays{T,N}, B::VectorOfArrays{U,N}) where {T,N,U} if !isempty(B) - append!(A.data, B.data) + # append_elemptr! places the elements of B directly after the last + # element of A, so A must not contain unused trailing data and only + # the data covered by the elements of B may be appended: + last(A.elem_ptr) == lastindex(A.data) + 1 || throw(ArgumentError("Cannot append to a VectorOfArrays that has unused trailing data")) + append!(A.data, vecflattened(B)) append_elemptr!(A.elem_ptr, B.elem_ptr) append!(A.kernel_size, B.kernel_size) end @@ -337,22 +522,52 @@ function Base.append!(A::VectorOfArrays{T,N}, B::AbstractVector{<:AbstractArray{ end -Base.vcat(V::VectorOfArrays) = V +# Concatenating vectors of arrays concatenates the data covered by their +# elements, in a single pass: -function Base.vcat(Vs::(VectorOfArrays{U,N} where U)...) where {N} - data = vcat(map(x -> x.data, Vs)...) +function _vcat_voas(Vs) + isempty(Vs) && throw(ArgumentError("reducing over an empty collection is not allowed")) + V1 = first(Vs) - elem_ptr = similar(Vs[1].elem_ptr, 1) - elem_ptr[1] = firstindex(data) - @inbounds for i in eachindex(Vs) - append_elemptr!(elem_ptr, Vs[i].elem_ptr) - end + T = mapreduce(V -> eltype(V.data), promote_type, Vs) + n_parts = sum(length, Vs) + n_data = sum(V -> Int(last(V.elem_ptr) - first(V.elem_ptr)), Vs) - kernel_size = vcat(map(x -> x.kernel_size, Vs)...) + data = similar(V1.data, T, n_data) + elem_ptr = similar(V1.elem_ptr, n_parts + 1) + kernel_size = similar(V1.kernel_size, n_parts) - VectorOfArrays(data, elem_ptr, kernel_size, no_consistency_checks) + i_ptr = firstindex(elem_ptr) + elem_ptr[i_ptr] = firstindex(data) + i_ksz = firstindex(kernel_size) + i_data = firstindex(data) + + for V in Vs + ep = V.elem_ptr + covered_len = last(ep) - first(ep) + copyto!(data, i_data, V.data, first(ep), covered_len) + i_data += covered_len + + @inbounds for j in firstindex(ep):(lastindex(ep) - 1) + elem_ptr[i_ptr + 1] = elem_ptr[i_ptr] + (ep[j + 1] - ep[j]) + i_ptr += 1 + end + + ksz = V.kernel_size + copyto!(kernel_size, i_ksz, ksz, firstindex(ksz), length(ksz)) + i_ksz += length(ksz) + end + + return VectorOfArrays(data, elem_ptr, kernel_size, no_consistency_checks) end +Base.vcat(V1::VectorOfArrays{<:Any,N}, Vs::VectorOfArrays{<:Any,N}...) where {N} = _vcat_voas((V1, Vs...)) + +Base.reduce(::typeof(vcat), Vs::AbstractVector{<:VectorOfArrays}) = _vcat_voas(Vs) +# Disambiguation: +Base.reduce(::typeof(vcat), Vs::AbstractVector{<:VectorOfArrays{<:Any,1}}) = _vcat_voas(Vs) +Base.reduce(::typeof(vcat), Vs::AbstractVector{<:VectorOfArrays{<:Any,2}}) = _vcat_voas(Vs) + function Base.copy(V::VectorOfArrays) VectorOfArrays(copy(V.data), copy(V.elem_ptr), copy(V.kernel_size), no_consistency_checks) @@ -388,7 +603,7 @@ function Base.push!(A::VectorOfArrays{T,N}, x::AbstractArray{U,N}) where {T,N,U} end -function Base.empty(A::VectorOfArrays{T,N}, ::Type{<:DenseArray{U,N}}) where {T,N,U} +function Base.empty(A::VectorOfArrays{T,N}, ::Type{<:AbstractArray{U,N}}) where {T,N,U} empty_data = empty(A.data, U) empty_elem_ptr = push!(empty(A.elem_ptr), firstindex(empty_data)) empty_kernel_size = empty(A.kernel_size) @@ -421,37 +636,41 @@ Base.Broadcast.broadcasted(::typeof(identity), A::VectorOfArrays) = A """ - VectorOfVectors{T,...} = VectorOfArrays{T,1,...} + PartsView{T,...} = VectorOfArrays{T,1,0,...} + +A vector of vectors (that may differ in length), stored in contiguous, +partitioned form. See [`VectorOfArrays`](@ref) for details. Constructors: ```julia -VectorOfVectors(A::AbstractVector{<:AbstractVector}) -VectorOfVectors{T}(A::AbstractVector{<:AbstractVector}) where {T} +PartsView(A::AbstractVector{<:AbstractVector}) +PartsView{T}(A::AbstractVector{<:AbstractVector}) where {T} -VectorOfVectors( +PartsView( data::AbstractVector, elem_ptr::AbstractVector{<:Integer}, checks::Function = full_consistency_checks ) - -See also [VectorOfArrays](@ref). ``` + +See also [`VectorOfArrays`](@ref). """ -const VectorOfVectors{ +const PartsView{ T, VT<:AbstractVector{T}, VI<:AbstractVector{<:Integer}, - VD<:AbstractVector{Dims{0}} -} = VectorOfArrays{T,1,0,VT,VI,VD} + VD<:AbstractVector{Dims{0}}, + ET<:AbstractVector{T} +} = VectorOfArrays{T,1,0,VT,VI,VD,ET} -export VectorOfVectors +export PartsView -VectorOfVectors{T}() where {T} = VectorOfArrays{T,1}() +PartsView{T}() where {T} = VectorOfArrays{T,1}() -VectorOfVectors{T}(A::AbstractVector{<:AbstractVector}) where {T} = VectorOfArrays{T,1}(A) -VectorOfVectors(A::AbstractVector{<:AbstractVector}) = VectorOfArrays(A) +PartsView{T}(A::AbstractVector{<:AbstractVector}) where {T} = VectorOfArrays{T,1}(A) +PartsView(A::AbstractVector{<:AbstractVector}) = VectorOfArrays(A) -VectorOfVectors( +PartsView( data::AbstractVector, elem_ptr::AbstractVector{I}, checks::Function = full_consistency_checks @@ -468,7 +687,7 @@ VectorOfVectors( consgrouped_ptrs(A::AbstractVector) Compute an element pointer vector, suitable for creation of a -`VectorOfVectors` that implies grouping equal consecutive entries of +`PartsView` that implies grouping equal consecutive entries of `A`. Example: @@ -476,15 +695,15 @@ Example: ```julia A = [1, 1, 2, 3, 3, 2, 2, 2] elem_ptr = consgrouped_ptrs(A) - first.(VectorOfVectors(A, elem_ptr)) == [1, 2, 3, 2] + first.(PartsView(A, elem_ptr)) == [1, 2, 3, 2] ``` -consgrouped_ptrs + Typically, `elem_ptr` will be used to apply the computed grouping to other data: ```julia B = [1, 2, 3, 4, 5, 6, 7, 8] - VectorOfVectors(B, elem_ptr) == [[1, 2], [3], [4, 5], [6, 7, 8]] + PartsView(B, elem_ptr) == [[1, 2], [3], [4, 5], [6, 7, 8]] ``` """ function consgrouped_ptrs end @@ -515,8 +734,8 @@ end Compute a grouping of equal consecutive elements on `source` via [`consgrouped_ptrs`](@ref) and apply the grouping to target, resp. each -element of `target`. `target` may be an vector or a named or unnamed tuple of -vectors. The result is a `VectorOfVectors`, resp. a tuple of such. +element of `target`. `target` may be a vector or a named or unnamed tuple of +vectors. The result is a `PartsView`, resp. a tuple of such. Example: @@ -561,15 +780,20 @@ export consgroupedview function consgroupedview(source::AbstractVector, target::AbstractVector) elem_ptr = consgrouped_ptrs(source) - VectorOfVectors(target, elem_ptr) + PartsView(target, elem_ptr) end function consgroupedview(source::AbstractVector, target::NTuple{N,AbstractVector}) where {N} elem_ptr = consgrouped_ptrs(source) - map(X -> VectorOfVectors(X, elem_ptr), target) + map(X -> PartsView(X, elem_ptr), target) end function consgroupedview(source::AbstractVector, target::NamedTuple{syms,<:NTuple{N,AbstractVector}}) where {syms,N} elem_ptr = consgrouped_ptrs(source) - map(X -> VectorOfVectors(X, elem_ptr), target) + map(X -> PartsView(X, elem_ptr), target) end + + +# Deprecated: + +Base.@deprecate_binding VectorOfVectors PartsView diff --git a/test/Project.toml b/test/Project.toml index 225e375..14704e4 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,9 +1,16 @@ [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4" ChainRulesTestUtils = "cdddcdb0-9152-4a09-a978-84456f9df70a" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" ElasticArrays = "fdbdab4c-e67f-52f5-8c3f-e7b388dad3d4" +FixedSizeArrays = "3821ddf9-e5b5-40d5-8e25-6813ab96b5e2" +GPUArraysCore = "46192b85-c4d5-4398-a991-12ede77f4527" +JLArrays = "27aeb0d3-9eb9-45fb-866b-73c2ecf80fcb" +KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" +Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91" @@ -11,3 +18,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [compat] Documenter = "1" +Mooncake = "0.5" diff --git a/test/array_of_similar_arrays.jl b/test/array_of_similar_arrays.jl index 4c2ca84..96ba40d 100644 --- a/test/array_of_similar_arrays.jl +++ b/test/array_of_similar_arrays.jl @@ -10,6 +10,22 @@ using ChainRulesTestUtils using Statistics using StatsBase: cov2cor +include("testdefs.jl") + +# Minimal third-party subtype of AbstractArrayOfSimilarArrays, only +# implements the standard array API and `fused`: +if !isdefined(Main, :TestSimilarVectors) + struct TestSimilarVectors{T,ET} <: AbstractArrayOfSimilarArrays{T,1,1,ET} + data::Matrix{T} + + TestSimilarVectors(data::Matrix{T}) where {T} = + new{T,typeof(view(data, :, firstindex(data, 2)))}(data) + end + Base.size(A::TestSimilarVectors) = (size(A.data, 2),) + Base.getindex(A::TestSimilarVectors, i::Int) = view(A.data, :, i) + ArraysOfArrays.fused(A::TestSimilarVectors) = A.data +end + @testset "array_of_similar_arrays" begin function rand_flat_array(Val_N::Val{N}) where {N} sz_max = (2,3,2,4,5) @@ -38,8 +54,22 @@ using StatsBase: cov2cor } @testset "$TT from Array{Float64,$L}" begin A = rand_flat_array(Val_L) - @test typeof(@inferred TT(A)) == RT - @test typeof(@inferred convert(TT, A)) == RT + @test typeof(@inferred TT(A)) <: RT + + AosA = TT(A) + @test typeof(AosA) == typeof(TT(A)) + @test @inferred(eltype(AosA)) == typeof(AosA[1]) + @test @inferred(parent(AosA)) === @inferred(flatview(AosA)) + if eltype(eltype(AosA)) == eltype(A) + @test @inferred(flatview(AosA)) === A + else + @test @inferred(flatview(AosA)) ≈ A + end + @test @inferred(stacked(AosA)) === flatview(AosA) + @test @inferred(stack(AosA)) == flatview(AosA) + @test stack(AosA) !== flatview(AosA) + + @test_deprecated convert(TT, A) == TT(A) end end @@ -52,38 +82,39 @@ using StatsBase: cov2cor @testset "$TT from Array{Array{Float64,$M},$N}" begin A = rand_nested_similar_arrays(Val_M, Val_N) - A2_ctor = @inferred TT(A) - U = eltype(flatview(A2_ctor)) + A2_conv = @inferred convert(TT, A) + U = eltype(flatview(A2_conv)) A_U = Array{Array{U,M},N}(A) - @test typeof(A2_ctor) == RT - @test A2_ctor == A_U + @test typeof(A2_conv) <: RT + @test A2_conv == A_U A2_conv = @inferred convert(TT, A) - @test typeof(A2_conv) == RT - @test A2_conv == A2_ctor + @test typeof(A2_conv) <: RT - U = eltype(flatview(A2_ctor)) - A3 = @inferred Array(A2_ctor) + U = eltype(flatview(A2_conv)) + A3 = @inferred Array(A2_conv) @test typeof(A3) == Array{Array{U,M},N} @test A3 == A_U + + @test_deprecated TT(A) == A2_conv end end @testset "construct/convert from flat array" begin - test_from_flat(ArrayOfSimilarArrays{Float64,2,3}, ArrayOfSimilarArrays{Float64,2,3,5,Array{Float64,5}}, Val(5)) - test_from_flat(ArrayOfSimilarArrays{Float64,2}, ArrayOfSimilarArrays{Float64,2,3,5,Array{Float64,5}}, Val(5)) - test_from_flat(ArrayOfSimilarVectors{Float64}, ArrayOfSimilarVectors{Float64,2,3,Array{Float64,3}}, Val(3)) - test_from_flat(VectorOfSimilarArrays{Float64}, VectorOfSimilarArrays{Float64,2,3,Array{Float64,3}}, Val(3)) + test_from_flat(ArrayOfSimilarArrays{Float64,2,3}, ArrayOfSimilarArrays{Float64,2,3,Array{Float64,5}}, Val(5)) + test_from_flat(ArrayOfSimilarArrays{Float64,2}, ArrayOfSimilarArrays{Float64,2,3,Array{Float64,5}}, Val(5)) + test_from_flat(ArrayOfSimilarVectors{Float64}, ArrayOfSimilarVectors{Float64,2,Array{Float64,3}}, Val(3)) + test_from_flat(VectorOfSimilarArrays{Float64}, VectorOfSimilarArrays{Float64,2,Array{Float64,3}}, Val(3)) test_from_flat(VectorOfSimilarVectors{Float64}, VectorOfSimilarVectors{Float64,Array{Float64,2}}, Val(2)) test_from_flat(VectorOfSimilarVectors, VectorOfSimilarVectors{Float64,Array{Float64,2}}, Val(2)) - test_from_flat(ArrayOfSimilarArrays{Float32,2,3}, ArrayOfSimilarArrays{Float32,2,3,5,Array{Float32,5}}, Val(5)) - test_from_flat(ArrayOfSimilarArrays{Float32,2}, ArrayOfSimilarArrays{Float32,2,3,5,Array{Float32,5}}, Val(5)) - test_from_flat(ArrayOfSimilarVectors{Float32}, ArrayOfSimilarVectors{Float32,2,3,Array{Float32,3}}, Val(3)) - test_from_flat(VectorOfSimilarArrays{Float32}, VectorOfSimilarArrays{Float32,2,3,Array{Float32,3}}, Val(3)) + test_from_flat(ArrayOfSimilarArrays{Float32,2,3}, ArrayOfSimilarArrays{Float32,2,3,Array{Float32,5}}, Val(5)) + test_from_flat(ArrayOfSimilarArrays{Float32,2}, ArrayOfSimilarArrays{Float32,2,3,Array{Float32,5}}, Val(5)) + test_from_flat(ArrayOfSimilarVectors{Float32}, ArrayOfSimilarVectors{Float32,2,Array{Float32,3}}, Val(3)) + test_from_flat(VectorOfSimilarArrays{Float32}, VectorOfSimilarArrays{Float32,2,Array{Float32,3}}, Val(3)) test_from_flat(VectorOfSimilarVectors{Float32}, VectorOfSimilarVectors{Float32,Array{Float32,2}}, Val(2)) test_from_flat(VectorOfSimilarVectors{Float32}, VectorOfSimilarVectors{Float32,Array{Float32,2}}, Val(2)) @@ -92,24 +123,24 @@ using StatsBase: cov2cor @testset "construct/convert from nested arrays" begin - test_from_nested(ArrayOfSimilarArrays, ArrayOfSimilarArrays{Float64,2,3,5,Array{Float64,5}}, Val(2), Val(3)) - test_from_nested(ArrayOfSimilarArrays{Float64,2,3}, ArrayOfSimilarArrays{Float64,2,3,5,Array{Float64,5}}, Val(2), Val(3)) - test_from_nested(ArrayOfSimilarArrays{Float64}, ArrayOfSimilarArrays{Float64,2,3,5,Array{Float64,5}}, Val(2), Val(3)) - test_from_nested(ArrayOfSimilarArrays{Float32}, ArrayOfSimilarArrays{Float32,2,3,5,Array{Float32,5}}, Val(2), Val(3)) + test_from_nested(ArrayOfSimilarArrays, ArrayOfSimilarArrays{Float64,2,3,Array{Float64,5}}, Val(2), Val(3)) + test_from_nested(ArrayOfSimilarArrays{Float64,2,3}, ArrayOfSimilarArrays{Float64,2,3,Array{Float64,5}}, Val(2), Val(3)) + test_from_nested(ArrayOfSimilarArrays{Float64}, ArrayOfSimilarArrays{Float64,2,3,Array{Float64,5}}, Val(2), Val(3)) + test_from_nested(ArrayOfSimilarArrays{Float32}, ArrayOfSimilarArrays{Float32,2,3,Array{Float32,5}}, Val(2), Val(3)) - test_from_nested(ArrayOfSimilarArrays, VectorOfSimilarArrays{Float64,4,5,Array{Float64,5}}, Val(4), Val(1)) - test_from_nested(ArrayOfSimilarArrays, ArrayOfSimilarVectors{Float64,4,5,Array{Float64,5}}, Val(1), Val(4)) + test_from_nested(ArrayOfSimilarArrays, VectorOfSimilarArrays{Float64,4,Array{Float64,5}}, Val(4), Val(1)) + test_from_nested(ArrayOfSimilarArrays, ArrayOfSimilarVectors{Float64,4,Array{Float64,5}}, Val(1), Val(4)) test_from_nested(ArrayOfSimilarArrays, VectorOfSimilarVectors{Float64,Array{Float64,2}}, Val(1), Val(1)) - test_from_nested(VectorOfSimilarArrays{Float64,4}, VectorOfSimilarArrays{Float64,4,5,Array{Float64,5}}, Val(4), Val(1)) - test_from_nested(VectorOfSimilarArrays{Float64}, VectorOfSimilarArrays{Float64,4,5,Array{Float64,5}}, Val(4), Val(1)) - test_from_nested(VectorOfSimilarArrays{Float32}, VectorOfSimilarArrays{Float32,4,5,Array{Float32,5}}, Val(4), Val(1)) - test_from_nested(VectorOfSimilarArrays, VectorOfSimilarArrays{Float64,4,5,Array{Float64,5}}, Val(4), Val(1)) + test_from_nested(VectorOfSimilarArrays{Float64,4}, VectorOfSimilarArrays{Float64,4,Array{Float64,5}}, Val(4), Val(1)) + test_from_nested(VectorOfSimilarArrays{Float64}, VectorOfSimilarArrays{Float64,4,Array{Float64,5}}, Val(4), Val(1)) + test_from_nested(VectorOfSimilarArrays{Float32}, VectorOfSimilarArrays{Float32,4,Array{Float32,5}}, Val(4), Val(1)) + test_from_nested(VectorOfSimilarArrays, VectorOfSimilarArrays{Float64,4,Array{Float64,5}}, Val(4), Val(1)) - test_from_nested(ArrayOfSimilarVectors{Float64,4}, ArrayOfSimilarVectors{Float64,4,5,Array{Float64,5}}, Val(1), Val(4)) - test_from_nested(ArrayOfSimilarVectors{Float64}, ArrayOfSimilarVectors{Float64,4,5,Array{Float64,5}}, Val(1), Val(4)) - test_from_nested(ArrayOfSimilarVectors{Float32}, ArrayOfSimilarVectors{Float32,4,5,Array{Float32,5}}, Val(1), Val(4)) - test_from_nested(ArrayOfSimilarVectors, ArrayOfSimilarVectors{Float64,4,5,Array{Float64,5}}, Val(1), Val(4)) + test_from_nested(ArrayOfSimilarVectors{Float64,4}, ArrayOfSimilarVectors{Float64,4,Array{Float64,5}}, Val(1), Val(4)) + test_from_nested(ArrayOfSimilarVectors{Float64}, ArrayOfSimilarVectors{Float64,4,Array{Float64,5}}, Val(1), Val(4)) + test_from_nested(ArrayOfSimilarVectors{Float32}, ArrayOfSimilarVectors{Float32,4,Array{Float32,5}}, Val(1), Val(4)) + test_from_nested(ArrayOfSimilarVectors, ArrayOfSimilarVectors{Float64,4,Array{Float64,5}}, Val(1), Val(4)) test_from_nested(VectorOfSimilarVectors{Float64}, VectorOfSimilarVectors{Float64,Array{Float64,2}}, Val(1), Val(1)) test_from_nested(VectorOfSimilarVectors{Float32}, VectorOfSimilarVectors{Float32,Array{Float32,2}}, Val(1), Val(1)) @@ -121,6 +152,79 @@ using StatsBase: cov2cor test_rrule(ArrayOfSimilarArrays{Float64,2,2}, [rand(2,3) for i in 1:5, j in 1:6]) end + @testset "AbstractSlices interface" begin + @test @inferred(ArrayOfSimilarArrays{Float32,2,3}(rand(3,4,5,6,7))) isa AbstractSlices + let A = ArrayOfSimilarArrays{Float32,2,3}(rand(3,4,5,6,7)) + @test @inferred(A[2, 3, 4]) isa AbstractArray{Float32,2} + ref_ET = typeof(A[2, 3, 4]) + @test A isa AbstractSlices{ref_ET, 3} + @test @inferred(eltype(A)) == ref_ET + @test @inferred(innersize(A)) == (3, 4) + @test @inferred(getslicemap(A)) == (:, :, 1, 2, 3) + @test @inferred(parent(A)) === A.data + end + end + + @testset "split mode API" begin + A_flat = rand_flat_array(Val(5)) + A = ArrayOfSimilarArrays{Float64,2,3}(A_flat) + test_api(A, Array(A), A_flat) + + V_flat = rand_flat_array(Val(2)) + V = VectorOfSimilarVectors(V_flat) + test_api(V, Array(V), V_flat) + + # mapat operates on the flat data and preserves structure: + @test @inferred(mapat(abs2, Val(2), A)) == innermap(abs2, A) + @test typeof(mapat(abs2, Val(2), A)) == typeof(A) + @test fused(@inferred(mapat(+, Val(2), A, A))) == 2 .* A_flat + @test_throws DimensionMismatch mapat(+, Val(2), A, ArrayOfSimilarArrays{Float64,3,2}(A_flat)) + + @test @inferred(innersizes(A)) == fill(innersize(A), size(A)) + @test @inferred(innerlengths(A)) == fill(prod(innersize(A)), size(A)) + + # Per-element reductions over the inner dimensions of the flat data: + @test @inferred(innersum(A)) == [sum(x) for x in A] + @test @inferred(innermapreduce(abs2, +, A)) ≈ [sum(abs2, x) for x in A] + @test @inferred(innerreduce(max, A)) == [maximum(x) for x in A] + + # bcastat with outer-aligned, scalar and flat-matching arguments: + w = rand(size(A)...) + r_bc = @inferred bcastat(+, Val(2), A, w) + @test r_bc isa ArrayOfSimilarArrays + @test collect(r_bc) == [A[i] .+ w[i] for i in eachindex(A, w)] + @test fused(bcastat(muladd, Val(2), A, 2.0, A_flat)) == muladd.(A_flat, 2.0, A_flat) + + # vecflattened rrule: + A_rr = ArrayOfSimilarArrays{Float64,1,1}(rand(3, 4)) + y, pb = rrule(vecflattened, A_rr) + @test y == vec(A_rr.data) + t = pb(collect(1.0:12.0)) + @test t[1] == NoTangent() + @test t[2] isa ArrayOfSimilarArrays + @test fused(t[2]) == reshape(1.0:12.0, 3, 4) + end + + @testset "custom subtypes" begin + data = rand(3, 5) + B = TestSimilarVectors(data) + + @test @inferred(getsplitmode(B)) === SplitSlices{1,1}() + @test @inferred(unstackmode(B)) === SplitSlices{1,1}() + @test @inferred(fused(B)) === data + @test @inferred(stacked(B)) === data + @test @inferred(flatview(B)) === data + @test @inferred(parent(B)) === data + @test @inferred(vecflattened(B)) == vec(data) + @test @inferred(innersize(B)) == (3,) + @test @inferred(getslicemap(B)) == (:, 1) + @test @inferred(stack(B)) == data + @test stack(B) !== data + @test B == VectorOfSimilarVectors(data) + @test isapprox(B, VectorOfSimilarVectors(data .+ 1e-14)) + @test splitup(fused(B), getsplitmode(B)) == B + end + @testset "add remove" begin EA_ref1 = rand_flat_array(Val(3)) EA_ref2 = rand_flat_array(Val(3)) @@ -161,13 +265,21 @@ using StatsBase: cov2cor @test @inferred(size(A)) == @inferred(size(A_similar)) @test @inferred(size(A.data)) == @inferred(size(A_similar.data)) @test typeof(A_similar.data) == typeof(A_data) + + # similar must respect the requested element type and outer dims: + A_similar32 = @inferred similar(A, Array{Float32, 2}, size(A)) + @test eltype(eltype(A_similar32)) == Float32 + @test eltype(A_similar32.data) == Float32 + A_similar_1d = @inferred similar(A, Array{Float64, 2}, (3,)) + @test size(A_similar_1d) == (3,) + @test innersize(A_similar_1d) == innersize(A) @test typeof(A_similar) == typeof(A) end @testset "adapt" begin A_flat = rand(2,3,4,5,6) - A_nested = nestedview(A_flat, 2) + A_nested = sliced(A_flat, 2) @test @inferred(adapt(identity, A_nested)) == A_nested @test typeof(adapt(identity, A_nested)) == typeof(A_nested) end @@ -182,64 +294,10 @@ using StatsBase: cov2cor @testset "flatview" begin A = rand_nested_similar_arrays(Val(3), Val(2)) - B = ArrayOfSimilarArrays(A) - @inferred(flatview(B))[:] == collect(flatview(A)) - end - - - @testset "deepgetindex" begin - A = rand_nested_similar_arrays(Val(3), Val(2)) - B = ArrayOfSimilarArrays(A) - - @test deepgetindex(A, 3, 4, 2, 1, 2) == @inferred deepgetindex(B, 3, 4, 2, 1, 2) - @test deepgetindex(A, 2:3, 4, 2, 1, 2) == @inferred deepgetindex(B, 2:3, 4, 2, 1, 2) - @test deepgetindex(A, 2:3, 2:4, 2, 1, 2) == @inferred deepgetindex(B, 2:3, 2:4, 2, 1, 2) - @test deepgetindex(A, 2, 4, :, 1, 2) == @inferred deepgetindex(B, 2, 4, :, 1, 2) - @test deepgetindex(A, 2, 4, :, 1, 1:2) == @inferred deepgetindex(B, 2, 4, :, 1, 1:2) - @test deepgetindex(A, 2:3, 4, :, 1, 2) == @inferred deepgetindex(B, 2:3, 4, :, 1, 2) - @test deepgetindex(A, 2:3, 4, :, 1, 1:2) == @inferred deepgetindex(B, 2:3, 4, :, 1, 1:2) - end - - - @testset "deepsetindex!" begin - function testdata() - A = rand_nested_similar_arrays(Val(3), Val(2)) - B = ArrayOfSimilarArrays(A) - A, B - end - - A, B = testdata() - @test deepsetindex!(A, 42, 3, 4, 2, 1, 2) == @inferred deepsetindex!(B, 42, 3, 4, 2, 1, 2) - @test deepgetindex(B, 3, 4, 2, 1, 2) == 42 - - A, B = testdata() - X1 = rand(2) - @test deepsetindex!(A, X1, 2:3, 4, 2, 1, 2) == @inferred deepsetindex!(B, X1, 2:3, 4, 2, 1, 2) - @test deepgetindex(B, 2:3, 4, 2, 1, 2) == X1 - - A, B = testdata() - X2 = rand(2,2) - @test deepsetindex!(A, X2, 2, 4, :, 1, 1:2) == @inferred deepsetindex!(B, X2, 2, 4, :, 1, 1:2) - @test deepgetindex(B, 2, 4, :, 1, 1:2) == X2 - - A, B = testdata() - X3 = [rand(2,2), rand(2,2)] - @test deepsetindex!(A, X3, 2:3, 4, :, 1, 1:2) == @inferred deepsetindex!(B, X3, 2:3, 4, :, 1, 1:2) - @test deepgetindex(B, 2:3, 4, :, 1, 1:2) == X3 - end - - - @testset "deepview" begin - A = rand_nested_similar_arrays(Val(3), Val(2)) - B = ArrayOfSimilarArrays(A) - - @test deepview(A, 3, 4, 2, 1, 2) == @inferred deepview(B, 3, 4, 2, 1, 2) - @test deepgetindex(A, 2:3, 4, 2, 1, 2) == @inferred deepview(B, 2:3, 4, 2, 1, 2) - @test deepgetindex(A, 2:3, 2:4, 2, 1, 2) == @inferred deepview(B, 2:3, 2:4, 2, 1, 2) - @test deepview(A, 2, 4, :, 1, 2) == @inferred deepview(B, 2, 4, :, 1, 2) - @test deepview(A, 2, 4, :, 1, 1:2) == @inferred deepview(B, 2, 4, :, 1, 1:2) - @test deepview(A, 2:3, 4, :, 1, 2) == @inferred deepview(B, 2:3, 4, :, 1, 2) - @test deepview(A, 2:3, 4, :, 1, 1:2) == @inferred deepview(B, 2:3, 4, :, 1, 1:2) + B = convert(ArrayOfSimilarArrays, A) + @test @inferred(flatview(B)) === B.data + @test flatview(B) == stacked(A) + @test_throws ArgumentError flatview(A) end @@ -264,6 +322,12 @@ using StatsBase: cov2cor VA = [rand(2,3,3) for i in 1:10] VA_aosa = ArrayOfSimilarArrays(VA) + # Non-Colon dims must forward to the generic implementations: + @test sum(VV_aosa; dims = 1) == sum(collect(VV_aosa); dims = 1) + @test mean(VV_aosa; dims = 1) == mean(collect(VV_aosa); dims = 1) + @test @inferred(sum(VV_aosa)) ≈ sum(collect(VV_aosa)) + @test @inferred(mean(VV_aosa)) ≈ mean(collect(VV_aosa)) + array_cmp(A, B) = (A ≈ B) && (size(A) == size(B)) function test_statistics_op(op::Function) @@ -296,13 +360,13 @@ using StatsBase: cov2cor @testset "examples" begin A_flat = rand(2,3,4,5,6) - A_nested = nestedview(A_flat, 2) + A_nested = sliced(A_flat, 2) @test A_nested isa AbstractArray{<:AbstractArray{T,2},3} where T @test flatview(A_nested) === A_flat A_flat = rand(4,4) - A_nested = @inferred(nestedview(A_flat)) + A_nested = @inferred(sliced(A_flat)) @test A_nested.data == A_flat @test @inferred(size(A_nested))[1] == @inferred(size(A_flat))[1] @@ -315,7 +379,7 @@ using StatsBase: cov2cor @test ASA.data == A_flat # ------------------------------------------------------------------- - A_nested = nestedview(ElasticArray{Float64}(undef, 2, 3, 0), 2) + A_nested = sliced(ElasticArray{Float64}(undef, 2, 3, 0), 2) A_nested_copy = deepcopy(A_nested) for i in 1:4 @@ -362,12 +426,12 @@ using StatsBase: cov2cor @test @inferred(deepmap(f, ASA)).data == ASA.data.*2 @test @inferred(innermap(f, ASA)).data == ASA.data.*2 - @test @inferred(ArraysOfArrays._innerlength(VSV)) == N + @test @inferred(prod(ArraysOfArrays.innersize(VSV))) == N end @testset "map and broadcast" begin A_flat = rand(2,3,4,5,6) - A = nestedview(A_flat, 2) + A = sliced(A_flat, 2) for do_map in (map, broadcast) @test @inferred(do_map(identity, A)) === A diff --git a/test/base_slices.jl b/test/base_slices.jl new file mode 100644 index 0000000..ac8a54b --- /dev/null +++ b/test/base_slices.jl @@ -0,0 +1,46 @@ +# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). + +using ArraysOfArrays +using Test + +using ArraysOfArrays: getinnerdims, getouterdims + +include("testdefs.jl") + +@testset "base_slices" begin + A_orig = rand(5,6,7,8,9) + A_orig_mat = rand(5,6) + + Aes1 = eachslice(A_orig; dims = (4,5)) + @test @inferred(getsplitmode(Aes1)) isa BaseSlicing{3,2,Tuple{Colon,Colon,Colon,Int,Int}} + @test @inferred(getinnerdims((1,2,3,4,5), getsplitmode(Aes1))) == (1,2,3) + @test @inferred(getouterdims((1,2,3,4,5), getsplitmode(Aes1))) == (4,5) + + Aes2 = eachslice(A_orig; dims = (3,1,5)) + @test @inferred(getsplitmode(Aes2)) isa BaseSlicing{2,3,Tuple{Int,Colon,Int,Colon,Int}} + @test @inferred(getinnerdims((1,2,3,4,5), getsplitmode(Aes2))) == (2,4) + @test @inferred(getouterdims((1,2,3,4,5), getsplitmode(Aes2))) == (3,1,5) + + # Non-involutive outer dimension order (a 3-cycle), regression test for + # getouterdims applying the inverse of the slicemap permutation: + Aes3 = eachslice(A_orig; dims = (2,3,1)) + @test @inferred(getsplitmode(Aes3)) isa BaseSlicing{2,3,Tuple{Int,Int,Int,Colon,Colon}} + @test @inferred(getinnerdims((1,2,3,4,5), getsplitmode(Aes3))) == (4,5) + @test @inferred(getouterdims((1,2,3,4,5), getsplitmode(Aes3))) == (2,3,1) + + Aec = eachcol(A_orig_mat) + @test @inferred(getsplitmode(Aec)) isa BaseSlicing{1,1,Tuple{Colon,Int}} + @test @inferred(getinnerdims((1,2), getsplitmode(Aec))) == (1,) + @test @inferred(getouterdims((1,2), getsplitmode(Aec))) == (2,) + + Aer = eachrow(A_orig_mat) + @test @inferred(getsplitmode(Aer)) isa BaseSlicing{1,1,Tuple{Int,Colon}} + @test @inferred(getinnerdims((1,2), getsplitmode(Aer))) == (2,) + @test @inferred(getouterdims((1,2), getsplitmode(Aer))) == (1,) + + test_api(Aes1, Array(Aes1), A_orig) + test_api(Aes2, Array(Aes2), A_orig) + test_api(Aes3, Array(Aes3), A_orig) + test_api(Aec, Array(Aec), A_orig_mat) + test_api(Aer, Array(Aer), A_orig_mat) +end diff --git a/test/broadcasting.jl b/test/broadcasting.jl index a2823b3..19cce78 100644 --- a/test/broadcasting.jl +++ b/test/broadcasting.jl @@ -18,7 +18,7 @@ using Test for Idxs in [ ([rand(eachindex(a), rand(1:length(a))) for a in A],), - (VectorOfVectors([rand(eachindex(a), rand(1:length(a))) for a in A]),), + (PartsView([rand(eachindex(a), rand(1:length(a))) for a in A]),), (tuple(1:1),), (tuple([1, 1, 1]),), (tuple(:),), (Ref(1:1),), (Ref([1, 1, 1]),), (Ref(:),), ] @@ -32,7 +32,7 @@ using Test for Idxs in [ ([rand(eachindex(a), rand(1:length(a))) for a in A],), - (VectorOfVectors([rand(eachindex(a), rand(1:length(a))) for a in A]),), + (PartsView([rand(eachindex(a), rand(1:length(a))) for a in A]),), ] @test @inferred(broadcast(getindex, A, Idxs...)) isa VectorOfArrays{eltype(eltype(A))} @test getindex.(A, Idxs...) == getindex.(refA, Idxs...) diff --git a/test/fixed_size_arrays.jl b/test/fixed_size_arrays.jl new file mode 100644 index 0000000..41d611b --- /dev/null +++ b/test/fixed_size_arrays.jl @@ -0,0 +1,35 @@ +# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). + +using ArraysOfArrays +using Test + +using FixedSizeArrays: FixedSizeVector + +include("testdefs.jl") + +@testset "FixedSizeArrays extension" begin + A = VectorOfArrays([[1, 2], [3, 4, 5]]) + + # Fixed-size shape information cannot be resized, so getsplitmode can + # share it without a defensive copy: + B = VectorOfArrays(A.data, FixedSizeVector(A.elem_ptr), FixedSizeVector(A.kernel_size)) + sm = getsplitmode(B) + @test sm.elem_ptr === B.elem_ptr + @test sm.kernel_size === B.kernel_size + @test splitup(fused(B), sm) == B + @test typeof(splitup(fused(B), sm)) == typeof(B) + @test B == A + test_api(B, Array(B), B.data) + + # Resizable shape information still gets copied: + sm_A = getsplitmode(A) + @test sm_A.elem_ptr !== A.elem_ptr + @test sm_A.kernel_size !== A.kernel_size + + # partitioned with fixed-size lengths results in fixed-size shape + # information: + p = partitioned(collect(1:10), FixedSizeVector([2, 3, 5])) + @test p == [[1, 2], [3, 4, 5], [6, 7, 8, 9, 10]] + @test p.elem_ptr isa FixedSizeVector + @test getsplitmode(p).elem_ptr === p.elem_ptr +end diff --git a/test/functions.jl b/test/functions.jl index 6b9f2d9..85612c9 100644 --- a/test/functions.jl +++ b/test/functions.jl @@ -3,92 +3,111 @@ using ArraysOfArrays using Test -using StaticArrays +using ArraysOfArrays: getinnerdims, getouterdims +include("testdefs.jl") @testset "functions" begin - function gen_nested() - A11 = [1 2; 3 4] - A21 = [4 5 6; 7 8 9] - A12 = [10 11; 12 13; 14 15] - A22 = [16 17; 18 19] - hcat([A11, A21], [A12, A22]) - end - + A_0 = fill(Float32(42)) + A_0_flat = fill(Float32(42)) - @testset "flatview and nestedview" begin - A = [(@SArray randn(3, 2, 4)) for i in 1:2, j in 1:2] - @test @inferred(nestedview(flatview(A), Val(3))) == A + A_1 = Float32[3, 8, 4, 6] + A_1_flat = Float32[3, 8, 4, 6] - B = rand(3, 2, 4) - @test @inferred(nestedview(flatview(B), SVector{3})) == @inferred(nestedview(B, Val(1))) - end + A_1e = Float32[] + A_1e_flat = Float32[] + A_2 = [Float32[3 8; 4 6], Float32[1 2; 5 7; 9 0]] + A_2_flat = Float32[3, 8, 4, 6, 1, 2, 5, 7, 9, 0] - @testset "innersize" begin - @test @inferred(innersize([[1, 2, 3], [4, 5, 6]])) == (3,) - @test @inferred(innersize([[]])) == (0,) - @test @inferred(innersize([2:5])) == (4,) - @test @inferred(innersize((2:5,))) == (4,) - @test @inferred(innersize(Ref(2:5))) == (4,) - @test_throws DimensionMismatch @inferred(innersize([[1, 2, 3], [4, 5]])) - end + A_2e = Matrix{Float32}[] + A_2e_flat = Matrix{Float32}(undef, 0, 0) + A_2b = [Float32[3 8; 4 6], Float32[1 2; 9 0]] + A_2b_flat = stack(A_2b) - @testset "deepgetindex" begin - A = gen_nested() - @test @inferred(deepgetindex(A, 1, 2)) === A[1, 2] - @test @inferred(deepgetindex(A, 1, 2, 2, 1)) == A[1, 2][2, 1] - @test @inferred(deepgetindex(A, 1, 2, 1:2, 1)) == A[1, 2][1:2, 1] - @test @inferred(deepgetindex(A, 1, 1:2, 2, 1)) == [A[1, 1][2, 1], A[1, 2][2, 1]] - @test @inferred(deepgetindex(A, 1, 1:2, 1:2, 2)) == [A[1, 1][1:2, 2], A[1, 2][1:2, 2]] - @test_throws MethodError deepgetindex(A, 1, 2, 2, 1, 2) + A_3 = [[Float32[3, 8, 4], Float32[1, 2, 5]], [Float32[6, 7], Float32[9, 0]]] + A_3_flat = reduce(vcat, A_3) - B = rand(3, 4, 5) - @test @inferred(deepgetindex(B, 2, 3, 4)) === getindex(B, 2, 3, 4) - @test_throws MethodError deepgetindex(B, 6) - @test_throws MethodError deepgetindex(B, 2, 3) + @testset "getsplitmode" begin + @test @inferred(getsplitmode(A_0)) isa NonSplitMode{0} + @test @inferred(getsplitmode(A_1)) isa NonSplitMode{1} + @test @inferred(getsplitmode(A_1e)) isa NonSplitMode{1} + @test @inferred(getsplitmode(A_2)) isa UnknownSplitMode{typeof(A_2)} + @test @inferred(getsplitmode(A_2e)) isa UnknownSplitMode{typeof(A_2e)} + @test @inferred(getsplitmode(A_2b)) isa UnknownSplitMode{typeof(A_2b)} + @test @inferred(getsplitmode(A_3)) isa UnknownSplitMode{typeof(A_3)} end + @testset "innersize" begin + @test @inferred(innersize(A_0)) == () + @test @inferred(innersize(A_1)) == () + @test @inferred(innersize(A_1e)) == () + @test_throws DimensionMismatch innersize(A_2) + @test @inferred(innersize(A_2e)) == (0, 0) + @test @inferred(innersize(A_2b)) == (2, 2) + @test @inferred(innersize(A_3)) == (2,) + end - @testset "deepsetindex!" begin - B12 = [20 21 22; 23 24 25] - X = [[41, 42], [43, 44]] - - A = gen_nested() - @test @inferred(deepsetindex!(A, B12, 1, 2)) === A - @test A[1, 2] == B12 - - A = gen_nested() - @test @inferred(deepsetindex!(A, 42, 1, 2, 2, 1)) === A - @test A[1, 2][2, 1] == 42 + @testset "innermap" begin + f = x -> x^2 + @test @inferred(innermap(f, A_0)) == fill(1764) + @test @inferred(innermap(f, A_1)) == Float32[9, 64, 16, 36] + @test @inferred(innermap(f, A_2)) == [Float32[9 64; 16 36], Float32[1 4; 25 49; 81 0]] + @test_throws MethodError innermap(f, A_3) + + @test @inferred(innermap(length, A_0)) == fill(1) + @test @inferred(innermap(length, A_1)) == fill(1, 4) + @test @inferred(innermap(length, A_2)) == [fill(1, 2, 2), fill(1, 3, 2)] + @test @inferred(innermap(length, A_3)) == [fill(3, 2), fill(2, 2)] + end - A = gen_nested() - @test @inferred(deepsetindex!(A, X, 1, 1:2, 1:2, 2)) === A - @test deepgetindex(A, 1, 1:2, 1:2, 2) == X + @testset "deepmap" begin + f = x -> x^2 + @test @inferred(deepmap(f, A_0)) == fill(1764) + @test @inferred(deepmap(f, A_1)) == Float32[9, 64, 16, 36] + @test @inferred(deepmap(f, A_2)) == [Float32[9 64; 16 36], Float32[1 4; 25 49; 81 0]] + @test @inferred(deepmap(f, A_3)) == [[Float32[9, 64, 16], Float32[1, 4, 25]], [Float32[36, 49], Float32[81, 0]]] + + @test @inferred(deepmap(length, A_0)) == fill(1) + @test @inferred(deepmap(length, A_1)) == fill(1, 4) + @test @inferred(deepmap(length, A_2)) == [fill(1, 2, 2), fill(1, 3, 2)] + @test @inferred(deepmap(length, A_3)) == [[fill(1, 3), fill(1, 3)], [fill(1, 2), fill(1, 2)]] end + @testset "mapat" begin + f = x -> x^2 + @test @inferred(mapat(f, Val(1), A_1)) == map(f, A_1) + @test @inferred(mapat(f, Val(2), A_2)) == innermap(f, A_2) + @test @inferred(mapat(f, Val(3), A_3)) == deepmap(f, A_3) + @test mapat(length, Val(1), A_2) == length.(A_2) + + # Depth exceeding the nesting depth applies at the innermost level: + @test mapat(f, Val(5), A_1) == map(f, A_1) - @testset "deepview" begin - A = gen_nested() - @test @inferred(deepview(A, 1, 2)) == view(A, 1, 2) - @test @inferred(deepview(A, 1, 2, 2, 1)) == view(A[1, 2], 2, 1) - @test @inferred(deepview(A, 1, 2, 1:2, 1)) == view(A[1, 2], 1:2, 1) - @test_throws ArgumentError deepview(A, 1, 1:2, 2, 1) - @test @inferred(deepview(A, 1, 1:2, 1:2, 2)) == [A[1, 1][1:2, 2], A[1, 2][1:2, 2]] - @test_throws MethodError deepview(A, 1, 2, 2, 1, 2) + # Multiple inputs, zipped like map: + @test @inferred(mapat(+, Val(1), A_1, A_1)) == 2 .* A_1 + @test mapat(+, Val(2), A_2, A_2) == innermap(x -> 2 * x, A_2) - B = rand(3, 4, 5) - @test @inferred(deepview(B, 2, 3, 4)) === view(B, 2, 3, 4) - @test_throws MethodError deepview(B, 6) - @test_throws MethodError deepview(B, 2, 3) + # Integer depth relies on constant propagation for type stability: + mapat_intdepth(g, A) = mapat(g, 2, A) + @test @inferred(mapat_intdepth(f, A_2)) == innermap(f, A_2) end + @testset "innersizes and innerlengths" begin + @test @inferred(innersizes(A_2)) == size.(A_2) + @test @inferred(innerlengths(A_2)) == length.(A_2) + @test innersizes(A_3) == size.(A_3) + @test innerlengths(A_2e) == Int[] + end - @testset "abstract_nestedarray_type" begin - @test @inferred(abstract_nestedarray_type(Int, Val(()))) == Int - @test @inferred(abstract_nestedarray_type(Int, Val((2,)))) == AbstractArray{Int, 2} - @test @inferred(abstract_nestedarray_type(Float32, Val((2,3,4)))) == - AbstractArray{<:AbstractArray{<:AbstractArray{Float32, 4}, 3}, 2} + @testset "Nested array API" begin + test_api(A_0, A_0, A_0_flat) + test_api(A_1, A_1, A_1_flat) + test_api(A_1e, A_1e, A_1e_flat) + test_api(A_2, A_2, A_2_flat) + test_api(A_2e, A_2e, A_2e_flat) + test_api(A_2b, A_2b, A_2b_flat) + test_api(A_3, A_3, A_3_flat) end end diff --git a/test/gpu_arrays.jl b/test/gpu_arrays.jl new file mode 100644 index 0000000..8603ad4 --- /dev/null +++ b/test/gpu_arrays.jl @@ -0,0 +1,102 @@ +# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). + +using ArraysOfArrays +using Test + +using Adapt +using GPUArraysCore: AbstractGPUArray +using JLArrays +import KernelAbstractions as KA + +include("testdefs.jl") + +# JLArray is the reference AbstractGPUArray implementation, tests must not +# perform scalar indexing on it: +JLArrays.allowscalar(false) + +@testset "GPU arrays (via JLArrays)" begin + @testset "ArrayOfSimilarArrays on device data" begin + x = jl(rand(Float32, 3, 4, 5)) + A = @inferred ArrayOfSimilarArrays{Float32,1,2}(x) + + @test A[2, 3] isa AbstractGPUArray + @test @inferred(flatview(A)) === x + @test @inferred(fused(A)) === x + @test splitup(fused(A), getsplitmode(A)) == A + @test @inferred(innersize(A)) == (3,) + @test collect(fused(innermap(abs, A))) == abs.(collect(x)) + @test stack(A) isa AbstractGPUArray + A[1, 1] = jl(rand(Float32, 3)) + + V = VectorOfSimilarVectors(jl(rand(Float32, 3, 5))) + @test sum(V) isa AbstractGPUArray + @test ArraysOfArrays.Statistics.mean(V) isa AbstractGPUArray + @test vcat(V, V) isa VectorOfSimilarVectors + @test fused(vcat(V, V)) isa AbstractGPUArray + + @test KA.get_backend(A) == KA.get_backend(x) + end + + @testset "VectorOfArrays, device data and host shape info" begin + xv = jl(rand(Float32, 10)) + V = VectorOfArrays(xv, [1, 3, 6, 11], [(), (), ()]) + @test V[2] isa AbstractGPUArray + @test @inferred(vecflattened(V)) isa AbstractGPUArray + @test reduce(vcat, [V, V]) isa VectorOfArrays + @test partitioned(xv, [2, 3, 5]) isa VectorOfArrays + @test KA.get_backend(V) == KA.get_backend(xv) + end + + @testset "VectorOfArrays, everything on device" begin + xv = jl(rand(Float32, 10)) + + # Consistency checks must work without scalar indexing: + V = VectorOfArrays(xv, jl([1, 3, 6, 11]), jl([(), (), ()])) + @test V isa VectorOfArrays + @test length(V) == 3 + @test fused(V) === xv + + V2 = VectorOfArrays(jl(rand(Float32, 12)), jl([1, 5, 13]), jl([(2,), (2,)])) + @test length(V2) == 2 + + # Inconsistent shape info must still be detected: + @test_throws ArgumentError VectorOfArrays(xv, jl([1, 3, 12]), jl([(), ()])) + @test_throws ArgumentError VectorOfArrays(xv, jl([1, 6, 3]), jl([(), ()])) + @test_throws ArgumentError VectorOfArrays(xv, jl([1, 4, 11]), jl([(2,), (2,)])) + + # mapat and innerlengths only touch device data and shape info: + @test collect(fused(mapat(abs2, Val(2), V))) == abs2.(collect(xv)) + @test innerlengths(V) isa AbstractGPUArray + @test collect(innerlengths(V)) == [2, 3, 5] + @test collect(innersizes(V2)) == [(2, 2), (2, 4)] + + # bcastat compiles to flat device broadcasts: + v_outer = jl(Float32[10, 20, 30]) + r_bc = bcastat(+, Val(2), V, v_outer) + @test r_bc isa VectorOfArrays + @test fused(r_bc) isa AbstractGPUArray + @test collect(fused(r_bc)) == collect(xv) .+ [10, 10, 20, 20, 20, 30, 30, 30, 30, 30] + + # Segmented reductions run as a single device kernel: + xv_h = collect(xv) + parts_h = [xv_h[1:2], xv_h[3:5], xv_h[6:10]] + @test innersum(V) isa AbstractGPUArray + @test collect(innersum(V)) ≈ sum.(parts_h) + @test collect(innermapreduce(abs2, +, V)) ≈ [sum(abs2, p) for p in parts_h] + @test collect(innerreduce(max, V; init = -Inf32)) ≈ maximum.(parts_h) + @test_throws ArgumentError innerreduce(max, VectorOfArrays(xv, jl([1, 3, 3, 11]), jl([(), (), ()]))) + + # Split mode round trip on device shape info: + sm = getsplitmode(V) + @test splitup(fused(V), sm) isa VectorOfArrays + @test typeof(splitup(fused(V), sm)) == typeof(V) + + # adapt moves data and shape info: + V_host = VectorOfArrays([Float32[1, 2], Float32[3, 4, 5]]) + V_dev = adapt(JLArray, V_host) + @test V_dev.data isa AbstractGPUArray + @test V_dev.elem_ptr isa AbstractGPUArray + @test adapt(Array, V_dev) == V_host + @test adapt(Array, getsplitmode(V_dev)).elem_ptr == getsplitmode(V_host).elem_ptr + end +end diff --git a/test/mooncake.jl b/test/mooncake.jl new file mode 100644 index 0000000..0ba665e --- /dev/null +++ b/test/mooncake.jl @@ -0,0 +1,49 @@ +# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). + +using ArraysOfArrays +using Test + +import Mooncake +import Random + +@testset "Mooncake extension" begin + rng = Random.Xoshiro(0x2b6ee80184a75b47) + + x_flat = rand(rng, 3, 4) + x_vec = rand(rng, 10) + lengths = [2, 3, 5] + + @testset "zero-derivative primitives" begin + A = sliced(x_flat, Val(1)) + p = partitioned(x_vec, lengths) + + Mooncake.TestUtils.test_rule(rng, getsplitmode, A; is_primitive = true, unsafe_perturb = true) + Mooncake.TestUtils.test_rule(rng, getsplitmode, p; is_primitive = true, unsafe_perturb = true) + Mooncake.TestUtils.test_rule(rng, unstackmode, A; is_primitive = true, unsafe_perturb = true) + Mooncake.TestUtils.test_rule(rng, is_memordered_splitmode, getsplitmode(A); is_primitive = true, unsafe_perturb = true) + Mooncake.TestUtils.test_rule(rng, innersize, A; is_primitive = true, unsafe_perturb = true) + Mooncake.TestUtils.test_rule(rng, getslicemap, A; is_primitive = true, unsafe_perturb = true) + Mooncake.TestUtils.test_rule(rng, consgrouped_ptrs, [1, 1, 2, 3, 3]; is_primitive = true, unsafe_perturb = true) + end + + # Mooncake differentiates the ArraysOfArrays types and operations without + # custom rules, verify gradient correctness end-to-end: + @testset "derived rules" begin + for (label, f, x) in [ + ("sliced and fused", x -> sum(abs2, fused(sliced(x, Val(1)))), x_flat), + ("stacked", x -> sum(stacked(ArrayOfSimilarArrays{Float64,1,1}(x))), x_flat), + ("vecflattened of ArrayOfSimilarArrays", x -> sum(abs2, vecflattened(ArrayOfSimilarArrays{Float64,1,1}(x))), x_flat), + ("partitioned", x -> sum(sum, partitioned(x, lengths)), x_vec), + ("VectorOfArrays round trip", x -> begin + p = partitioned(x, lengths) + sum(abs2, fused(splitup(fused(p), getsplitmode(p)))) + end, x_vec), + ("reduce vcat", x -> sum(sum, reduce(vcat, [partitioned(x, [2, 3]), partitioned(x, [4])])), x_vec), + ("innermap", x -> sum(fused(innermap(abs2, sliced(x, Val(1))))), x_flat), + ] + @testset "$label" begin + Mooncake.TestUtils.test_rule(rng, f, x; is_primitive = false, unsafe_perturb = true) + end + end + end +end diff --git a/test/runtests.jl b/test/runtests.jl index fa08bb9..2ff6a88 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -5,8 +5,13 @@ import Test Test.@testset "Package ArraysOfArrays" begin include("test_aqua.jl") include("functions.jl") + include("base_slices.jl") include("array_of_similar_arrays.jl") include("vector_of_arrays.jl") include("broadcasting.jl") + include("static_arrays.jl") + include("fixed_size_arrays.jl") + include("gpu_arrays.jl") + include("mooncake.jl") include("test_docs.jl") end # testset diff --git a/test/static_arrays.jl b/test/static_arrays.jl new file mode 100644 index 0000000..77fd454 --- /dev/null +++ b/test/static_arrays.jl @@ -0,0 +1,18 @@ +# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). + +using ArraysOfArrays +using Test + +using StaticArrays + + +@testset "StaticArray Extension" begin + @testset "flatview and sliced" begin + A = [(@SArray randn(3, 2, 4)) for i in 1:2, j in 1:2] + @test @inferred(sliced(flatview(A), Val(3))) == A + + B = rand(3, 2, 4) + @test @inferred(sliced(flatview(B), SVector{3})) == @inferred(sliced(B, Val(1))) + @test_deprecated nestedview(flatview(B), SVector{3}) + end +end diff --git a/test/testdefs.jl b/test/testdefs.jl new file mode 100644 index 0000000..561f168 --- /dev/null +++ b/test/testdefs.jl @@ -0,0 +1,175 @@ +# This file is a part of ArraysOfArrays.jl, licensed under the MIT License (MIT). + +using ChainRulesCore: rrule, unthunk, AbstractThunk, Thunk, NoTangent + +using ArraysOfArrays: getinnerdims, getouterdims + +if !isdefined(Main, :test_api) + maptest_f(x::Number) = x^2 + maptest_f(x::AbstractArray{<:Number}) = sum(x)^2 + maptest_f(x::AbstractArray) = length(x)^2 + + function test_notangent_rrule(f, args::Vararg{Any,N}) where {N} + y, f_pullback = @inferred(rrule(f, args...)) + @test y == f(args...) + dy = NoTangent() + @test @inferred(f_pullback(dy)) == ntuple(_ -> NoTangent(), Val(N+1)) + end + + function test_fused_rrule(A) + @testset "Test fused rrule for $(nameof(typeof(A)))" begin + A_joined, f_pullback = @inferred(rrule(fused, A)) + @test A_joined == fused(A) + dy = Thunk(Returns(A_joined)) + tangents = @inferred(f_pullback(dy)) + @test tangents isa Tuple{NoTangent,<:AbstractThunk} + @test unthunk(tangents[2]) == A + end + end + + function test_splitview_rrule(A, smode) + @testset "Test splitup rrule for $(nameof(typeof(A)))" begin + A_split, f_pullback = @inferred(rrule(splitup, A, smode)) + @test A_split == splitup(A, smode) + dy = Thunk(Returns(A_split)) + tangents = @inferred(f_pullback(dy)) + @test tangents isa Tuple{NoTangent,<:AbstractThunk,NoTangent} + @test unthunk(tangents[2]) == A + end + end + + function test_api(A, A_array_ref, A_unsplit_ref) + @testset "Test API for $(nameof(typeof(A)))" begin + global g_state = (;A, A_array_ref, A_unsplit_ref) + + @test Array(A) isa Array{<:Any,ndims(A)} + A_array = Array(A) + @test A == A_array + @test isequal(A, A_array) + @test isequal(A_array, A_array_ref) + + @test @inferred(getsplitmode(A)) isa AbstractSplitMode + smode = getsplitmode(A) + test_notangent_rrule(getsplitmode, A) + + @test @inferred(is_memordered_splitmode(smode)) isa Bool + test_notangent_rrule(is_memordered_splitmode, smode) + + if A isa AbstractSlices + let M = ndims(eltype(A)), N = ndims(A) + @test smode isa AbstractSlicingMode{M,N} + end + elseif eltype(A) <: Number + @test smode isa NonSplitMode{ndims(A)} + end + + @test @inferred(eltype(A)) <: Union{Number,AbstractArray} + T_elem = eltype(A) + + innersz = if smode isa NonSplitMode + @inferred(innersize(A)) + elseif smode isa AbstractSlicingMode + @inferred(innersize(A)) + else + let sz = (try innersize(A); catch err; err; end) + sz isa Exception ? sz : @inferred(innersize(A)) + end + end + + if !(innersz isa Exception) + test_notangent_rrule(innersize, A) + end + + @test innersz isa Union{Exception, Dims} + + if !isempty(A) + A_elem_1 = @inferred(A[first(eachindex(A))]) + @test typeof(A_elem_1) == T_elem + if !(innersz isa Exception) + @test all(size.(A) .== Ref(innersz)) + end + + @inferred(innermap(maptest_f, A)) == innermap(maptest_f, Array(A)) + @inferred(deepmap(maptest_f, A)) == deepmap(maptest_f, Array(A)) + end + + _smode_M(::AbstractSlicingMode{M,N}) where {M,N} = M + _smode_N(::AbstractSlicingMode{M,N}) where {M,N} = N + + dimstpl = ntuple(identity, Val(ndims(A_unsplit_ref))) + + if smode isa AbstractSlicingMode + M, N = _smode_M(smode), _smode_N(smode) + A_array_stacked = stack(A_array) + @test M == ndims(eltype(A)) + @test N == ndims(A) + + @test Array(stack(A)) == A_array_stacked + + @test @inferred(unstackmode(A)) isa AbstractSlicingMode{M,N} + umode = unstackmode(A) + @test @inferred(splitup(stacked(A), umode)) == A + + if is_memordered_splitmode(smode) + # stack(A) must return an independent array, even if the + # parent array could be returned as-is: + @test @inferred(stack(A)) == A_unsplit_ref + @test stack(A) !== A_unsplit_ref + @test @inferred(stacked(A)) === A_unsplit_ref + @test @inferred(flatview(A)) === A_unsplit_ref + @test stacked(splitup(stacked(A), umode)) === stacked(A) + else + @test Array(@inferred(stack(A))) == A_array_stacked + @test Array(@inferred(stacked(A))) == A_array_stacked + @test_throws ArgumentError flatview(A) + end + + @test @inferred(getinnerdims(dimstpl, smode)) isa NTuple{M,Int} + @test @inferred(getouterdims(dimstpl, smode)) isa NTuple{N,Int} + innerdims = getinnerdims(dimstpl, smode) + outerdims = getouterdims(dimstpl, smode) + @test permutedims(A_unsplit_ref, (innerdims..., outerdims...)) == A_array_stacked + elseif smode isa NonSplitMode + @test @inferred(stacked(A)) === A + else + stacked_A = try stack(A); catch err; err; end + if stacked_A isa Exception + @test_throws typeof(stacked_A) stacked(A) + else + @inferred(stacked(A)) == stacked_A + end + end + + if smode isa UnknownSplitMode + @test @inferred(is_memordered_splitmode(smode)) == false + @test_throws ArgumentError fused(A) + @test_throws ArgumentError flatview(A) + @test_throws ArgumentError splitup(A_unsplit_ref, smode) + @test_throws ArgumentError getinnerdims(dimstpl, smode) + @test_throws ArgumentError getouterdims(dimstpl, smode) + elseif smode isa NonSplitMode + @test @inferred(is_memordered_splitmode(smode)) == true + @test @inferred(fused(A)) === A + @test @inferred(flatview(A)) === A + @test @inferred(splitup(A, smode)) === A + @test @inferred(getinnerdims(dimstpl, smode)) == () + @test @inferred(getouterdims(dimstpl, smode)) == dimstpl + + test_fused_rrule(A) + test_splitview_rrule(A, smode) + else + if A isa Slices + @test fused(A) === parent(A) + end + @test @inferred(fused(A)) == A_unsplit_ref + A_unsplit = fused(A) + @test typeof(A_unsplit) == typeof(A_unsplit_ref) + @test typeof(splitup(A_unsplit, smode)) == typeof(A) + @test splitup(A_unsplit, smode) == A + + test_fused_rrule(A) + test_splitview_rrule(A_unsplit, smode) + end + end + end +end diff --git a/test/vector_of_arrays.jl b/test/vector_of_arrays.jl index 8ea4d00..6862289 100644 --- a/test/vector_of_arrays.jl +++ b/test/vector_of_arrays.jl @@ -8,6 +8,8 @@ using Adapt using ArraysOfArrays: full_consistency_checks, append_elemptr!, element_ptr +include("testdefs.jl") + @testset "vector_of_arrays" begin ref_flatview(A::AbstractVector{<:AbstractArray}) = vcat(map(vec, Array(A))...) @@ -32,10 +34,10 @@ using ArraysOfArrays: full_consistency_checks, append_elemptr!, element_ptr @test @inferred(VectorOfArrays{Float64,1}(deepcopy(A1))) isa VectorOfArrays{Float64,1,0,Array{Float64,1},Array{Int,1},Array{Tuple{},1}} @test VectorOfArrays{Float64,1}(deepcopy(A1)) == A1 - @test @inferred(VectorOfVectors(deepcopy(A1))) isa VectorOfArrays{Float32,1,0,Array{Float32,1},Array{Int,1},Array{Tuple{},1}} - @test VectorOfVectors(deepcopy(A1)) == A1 - @test @inferred(VectorOfVectors{Float64}(deepcopy(A1))) isa VectorOfArrays{Float64,1,0,Array{Float64,1},Array{Int,1},Array{Tuple{},1}} - @test VectorOfVectors{Float64}(deepcopy(A1)) == A1 + @test @inferred(PartsView(deepcopy(A1))) isa VectorOfArrays{Float32,1,0,Array{Float32,1},Array{Int,1},Array{Tuple{},1}} + @test PartsView(deepcopy(A1)) == A1 + @test @inferred(PartsView{Float64}(deepcopy(A1))) isa VectorOfArrays{Float64,1,0,Array{Float64,1},Array{Int,1},Array{Tuple{},1}} + @test PartsView{Float64}(deepcopy(A1)) == A1 A1_empty = ref_AoA1(Float32, 0) @test @inferred(VectorOfArrays(deepcopy(A1_empty))) isa VectorOfArrays{Float32,1,0,Array{Float32,1},Array{Int,1},Array{Tuple{},1}} @@ -45,10 +47,10 @@ using ArraysOfArrays: full_consistency_checks, append_elemptr!, element_ptr @test @inferred(VectorOfArrays{Float64,1}(deepcopy(A1_empty))) isa VectorOfArrays{Float64,1,0,Array{Float64,1},Array{Int,1},Array{Tuple{},1}} @test VectorOfArrays{Float64,1}(deepcopy(A1_empty)) == A1_empty - @test @inferred(VectorOfVectors(deepcopy(A1_empty))) isa VectorOfArrays{Float32,1,0,Array{Float32,1},Array{Int,1},Array{Tuple{},1}} - @test VectorOfVectors(deepcopy(A1_empty)) == A1_empty - @test @inferred(VectorOfVectors{Float64}(deepcopy(A1_empty))) isa VectorOfArrays{Float64,1,0,Array{Float64,1},Array{Int,1},Array{Tuple{},1}} - @test VectorOfVectors{Float64}(deepcopy(A1_empty)) == A1_empty + @test @inferred(PartsView(deepcopy(A1_empty))) isa VectorOfArrays{Float32,1,0,Array{Float32,1},Array{Int,1},Array{Tuple{},1}} + @test PartsView(deepcopy(A1_empty)) == A1_empty + @test @inferred(PartsView{Float64}(deepcopy(A1_empty))) isa VectorOfArrays{Float64,1,0,Array{Float64,1},Array{Int,1},Array{Tuple{},1}} + @test PartsView{Float64}(deepcopy(A1_empty)) == A1_empty A2 = ref_AoA2(Float32, 4) @test @inferred(VectorOfArrays(deepcopy(A2))) isa VectorOfArrays{Float32,2,1,Array{Float32,1},Array{Int,1},Array{Tuple{Int},1}} @@ -83,27 +85,323 @@ using ArraysOfArrays: full_consistency_checks, append_elemptr!, element_ptr B1 = VectorOfArrays(A1); B2 = VectorOfArrays(A2); B3 = VectorOfArrays(A3); B4 = VectorOfArrays(A4); - @test @inferred(vcat(B1)) === B1 + # vcat must return an independent array: + @test @inferred(vcat(B1)) == B1 + @test vcat(B1) !== B1 + @test vcat(B1).data !== B1.data full_consistency_checks(vcat(B1)) @test @inferred(vcat(B1, B2)) isa VectorOfArrays @test vcat(B1, B2) == vcat(A1, A2) - @test eltype(vcat(B1, B2)) == Array{Float32,3} + @test eltype(vcat(B1, B2)) <: AbstractArray{Float32,3} + @test eltype(vcat(B1, B2)) == eltype(B1) full_consistency_checks(vcat(B1, B2)) @test @inferred(vcat(B1, B3)) isa VectorOfArrays @test vcat(B1, B3) == vcat(A1, A3) - @test eltype(vcat(B1, B3)) == Array{Float32,3} + @test eltype(vcat(B1, B3)) <: AbstractArray{Float32,3} full_consistency_checks(vcat(B1, B3)) @test @inferred(vcat(B1, B2, B3, B4)) isa VectorOfArrays @test vcat(B1, B2, B3, B4) == vcat(A1, A2, A3, A4) - @test eltype(vcat(B1, B2, B3, B4)) == Array{Float64,3} + @test eltype(vcat(B1, B2, B3, B4)) <: AbstractArray{Float64,3} full_consistency_checks(vcat(B1, B2, B3, B4)) B1_copy = @inferred(copy(B1)); B3_copy = @inferred(copy(B3)) append!(B1_copy, B3_copy) @test B1_copy.data == vcat(B1.data, B3.data) + + # reduce(vcat, ...) uses a single-pass implementation: + @test @inferred(reduce(vcat, [B1, B2, B3])) isa VectorOfArrays + @test reduce(vcat, [B1, B2, B3]) == vcat(A1, A2, A3) + @test reduce(vcat, [B1, B2, B3]) == vcat(B1, B2, B3) + full_consistency_checks(reduce(vcat, [B1, B2, B3])) + + # vcat and append! of vectors of arrays with unused data regions + # must only use the data covered by their elements: + p_partial = partitioned(collect(1:10), [2, 3]) + q = VectorOfArrays([[6], [7, 8]]) + @test vcat(p_partial, q) == vcat(collect(p_partial), collect(q)) + @test reduce(vcat, [p_partial, q]) == vcat(collect(p_partial), collect(q)) + full_consistency_checks(vcat(p_partial, q)) + + q_copy = copy(q) + append!(q_copy, p_partial) + @test q_copy == vcat(collect(q), collect(p_partial)) + full_consistency_checks(q_copy) + + # Appending to a VectorOfArrays with unused trailing data would + # corrupt it: + @test_throws ArgumentError append!(partitioned(collect(1:10), [2, 3]), q) + end + + + @testset "vcat of VectorOfSimilarArrays" begin + VsoA = [VectorOfSimilarArrays(rand(Float32, 2, 3, n)) for n in (2, 0, 4)] + Vs_ref = vcat(map(collect, VsoA)...) + + @test @inferred(vcat(VsoA[1])) isa VectorOfSimilarArrays + @test vcat(VsoA[1]) == VsoA[1] + @test vcat(VsoA[1]).data !== VsoA[1].data + + @test @inferred(vcat(VsoA...)) isa VectorOfSimilarArrays + @test vcat(VsoA...) == Vs_ref + @test @inferred(reduce(vcat, VsoA)) isa VectorOfSimilarArrays + @test reduce(vcat, VsoA) == Vs_ref + + # Element type promotion like Base.vcat: + V64 = VectorOfSimilarArrays(rand(Float64, 2, 3, 2)) + @test eltype(eltype(vcat(VsoA[1], V64))) == Float64 + @test vcat(VsoA[1], V64) == vcat(collect(VsoA[1]), collect(V64)) + + # Inner size mismatch: + V_bad = VectorOfSimilarArrays(rand(Float32, 3, 2, 2)) + @test_throws DimensionMismatch vcat(VsoA[1], V_bad) + + # Vectors of similar vectors: + VsoV = [VectorOfSimilarVectors(rand(3, n)) for n in (2, 4)] + @test @inferred(reduce(vcat, VsoV)) isa VectorOfSimilarVectors + @test reduce(vcat, VsoV) == vcat(map(collect, VsoV)...) + end + + + @testset "equality" begin + # == and isequal must be equivalent to elementwise comparison, + # independent of the underlying data layout: + p = partitioned(collect(1:10), [2, 3]) # unused data after covered range + q = VectorOfArrays([[1, 2], [3, 4, 5]]) + @test collect(p) == collect(q) + @test p == q + @test q == p + @test isequal(p, q) + + @test p != VectorOfArrays([[1, 2], [3, 4, 6]]) + @test p != VectorOfArrays([[1, 2, 3], [4, 5]]) + @test p != VectorOfArrays([[1, 2], [3, 4, 5], [6]]) + + # Equal-length parts with different shapes are not equal: + r1 = VectorOfArrays([[1 2; 3 4]]) + r2 = VectorOfArrays([[1 3; 2 4]]) + r3 = partitioned(collect([1, 3, 2, 4]), [(2, 2)]) + @test r1 != r2 + @test r1 == permutedims.(r2) + @test r1 == r3 && r2 != r3 + + # missing propagates through ==, but not isequal: + m1 = VectorOfArrays([[1, missing], [3]]) + m2 = VectorOfArrays([[1, missing], [3]]) + @test ismissing(m1 == m2) + @test isequal(m1, m2) + end + + + @testset "split mode API" begin + B1 = VectorOfArrays(ref_AoA1(Float32, 5)) + B1e = VectorOfArrays(ref_AoA1(Float32, 0)) + B3 = VectorOfArrays(ref_AoA3(Float32, 3)) + Bu = VectorOfArrays([rand(Float32, 2, 3) for i in 1:4]) + + for B in [B1, B1e, B3, Bu] + @test @inferred(getsplitmode(B)) isa AbstractPartMode{ndims(eltype(B)),1} + @test @inferred(fused(B)) === B.data + @test @inferred(splitup(fused(B), getsplitmode(B))) == B + @test typeof(splitup(fused(B), getsplitmode(B))) == typeof(B) + test_api(B, Array(B), B.data) + end + + # mapat operates on the flat data and preserves structure: + @test @inferred(mapat(abs2, Val(2), B3)) == innermap(abs2, B3) + @test typeof(mapat(abs2, Val(2), B3)) == typeof(B3) + @test mapat(+, Val(2), B1, B1) == [x .+ x for x in B1] + @test_throws DimensionMismatch mapat(+, Val(2), B1, B3) + + @test @inferred(innerlengths(B1)) == length.(collect(B1)) + @test @inferred(innersizes(B3)) == size.(collect(B3)) + + # getsplitmode must not be affected by later resizing: + B_grow = VectorOfArrays([[1, 2], [3, 4, 5]]) + sm_grow = getsplitmode(B_grow) + push!(B_grow, [6]) + @test splitup(collect(1:5), sm_grow) == [[1, 2], [3, 4, 5]] + + # Uniform element size, so stackable: + @test @inferred(stacked(Bu)) == stack(Array(Bu)) + @test @inferred(splitup(stacked(Bu), unstackmode(Bu))) == Bu + + # flatview on view-backed VectorOfArrays returns a view of only the + # covered data range, in a type-stable fashion: + B3_view = view(B3, 2:3) + @test @inferred(flatview(B3_view)) isa SubArray + @test flatview(B3_view) == B3.data[B3.elem_ptr[2]:(B3.elem_ptr[4] - 1)] + @test @inferred(vecflattened(B3_view)) == flatview(B3_view) + @test @inferred(flatview(B3)) === B3.data + end + + + @testset "partitioned" begin + x = collect(1:10) + + p = @inferred(partitioned(x, [2, 3, 5])) + @test p isa PartsView{Int} + @test p == [[1, 2], [3, 4, 5], [6, 7, 8, 9, 10]] + @test @inferred(fused(p)) === x + @test flatview(p) === x + test_api(p, Array(p), x) + + # vecflattened may share memory, reduce/mapreduce must not: + @test @inferred(vecflattened(p)) == x + @test parent(vecflattened(p)) === x + @test @inferred(reduce(vcat, p)) == x + @test reduce(vcat, p) !== x + @test typeof(reduce(vcat, p)) == typeof(x) + @test @inferred(mapreduce(vec, vcat, p)) == x + @test mapreduce(vec, vcat, p) !== x + + # Partial partitions are allowed: + p_partial = @inferred(partitioned(x, [2, 3])) + @test p_partial == [[1, 2], [3, 4, 5]] + + @test_throws ArgumentError partitioned(x, [2, 3, 6]) + @test_throws ArgumentError partitioned(x, [2, -1, 5]) + + p2 = @inferred(partitioned(x, [(1, 2), (2, 2)])) + @test p2 isa VectorOfArrays{Int,2} + @test p2[1] == [1 2] + @test p2[2] == [3 5; 4 6] + @test @inferred(fused(p2)) === x + end + + + @testset "bcastat" begin + x = collect(Float32, 1:10) + p = partitioned(x, [2, 3, 5]) + v = Float32[10, 20, 30] + + # One value per element, broadcast over the element contents: + r = @inferred bcastat(+, Val(2), p, v) + @test r isa PartsView + @test r == [xi .+ vi for (xi, vi) in zip(p, v)] + + # Scalars broadcast over everything: + @test bcastat(+, Val(2), p, 1) == [xi .+ 1 for xi in p] + + # Aligned nested arguments and flat-matching arguments: + @test bcastat(+, Val(2), p, p) == [xi .+ xi for xi in p] + @test bcastat(+, Val(2), p, x) == [xi .+ xi for xi in p] + + # Mixed argument kinds in one call: + @test bcastat(muladd, Val(2), p, v, 2) == [muladd.(xi, vi, 2) for (xi, vi) in zip(p, v)] + + # Data not covered by any part gets no contribution: + p_part = partitioned(x, [2, 3]) + @test bcastat(+, Val(2), p_part, Float32[10, 20]) == [Float32[11, 12], Float32[23, 24, 25]] + + # Depth exceeding the nesting depth applies at the innermost level: + @test bcastat(+, Val(3), p, v) == bcastat(+, Val(2), p, v) + + # Integer depth relies on constant propagation for type stability: + bcastat_intdepth(g, A, y) = bcastat(g, 2, A, y) + @test @inferred(bcastat_intdepth(+, p, v)) == bcastat(+, Val(2), p, v) + + # Two nesting levels over a single flat buffer: + VV = VectorOfArrays(partitioned(collect(1.0:10.0), [2, 3, 5]), [1, 3, 4], [(), ()]) + w = [100.0, 200.0] + r2 = bcastat(+, Val(3), VV, w) + @test r2 isa VectorOfArrays + @test fused(fused(r2)) == [101, 102, 103, 104, 105, 206, 207, 208, 209, 210] + + @test_throws DimensionMismatch bcastat(+, Val(2), p, Float32[1, 2]) + @test_throws DimensionMismatch bcastat(+, Val(2), p, partitioned(x, [4, 6])) + @test_throws ArgumentError bcastat(+, Val(2), [[1, 2], [3]], 1) + end + + @testset "outer broadcast" begin + A = VectorOfArrays([[1.0, 2.0], [3.0, 4.0, 5.0], [6.0]]) + A_ref = collect(A) + + # Array-valued results at the outer level yield a VectorOfArrays: + r = (x -> 2 .* x).(A) + @test r isa VectorOfArrays{Float64,1} + @test r == [2 .* x for x in A_ref] + full_consistency_checks(r) + + # Results may be ragged in new ways: + r2 = (x -> vcat(x, sum(x))).(A) + @test r2 isa VectorOfArrays{Float64,1} + @test r2 == [vcat(x, sum(x)) for x in A_ref] + + # Multiple and mixed arguments: + @test broadcast((x, y) -> x .+ y, A, A) == [x .+ x for x in A_ref] + @test broadcast((x, y) -> x .+ y, A, A) isa VectorOfArrays + @test broadcast((x, s) -> x .* s, A, 2.0) isa VectorOfArrays + @test broadcast((x, s) -> x .+ s, A, [10.0, 20.0, 30.0]) == [x .+ s for (x, s) in zip(A_ref, [10.0, 20.0, 30.0])] + + # Scalar-valued results stay plain arrays: + @test sum.(A) isa Vector{Float64} + @test sum.(A) == sum.(A_ref) + + # Matrix elements: + B = VectorOfArrays([rand(2, 2), rand(3, 2)]) + rB = (x -> 2 .* x).(B) + @test rB isa VectorOfArrays{Float64,2} + @test rB == [2 .* x for x in collect(B)] + + # VectorOfSimilarArrays broadcasts to VectorOfArrays too: + C = VectorOfSimilarVectors(rand(3, 4)) + rC = (x -> x .+ 1).(C) + @test rC isa VectorOfArrays{Float64,1} + @test rC == [x .+ 1 for x in collect(C)] + + # Multi-dimensional outer structure falls back to the default + # behavior: + D = ArrayOfSimilarArrays{Float64,1,2}(rand(2, 3, 4)) + rD = (x -> 2 .* x).(D) + @test rD isa Matrix{Vector{Float64}} + @test rD == [2 .* D[i, j] for i in axes(D, 1), j in axes(D, 2)] + end + + @testset "inner reductions" begin + x = collect(Float32, 1:10) + p = partitioned(x, [2, 3, 5]) + + @test @inferred(innermapreduce(abs2, +, p)) == [sum(abs2, xi) for xi in p] + @test @inferred(innerreduce(max, p)) == [maximum(xi) for xi in p] + @test @inferred(innersum(p)) == [sum(xi) for xi in p] + + # Empty element arrays require an init value, except for innersum: + pe = partitioned(collect(Float32, 1:5), [2, 0, 3]) + @test innersum(pe) == Float32[3, 0, 12] + @test_throws ArgumentError innerreduce(max, pe) + @test innerreduce(max, pe; init = -Inf32) == Float32[2, -Inf32, 5] + end + + @testset "rrules" begin + x = collect(Float64, 1:10) + + for lengths in ([2, 3, 5], [2, 3]) + Y, pb = rrule(partitioned, x, lengths) + @test Y == partitioned(x, lengths) + ΔY = [fill(1.0, l) for l in lengths] + ΔA = pb(ΔY)[2] + want = zero(x) + want[1:sum(lengths)] .= 1.0 + @test pb(ΔY)[1] == NoTangent() && pb(ΔY)[3] == NoTangent() + @test ΔA == want + end + + Y2, pb2 = rrule(partitioned, x, [(1, 2), (2, 2)]) + @test Y2 == partitioned(x, [(1, 2), (2, 2)]) + ΔA2 = pb2([fill(1.0, 1, 2), fill(1.0, 2, 2)])[2] + @test ΔA2 == [ones(6); zeros(4)] + + p = partitioned(x, [2, 3]) + y, pb3 = rrule(vecflattened, p) + @test y == vecflattened(p) + t = pb3(collect(1.0:5.0)) + @test t[1] == NoTangent() + @test t[2] isa VectorOfArrays + @test t[2] == [[1.0, 2.0], [3.0, 4.0, 5.0]] + @test fused(t[2])[6:10] == zeros(5) end @@ -128,7 +426,7 @@ using ArraysOfArrays: full_consistency_checks, append_elemptr!, element_ptr @test getindex_of_UR == V12 @test getindex_of_vector == getindex_of_UR - VV = @inferred(VectorOfVectors{Float64}()) + VV = @inferred(PartsView{Float64}()) data = @inferred(rand(5)) @inferred(push!(VV, data)) @test @inferred(getindex(VV, 1)) == data @@ -239,6 +537,10 @@ using ArraysOfArrays: full_consistency_checks, append_elemptr!, element_ptr A3 = VectorOfArrays(ref_AoA3(Float32, 3)) @test @inferred(adapt(identity, A3)) == A3 @test typeof(adapt(identity, A3)) == typeof(A3) + + sm = getsplitmode(A3) + @test @inferred(adapt(identity, sm)) == sm + @test typeof(adapt(identity, sm)) == typeof(sm) end @@ -253,7 +555,7 @@ using ArraysOfArrays: full_consistency_checks, append_elemptr!, element_ptr # ------------------------------------------------------------------- - VV = @inferred(VectorOfVectors{Float64}()) + VV = @inferred(PartsView{Float64}()) d1 = @inferred(rand(5)) d2 = @inferred(rand(4)) @@ -290,21 +592,21 @@ using ArraysOfArrays: full_consistency_checks, append_elemptr!, element_ptr A_grouped_ref = [[1, 1], [2], [3, 3], [2, 2, 2]] elem_ptr = consgrouped_ptrs(A) elem_ptr32 = Int32.(consgrouped_ptrs(A)) - @test first.(@inferred(VectorOfVectors(A, elem_ptr))) == [1, 2, 3, 2] - @test first.(@inferred(VectorOfVectors(A, elem_ptr32))) == [1, 2, 3, 2] + @test first.(@inferred(PartsView(A, elem_ptr))) == [1, 2, 3, 2] + @test first.(@inferred(PartsView(A, elem_ptr32))) == [1, 2, 3, 2] B = [1, 2, 3, 4, 5, 6, 7, 8] B_grouped_ref = [[1, 2], [3], [4, 5], [6, 7, 8]] - @test @inferred(VectorOfVectors(B, elem_ptr)) == B_grouped_ref - @test @inferred(VectorOfVectors(B, elem_ptr32)) == B_grouped_ref + @test @inferred(PartsView(B, elem_ptr)) == B_grouped_ref + @test @inferred(PartsView(B, elem_ptr32)) == B_grouped_ref C = [1.1, 2.2, 3.3, 4.4, 5.5, 6.6, 7.7, 8.8] C_grouped_ref = [[1.1, 2.2], [3.3], [4.4, 5.5], [6.6, 7.7, 8.8]] - @test @inferred(consgroupedview(A, B)) isa VectorOfVectors + @test @inferred(consgroupedview(A, B)) isa PartsView @test consgroupedview(A, B) == B_grouped_ref - @test @inferred(consgroupedview(A, (B, C))) isa NTuple{2, VectorOfVectors} + @test @inferred(consgroupedview(A, (B, C))) isa NTuple{2, PartsView} @test consgroupedview(A, (B, C)) == (B_grouped_ref, C_grouped_ref) data = (a = A, b = B, c = C)