How to Sensibly Implement Generic Slice Indexing for Custom Array Type

General Usage
, , ,
1

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?

2

At least for a specific number of indices I can write


"""
    catindex(rowindex::Union{CartesianIndex{M},NTuple{M,T}}, columnindex::Union{CartesianIndex{N},NTuple{N,T}}) where {M,N,T}

Concatenate various indexing types.
"""
function catindex(rowindex::Union{CartesianIndex{M},NTuple{M,T}}, columnindex::Union{CartesianIndex{N},NTuple{N,T}}) where {M,N,T}
    ntuple(function (k)
            if k <= M
                rowindex[k]
            else
                columnindex[k-M]
            end
        end,
        M + N)
end

function catindex(rowindex::Union{CartesianIndex{M},NTuple{M,T}}, columnindex::T) where {M,T}
    ntuple(function (k)
            if k <= M
                rowindex[k]
            else
                columnindex
            end
        end,
        M + 1)
end

function catindex(rowindex::T, columnindex::Union{CartesianIndex{N},NTuple{N,T}}) where {N,T}
    ntuple(function (k)
            if k == 1
                rowindex
            else
                columnindex[k-1]
            end
        end,
        1 + N)
end

function catindex(rowindex::T, colindex::T) where {T<:Integer}
    NTuple((rowindex, colindex))
end


function Base.getindex(A::MyReadOnlyArray{T,N}, rows, columns) where {T,N}
    rows = collect(rows)
    columns = collect(columns)
    # map over index combinations in parallel
    ....
end

This works for combinations of Int, Iterator, and Vec of Int, Tuple, and CartesianIndex.

3

A tuple isn’t usually a supported index for AbstractArrays. Maybe to_indices can help simplify the indices (it’s how : and multiple CartesianIndex are handled) for dispatch and processing, but it’s beyond me how this would work into GPU-accelerated slicing of a seemingly non-GPU array with non-GPU-array indices.

1 Like
4

The underlying array doesn’t really exist. I am working on an accelerated version of this: Function as Type Parameter - #6 by AlexanderNenninger. Can’t post code due to IP contracts with my employer, but once I have the indices as arrays of numbers/tuples/Cartesian indices I can work with them.