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    for (s, new_values) in state
415        .borrow_mut()
416        .updatable_states_mut()
417        .into_iter()
418        .zip(result_plus_state_output[suffix_start..].iter())
419    {
420        update_by_replace_with_ref_to_new_array(s, new_values);
421    }
422
423    let result_len = result_plus_state_output.len() - state_params_len;
424    Ok(result_plus_state_output
425        .into_iter()
426        .take(result_len)
427        .collect())
428}
429
430impl<F> CompiledState<F> {
431    fn retry_call_mut_with_state<U>(
432        &mut self,
433        state: &mut U,
434        args: &[impl AsRef<Array>],
435    ) -> Result<Vec<Array>, Exception>
436    where
437        F: FnMut(&mut U, &[Array]) -> Vec<Array>,
438        U: Updatable,
439    {
440        self.call_mut_with_state(state, args).or_else(|_e| {
441            // Somehow the mlx_closure_apply may fail on the first call for
442            // certain types of state with the error message:
443            // "unordered_map::at: key not found", so we just try again.
444            //
445            // One type that is known to cause this is a tuple of
446            // `Module` and `Optimizer` eg. `(<Module>, <Optimizer>)`
447            self.call_mut_with_state(state, args)
448        })
449    }
450
451    fn retry_fallible_call_mut_with_state<U>(
452        &mut self,
453        state: &mut U,
454        args: &[impl AsRef<Array>],
455    ) -> Result<Vec<Array>, Exception>
456    where
457        F: FnMut(&mut U, &[Array]) -> Result<Vec<Array>, Exception>,
458        U: Updatable,
459    {
460        self.fallible_call_mut_with_state(state, args)
461            .or_else(|_e| {
462                // Somehow the mlx_closure_apply may fail on the first call for
463                // certain types of state with the error message:
464                // "unordered_map::at: key not found", so we just try again.
465                //
466                // One type that is known to cause this is a tuple of
467                // `Module` and `Optimizer` eg. `(<Module>, <Optimizer>)`
468                self.fallible_call_mut_with_state(state, args)
469            })
470    }
471
472    fn call_mut_with_state<U>(
473        &mut self,
474        state: &mut U,
475        args: &[impl AsRef<Array>],
476    ) -> Result<Vec<Array>, Exception>
477    where
478        F: FnMut(&mut U, &[Array]) -> Vec<Array>,
479        U: Updatable,
480    {
481        let args_len = args.len();
482        let state = Rc::new(RefCell::new(state));
483        let f = &mut self.f;
484
485        let state_clone = Rc::clone(&state);
486        let inner = move |tracers: &[Array]| -> Vec<Array> {
487            // put the tracers in their appropriate places:
488            // - arguments to the function
489            // - inner state
490
491            let tracer_args = &tracers[..args_len];
492
493            // save a snapshot of the inner state
494            let saved_state_inputs = state_clone
495                .borrow()
496                .updatable_states()
497                .into_iter()
498                .map(|array| (*array).clone())
499                .collect::<Vec<Array>>();
500
501            // replace the inner state with the tracers
502            for (s, tracer) in state_clone
503                .borrow_mut()
504                .updatable_states_mut()
505                .into_iter()
506                .zip(tracers.iter().skip(args_len))
507            {
508                update_by_replace_with_ref_to_new_array(s, tracer);
509            }
510
511            // call the function with the tracer arguments and the state holding tracers
512            let mut result = (f)(*state_clone.borrow_mut(), tracer_args);
513
514            // recapture the state as it may have changed
515            let mut state_output_tracers = state_clone
516                .borrow()
517                .updatable_states()
518                .into_iter()
519                .map(|array| (*array).clone())
520                .collect::<Vec<Array>>();
521
522            // put the original values back in the state
523            for (s, saved) in state_clone
524                .borrow_mut()
525                .updatable_states_mut()
526                .into_iter()
527                .zip(saved_state_inputs)
528            {
529                update_by_replace_with_ref_to_new_array(s, &saved);
530            }
531
532            // return the result of the function and the state
533            result.append(&mut state_output_tracers);
534
535            result
536        };
537
538        let inner_closure = Closure::new(inner);
539        call_mut_with_state_inner(inner_closure, self.id, self.shapeless, state, args)
540    }
541
542    fn fallible_call_mut_with_state<U>(
543        &mut self,
544        state: &mut U,
545        args: &[impl AsRef<Array>],
546    ) -> Result<Vec<Array>, Exception>
547    where
548        F: FnMut(&mut U, &[Array]) -> Result<Vec<Array>, Exception>,
549        U: Updatable,
550    {
551        let args_len = args.len();
552        let state = Rc::new(RefCell::new(state));
553        let f = &mut self.f;
554
555        let state_clone = Rc::clone(&state);
556        let inner = move |tracers: &[Array]| -> Result<Vec<Array>, Exception> {
557            // put the tracers in their appropriate places:
558            // - arguments to the function
559            // - inner state
560
561            let tracer_args = &tracers[..args_len];
562
563            // save a snapshot of the inner state
564            let saved_state_inputs = state_clone
565                .borrow()
566                .updatable_states()
567                .into_iter()
568                .map(|array| (*array).clone())
569                .collect::<Vec<Array>>();
570
571            // replace the inner state with the tracers
572            for (s, tracer) in state_clone
573                .borrow_mut()
574                .updatable_states_mut()
575                .into_iter()
576                .zip(tracers.iter().skip(args_len))
577            {
578                update_by_replace_with_ref_to_new_array(s, tracer);
579            }
580
581            // call the function with the tracer arguments and the state holding tracers
582            let mut result = (f)(*state_clone.borrow_mut(), tracer_args)?;
583
584            // recapture the state as it may have changed
585            let mut state_output_tracers = state_clone
586                .borrow()
587                .updatable_states()
588                .into_iter()
589                .map(|array| (*array).clone())
590                .collect::<Vec<Array>>();
591
592            // put the original values back in the state
593            for (s, saved) in state_clone
594                .borrow_mut()
595                .updatable_states_mut()
596                .into_iter()
597                .zip(saved_state_inputs)
598            {
599                update_by_replace_with_ref_to_new_array(s, &saved);
600            }
601
602            // return the result of the function and the state
603            result.append(&mut state_output_tracers);
604
605            Ok(result)
606        };
607
608        let inner_closure = Closure::new_fallible(inner);
609        call_mut_with_state_inner(inner_closure, self.id, self.shapeless, state, args)
610    }
611}
mlx_rs/transforms/compile/compile_with_state.rs

mlx_rs/transforms/compile/
compile_with_state.rs
