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

nuvec_core/
kernel.rs
