Enum poc_kokkos_rs::functor::KernelArgs

pub enum KernelArgs<const N: usize> {
    Index1D(usize),
    IndexND([usize; N]),
    Handle,
}

Kernel argument enum

In the Kokkos library, there is a finite number of kernel signatures. Each is associated to/determined by a given execution policy. In order to have kernel genericity in Rust, without introducing overhead due to downcasting, the solution was to define kernel arguments as a struct-like enum.

Example

One-dimensional kernel:

// Range is defined in the execution policy
use poc_kokkos_rs::functor::KernelArgs;

let kern = |arg: KernelArgs<1>| match arg {
        KernelArgs::Index1D(i) => {
            // body of the kernel
            println!("Hello from iteration {i}")
        },
        KernelArgs::IndexND(_) => unimplemented!(),
        KernelArgs::Handle => unimplemented!(),
    };

3D kernel:

use poc_kokkos_rs::functor::KernelArgs;

// Use the array
let kern = |arg: KernelArgs<3>| match arg {
        KernelArgs::Index1D(_) => unimplemented!(),
        KernelArgs::IndexND(idx) => { // idx: [usize; 3]
            // body of the kernel
            println!("Hello from iteration {idx:?}")
        },
        KernelArgs::Handle => unimplemented!(),
    };

// Decompose the array
let kern = |arg: KernelArgs<3>| match arg {
        KernelArgs::Index1D(_) => unimplemented!(),
        KernelArgs::IndexND([i, j, k]) => { // i,j,k: usize
            // body of the kernel
            println!("Hello from iteration {i},{j},{k}");
        },
        KernelArgs::Handle => unimplemented!(),
    };

Variants

Index1D(usize)

Arguments of a one-dimensionnal kernel (e.g. a RangePolicy).

IndexND([usize; N])

Arguments of a N-dimensionnal kernel (e.g. a MDRangePolicy).

Handle

Arguments of a team-based kernel.