mlx_rs/transforms/compile/
compile_with_state.rs

1//! Compilation of functions with state.
2//!
3//! # Unit tests
4//!
5//! See `mlx-rs/mlx-tests/tests/test_compile.rs` for unit tests.
6
7// TODO: there's plenty boilerplate code here but it's not clear how to reduce it
8
9use std::{cell::RefCell, marker::PhantomData, rc::Rc};
10
11use crate::{
12    error::Exception,
13    transforms::compile::{type_id_to_usize, CompiledState},
14    utils::Updatable,
15    Array,
16};
17
18use super::{update_by_replace_with_ref_to_new_array, Closure, Compiled, Guarded, VectorArray};
19
20/// Similar to [`crate::transforms::compile`] but allows for functions that take
21/// a mutable reference to a state `U`.
22pub fn compile_with_state<F, U, A, O, E>(
23    f: F,
24    shapeless: impl Into<Option<bool>>,
25) -> impl for<'a> FnMut(&mut U, F::Args<'a>) -> Result<O, Exception>
26where
27    F: CompileWithState<U, A, O, E> + Copy + 'static,
28    U: Updatable,
29{
30    let shapeless = shapeless.into().unwrap_or(false);
31    move |state, args| {
32        let mut compiled = f.compile(shapeless);
33        compiled.call_mut(state, args)
34    }
35}
36
37/// A trait for functions that can be compiled with state.
38///
39/// This trait is used to compile a function that takes a mutable reference to a state
40/// and some arguments and returns a result.
41///
42/// # Generic parameters
43///
44/// - `U`: The type of the state.
45/// - `A`: The type of the arguments.
46/// - `O`: The type of the output.
47/// - `E`: The type of the exception.
48pub trait CompileWithState<U, A, O, E> {
49    /// The type of the arguments that the returned closure takes.
50    ///
51    /// This is needed to relax the lifetime requirements of the returned
52    /// closure. Otherwise, the arguments to the returned closure would have to
53    /// live longer than the closure itself.
54    type Args<'a>;
55
56    /// Compile the function.
57    fn compile<'args>(self, shapeless: bool) -> impl CallMutWithState<U, Self::Args<'args>, O, E>;
58}
59
60impl<F, U> CompileWithState<U, &[Array], Vec<Array>, ()> for F
61where
62    F: FnMut(&mut U, &[Array]) -> Vec<Array> + 'static,
63    U: Updatable,
64{
65    type Args<'a> = &'a [Array];
66
67    fn compile<'args>(
68        self,
69        shapeless: bool,
70    ) -> impl CallMutWithState<U, Self::Args<'args>, Vec<Array>, ()> {
71        let id = type_id_to_usize(&self);
72        let state = CompiledState {
73            f: self,
74            shapeless,
75            id,
76        };
77        Compiled {
78            f_marker: PhantomData::<F>,
79            state,
80        }
81    }
82}
83
84impl<F, U> CompileWithState<U, &Array, Array, ()> for F
85where
86    F: FnMut(&mut U, &Array) -> Array + 'static,
87    U: Updatable,
88{
89    type Args<'a> = &'a Array;
90
91    fn compile<'args>(
92        mut self,
93        shapeless: bool,
94    ) -> impl CallMutWithState<U, Self::Args<'args>, Array, ()> {
95        let id = type_id_to_usize(&self);
96        let f = move |state: &mut U, args: &[Array]| -> Vec<Array> {
97            let result = (self)(state, &args[0]);
98            vec![result]
99        };
100        let state = CompiledState { f, shapeless, id };
101        Compiled {
102            f_marker: PhantomData::<F>,
103            state,
104        }
105    }
106}
107
108impl<F, U> CompileWithState<U, (&Array, &Array), Array, ()> for F
109where
110    F: FnMut(&mut U, (&Array, &Array)) -> Array + 'static,
111    U: Updatable,
112{
113    type Args<'a> = (&'a Array, &'a Array);
114
115    fn compile<'args>(
116        mut self,
117        shapeless: bool,
118    ) -> impl CallMutWithState<U, Self::Args<'args>, Array, ()> {
119        let id = type_id_to_usize(&self);
120        let f = move |state: &mut U, args: &[Array]| -> Vec<Array> {
121            let result = (self)(state, (&args[0], &args[1]));
122            vec![result]
123        };
124        let state = CompiledState { f, shapeless, id };
125        Compiled {
126            f_marker: PhantomData::<F>,
127            state,
128        }
129    }
130}
131
132impl<F, U> CompileWithState<U, (&Array, &Array, &Array), Array, ()> for F
133where
134    F: FnMut(&mut U, (&Array, &Array, &Array)) -> Array + 'static,
135    U: Updatable,
136{
137    type Args<'a> = (&'a Array, &'a Array, &'a Array);
138
139    fn compile<'args>(
140        mut self,
141        shapeless: bool,
142    ) -> impl CallMutWithState<U, Self::Args<'args>, Array, ()> {
143        let id = type_id_to_usize(&self);
144        let f = move |state: &mut U, args: &[Array]| -> Vec<Array> {
145            let result = (self)(state, (&args[0], &args[1], &args[2]));
146            vec![result]
147        };
148        let state = CompiledState { f, shapeless, id };
149        Compiled {
150            f_marker: PhantomData::<F>,
151            state,
152        }
153    }
154}
155
156impl<F, U> CompileWithState<U, &[Array], Vec<Array>, Exception> for F
157where
158    F: FnMut(&mut U, &[Array]) -> Result<Vec<Array>, Exception> + 'static,
159    U: Updatable,
160{
161    type Args<'a> = &'a [Array];
162
163    fn compile<'args>(
164        self,
165        shapeless: bool,
166    ) -> impl CallMutWithState<U, Self::Args<'args>, Vec<Array>, Exception> {
167        let id = type_id_to_usize(&self);
168        let state = CompiledState {
169            f: self,
170            shapeless,
171            id,
172        };
173        Compiled {
174            f_marker: PhantomData::<F>,
175            state,
176        }
177    }
178}
179
180impl<F, U> CompileWithState<U, &Array, Array, Exception> for F
181where
182    F: FnMut(&mut U, &Array) -> Result<Array, Exception> + 'static,
183    U: Updatable,
184{
185    type Args<'a> = &'a Array;
186
187    fn compile<'args>(
188        mut self,
189        shapeless: bool,
190    ) -> impl CallMutWithState<U, Self::Args<'args>, Array, Exception> {
191        let id = type_id_to_usize(&self);
192        let f = move |state: &mut U, args: &[Array]| -> Result<Vec<Array>, Exception> {
193            let result = (self)(state, &args[0])?;
194            Ok(vec![result])
195        };
196        let state = CompiledState { f, shapeless, id };
197        Compiled {
198            f_marker: PhantomData::<F>,
199            state,
200        }
201    }
202}
203
204impl<F, U> CompileWithState<U, (&Array, &Array), Array, Exception> for F
205where
206    F: FnMut(&mut U, (&Array, &Array)) -> Result<Array, Exception> + 'static,
207    U: Updatable,
208{
209    type Args<'a> = (&'a Array, &'a Array);
210
211    fn compile<'args>(
212        mut self,
213        shapeless: bool,
214    ) -> impl CallMutWithState<U, Self::Args<'args>, Array, Exception> {
215        let id = type_id_to_usize(&self);
216        let f = move |state: &mut U, args: &[Array]| -> Result<Vec<Array>, Exception> {
217            let result = (self)(state, (&args[0], &args[1]))?;
218            Ok(vec![result])
219        };
220        let state = CompiledState { f, shapeless, id };
221        Compiled {
222            f_marker: PhantomData::<F>,
223            state,
224        }
225    }
226}
227
228impl<F, U> CompileWithState<U, (&Array, &Array, &Array), Array, Exception> for F
229where
230    F: FnMut(&mut U, (&Array, &Array, &Array)) -> Result<Array, Exception> + 'static,
231    U: Updatable,
232{
233    type Args<'a> = (&'a Array, &'a Array, &'a Array);
234
235    fn compile<'args>(
236        mut self,
237        shapeless: bool,
238    ) -> impl CallMutWithState<U, Self::Args<'args>, Array, Exception> {
239        let id = type_id_to_usize(&self);
240        let f = move |state: &mut U, args: &[Array]| -> Result<Vec<Array>, Exception> {
241            let result = (self)(state, (&args[0], &args[1], &args[2]))?;
242            Ok(vec![result])
243        };
244        let state = CompiledState { f, shapeless, id };
245        Compiled {
246            f_marker: PhantomData::<F>,
247            state,
248        }
249    }
250}
251
252/// A trait for functions that can be called with state.
253pub trait CallMutWithState<U, A, O, E> {
254    /// Call the function with the given state and arguments.
255    fn call_mut(&mut self, state: &mut U, args: A) -> Result<O, Exception>;
256}
257
258impl<U, F, G> CallMutWithState<U, &[Array], Vec<Array>, ()> for Compiled<F, G>
259where
260    F: FnMut(&mut U, &[Array]) -> Vec<Array>,
261    G: FnMut(&mut U, &[Array]) -> Vec<Array>,
262    U: Updatable,
263{
264    fn call_mut(&mut self, state: &mut U, args: &[Array]) -> Result<Vec<Array>, Exception> {
265        self.state.retry_call_mut_with_state(state, args)
266    }
267}
268
269impl<U, F, G> CallMutWithState<U, &Array, Array, ()> for Compiled<F, G>
270where
271    F: FnMut(&mut U, &Array) -> Array,
272    G: FnMut(&mut U, &[Array]) -> Vec<Array>,
273    U: Updatable,
274{
275    fn call_mut(&mut self, state: &mut U, args: &Array) -> Result<Array, Exception> {
276        let args = std::slice::from_ref(args);
277        let result = self.state.retry_call_mut_with_state(state, args)?;
278        Ok(result.into_iter().next().unwrap())
279    }
280}
281
282impl<U, F, G> CallMutWithState<U, (&Array, &Array), Array, ()> for Compiled<F, G>
283where
284    F: FnMut(&mut U, (&Array, &Array)) -> Array,
285    G: FnMut(&mut U, &[Array]) -> Vec<Array>,
286    U: Updatable,
287{
288    fn call_mut(&mut self, state: &mut U, args: (&Array, &Array)) -> Result<Array, Exception> {
289        let args = &[args.0, args.1];
290        let result = self.state.retry_call_mut_with_state(state, args)?;
291        Ok(result.into_iter().next().unwrap())
292    }
293}
294
295impl<U, F, G> CallMutWithState<U, (&Array, &Array, &Array), Array, ()> for Compiled<F, G>
296where
297    F: FnMut(&mut U, (&Array, &Array, &Array)) -> Array,
298    G: FnMut(&mut U, &[Array]) -> Vec<Array>,
299    U: Updatable,
300{
301    fn call_mut(
302        &mut self,
303        state: &mut U,
304        args: (&Array, &Array, &Array),
305    ) -> Result<Array, Exception> {
306        let args = &[args.0, args.1, args.2];
307        let result = self.state.retry_call_mut_with_state(state, args)?;
308        Ok(result.into_iter().next().unwrap())
309    }
310}
311
312impl<U, F, G> CallMutWithState<U, &[Array], Vec<Array>, Exception> for Compiled<F, G>
313where
314    F: FnMut(&mut U, &[Array]) -> Result<Vec<Array>, Exception>,
315    G: FnMut(&mut U, &[Array]) -> Result<Vec<Array>, Exception>,
316    U: Updatable,
317{
318    fn call_mut(&mut self, state: &mut U, args: &[Array]) -> Result<Vec<Array>, Exception> {
319        self.state.retry_fallible_call_mut_with_state(state, args)
320    }
321}
322
323impl<U, F, G> CallMutWithState<U, &Array, Array, Exception> for Compiled<F, G>
324where
325    F: FnMut(&mut U, &Array) -> Result<Array, Exception>,
326    G: FnMut(&mut U, &[Array]) -> Result<Vec<Array>, Exception>,
327    U: Updatable,
328{
329    fn call_mut(&mut self, state: &mut U, args: &Array) -> Result<Array, Exception> {
330        let args = std::slice::from_ref(args);
331        let result = self.state.retry_fallible_call_mut_with_state(state, args)?;
332        Ok(result.into_iter().next().unwrap())
333    }
334}
335
336impl<U, F, G> CallMutWithState<U, (&Array, &Array), Array, Exception> for Compiled<F, G>
337where
338    F: FnMut(&mut U, (&Array, &Array)) -> Result<Array, Exception>,
339    G: FnMut(&mut U, &[Array]) -> Result<Vec<Array>, Exception>,
340    U: Updatable,
341{
342    fn call_mut(&mut self, state: &mut U, args: (&Array, &Array)) -> Result<Array, Exception> {
343        let args = &[args.0, args.1];
344        let result = self.state.retry_fallible_call_mut_with_state(state, args)?;
345        Ok(result.into_iter().next().unwrap())
346    }
347}
348
349impl<U, F, G> CallMutWithState<U, (&Array, &Array, &Array), Array, Exception> for Compiled<F, G>
350where
351    F: FnMut(&mut U, (&Array, &Array, &Array)) -> Result<Array, Exception>,
352    G: FnMut(&mut U, &[Array]) -> Result<Vec<Array>, Exception>,
353    U: Updatable,
354{
355    fn call_mut(
356        &mut self,
357        state: &mut U,
358        args: (&Array, &Array, &Array),
359    ) -> Result<Array, Exception> {
360        let args = &[args.0, args.1, args.2];
361        let result = self.state.retry_fallible_call_mut_with_state(state, args)?;
362        Ok(result.into_iter().next().unwrap())
363    }
364}
365
366#[inline]
367fn call_mut_with_state_inner<U>(
368    inner_closure: Closure,
369    fun_id: usize,
370    shapeless: bool,
371    state: Rc<RefCell<&mut U>>,
372    args: &[impl AsRef<Array>],
373) -> crate::error::Result<Vec<Array>>
374where
375    U: Updatable,
376{
377    // note: this will use the cached compile (via the id)
378    // but will be able to re-evaluate with fresh state if needed
379    let compiled = Closure::try_from_op(|res| unsafe {
380        let constants = &[];
381        mlx_sys::mlx_detail_compile(
382            res,
383            inner_closure.as_ptr(),
384            fun_id,
385            shapeless,
386            constants.as_ptr(),
387            0,
388        )
389    })?;
390
391    let inner_inputs_vector = {
392        let borrow = state.borrow();
393        VectorArray::try_from_iter(
394            args.iter()
395                .map(AsRef::as_ref)
396                .chain(borrow.updatable_states()),
397        )?
398    };
399
400    // will compile the function (if needed) and evaluate the
401    // compiled graph
402    let result_vector = VectorArray::try_from_op(|res| unsafe {
403        mlx_sys::mlx_closure_apply(res, compiled.as_ptr(), inner_inputs_vector.as_ptr())
404    })?;
405
406    // number of states may change during the call
407    let state_params_len = state.borrow().updatable_states_len();
408
409    let result_plus_state_output: Vec<Array> = result_vector.try_into_values()?;
410
411    // push the stateOutput into the state
412    let result_plus_state_output_len = result_plus_state_output.len();
413    let suffix_start = result_plus_state_output_len - state_params_len;
414
415    for (s, new_values) in state
416        .borrow_mut()
417        .updatable_states_mut()
418        .into_iter()
419        .zip(result_plus_state_output[suffix_start..].iter())
420    {
421        update_by_replace_with_ref_to_new_array(s, new_values);
422    }
423
424    let result_len = result_plus_state_output.len() - state_params_len;
425    Ok(result_plus_state_output
426        .into_iter()
427        .take(result_len)
428        .collect())
429}
430
431impl<F> CompiledState<F> {
432    fn retry_call_mut_with_state<U>(
433        &mut self,
434        state: &mut U,
435        args: &[impl AsRef<Array>],
436    ) -> Result<Vec<Array>, Exception>
437    where
438        F: FnMut(&mut U, &[Array]) -> Vec<Array>,
439        U: Updatable,
440    {
441        self.call_mut_with_state(state, args).or_else(|_e| {
442            // Somehow the mlx_closure_apply may fail on the first call for
443            // certain types of state with the error message:
444            // "unordered_map::at: key not found", so we just try again.
445            //
446            // One type that is known to cause this is a tuple of
447            // `Module` and `Optimizer` eg. `(<Module>, <Optimizer>)`
448            self.call_mut_with_state(state, args)
449        })
450    }
451
452    fn retry_fallible_call_mut_with_state<U>(
453        &mut self,
454        state: &mut U,
455        args: &[impl AsRef<Array>],
456    ) -> Result<Vec<Array>, Exception>
457    where
458        F: FnMut(&mut U, &[Array]) -> Result<Vec<Array>, Exception>,
459        U: Updatable,
460    {
461        self.fallible_call_mut_with_state(state, args)
462            .or_else(|_e| {
463                // Somehow the mlx_closure_apply may fail on the first call for
464                // certain types of state with the error message:
465                // "unordered_map::at: key not found", so we just try again.
466                //
467                // One type that is known to cause this is a tuple of
468                // `Module` and `Optimizer` eg. `(<Module>, <Optimizer>)`
469                self.fallible_call_mut_with_state(state, args)
470            })
471    }
472
473    fn call_mut_with_state<U>(
474        &mut self,
475        state: &mut U,
476        args: &[impl AsRef<Array>],
477    ) -> Result<Vec<Array>, Exception>
478    where
479        F: FnMut(&mut U, &[Array]) -> Vec<Array>,
480        U: Updatable,
481    {
482        let args_len = args.len();
483        let state = Rc::new(RefCell::new(state));
484        let f = &mut self.f;
485
486        let state_clone = Rc::clone(&state);
487        let inner = move |tracers: &[Array]| -> Vec<Array> {
488            // put the tracers in their appropriate places:
489            // - arguments to the function
490            // - inner state
491
492            let tracer_args = &tracers[..args_len];
493
494            // save a snapshot of the inner state
495            let saved_state_inputs = state_clone
496                .borrow()
497                .updatable_states()
498                .into_iter()
499                .map(|array| (*array).clone())
500                .collect::<Vec<Array>>();
501
502            // replace the inner state with the tracers
503            for (s, tracer) in state_clone
504                .borrow_mut()
505                .updatable_states_mut()
506                .into_iter()
507                .zip(tracers.iter().skip(args_len))
508            {
509                update_by_replace_with_ref_to_new_array(s, tracer);
510            }
511
512            // call the function with the tracer arguments and the state holding tracers
513            let mut result = (f)(*state_clone.borrow_mut(), tracer_args);
514
515            // recapture the state as it may have changed
516            let mut state_output_tracers = state_clone
517                .borrow()
518                .updatable_states()
519                .into_iter()
520                .map(|array| (*array).clone())
521                .collect::<Vec<Array>>();
522
523            // put the original values back in the state
524            for (s, saved) in state_clone
525                .borrow_mut()
526                .updatable_states_mut()
527                .into_iter()
528                .zip(saved_state_inputs)
529            {
530                update_by_replace_with_ref_to_new_array(s, &saved);
531            }
532
533            // return the result of the function and the state
534            result.append(&mut state_output_tracers);
535
536            result
537        };
538
539        let inner_closure = Closure::new(inner);
540        call_mut_with_state_inner(inner_closure, self.id, self.shapeless, state, args)
541    }
542
543    fn fallible_call_mut_with_state<U>(
544        &mut self,
545        state: &mut U,
546        args: &[impl AsRef<Array>],
547    ) -> Result<Vec<Array>, Exception>
548    where
549        F: FnMut(&mut U, &[Array]) -> Result<Vec<Array>, Exception>,
550        U: Updatable,
551    {
552        let args_len = args.len();
553        let state = Rc::new(RefCell::new(state));
554        let f = &mut self.f;
555
556        let state_clone = Rc::clone(&state);
557        let inner = move |tracers: &[Array]| -> Result<Vec<Array>, Exception> {
558            // put the tracers in their appropriate places:
559            // - arguments to the function
560            // - inner state
561
562            let tracer_args = &tracers[..args_len];
563
564            // save a snapshot of the inner state
565            let saved_state_inputs = state_clone
566                .borrow()
567                .updatable_states()
568                .into_iter()
569                .map(|array| (*array).clone())
570                .collect::<Vec<Array>>();
571
572            // replace the inner state with the tracers
573            for (s, tracer) in state_clone
574                .borrow_mut()
575                .updatable_states_mut()
576                .into_iter()
577                .zip(tracers.iter().skip(args_len))
578            {
579                update_by_replace_with_ref_to_new_array(s, tracer);
580            }
581
582            // call the function with the tracer arguments and the state holding tracers
583            let mut result = (f)(*state_clone.borrow_mut(), tracer_args)?;
584
585            // recapture the state as it may have changed
586            let mut state_output_tracers = state_clone
587                .borrow()
588                .updatable_states()
589                .into_iter()
590                .map(|array| (*array).clone())
591                .collect::<Vec<Array>>();
592
593            // put the original values back in the state
594            for (s, saved) in state_clone
595                .borrow_mut()
596                .updatable_states_mut()
597                .into_iter()
598                .zip(saved_state_inputs)
599            {
600                update_by_replace_with_ref_to_new_array(s, &saved);
601            }
602
603            // return the result of the function and the state
604            result.append(&mut state_output_tracers);
605
606            Ok(result)
607        };
608
609        let inner_closure = Closure::new_fallible(inner);
610        call_mut_with_state_inner(inner_closure, self.id, self.shapeless, state, args)
611    }
612}
mlx_rs/transforms/compile/compile_with_state.rs

mlx_rs/transforms/compile/
compile_with_state.rs
