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

nuvec_core/
kernel.rs
