How to Sensibly Implement Generic Slice Indexing for Custom Array Type

AlexanderNenninger
Fri, 26 Jan 2024 08:53:00 +0000
General Usage

I have a custom array type that I want to implement custom Base.getindex for:

struct MyReadOnlyArray{T, N} <: AbstractArray{T, N}
    ...
    size::NTuple{N, Int}
end

Not it needs to be n-dimensional. For single entries, I can just implement

function Base.getindex(A::MyReadOnlyArray{T,N}, I::Vararg{Int,N}) where {T,N}
    ... # go reasonably fast sequentially
end

and the standard library gives me all of the fancy indexing Julia has to offer.

Now this custom array type can really profit from offloading the calculation of large slices to the GPU, e.g.

function Base.getindex(A::MyReadOnlyArray{T,N}, rows::Vector{CartesianIndex}, columns::Vector{CartesianIndex}) where {T,N}
... # go fast in parallel
end

I could of course provide methods for all kinds of indexing operations, but that’s not scalable. Continuing the example, I’d have to write

function Base.getindex(A::MyReadOnlyArray{T,N}, rows::Vector{CartesianIndex}, columns::UnitRange{Int}) where {T,N}
... # go fast in parallel too
end

function Base.getindex(A::MyReadOnlyArray{T,N}, rows::Vector, columns::UnitRange{Int}, stack::Int) where {T,N}
... # go fast in parallel too
end

Ignoring implementation details, What type signature do I give specialized methods such that getindex will dispatch to the fast implementation for slices without writing one for each index combination separately?

4 posts - 2 participants