Skip to content

Node to source edge mapping has ragged array; not GPU compatible #648

@BSchilperoort

Description

@BSchilperoort

The EdgesAtNode struct contains a mapping from nodes to source edges. This is represented in a ragged array. For example;

mapping = Vector([
    Vector([1]),
    Vector([2, 3]),
    Vector([4, 5, 6]),
])

This nested vector structure is not compatible with (julia's/KernelAbstractions) GPU arrays*. If the array cannot be copied to GPU we cannot use it inside a kernel.
To make it GPU compatible, we can use a 2D array, with 0 as padding value (as julia has no 0 index);

mapping = Array([
    1 0 0
    2 3 0
    4 5 6
])

Seeing as this will be an array of size (nodes, <node with maximum number of source edges>) the array shouldn't become much bigger than its current size.

The sum_at function needs to be modified to accommodate this, by skipping any 0 values.

function sum_at(A, inds)
    v = zero(eltype(A))
    for i in inds
        if i != 0
            v += A[i]
        end
    end
    return v
end

* Technically CUDA and ROCm support ragged tensors, but this is not available in Julia's generic GPU packages.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions