nuvec_core/
kernel.rs

1//! Custom Metal kernels from source, as MLX's `mx.fast.metal_kernel`.
2//!
3//! A kernel is a name, the names of its inputs and outputs, and the body of a Metal kernel
4//! function. Calling it with arrays generates the function's signature (as MLX does):
5//!
6//! - an input of at least one dimension is `const device T* name`, a 0-d input
7//!   `const constant T& name`; `name_shape`, `name_strides` and `name_ndim` are added when
8//!   the body uses them;
9//! - an output is `device T* name` (`device atomic<T>* name` with atomic outputs);
10//! - thread attributes the body uses (`thread_position_in_grid`, `thread_index_in_simdgroup`,
11//!   ...) become parameters;
12//! - template arguments become `constexpr` values and `typedef`s at the top of the body.
13//!
14//! The kernel runs on the GPU only, with one thread per grid position (`dispatchThreads`).
15//! It has no derivative or batching rule; wrap it in [`crate::custom_function`] (or [`crate::custom_function_flat`]) to give it
16//! one.
17
18use std::hash::{DefaultHasher, Hash, Hasher};
19use std::sync::Arc;
20
21use crate::primitive::Primitive;
22use crate::{Array, DType, DeviceKind, Error, Result, default_stream};
23
24/// A template argument of a custom kernel.
25#[derive(Debug, Clone, PartialEq, Eq, Hash)]
26pub enum TemplateArg {
27    /// `constexpr int name = value;`
28    Int(i64),
29    /// `constexpr bool name = value;`
30    Bool(bool),
31    /// `typedef type name;`
32    DType(DType),
33}
34
35#[derive(Debug, Clone, PartialEq, Eq, Hash)]
36pub(crate) struct KernelSpec {
37    name: String,
38    inputs: Vec<String>,
39    outputs: Vec<String>,
40    source: String,
41    header: String,
42    ensure_row_contiguous: bool,
43    atomic_outputs: bool,
44}
45
46/// A Metal kernel built from source, as MLX's `mx.fast.metal_kernel`: a name, the names
47/// of its inputs and outputs, and the body of a kernel function. Calling it generates the
48/// function's signature as MLX does:
49///
50/// - an input of at least one dimension is `const device T* name`, a 0-d input
51///   `const constant T& name`; `name_shape`, `name_strides` and `name_ndim` are added when
52///   the body uses them;
53/// - an output is `device T* name` (`device atomic<T>* name` with atomic outputs);
54/// - thread attributes the body uses (`thread_position_in_grid`,
55///   `thread_index_in_simdgroup`, ...) become parameters;
56/// - template arguments become `constexpr` values and `typedef`s at the top of the body;
57///   MLX's type names (`bfloat16_t`, `float16_t`) and `elem_to_loc` are defined.
58///
59/// It runs on the GPU only, one thread per grid position (`dispatchThreads`), and has no
60/// derivative or batching rule; wrap it in [`crate::custom_function`] (or [`crate::custom_function_flat`]) to give it one.
61///
62/// The source is trusted code, run as written: nuvec checks the inputs, outputs and
63/// launch configuration, but not what the body does with them. A body that reads or
64/// writes outside its inputs' and outputs' elements accesses other GPU memory, which is
65/// undefined behaviour (as with any `unsafe` code); pass only source you would compile
66/// into your program.
67///
68/// ```
69/// use nuvec_core::fast::{KernelConfig, MetalKernel};
70/// use nuvec_core::{Array, DType, Device, Result, array, default_stream};
71///
72/// # fn main() -> Result<()> {
73/// # if default_stream().device != Device::gpu(0) || (Array::from(1.0f32) * 2.0).eval().is_err() { return Ok(()); }
74/// let exp = MetalKernel::new(
75///     "myexp",
76///     &["inp"],
77///     &["out"],
78///     "uint elem = thread_position_in_grid.x;\nout[elem] = metal::exp(inp[elem]);",
79/// );
80/// let x = array![0.0, 1.0];
81/// let config = KernelConfig { grid: [2, 1, 1], threadgroup: [2, 1, 1], ..Default::default() };
82/// let y = exp.call(&[x], &[(&[2], DType::F32)], &config).remove(0);
83/// assert!((y.to_vec::<f32>()?[1] - std::f32::consts::E).abs() < 1e-6);
84/// # Ok(())
85/// # }
86/// ```
87#[derive(Debug, Clone)]
88pub struct MetalKernel {
89    spec: KernelSpec,
90}
91
92/// How a [`MetalKernel`] is launched.
93#[derive(Debug, Clone, Default, PartialEq)]
94pub struct KernelConfig {
95    /// Threads in the grid, `[x, y, z]`.
96    pub grid: [usize; 3],
97    /// Threads per threadgroup, `[x, y, z]` (at most 1,024 in all).
98    pub threadgroup: [usize; 3],
99    /// Template arguments, by name.
100    pub template: Vec<(String, TemplateArg)>,
101    /// A value every output holds before the kernel runs (else they are uninitialised).
102    pub init_value: Option<f64>,
103}
104
105/// A recorded call of a custom kernel.
106#[derive(Debug, Clone)]
107pub struct KernelCall {
108    pub(crate) spec: Arc<KernelSpec>,
109    pub(crate) template: Vec<(String, TemplateArg)>,
110    pub(crate) grid: [usize; 3],
111    pub(crate) threadgroup: [usize; 3],
112    pub(crate) outputs: Vec<(DType, Vec<usize>)>,
113    /// The outputs' initial values are the last inputs.
114    pub(crate) initialised: bool,
115}
116
117impl PartialEq for KernelCall {
118    fn eq(&self, other: &Self) -> bool {
119        self.spec == other.spec
120            && self.template == other.template
121            && self.grid == other.grid
122            && self.threadgroup == other.threadgroup
123            && self.outputs == other.outputs
124            && self.initialised == other.initialised
125    }
126}
127
128impl Eq for KernelCall {}
129
130impl Hash for KernelCall {
131    fn hash<H: Hasher>(&self, state: &mut H) {
132        self.spec.hash(state);
133        self.template.hash(state);
134        self.grid.hash(state);
135        self.threadgroup.hash(state);
136        self.outputs.hash(state);
137        self.initialised.hash(state);
138    }
139}
140
141impl MetalKernel {
142    /// A kernel named `name` (letters, digits and underscores) whose body `source` reads
143    /// `inputs` and writes `outputs` (in this order).
144    ///
145    /// Each distinct instantiation (source, template arguments, input and output types and
146    /// layouts) compiles one pipeline, cached for the life of the process: call a kernel
147    /// with a bounded set of template arguments (not, say, a sequence length) to keep the
148    /// cache bounded.
149    pub fn new(name: &str, inputs: &[&str], outputs: &[&str], source: &str) -> Self {
150        Self {
151            spec: KernelSpec {
152                name: name.to_owned(),
153                inputs: inputs.iter().map(|s| (*s).to_owned()).collect(),
154                outputs: outputs.iter().map(|s| (*s).to_owned()).collect(),
155                source: source.to_owned(),
156                header: String::new(),
157                ensure_row_contiguous: true,
158                atomic_outputs: false,
159            },
160        }
161    }
162
163    /// Code placed before the kernel function (helper functions, includes).
164    pub fn header(mut self, header: &str) -> Self {
165        self.spec.header = header.to_owned();
166        self
167    }
168
169    /// Whether inputs are copied to row-major order first (the default, as MLX); otherwise
170    /// the body indexes them through `name_strides`.
171    pub fn ensure_row_contiguous(mut self, yes: bool) -> Self {
172        self.spec.ensure_row_contiguous = yes;
173        self
174    }
175
176    /// Whether outputs are `device atomic<T>*` (to be updated with `atomic_*` functions).
177    pub fn atomic_outputs(mut self, yes: bool) -> Self {
178        self.spec.atomic_outputs = yes;
179        self
180    }
181
182    /// Runs the kernel on `inputs`, producing one array per output of the given shape and
183    /// dtype. Invalid arguments (names, counts, dtypes, sizes) give failed arrays.
184    #[track_caller]
185    pub fn call(
186        &self,
187        inputs: &[Array],
188        outputs: &[(&[i32], DType)],
189        config: &KernelConfig,
190    ) -> Vec<Array> {
191        let ins: Vec<&Array> = inputs.iter().collect();
192        Array::sticky_vec(
193            "metal_kernel",
194            std::panic::Location::caller(),
195            &ins,
196            outputs.len().max(1),
197            || {
198                let shapes = outputs
199                    .iter()
200                    .map(|(shape, _)| {
201                        let spec: Vec<i64> = shape.iter().map(|&d| i64::from(d)).collect();
202                        crate::axes::dims(&spec)
203                    })
204                    .collect::<Result<Vec<_>>>()?;
205                let outs: Vec<(&[usize], DType)> = shapes
206                    .iter()
207                    .zip(outputs)
208                    .map(|(s, (_, d))| (s.as_slice(), *d))
209                    .collect();
210                self.call_impl(inputs, &outs, config)
211            },
212        )
213    }
214
215    /// Runs the kernel on `inputs`, producing arrays of the given shapes and dtypes (`usize`
216    /// shapes; the implementation behind [`MetalKernel::call`]).
217    pub(crate) fn call_impl(
218        &self,
219        inputs: &[Array],
220        outputs: &[(&[usize], DType)],
221        config: &KernelConfig,
222    ) -> Result<Vec<Array>> {
223        let invalid = |msg: String| Err(Error::InvalidArgument(format!("metal_kernel: {msg}")));
224        let spec = &self.spec;
225        let word = |s: &String| {
226            !s.is_empty()
227                && !s.starts_with(|c: char| c.is_ascii_digit())
228                && s.chars().all(|c| c.is_ascii_alphanumeric() || c == '_')
229        };
230        if !word(&spec.name) || !spec.inputs.iter().chain(&spec.outputs).all(word) {
231            return invalid("names must be identifiers".into());
232        }
233        if inputs.len() != spec.inputs.len() || outputs.len() != spec.outputs.len() {
234            return invalid(format!(
235                "{} inputs and {} outputs expected, got {} and {}",
236                spec.inputs.len(),
237                spec.outputs.len(),
238                inputs.len(),
239                outputs.len()
240            ));
241        }
242        if outputs.is_empty() {
243            return invalid("a kernel needs at least one output".into());
244        }
245        if config.grid.contains(&0) || config.threadgroup.contains(&0) {
246            return invalid("grid and threadgroup sizes must be positive".into());
247        }
248        if config.threadgroup.iter().product::<usize>() > 1024 {
249            return invalid("at most 1,024 threads per threadgroup".into());
250        }
251        let all = inputs
252            .iter()
253            .map(Array::dtype)
254            .chain(outputs.iter().map(|o| o.1));
255        if let Some(dtype) = all
256            .chain(config.template.iter().filter_map(|(_, a)| match a {
257                TemplateArg::DType(d) => Some(*d),
258                _ => None,
259            }))
260            .find(|&d| msl_type(d).is_none())
261        {
262            return Err(Error::unsupported_dtype("metal_kernel", dtype));
263        }
264        if default_stream().device.kind != DeviceKind::Gpu {
265            return Err(Error::UnsupportedDevice(default_stream().device));
266        }
267        let mut operands = inputs.to_vec();
268        if let Some(v) = config.init_value {
269            for &(shape, dtype) in outputs {
270                operands.push(crate::ops::imp::broadcast_to(
271                    &Array::from(v).try_astype(dtype)?,
272                    shape,
273                )?);
274            }
275        }
276        let call = KernelCall {
277            spec: Arc::new(KernelSpec {
278                name: spec.name.clone(),
279                inputs: spec.inputs.clone(),
280                outputs: spec.outputs.clone(),
281                source: spec.source.clone(),
282                header: spec.header.clone(),
283                ensure_row_contiguous: spec.ensure_row_contiguous,
284                atomic_outputs: spec.atomic_outputs,
285            }),
286            template: config.template.clone(),
287            grid: config.grid,
288            threadgroup: config.threadgroup,
289            outputs: outputs.iter().map(|&(s, d)| (d, s.to_vec())).collect(),
290            initialised: config.init_value.is_some(),
291        };
292        Array::from_op(Primitive::MetalKernel(call), operands)
293    }
294}
295
296/// MLX's name of a dtype in kernel source.
297pub(crate) fn msl_type(dtype: DType) -> Option<&'static str> {
298    Some(match dtype {
299        DType::Bool => "bool",
300        DType::U8 => "uint8_t",
301        DType::U16 => "uint16_t",
302        DType::U32 => "uint32_t",
303        DType::U64 => "uint64_t",
304        DType::I8 => "int8_t",
305        DType::I16 => "int16_t",
306        DType::I32 => "int32_t",
307        DType::I64 => "int64_t",
308        DType::F16 => "float16_t",
309        DType::BF16 => "bfloat16_t",
310        DType::F32 => "float",
311        DType::F64 | DType::Complex64 => return None,
312    })
313}
314
315/// Thread attributes a body may use, with their types.
316const ATTRIBUTES: [(&str, &str); 13] = [
317    ("thread_position_in_grid", "uint3"),
318    ("threads_per_grid", "uint3"),
319    ("threadgroup_position_in_grid", "uint3"),
320    ("threadgroups_per_grid", "uint3"),
321    ("thread_position_in_threadgroup", "uint3"),
322    ("threads_per_threadgroup", "uint3"),
323    ("thread_index_in_threadgroup", "uint"),
324    ("simdgroup_index_in_threadgroup", "uint"),
325    ("thread_index_in_simdgroup", "uint"),
326    ("simdgroups_per_threadgroup", "uint"),
327    ("threads_per_simdgroup", "uint"),
328    ("quadgroup_index_in_threadgroup", "uint"),
329    ("thread_index_in_quadgroup", "uint"),
330];
331
332/// Whether `word` occurs in `text` as an identifier.
333fn uses(text: &str, word: &str) -> bool {
334    let ident = |c: char| c.is_ascii_alphanumeric() || c == '_';
335    text.match_indices(word).any(|(i, _)| {
336        let before = text[..i].chars().next_back();
337        let after = text[i + word.len()..].chars().next();
338        !before.is_some_and(ident) && !after.is_some_and(ident)
339    })
340}
341
342/// What a kernel's signature needs about each input.
343pub(crate) struct InputLayout {
344    pub(crate) dtype: DType,
345    pub(crate) scalar: bool,
346}
347
348/// Buffer roles in a generated signature, in order.
349#[derive(Debug, Clone, Copy, PartialEq, Eq)]
350pub(crate) enum Binding {
351    /// Input `i`'s data.
352    Input(usize),
353    /// Input `i`'s shape (`int` per dimension).
354    Shape(usize),
355    /// Input `i`'s strides (`int64_t` per dimension).
356    Strides(usize),
357    /// Input `i`'s number of dimensions (`int`).
358    Ndim(usize),
359    /// Output `i`.
360    Output(usize),
361}
362
363impl KernelCall {
364    /// The kernel's own inputs (the rest of a call's operands are initial outputs).
365    pub(crate) fn input_count(&self) -> usize {
366        self.spec.inputs.len()
367    }
368
369    /// Whether inputs are bound row-major.
370    pub(crate) fn row_contiguous_inputs(&self) -> bool {
371        self.spec.ensure_row_contiguous
372    }
373
374    /// The buffers of the generated signature, in order.
375    pub(crate) fn bindings(&self, inputs: &[InputLayout]) -> Vec<Binding> {
376        let spec = &self.spec;
377        let mut out = Vec::new();
378        for (i, (name, layout)) in spec.inputs.iter().zip(inputs).enumerate() {
379            out.push(Binding::Input(i));
380            if layout.scalar {
381                continue;
382            }
383            if uses(&spec.source, &format!("{name}_shape")) {
384                out.push(Binding::Shape(i));
385            }
386            if uses(&spec.source, &format!("{name}_strides")) {
387                out.push(Binding::Strides(i));
388            }
389            if uses(&spec.source, &format!("{name}_ndim")) {
390                out.push(Binding::Ndim(i));
391            }
392        }
393        out.extend((0..spec.outputs.len()).map(Binding::Output));
394        out
395    }
396
397    /// The pipeline name and the Metal source of the kernel for inputs laid out as given.
398    pub(crate) fn source(&self, inputs: &[InputLayout]) -> (String, String) {
399        let spec = &self.spec;
400        let ty = |d: DType| msl_type(d).expect("checked when called");
401        let mut params = Vec::new();
402        for (k, binding) in self.bindings(inputs).into_iter().enumerate() {
403            params.push(match binding {
404                Binding::Input(i) => {
405                    let (name, layout) = (&spec.inputs[i], &inputs[i]);
406                    match layout.scalar {
407                        true => format!(
408                            "const constant {}& {name} [[buffer({k})]]",
409                            ty(layout.dtype)
410                        ),
411                        false => {
412                            format!("const device {}* {name} [[buffer({k})]]", ty(layout.dtype))
413                        }
414                    }
415                }
416                Binding::Shape(i) => {
417                    format!(
418                        "const constant int* {}_shape [[buffer({k})]]",
419                        spec.inputs[i]
420                    )
421                }
422                Binding::Strides(i) => {
423                    format!(
424                        "const constant int64_t* {}_strides [[buffer({k})]]",
425                        spec.inputs[i]
426                    )
427                }
428                Binding::Ndim(i) => {
429                    format!(
430                        "const constant int& {}_ndim [[buffer({k})]]",
431                        spec.inputs[i]
432                    )
433                }
434                Binding::Output(i) => {
435                    let t = ty(self.outputs[i].0);
436                    match spec.atomic_outputs {
437                        true => format!("device atomic<{t}>* {} [[buffer({k})]]", spec.outputs[i]),
438                        false => format!("device {t}* {} [[buffer({k})]]", spec.outputs[i]),
439                    }
440                }
441            });
442        }
443        for (attr, t) in ATTRIBUTES {
444            if uses(&spec.source, attr) {
445                params.push(format!("{t} {attr} [[{attr}]]"));
446            }
447        }
448        let template: String = self
449            .template
450            .iter()
451            .map(|(name, arg)| match arg {
452                TemplateArg::Int(v) => format!("  constexpr int {name} = {v};\n"),
453                TemplateArg::Bool(v) => format!("  constexpr bool {name} = {v};\n"),
454                TemplateArg::DType(d) => format!("  typedef {} {name};\n", ty(*d)),
455            })
456            .collect();
457        let body = format!(
458            "{PRELUDE}{}\n[[kernel]] void KERNEL_NAME(\n  {}) {{\n{template}{}\n}}\n",
459            spec.header,
460            params.join(",\n  "),
461            spec.source
462        );
463        let mut h = DefaultHasher::new();
464        body.hash(&mut h);
465        let name = format!("nuvec_custom_{}_{:016x}", spec.name, h.finish());
466        (name.clone(), body.replace("KERNEL_NAME", &name))
467    }
468}
469
470/// Definitions every custom kernel sees: Metal's standard library and MLX's type names and
471/// `elem_to_loc`.
472const PRELUDE: &str = r#"#include <metal_stdlib>
473#include <metal_simdgroup>
474#include <metal_atomic>
475using namespace metal;
476typedef bfloat bfloat16_t;
477typedef half float16_t;
478
479// The element of a strided array at row-major position `elem` (MLX's `elem_to_loc`).
480inline int64_t elem_to_loc(uint elem, constant const int* shape, constant const int64_t* strides, int ndim) {
481  int64_t loc = 0;
482  for (int i = ndim - 1; i >= 0 && elem > 0; --i) {
483    loc += int64_t(elem % shape[i]) * strides[i];
484    elem /= shape[i];
485  }
486  return loc;
487}
488"#;
nuvec_core/kernel.rs

nuvec_core/
kernel.rs
