Stable element types and new API features, major changes

.github/workflows/ci.yml → .github/workflows/CI.yml

@@ -26,7 +26,7 @@ jobs:
       fail-fast: false
       matrix:
         version:
-          - '1.6'
+          - '1.10'
           - '1'
           - 'pre'
         os:

.gitignore

@@ -5,3 +5,5 @@
 .ipynb_checkpoints
 .vscode
 Manifest.toml
+Manifest-v*.toml
+LocalPreferences.toml

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"

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
 

docs/src/index.md

@@ -4,21 +4,47 @@ 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)
 
 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:
 
 ```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
 

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

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
 

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

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