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 to give it one.

use nuvec_core::fast::{KernelConfig, MetalKernel};
use nuvec_core::{Array, DType, Device, Result, Stream, 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::from_slice(&[0.0f32, 1.0], &[2])?;
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).

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: &[(&[usize], DType)], config: &KernelConfig, ) -> Result<Vec<Array>>

Runs the kernel on inputs, producing arrays of the given shapes and dtypes.

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.