Struct MetalKernel 

pub struct MetalKernel { /* private fields */ }

A custom Metal kernel, compiled from source on first use and cached by MLX.

The kernel body is written against the names given at construction: each input a is a const device T* a, or const constant T* a when it has fewer than eight elements. Non-scalar inputs also expose a_shape, a_strides and a_ndim when referenced in the source. A 0-d input is a const constant T& a and has no shape, strides or ndim arguments. Each output is a device T* (or an atomic pointer with atomic_outputs), and Metal attributes such as thread_position_in_grid are declared automatically when referenced.

Kernels only run on the GPU. MetalKernel::apply uses the stream selected by crate::with_stream or crate::with_device, falling back to the GPU outside a scope, even when the default device is the CPU. Selecting a CPU scope returns an error.

use mlx_rs::{fast::{MetalKernel, MetalKernelApplyOptions, MetalKernelOptions, MetalKernelOutput}, Array, Dtype};

let kernel = MetalKernel::new(
    "double",
    &["a"],
    &["out"],
    "uint i = thread_position_in_grid.x; out[i] = 2 * a[i];",
    MetalKernelOptions::default(),
)?;
let a = Array::from_slice(&[1.0_f32, 2.0, 3.0, 4.0], &[4]);
// SAFETY: Four threads read four f32 inputs and initialize distinct f32 outputs.
let out = unsafe {
    kernel.apply(
        &[&a],
        &[MetalKernelOutput { shape: vec![4], dtype: Dtype::Float32 }],
        (4, 1, 1),
        (4, 1, 1),
        MetalKernelApplyOptions::default(),
    )?
};
assert_eq!(out[0].as_slice::<f32>(), &[2.0, 4.0, 6.0, 8.0]);

Implementations

impl MetalKernel

pub fn new( name: &str, input_names: &[&str], output_names: &[&str], source: &str, options: MetalKernelOptions, ) -> Result<Self>

Builds a kernel from its name, input and output names, and Metal body source.

At least one output name is required. Compilation is deferred until the outputs are evaluated, so source errors surface from Array::eval, not apply.

pub unsafe fn apply( &self, inputs: &[impl AsRef<Array>], outputs: &[MetalKernelOutput], grid: (i32, i32, i32), thread_group: (i32, i32, i32), options: MetalKernelApplyOptions, ) -> Result<Vec<Array>>

Runs the kernel over grid threads in groups of thread_group, returning one array per entry of outputs.

inputs and outputs must match the counts given at construction. Like other operations the result is lazy; the kernel compiles and runs when the outputs are evaluated. Source and threadgroup errors surface from Array::eval. The kernel and option values may be dropped before evaluating the outputs.

Safety

For the given inputs, outputs, template arguments and launch configuration, the kernel must only access memory inside the input and output buffers and must never write an input. It must write every output element before it is read unless init_value is set, only write values valid for the output dtype, and have no unsynchronized conflicting accesses. These requirements apply throughout deferred execution.

Trait Implementations

impl Debug for MetalKernel

fn fmt(&self, f: &mut Formatter<'_>) -> Result

Formats the value using the given formatter.

impl Drop for MetalKernel

fn drop(&mut self)

Executes the destructor for this type.

fn pin_drop(self: Pin<&mut Self>)

🔬This is a nightly-only experimental API. (pin_ergonomics)

Execute the destructor for this type, but different to Drop::drop, it requires self to be pinned.