Struct MetalKernel 

pub struct MetalKernel { /* private fields */ }

A Metal kernel built from source, as MLX’s mx.fast.metal_kernel: a name, the names of its inputs and outputs, and the body of a kernel function. Calling it generates the function’s signature as MLX does:

• an input of at least one dimension is const device T* name, a 0-d input const constant T& name; name_shape, name_strides and name_ndim are added when the body uses them;
• an output is device T* name (device atomic<T>* name with atomic outputs);
• thread attributes the body uses (thread_position_in_grid, thread_index_in_simdgroup, …) become parameters;
• template arguments become constexpr values and typedefs at the top of the body; MLX’s type names (bfloat16_t, float16_t) and elem_to_loc are defined.

It runs on the GPU only, one thread per grid position (dispatchThreads), and has no derivative or batching rule; wrap it in crate::custom_function (or crate::custom_function_flat) to give it one.

The source is trusted code, run as written: nuvec checks the inputs, outputs and launch configuration, but not what the body does with them. A body that reads or writes outside its inputs’ and outputs’ elements accesses other GPU memory, which is undefined behaviour (as with any unsafe code); pass only source you would compile into your program.

use nuvec_core::fast::{KernelConfig, MetalKernel};
use nuvec_core::{Array, DType, Device, Result, array, default_stream};

let exp = MetalKernel::new(
    "myexp",
    &["inp"],
    &["out"],
    "uint elem = thread_position_in_grid.x;\nout[elem] = metal::exp(inp[elem]);",
);
let x = array![0.0, 1.0];
let config = KernelConfig { grid: [2, 1, 1], threadgroup: [2, 1, 1], ..Default::default() };
let y = exp.call(&[x], &[(&[2], DType::F32)], &config).remove(0);
assert!((y.to_vec::<f32>()?[1] - std::f32::consts::E).abs() < 1e-6);

Implementations

impl MetalKernel

pub fn new(name: &str, inputs: &[&str], outputs: &[&str], source: &str) -> Self

A kernel named name (letters, digits and underscores) whose body source reads inputs and writes outputs (in this order).

Each distinct instantiation (source, template arguments, input and output types and layouts) compiles one pipeline, cached for the life of the process: call a kernel with a bounded set of template arguments (not, say, a sequence length) to keep the cache bounded.

pub fn header(self, header: &str) -> Self

Code placed before the kernel function (helper functions, includes).

pub fn ensure_row_contiguous(self, yes: bool) -> Self

Whether inputs are copied to row-major order first (the default, as MLX); otherwise the body indexes them through name_strides.

pub fn atomic_outputs(self, yes: bool) -> Self

Whether outputs are device atomic<T>* (to be updated with atomic_* functions).

pub fn call( &self, inputs: &[Array], outputs: &[(&[i32], DType)], config: &KernelConfig, ) -> Vec<Array>

Runs the kernel on inputs, producing one array per output of the given shape and dtype. Invalid arguments (names, counts, dtypes, sizes) give failed arrays.

Trait Implementations

impl Clone for MetalKernel

fn clone(&self) -> Self

Returns a duplicate of the value.

fn clone_from(&mut self, source: &Self)

Performs copy-assignment from source.

impl Debug for MetalKernel

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

Formats the value using the given formatter.