Keyboard shortcuts

Press ← or → to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

Your first kernel

This is the complete vec-add executable:

use metal_oxide::device::{DisjointBuffer, ReadBuffer, thread_index};

/// Adds two buffers of equal length element by element.
#[metal_oxide::kernel(requires(a.len() == out.len(), b.len() == out.len()))]
pub fn vec_add(a: ReadBuffer<f32>, b: ReadBuffer<f32>, mut out: DisjointBuffer<f32>) {
    let index = thread_index();
    if let Some(out) = out.get_mut(&index) {
        let i = index.get();
        *out = a.as_slice()[i] + b.as_slice()[i];
    }
}

#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
fn main() -> Result<(), Box<dyn std::error::Error>> {
    use metal_oxide::{Device, LaunchConfig, Operation};

    let device = Device::system_default()?;
    println!("Device: {}", device.name());
    for n in [0_u32, 1, 255, 256, 257, 1_000_003] {
        let a = (0..n).map(|i| (i % 1024) as f32 * 0.25).collect::<Vec<_>>();
        let b = (0..n).map(|i| (i % 511) as f32 * -0.5).collect::<Vec<_>>();
        let input_a = device.buffer_from_slice(&a)?;
        let input_b = device.buffer_from_slice(&b)?;
        let mut output = device.buffer_zeroed::<f32>(n as usize)?;
        let config = LaunchConfig::<256>::for_elements(n)?;
        vec_add(&device, config, &input_a, &input_b, &mut output).wait()?;
        for (i, actual) in output.read().wait().iter().enumerate() {
            if *actual != a[i] + b[i] {
                return Err(format!("vec_add mismatch at n={n}, index={i}").into());
            }
        }
        println!(
            "vec_add n={n}: verified ({} blocks x {} threads)",
            config.grid.x,
            LaunchConfig::<256>::BLOCK.x
        );
    }
    Ok(())
}

#[cfg(not(all(target_os = "macos", target_arch = "aarch64")))]
fn main() {
    eprintln!("vec-add runs on macOS with Apple Silicon");
}

#[cfg(all(test, target_os = "macos", target_arch = "aarch64"))]
mod tests;

The #[kernel] macro marks a Rust entrypoint for the compiler and produces a host launch function of the same name. Its arguments begin with the device and launch configuration, followed by the kernel’s arguments. requires relations are checked before encoding; a violation returns Error::Contract.

ReadBuffer exposes a checked shared slice. DisjointBuffer::get_mut uses a thread_index() witness to give each thread exclusive access to its own element. Extra threads get None. Here the length contracts prove the input indexing safe.

The host allocates initialized buffers, constructs a configuration, waits for the launch and reads the output. Import Operation to use .wait(). The hardware tests run only when explicitly requested; ordinary workspace tests do not run them.

For an in-place update, the tested SAXPY example uses the same ownership rules:

use metal_oxide::device::{DisjointBuffer, ReadBuffer, thread_index};

/// `y = a * x + y`: each thread reads and writes its own element of `y`.
#[metal_oxide::kernel(requires(x.len() == y.len()))]
pub fn saxpy(a: f32, x: ReadBuffer<f32>, mut y: DisjointBuffer<f32>) {
    let index = thread_index();
    if let Some(y) = y.get_mut(&index) {
        *y += a * x.as_slice()[index.get()];
    }
}

#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
fn main() -> Result<(), Box<dyn std::error::Error>> {
    use metal_oxide::{Device, LaunchConfig, Operation};

    let device = Device::system_default()?;
    println!("Device: {}", device.name());
    let a = 2.5_f32;
    for n in [0_u32, 1, 255, 256, 257, 1_000_003] {
        let x = (0..n).map(|i| (i % 1024) as f32 * 0.25).collect::<Vec<_>>();
        let y = (0..n).map(|i| (i % 511) as f32 * -0.5).collect::<Vec<_>>();
        let input = device.buffer_from_slice(&x)?;
        let mut output = device.buffer_from_slice(&y)?;
        let config = LaunchConfig::<256>::for_elements(n)?;
        saxpy(&device, config, a, &input, &mut output).wait()?;
        for (i, actual) in output.read().wait().iter().enumerate() {
            if *actual != a * x[i] + y[i] {
                return Err(format!("saxpy mismatch at n={n}, index={i}").into());
            }
        }
        println!("saxpy n={n}: verified");
    }
    Ok(())
}

#[cfg(not(all(target_os = "macos", target_arch = "aarch64")))]
fn main() {
    eprintln!("saxpy runs on macOS with Apple Silicon");
}

#[cfg(all(test, target_os = "macos", target_arch = "aarch64"))]
mod tests;

Run it from the checkout with cargo run -p saxpy --locked. For larger programs, #[metal_oxide::kernels] shares types and helpers in an inline module; #[metal_oxide::kernels(path = "../kernels")] imports a separate no_std kernel crate. See compiler internals for their build boundary.