mlx_rs/fast/
metal_kernel.rs

1use std::ffi::CString;
2
3use crate::error::{Exception, Result};
4use crate::utils::{guard::Guarded, SUCCESS};
5use crate::{Array, Dtype, Stream};
6
7/// A compile-time template parameter for a [`MetalKernel`].
8///
9/// Each argument is emitted as a template parameter of the generated `[[kernel]]` function, so
10/// the kernel source can refer to it by name (for example `T`, `N` or `USE_BIAS`).
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
12pub enum TemplateArg {
13    /// A type parameter, emitted as the matching Metal type (for example `float`).
14    Dtype(Dtype),
15    /// An `int` parameter.
16    Int(i32),
17    /// A `bool` parameter.
18    Bool(bool),
19}
20
21/// Construction options for [`MetalKernel::new`].
22#[derive(Debug, Clone, Default, PartialEq, Eq)]
23pub struct MetalKernelOptions {
24    /// Metal source placed before the generated kernel, such as helper functions or includes.
25    pub header: String,
26    /// Copy inputs to row-contiguous layout only when `true`. Defaults to `false`; Python
27    /// defaults to `true`. When `false`, the source must handle the inputs' actual layout
28    /// (for example through `_strides` and `elem_to_loc`).
29    pub ensure_row_contiguous: bool,
30    /// Declare outputs as `device atomic<T>*` so the kernel can use Metal atomics on them.
31    pub atomic_outputs: bool,
32}
33
34/// The shape and dtype of one [`MetalKernel`] output.
35#[derive(Debug, Clone, PartialEq, Eq)]
36pub struct MetalKernelOutput {
37    /// Output shape.
38    pub shape: Vec<i32>,
39    /// Output element type.
40    pub dtype: Dtype,
41}
42
43/// Independent per-call options for [`MetalKernel::apply`].
44#[derive(Debug, Clone, Default, PartialEq)]
45pub struct MetalKernelApplyOptions {
46    /// Template parameters, in declaration order.
47    pub template_args: Vec<(String, TemplateArg)>,
48    /// Fill every output with this value before the kernel runs. When absent, outputs start
49    /// uninitialized and the kernel must write every element.
50    pub init_value: Option<f32>,
51    /// Print the generated Metal source on every call to `apply`.
52    pub verbose: bool,
53}
54
55/// A custom Metal kernel, compiled from source on first use and cached by MLX.
56///
57/// The kernel body is written against the names given at construction: each input `a` is a
58/// `const device T* a`, or `const constant T* a` when it has fewer than eight elements.
59/// Non-scalar inputs also expose `a_shape`, `a_strides` and `a_ndim` when referenced in the
60/// source. A 0-d input is a `const constant T& a` and has no shape, strides or ndim arguments.
61/// Each output is a `device T*` (or an atomic pointer with `atomic_outputs`), and Metal
62/// attributes such as `thread_position_in_grid` are declared automatically when referenced.
63///
64/// Kernels only run on the GPU. [`MetalKernel::apply`] uses the stream selected by
65/// [`crate::with_stream`] or [`crate::with_device`], falling back to the GPU outside a scope,
66/// even when the default device is the CPU. Selecting a CPU scope returns an error.
67///
68/// ```rust
69/// use mlx_rs::{fast::{MetalKernel, MetalKernelApplyOptions, MetalKernelOptions, MetalKernelOutput}, Array, Dtype};
70///
71/// let kernel = MetalKernel::new(
72///     "double",
73///     &["a"],
74///     &["out"],
75///     "uint i = thread_position_in_grid.x; out[i] = 2 * a[i];",
76///     MetalKernelOptions::default(),
77/// )?;
78/// let a = Array::from_slice(&[1.0_f32, 2.0, 3.0, 4.0], &[4]);
79/// // SAFETY: Four threads read four f32 inputs and initialize distinct f32 outputs.
80/// let out = unsafe {
81///     kernel.apply(
82///         &[&a],
83///         &[MetalKernelOutput { shape: vec![4], dtype: Dtype::Float32 }],
84///         (4, 1, 1),
85///         (4, 1, 1),
86///         MetalKernelApplyOptions::default(),
87///     )?
88/// };
89/// assert_eq!(out[0].as_slice::<f32>(), &[2.0, 4.0, 6.0, 8.0]);
90/// # Ok::<(), mlx_rs::error::Exception>(())
91/// ```
92pub struct MetalKernel {
93    raw: mlx_sys::mlx_fast_metal_kernel,
94}
95
96impl MetalKernel {
97    /// Builds a kernel from its name, input and output names, and Metal body `source`.
98    ///
99    /// At least one output name is required. Compilation is deferred until the outputs are
100    /// evaluated, so source errors surface from [`Array::eval`], not [`apply`](Self::apply).
101    pub fn new(
102        name: &str,
103        input_names: &[&str],
104        output_names: &[&str],
105        source: &str,
106        options: MetalKernelOptions,
107    ) -> Result<Self> {
108        crate::error::INIT_ERR_HANDLER.call_once(crate::error::setup_mlx_error_handler);
109        let name = c_string(name)?;
110        let source = c_string(source)?;
111        let header = c_string(&options.header)?;
112        let input_names = StringVector::new(input_names)?;
113        let output_names = StringVector::new(output_names)?;
114        let raw = unsafe {
115            mlx_sys::mlx_fast_metal_kernel_new(
116                name.as_ptr(),
117                input_names.raw,
118                output_names.raw,
119                source.as_ptr(),
120                header.as_ptr(),
121                options.ensure_row_contiguous,
122                options.atomic_outputs,
123            )
124        };
125        if raw.ctx.is_null() {
126            return Err(crate::error::exception_from_status(
127                SUCCESS + 1,
128                "creating a Metal kernel",
129            ));
130        }
131        Ok(Self { raw })
132    }
133
134    /// Runs the kernel over `grid` threads in groups of `thread_group`, returning one array per
135    /// entry of `outputs`.
136    ///
137    /// `inputs` and `outputs` must match the counts given at construction. Like other
138    /// operations the result is lazy; the kernel compiles and runs when the outputs are
139    /// evaluated. Source and threadgroup errors surface from [`Array::eval`]. The kernel and
140    /// option values may be dropped before evaluating the outputs.
141    ///
142    /// # Safety
143    ///
144    /// For the given inputs, outputs, template arguments and launch configuration, the kernel
145    /// must only access memory inside the input and output buffers and must never write an input.
146    /// It must write every output element before it is read unless `init_value` is set, only
147    /// write values valid for the output dtype, and have no unsynchronized conflicting accesses.
148    /// These requirements apply throughout deferred execution.
149    pub unsafe fn apply(
150        &self,
151        inputs: &[impl AsRef<Array>],
152        outputs: &[MetalKernelOutput],
153        grid: (i32, i32, i32),
154        thread_group: (i32, i32, i32),
155        options: MetalKernelApplyOptions,
156    ) -> Result<Vec<Array>> {
157        crate::error::INIT_ERR_HANDLER.call_once(crate::error::setup_mlx_error_handler);
158        let config = KernelConfig::new()?;
159        for output in outputs {
160            config.check(
161                unsafe {
162                    mlx_sys::mlx_fast_metal_kernel_config_add_output_arg(
163                        config.raw,
164                        output.shape.as_ptr(),
165                        output.shape.len(),
166                        output.dtype.into(),
167                    )
168                },
169                "adding a Metal kernel output",
170            )?;
171        }
172        config.check(
173            unsafe {
174                mlx_sys::mlx_fast_metal_kernel_config_set_grid(config.raw, grid.0, grid.1, grid.2)
175            },
176            "setting the Metal kernel grid",
177        )?;
178        config.check(
179            unsafe {
180                mlx_sys::mlx_fast_metal_kernel_config_set_thread_group(
181                    config.raw,
182                    thread_group.0,
183                    thread_group.1,
184                    thread_group.2,
185                )
186            },
187            "setting the Metal kernel thread group",
188        )?;
189        for (name, arg) in &options.template_args {
190            let name = c_string(name)?;
191            let status = unsafe {
192                match *arg {
193                    TemplateArg::Dtype(dtype) => {
194                        mlx_sys::mlx_fast_metal_kernel_config_add_template_arg_dtype(
195                            config.raw,
196                            name.as_ptr(),
197                            dtype.into(),
198                        )
199                    }
200                    TemplateArg::Int(value) => {
201                        mlx_sys::mlx_fast_metal_kernel_config_add_template_arg_int(
202                            config.raw,
203                            name.as_ptr(),
204                            value,
205                        )
206                    }
207                    TemplateArg::Bool(value) => {
208                        mlx_sys::mlx_fast_metal_kernel_config_add_template_arg_bool(
209                            config.raw,
210                            name.as_ptr(),
211                            value,
212                        )
213                    }
214                }
215            };
216            config.check(status, "adding a Metal kernel template argument")?;
217        }
218        if let Some(value) = options.init_value {
219            config.check(
220                unsafe { mlx_sys::mlx_fast_metal_kernel_config_set_init_value(config.raw, value) },
221                "setting the Metal kernel init value",
222            )?;
223        }
224        config.check(
225            unsafe {
226                mlx_sys::mlx_fast_metal_kernel_config_set_verbose(config.raw, options.verbose)
227            },
228            "setting Metal kernel verbosity",
229        )?;
230
231        let inputs = crate::utils::VectorArray::try_from_iter(inputs.iter())?;
232        let stream = Stream::thread_local_or_gpu();
233        Vec::<Array>::try_from_op(|res| unsafe {
234            mlx_sys::mlx_fast_metal_kernel_apply(
235                res,
236                self.raw,
237                inputs.as_ptr(),
238                config.raw,
239                stream.as_ref().as_ptr(),
240            )
241        })
242    }
243}
244
245impl Drop for MetalKernel {
246    fn drop(&mut self) {
247        unsafe { mlx_sys::mlx_fast_metal_kernel_free(self.raw) };
248    }
249}
250
251impl std::fmt::Debug for MetalKernel {
252    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
253        f.debug_struct("MetalKernel").finish_non_exhaustive()
254    }
255}
256
257struct KernelConfig {
258    raw: mlx_sys::mlx_fast_metal_kernel_config,
259}
260
261impl KernelConfig {
262    fn new() -> Result<Self> {
263        let raw = unsafe { mlx_sys::mlx_fast_metal_kernel_config_new() };
264        if raw.ctx.is_null() {
265            return Err(crate::error::exception_from_status(
266                SUCCESS + 1,
267                "creating a Metal kernel config",
268            ));
269        }
270        Ok(Self { raw })
271    }
272
273    fn check(&self, status: i32, operation: &str) -> Result<()> {
274        if status == SUCCESS {
275            Ok(())
276        } else {
277            Err(crate::error::exception_from_status(status, operation))
278        }
279    }
280}
281
282impl Drop for KernelConfig {
283    fn drop(&mut self) {
284        unsafe { mlx_sys::mlx_fast_metal_kernel_config_free(self.raw) };
285    }
286}
287
288struct StringVector {
289    raw: mlx_sys::mlx_vector_string,
290}
291
292impl StringVector {
293    fn new(values: &[&str]) -> Result<Self> {
294        let owned = values
295            .iter()
296            .map(|value| c_string(value))
297            .collect::<Result<Vec<_>>>()?;
298        let mut pointers: Vec<_> = owned.iter().map(|value| value.as_ptr()).collect();
299        // mlx-c copies each string into the vector, so `owned` only needs to outlive this call.
300        let raw =
301            unsafe { mlx_sys::mlx_vector_string_new_data(pointers.as_mut_ptr(), pointers.len()) };
302        if raw.ctx.is_null() {
303            return Err(crate::error::exception_from_status(
304                SUCCESS + 1,
305                "creating an MLX string vector",
306            ));
307        }
308        Ok(Self { raw })
309    }
310}
311
312impl Drop for StringVector {
313    fn drop(&mut self) {
314        let _ = unsafe { mlx_sys::mlx_vector_string_free(self.raw) };
315    }
316}
317
318fn c_string(value: &str) -> Result<CString> {
319    CString::new(value).map_err(|error| Exception::custom(error.to_string()))
320}
321
322#[cfg(test)]
323mod tests {
324    use super::*;
325    use pretty_assertions::assert_eq;
326
327    #[test]
328    fn test_metal_kernel_template_args_and_shapes() {
329        let kernel = MetalKernel::new(
330            "scale_shift",
331            &["a", "b"],
332            &["out1", "out2"],
333            "uint elem = thread_position_in_grid.x;
334             T v = a[elem] * T(SCALE);
335             out1[elem] = SHIFT ? v + b[0] : v;
336             out2[elem] = a_shape[0];",
337            MetalKernelOptions::default(),
338        )
339        .unwrap();
340        let a = &crate::ops::arange::<_, f32>(None, 8, None).unwrap() / 4.0;
341        let b = Array::from_slice(&[0.5_f32], &[1]);
342        // SAFETY: Eight threads read valid inputs and initialize distinct outputs of matching dtypes.
343        let outputs = crate::with_device(crate::Device::gpu(), || unsafe {
344            kernel.apply(
345                &[&a, &b],
346                &[
347                    MetalKernelOutput {
348                        shape: vec![8],
349                        dtype: Dtype::Float32,
350                    },
351                    MetalKernelOutput {
352                        shape: vec![8],
353                        dtype: Dtype::Int32,
354                    },
355                ],
356                (8, 1, 1),
357                (4, 1, 1),
358                MetalKernelApplyOptions {
359                    template_args: vec![
360                        ("T".into(), TemplateArg::Dtype(Dtype::Float32)),
361                        ("SCALE".into(), TemplateArg::Int(3)),
362                        ("SHIFT".into(), TemplateArg::Bool(true)),
363                    ],
364                    ..Default::default()
365                },
366            )
367        })
368        .unwrap();
369
370        assert_eq!(
371            outputs[0].to_vec_exact::<f32>().unwrap(),
372            [0.5, 1.25, 2.0, 2.75, 3.5, 4.25, 5.0, 5.75]
373        );
374        assert_eq!(outputs[1].to_vec_exact::<i32>().unwrap(), [8; 8]);
375    }
376
377    #[test]
378    fn test_metal_kernel_init_value_fills_unwritten_outputs() {
379        let kernel = MetalKernel::new(
380            "touch_half",
381            &["a"],
382            &["out"],
383            "uint e = thread_position_in_grid.x; if (e < 2) out[e] = a[e];",
384            MetalKernelOptions::default(),
385        )
386        .unwrap();
387        let a = Array::from_slice(&[0.0_f32, 0.25, 0.5, 0.75], &[4]);
388        // SAFETY: Threads 0 and 1 copy valid f32 elements; init_value initializes the rest.
389        let out = crate::with_device(crate::Device::gpu(), || unsafe {
390            kernel.apply(
391                &[&a],
392                &[MetalKernelOutput {
393                    shape: vec![4],
394                    dtype: Dtype::Float32,
395                }],
396                (4, 1, 1),
397                (4, 1, 1),
398                MetalKernelApplyOptions {
399                    init_value: Some(-1.0),
400                    ..Default::default()
401                },
402            )
403        })
404        .unwrap();
405        assert_eq!(
406            out[0].to_vec_exact::<f32>().unwrap(),
407            [0.0, 0.25, -1.0, -1.0]
408        );
409    }
410
411    #[test]
412    fn test_metal_kernel_rejects_invalid_arguments() {
413        let error = MetalKernel::new("none", &["a"], &[], "", MetalKernelOptions::default())
414            .expect_err("a kernel needs at least one output");
415        assert!(error.what().contains("at least one output"));
416        assert!(MetalKernel::new("nul", &["a\0"], &["out"], "", Default::default()).is_err());
417
418        let kernel = MetalKernel::new(
419            "copy",
420            &["a"],
421            &["out"],
422            "out[0] = a[0];",
423            Default::default(),
424        )
425        .unwrap();
426        let output = [MetalKernelOutput {
427            shape: vec![1],
428            dtype: Dtype::Float32,
429        }];
430        let a = Array::from_slice(&[1.0_f32], &[1]);
431        let no_inputs: [&Array; 0] = [];
432        // SAFETY: The input count mismatch is rejected before any kernel can execute.
433        let error = unsafe {
434            kernel.apply(
435                &no_inputs,
436                &output,
437                (1, 1, 1),
438                (1, 1, 1),
439                Default::default(),
440            )
441        }
442        .expect_err("input count must match construction");
443        assert!(error.what().contains("`inputs`"));
444
445        let error = crate::with_device(crate::Device::cpu(), || {
446            // SAFETY: One thread copies one f32 between single-element buffers.
447            unsafe { kernel.apply(&[&a], &output, (1, 1, 1), (1, 1, 1), Default::default()) }
448        })
449        .expect_err("custom Metal kernels only run on the GPU");
450        assert!(error.what().contains("[metal_kernel]"));
451    }
452
453    #[test]
454    fn test_metal_kernel_gpu_fallback_and_scopes() {
455        let kernel = MetalKernel::new(
456            "stream_copy",
457            &["a"],
458            &["out"],
459            "out[0] = a[0];",
460            Default::default(),
461        )
462        .unwrap();
463        let a = Array::from_slice(&[3.0_f32], &[1]);
464        let output = [MetalKernelOutput {
465            shape: vec![1],
466            dtype: Dtype::Float32,
467        }];
468        let apply = || {
469            // SAFETY: One thread copies one f32 between single-element buffers.
470            unsafe { kernel.apply(&[&a], &output, (1, 1, 1), (1, 1, 1), Default::default()) }
471        };
472        assert_eq!(apply().unwrap()[0].to_vec_exact::<f32>().unwrap(), [3.0]);
473        let previous = crate::Device::default();
474        crate::Device::set_default(&crate::Device::cpu());
475        let fallback = apply();
476        let gpu_scope = crate::with_device(crate::Device::gpu(), apply);
477        let cpu_scope = crate::with_device(crate::Device::cpu(), apply);
478        let gpu_stream = crate::with_stream(&Stream::gpu(), apply);
479        let cpu_stream = crate::with_stream(&Stream::cpu(), apply);
480        crate::Device::set_default(&previous);
481        for outputs in [fallback, gpu_scope, gpu_stream] {
482            assert_eq!(outputs.unwrap()[0].to_vec_exact::<f32>().unwrap(), [3.0]);
483        }
484        assert!(cpu_scope.is_err());
485        assert!(cpu_stream.is_err());
486    }
487
488    #[test]
489    fn test_metal_kernel_layout_copies_are_opt_in() {
490        let a = Array::from_slice(&[0.0_f32, 1.0, 2.0, 3.0, 4.0, 5.0], &[2, 3]);
491        let transposed = a.transpose().unwrap();
492        for (options, expected) in [
493            (
494                MetalKernelOptions::default(),
495                [0.0, 1.0, 2.0, 3.0, 4.0, 5.0],
496            ),
497            (
498                MetalKernelOptions {
499                    ensure_row_contiguous: true,
500                    ..Default::default()
501                },
502                [0.0, 3.0, 1.0, 4.0, 2.0, 5.0],
503            ),
504        ] {
505            let kernel = MetalKernel::new(
506                "physical_copy",
507                &["a"],
508                &["out"],
509                "uint i = thread_position_in_grid.x; out[i] = a[i];",
510                options,
511            )
512            .unwrap();
513            // SAFETY: Six threads copy distinct f32 elements within both six-element buffers.
514            let outputs = unsafe {
515                kernel.apply(
516                    &[&transposed],
517                    &[MetalKernelOutput {
518                        shape: vec![6],
519                        dtype: Dtype::Float32,
520                    }],
521                    (6, 1, 1),
522                    (6, 1, 1),
523                    Default::default(),
524                )
525            }
526            .unwrap();
527            assert_eq!(outputs[0].to_vec_exact::<f32>().unwrap(), expected);
528        }
529    }
530
531    #[test]
532    fn test_metal_kernel_errors_are_deferred() {
533        for (name, source, size, group) in [
534            ("bad_source", "out[0] = a[0] +;", 1, 1),
535            (
536                "bad_group",
537                "uint i = thread_position_in_grid.x; out[i] = a[i];",
538                4096,
539                4096,
540            ),
541        ] {
542            let kernel =
543                MetalKernel::new(name, &["a"], &["out"], source, Default::default()).unwrap();
544            let input = Array::from_slice(&vec![1.0_f32; size as usize], &[size]);
545            // SAFETY: Invalid source cannot execute; the copy case has one f32 per grid thread.
546            let outputs = unsafe {
547                kernel.apply(
548                    &[&input],
549                    &[MetalKernelOutput {
550                        shape: vec![size],
551                        dtype: Dtype::Float32,
552                    }],
553                    (size, 1, 1),
554                    (group, 1, 1),
555                    Default::default(),
556                )
557            }
558            .expect("source and threadgroup errors are deferred");
559            assert!(outputs[0].eval().is_err());
560        }
561        let kernel = MetalKernel::new(
562            "after_error",
563            &["a"],
564            &["out"],
565            "out[0] = a;",
566            Default::default(),
567        )
568        .unwrap();
569        let input = Array::from_f32(2.0);
570        // SAFETY: One thread reads a scalar and initializes the single f32 output.
571        let outputs = unsafe {
572            kernel.apply(
573                &[&input],
574                &[MetalKernelOutput {
575                    shape: vec![1],
576                    dtype: Dtype::Float32,
577                }],
578                (1, 1, 1),
579                (1, 1, 1),
580                Default::default(),
581            )
582        }
583        .unwrap();
584        assert_eq!(outputs[0].to_vec_exact::<f32>().unwrap(), [2.0]);
585    }
586}
mlx_rs/fast/metal_kernel.rs

mlx_rs/fast/
metal_kernel.rs
