mlx_rs/transforms/compile/
compile.rs

1//! Compilation of functions.
2
3// TODO: there's plenty boilerplate code here but it's not clear how to reduce it
4
5use std::marker::PhantomData;
6
7use crate::{error::Exception, Array};
8
9use super::{Closure, Compiled, CompiledState, Guarded, VectorArray};
10
11/// Returns a compiled function that produces the same output as `f`.
12///
13/// Please refer to the [swift binding
14/// documentation](https://swiftpackageindex.com/ml-explore/mlx-swift/main/documentation/mlx/compilation)
15/// for more information.
16pub fn compile<F, A, O, E>(
17    f: F,
18    shapeless: impl Into<Option<bool>>,
19) -> impl for<'a> FnMut(F::Args<'a>) -> Result<O, Exception>
20where
21    F: Compile<A, O, E> + 'static,
22{
23    let shapeless = shapeless.into().unwrap_or(false);
24    let mut compiled = f.compile(shapeless);
25    move |args| compiled.call_mut(args)
26}
27
28/// A trait for functions that can be compiled.
29///
30/// # Generic parameters
31///
32/// - `A`: The type of the array arguments
33/// - `O`: The type of the output
34/// - `E`: The type of the error
35pub trait Compile<A, O, E>: Sized {
36    /// The type of the arguments that the returned closure takes.
37    ///
38    /// This is needed to relax the lifetime requirements of the returned
39    /// closure. Otherwise, the arguments to the returned closure would have to
40    /// live longer than the closure itself.
41    type Args<'a>;
42
43    /// Compiles the function.
44    fn compile(self, shapeless: bool) -> impl for<'args> CallMut<Self::Args<'args>, O, E>;
45}
46
47impl<F> Compile<&[Array], Vec<Array>, ()> for F
48where
49    F: FnMut(&[Array]) -> Vec<Array> + 'static,
50{
51    type Args<'a> = &'a [Array];
52
53    fn compile(
54        self,
55        shapeless: bool,
56    ) -> impl for<'args> CallMut<Self::Args<'args>, Vec<Array>, ()> {
57        let state = CompiledState::new(self, shapeless);
58        Compiled::<F, _> {
59            f_marker: PhantomData,
60            state,
61        }
62    }
63}
64
65impl<F> Compile<&Array, Array, ()> for F
66where
67    F: FnMut(&Array) -> Array + 'static,
68{
69    type Args<'a> = &'a Array;
70
71    fn compile(mut self, shapeless: bool) -> impl for<'args> CallMut<Self::Args<'args>, Array, ()> {
72        let f = move |args: &[Array]| -> Vec<Array> {
73            let result = (self)(&args[0]);
74            vec![result]
75        };
76        let state = CompiledState::new(f, shapeless);
77        Compiled::<F, _> {
78            f_marker: PhantomData,
79            state,
80        }
81    }
82}
83
84impl<F> Compile<(&Array, &Array), Array, ()> for F
85where
86    F: FnMut((&Array, &Array)) -> Array + 'static,
87{
88    type Args<'a> = (&'a Array, &'a Array);
89
90    fn compile(mut self, shapeless: bool) -> impl for<'args> CallMut<Self::Args<'args>, Array, ()> {
91        let f = move |args: &[Array]| -> Vec<Array> {
92            let result = (self)((&args[0], &args[1]));
93            vec![result]
94        };
95        let state = CompiledState::new(f, shapeless);
96        Compiled::<F, _> {
97            f_marker: PhantomData,
98            state,
99        }
100    }
101}
102
103impl<F> Compile<(&Array, &Array, &Array), Array, ()> for F
104where
105    F: FnMut((&Array, &Array, &Array)) -> Array + 'static,
106{
107    type Args<'a> = (&'a Array, &'a Array, &'a Array);
108
109    fn compile(mut self, shapeless: bool) -> impl for<'args> CallMut<Self::Args<'args>, Array, ()> {
110        let f = move |args: &[Array]| -> Vec<Array> {
111            let result = (self)((&args[0], &args[1], &args[2]));
112            vec![result]
113        };
114        let state = CompiledState::new(f, shapeless);
115        Compiled::<F, _> {
116            f_marker: PhantomData,
117            state,
118        }
119    }
120}
121
122impl<F> Compile<&[Array], Vec<Array>, Exception> for F
123where
124    F: FnMut(&[Array]) -> Result<Vec<Array>, Exception> + 'static,
125{
126    type Args<'a> = &'a [Array];
127
128    fn compile(
129        self,
130        shapeless: bool,
131    ) -> impl for<'args> CallMut<Self::Args<'args>, Vec<Array>, Exception> {
132        let state = CompiledState::new(self, shapeless);
133        Compiled::<F, _> {
134            f_marker: PhantomData,
135            state,
136        }
137    }
138}
139
140impl<F> Compile<&Array, Array, Exception> for F
141where
142    F: FnMut(&Array) -> Result<Array, Exception> + 'static,
143{
144    type Args<'a> = &'a Array;
145
146    fn compile(
147        mut self,
148        shapeless: bool,
149    ) -> impl for<'args> CallMut<Self::Args<'args>, Array, Exception> {
150        let f = move |args: &[Array]| -> Result<Vec<Array>, Exception> {
151            let result = (self)(&args[0])?;
152            Ok(vec![result])
153        };
154        let state = CompiledState::new(f, shapeless);
155        Compiled::<F, _> {
156            f_marker: PhantomData,
157            state,
158        }
159    }
160}
161
162impl<F> Compile<(&Array, &Array), Array, Exception> for F
163where
164    F: FnMut((&Array, &Array)) -> Result<Array, Exception> + 'static,
165{
166    type Args<'a> = (&'a Array, &'a Array);
167
168    fn compile(
169        mut self,
170        shapeless: bool,
171    ) -> impl for<'args> CallMut<Self::Args<'args>, Array, Exception> {
172        let f = move |args: &[Array]| -> Result<Vec<Array>, Exception> {
173            let result = (self)((&args[0], &args[1]))?;
174            Ok(vec![result])
175        };
176        let state = CompiledState::new(f, shapeless);
177        Compiled::<F, _> {
178            f_marker: PhantomData,
179            state,
180        }
181    }
182}
183
184impl<F> Compile<(&Array, &Array, &Array), Array, Exception> for F
185where
186    F: FnMut((&Array, &Array, &Array)) -> Result<Array, Exception> + 'static,
187{
188    type Args<'a> = (&'a Array, &'a Array, &'a Array);
189
190    fn compile(
191        mut self,
192        shapeless: bool,
193    ) -> impl for<'args> CallMut<Self::Args<'args>, Array, Exception> {
194        let f = move |args: &[Array]| -> Result<Vec<Array>, Exception> {
195            let result = (self)((&args[0], &args[1], &args[2]))?;
196            Ok(vec![result])
197        };
198        let state = CompiledState::new(f, shapeless);
199        Compiled::<F, _> {
200            f_marker: PhantomData,
201            state,
202        }
203    }
204}
205
206/// A trait for a compiled function that can be called.
207pub trait CallMut<A, O, E> {
208    /// Calls the compiled function with the given arguments.
209    fn call_mut(&mut self, args: A) -> Result<O, Exception>;
210}
211
212impl<'a, F, G> CallMut<&'a [Array], Vec<Array>, ()> for Compiled<F, G>
213where
214    F: FnMut(&[Array]) -> Vec<Array> + 'a,
215    G: FnMut(&[Array]) -> Vec<Array> + 'a,
216{
217    fn call_mut(&mut self, args: &[Array]) -> Result<Vec<Array>, Exception> {
218        self.state.call_mut(args)
219    }
220}
221
222impl<'a, F, G> CallMut<&'a Array, Array, ()> for Compiled<F, G>
223where
224    F: FnMut(&Array) -> Array + 'a,
225    G: FnMut(&[Array]) -> Vec<Array> + 'a,
226{
227    fn call_mut(&mut self, args: &Array) -> Result<Array, Exception> {
228        let args = std::slice::from_ref(args);
229        let result = self.state.call_mut(args)?;
230        Ok(result.into_iter().next().unwrap())
231    }
232}
233
234impl<'a, F, G> CallMut<(&'a Array, &'a Array), Array, ()> for Compiled<F, G>
235where
236    F: FnMut((&Array, &Array)) -> Array + 'a,
237    G: FnMut(&[Array]) -> Vec<Array> + 'a,
238{
239    fn call_mut(&mut self, args: (&Array, &Array)) -> Result<Array, Exception> {
240        let args = &[args.0, args.1];
241        let result = self.state.call_mut(args)?;
242        Ok(result.into_iter().next().unwrap())
243    }
244}
245
246impl<'a, F, G> CallMut<(&'a Array, &'a Array, &'a Array), Array, ()> for Compiled<F, G>
247where
248    F: FnMut((&Array, &Array, &Array)) -> Array + 'a,
249    G: FnMut(&[Array]) -> Vec<Array> + 'a,
250{
251    fn call_mut(&mut self, args: (&Array, &Array, &Array)) -> Result<Array, Exception> {
252        // Is there any way to avoid this shallow clone?
253        let args = &[args.0, args.1, args.2];
254        let result = self.state.call_mut(args)?;
255        Ok(result.into_iter().next().unwrap())
256    }
257}
258
259impl<'a, F, G> CallMut<&'a [Array], Vec<Array>, Exception> for Compiled<F, G>
260where
261    F: FnMut(&[Array]) -> Result<Vec<Array>, Exception> + 'a,
262    G: FnMut(&[Array]) -> Result<Vec<Array>, Exception> + 'a,
263{
264    fn call_mut(&mut self, args: &[Array]) -> Result<Vec<Array>, Exception> {
265        self.state.fallible_call_mut(args)
266    }
267}
268
269impl<'a, F, G> CallMut<&'a Array, Array, Exception> for Compiled<F, G>
270where
271    F: FnMut(&Array) -> Result<Array, Exception> + 'a,
272    G: FnMut(&[Array]) -> Result<Vec<Array>, Exception> + 'a,
273{
274    fn call_mut(&mut self, args: &Array) -> Result<Array, Exception> {
275        let args = &[args];
276        let result = self.state.fallible_call_mut(args)?;
277        Ok(result.into_iter().next().unwrap())
278    }
279}
280
281impl<'a, F, G> CallMut<(&'a Array, &'a Array), Array, Exception> for Compiled<F, G>
282where
283    F: FnMut((&Array, &Array)) -> Result<Array, Exception> + 'a,
284    G: FnMut(&[Array]) -> Result<Vec<Array>, Exception> + 'a,
285{
286    fn call_mut(&mut self, args: (&Array, &Array)) -> Result<Array, Exception> {
287        let args = &[args.0, args.1];
288        let result = self.state.fallible_call_mut(args)?;
289        Ok(result.into_iter().next().unwrap())
290    }
291}
292
293impl<'a, F, G> CallMut<(&'a Array, &'a Array, &'a Array), Array, Exception> for Compiled<F, G>
294where
295    F: FnMut((&Array, &Array, &Array)) -> Result<Array, Exception> + 'a,
296    G: FnMut(&[Array]) -> Result<Vec<Array>, Exception> + 'a,
297{
298    fn call_mut(&mut self, args: (&Array, &Array, &Array)) -> Result<Array, Exception> {
299        let args = &[args.0, args.1, args.2];
300        let result = self.state.fallible_call_mut(args)?;
301        Ok(result.into_iter().next().unwrap())
302    }
303}
304
305#[inline]
306fn call_mut_inner(
307    inner_closure: Closure,
308    fun_id: usize,
309    shapeless: bool,
310    args: &[impl AsRef<Array>],
311) -> crate::error::Result<Vec<Array>> {
312    // note: this will use the cached compile (via the id)
313    // but will be able to re-evaluate with fresh state if needed
314    let compiled = Closure::try_from_op(|res| unsafe {
315        let constants = &[];
316        mlx_sys::mlx_detail_compile(
317            res,
318            inner_closure.as_ptr(),
319            fun_id,
320            shapeless,
321            constants.as_ptr(),
322            0,
323        )
324    })?;
325
326    let inner_inputs_vector = VectorArray::try_from_iter(args.iter())?;
327
328    // will compile the function (if needed) and evaluate the
329    // compiled graph
330    let result_vector = VectorArray::try_from_op(|res| unsafe {
331        mlx_sys::mlx_closure_apply(res, compiled.as_ptr(), inner_inputs_vector.as_ptr())
332    })?;
333    let result_plus_state_output: Vec<Array> = result_vector.try_into_values()?;
334
335    let result_len = result_plus_state_output.len();
336    Ok(result_plus_state_output
337        .into_iter()
338        .take(result_len)
339        .collect())
340}
341
342impl<F> CompiledState<F> {
343    fn call_mut(&mut self, args: &[impl AsRef<Array>]) -> Result<Vec<Array>, Exception>
344    where
345        F: FnMut(&[Array]) -> Vec<Array>,
346    {
347        let inner_closure = Closure::new(&mut self.f);
348
349        call_mut_inner(inner_closure, self.id, self.shapeless, args)
350    }
351
352    fn fallible_call_mut(&mut self, args: &[impl AsRef<Array>]) -> Result<Vec<Array>, Exception>
353    where
354        F: FnMut(&[Array]) -> Result<Vec<Array>, Exception>,
355    {
356        let inner_closure = Closure::new_fallible(&mut self.f);
357
358        call_mut_inner(inner_closure, self.id, self.shapeless, args)
359    }
360}
361
362#[cfg(test)]
363mod tests {
364    use crate::{
365        array,
366        error::Exception,
367        ops::{multiply, ones},
368        test_utils::{assert_array_eq, tolerances},
369        Array,
370    };
371
372    use super::compile;
373
374    fn example_fn_0(x: f32) -> f32 {
375        x + 1.0
376    }
377
378    #[test]
379    fn compile_ids_are_unique_for_live_instances_of_the_same_type() {
380        let first = super::CompiledState::new(example_fn_0, false);
381        let second = super::CompiledState::new(example_fn_0, false);
382
383        assert_ne!(first.id, second.id);
384    }
385
386    #[test]
387    fn cloned_compile_state_gets_a_fresh_monotonic_id() {
388        let first = super::CompiledState::new(example_fn_0, false);
389        let cloned = super::CompiledState::clone(&first);
390
391        assert!(cloned.id > first.id);
392    }
393
394    #[test]
395    fn test_compile() {
396        // This unit test is modified from the mlx-swift codebase
397
398        let f = |inputs: &[Array]| -> Vec<Array> { vec![&inputs[0] * &inputs[1]] };
399        let mut compiled = compile(f, None);
400
401        let i1 = ones::<f32>(&[20, 20]).unwrap();
402        let i2 = ones::<f32>(&[20, 20]).unwrap();
403
404        let args = [i1, i2];
405
406        // evaluate directly
407        let r1 = f(&args).drain(0..1).next().unwrap();
408        // evaluate compiled
409        let r2 = compiled(&args).unwrap().drain(0..1).next().unwrap();
410
411        assert_array_eq(&r1, &r2, tolerances::EXACT.rtol, tolerances::EXACT.atol);
412
413        let r3 = compiled(&args).unwrap().drain(0..1).next().unwrap();
414        assert_array_eq(&r1, &r3, tolerances::EXACT.rtol, tolerances::EXACT.atol);
415    }
416
417    #[test]
418    fn test_compile_with_error() {
419        let f = |inputs: &[Array]| -> Result<Vec<Array>, Exception> {
420            multiply(&inputs[0], &inputs[1]).map(|x| vec![x])
421        };
422
423        // Success case
424        let i1 = ones::<f32>(&[20, 20]).unwrap();
425        let i2 = ones::<f32>(&[20, 20]).unwrap();
426        let args = [i1, i2];
427
428        // evaluate directly
429        let r1 = f(&args).unwrap().drain(0..1).next().unwrap();
430
431        // evaluate compiled
432        let mut compiled = compile(f, None);
433        let r2 = compiled(&args).unwrap().drain(0..1).next().unwrap();
434
435        assert_array_eq(&r1, &r2, tolerances::EXACT.rtol, tolerances::EXACT.atol);
436
437        let r3 = compiled(&args).unwrap().drain(0..1).next().unwrap();
438        assert_array_eq(&r1, &r3, tolerances::EXACT.rtol, tolerances::EXACT.atol);
439
440        // Error case
441        let a = array!([1.0, 2.0, 3.0]);
442        let b = array!([4.0, 5.0]);
443        let args = [a, b];
444
445        // The cache is keyed by function pointer and argument shapes
446        let c = array!([4.0, 5.0, 6.0]);
447        let d = array!([7.0, 8.0]);
448        let another_args = [c, d];
449
450        // evaluate directly
451        let result = f(&args);
452        assert!(result.is_err());
453
454        // evaluate compiled
455        let mut compiled = compile(f, None);
456        let result = compiled(&args);
457        assert!(result.is_err());
458
459        let result = compiled(&args);
460        assert!(result.is_err());
461
462        let result = compiled(&another_args);
463        assert!(result.is_err());
464    }
465
466    #[test]
467    fn test_compile_with_one_arg() {
468        let f = |x: &Array| x * x;
469
470        let i = ones::<f32>(&[20, 20]).unwrap();
471
472        // evaluate directly
473        let r1 = f(&i);
474
475        // evaluate compiled
476        let mut compiled = compile(f, None);
477        let r2 = compiled(&i).unwrap();
478
479        assert_array_eq(&r1, &r2, tolerances::EXACT.rtol, tolerances::EXACT.atol);
480
481        let r3 = compiled(&i).unwrap();
482        assert_array_eq(&r1, &r3, tolerances::EXACT.rtol, tolerances::EXACT.atol);
483    }
484
485    #[test]
486    fn test_compile_with_two_args() {
487        let f = |(x, y): (&Array, &Array)| x * y;
488
489        let i1 = ones::<f32>(&[20, 20]).unwrap();
490        let i2 = ones::<f32>(&[20, 20]).unwrap();
491
492        // evaluate directly
493        let r1 = f((&i1, &i2));
494
495        // evaluate compiled
496        let mut compiled = compile(f, None);
497        let r2 = compiled((&i1, &i2)).unwrap();
498
499        assert_array_eq(&r1, &r2, tolerances::EXACT.rtol, tolerances::EXACT.atol);
500
501        let r3 = compiled((&i1, &i2)).unwrap();
502        assert_array_eq(&r1, &r3, tolerances::EXACT.rtol, tolerances::EXACT.atol);
503    }
504
505    #[test]
506    fn test_compile_with_three_args() {
507        let f = |(x, y, z): (&Array, &Array, &Array)| x * y * z;
508        let mut compiled = compile(f, None);
509
510        let i1 = ones::<f32>(&[20, 20]).unwrap();
511        let i2 = ones::<f32>(&[20, 20]).unwrap();
512        let i3 = ones::<f32>(&[20, 20]).unwrap();
513
514        // evaluate directly
515        let r1 = f((&i1, &i2, &i3));
516
517        // evaluate compiled
518        let r2 = compiled((&i1, &i2, &i3)).unwrap();
519
520        assert_array_eq(&r1, &r2, tolerances::EXACT.rtol, tolerances::EXACT.atol);
521
522        let r3 = compiled((&i1, &i2, &i3)).unwrap();
523        assert_array_eq(&r1, &r3, tolerances::EXACT.rtol, tolerances::EXACT.atol);
524    }
525}
mlx_rs/transforms/compile/compile.rs

mlx_rs/transforms/compile/
compile.rs
