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

Launch geometry

LaunchConfig<X, Y, Z> stores the threads per block in its type; omitted axes are one. Its runtime grid: Dim3 counts blocks. Each block is a Metal threadgroup. LaunchConfig::<256>::for_elements(257) launches two blocks of 256 threads. The kernel must guard the 255 padding threads.

DynamicLaunchConfig carries both grid and block dimensions at runtime. It shares the const configuration’s validation: nonzero block dimensions, device and pipeline limits, total threadgroup storage and nonoverflowing global dimensions. An empty grid executes no threads but still validates its bindings and geometry. Scalar kernel parameters do not set the grid automatically.

Device coordinates

Rust functionMeaningCUDA equivalent
thread_idx()Thread within its blockthreadIdx
block_idx()Block within the gridblockIdx
block_dim()Threads per blockblockDim
grid_dim()Blocks in the gridgridDim

All four return x/y/z coordinates. The usual global x index is block_idx().x * block_dim().x + thread_idx().x. thread_index() instead returns an ownership witness for the global x index. Kernels using it, including through helpers, require one-dimensional launches; other shapes are rejected.

Required shapes

#[kernel(block = (128, 1, 1))] records an exact block requirement. Its generated host function accepts the corresponding const configuration. Host const generics alone do not specialize handwritten MSL.

These tested pipeline kernels use two required shapes and a shape-independent final stage:

/// Two stages whose kernels share an affine helper and a generic scaling body.
#[metal_oxide::kernels]
pub mod stages {
    use metal_oxide::device::{ReadBuffer, WriteBuffer, block_dim, block_idx, thread_idx};

    struct Affine {
        scale: f32,
        offset: f32,
    }

    impl Affine {
        fn apply(self, value: f32) -> f32 {
            value * self.scale + self.offset
        }
    }

    unsafe fn scale<const BLOCK: u32, const SCALE: u32>(
        input: ReadBuffer<f32>,
        output: WriteBuffer<f32>,
        n: u32,
    ) {
        let i = block_idx().x * BLOCK + thread_idx().x;
        if i < n {
            let affine = Affine {
                scale: SCALE as f32,
                offset: 3.0,
            };
            // SAFETY: entrypoint callers supply n elements and one writer per index.
            unsafe { output.store_unchecked(i, affine.apply(input.load_unchecked(i))) };
        }
    }

    /// # Safety
    /// Buffers cover n elements, do not overlap, and the launch uses 128 threads per block.
    #[kernel(block = (128, 1, 1))]
    pub unsafe fn scale_128(input: ReadBuffer<f32>, output: WriteBuffer<f32>, n: u32) {
        // SAFETY: the entrypoint contract fixes the index calculation and buffer bounds.
        unsafe { scale::<128, 2>(input, output, n) };
    }

    /// # Safety
    /// Buffers cover n elements, do not overlap, and the launch uses 256 threads per block.
    #[kernel(block = (256, 1, 1))]
    pub unsafe fn scale_256(input: ReadBuffer<f32>, output: WriteBuffer<f32>, n: u32) {
        // SAFETY: the entrypoint contract fixes the index calculation and buffer bounds.
        unsafe { scale::<256, 4>(input, output, n) };
    }

    /// # Safety
    /// Buffers cover n elements, do not overlap, and the launch is one dimensional.
    #[kernel]
    pub unsafe fn finish(input: ReadBuffer<f32>, output: WriteBuffer<f32>, n: u32) {
        let i = block_idx().x * block_dim().x + thread_idx().x;
        if i < n {
            let affine = Affine {
                scale: 5.0,
                offset: 7.0,
            };
            // SAFETY: each active thread reads an initialized input and writes its own output.
            unsafe { output.store_unchecked(i, affine.apply(input.load_unchecked(i))) };
        }
    }
}

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

#[cfg(all(target_os = "macos", target_arch = "aarch64"))]
fn main() -> Result<(), Box<dyn std::error::Error>> {
    run::block_on(run::verify())
}

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

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

The async chapter includes their host launches. SIMD-group matrix kernels require a declared shape whose thread count is a multiple of 32. Barriers and collectives also impose participation rules; see memory and safety.