Support for Custom Data Types in Buffers and Kernel Parameters

[NoahGav]

As far as I'm aware it is not possible to use custom structs in buffers or as parameters for kernels. This is very limiting. For example, this is not currently possible:

use bytemuck::{Pod, Zeroable};

#[repr(C)]
#[derive(Clone, Copy, Pod, Zeroable)]
struct Pixel {
    r: u8,
    g: u8,
    b: u8,
}

#[kernel]
fn foo(#[global] pixels: UnsafeSlice<Pixel>) { ... }

[charles-r-earp]

You could try something like this:

#[repr(C)]
#[derive(Clone, Copy, Pod, Zeroable)]
struct Pixel {
    r: u8,
    g: u8,
    b: u8,
}

struct UnsafePixelSlice<'a> {
    inner: UnsafeSlice<'a, u8>,
}

impl<'a> From<UnsafeSlice<'a, u8>> for UnsafePixelSlice<'a> {
    fn from(slice: UnsafeSlice<'a, u8>) -> Self {
        Self { inner: slice }
    }
}

impl UnsafePixelSlice<'_> {
    fn len(&self) -> usize {
        self.inner.len() / size_of::<Pixel>()
    }
    unsafe fn unsafe_load(&self, index: usize) -> Pixel {
        let len = self.len();
        if index < len {
            unsafe {
                Pixel {
                    r: *self.inner.unsafe_index(index * 3),
                    g: *self.inner.unsafe_index(index * 3 + 1),
                    b: *self.inner.unsafe_index(index * 3 + 2),
                }
            }
        } else {
            panic!("index out of bounds: the len is {index} but the index is {len}")
        }
    }
    unsafe fn unsafe_store(&self, index: usize, pixel: Pixel) {
        let len = self.len();
        if index < len {
            unsafe {
                *self.inner.unsafe_index_mut(index * 3) = pixel.r;
                *self.inner.unsafe_index_mut(index * 3 + 1) = pixel.g;
                *self.inner.unsafe_index_mut(index * 3 + 2) = pixel.b;
            }
        } else {
            panic!("index out of bounds: the len is {index} but the index is {len}")
        }
    }
}

#[kernel]
fn foo(#[global] data: UnsafeSlice<u8>) {
    let pixel_index = kernel.global_id() * size_of::<Pixel>();
    let pixels = UnsafePixelSlice::from(data);
    if pixel_index < pixels.len() {
        let mut pixel = unsafe { pixels.unsafe_load(pixel_index) };
        /* .. */
        unsafe {
            pixels.unsafe_store(pixel_index, pixel);
        }
    }
}

Note that rust-gpu doesn't currently support casting between pointer types, which would allow implementing UnsafeIndex.

Currently buffers on the host and device require T: Scalar. This isn't actually necessary, as simply allocating and copying could just require Pod or whatever. But it probably isn't worth changing unless kernels also supported more types.

I would potentially like to support arrays and or vectors. Unfortunately the only vector library that supports spirv arch is glam, which isn't generic.

Custom types would add significant flexibility, but it's harder to ensure soundness. This is because krnlc uses a different toolchain, and a different set of dependencies, and potentially different versions, and isn't re-run when the host crate is compiled, potentially downstream. Stable rust types, and half, have a fixed definition, that cannot or is assumed not to change. Custom types, whether in a dependency or the current crate, might change, and can't be trusted.

It might be possible by requiring an unsafe trait, plus validate layouts and such, but it requires significant design.