tract_linalg/frame/reduce/
mod.rs

1pub mod max;
2pub mod min;
3pub mod softmax;
4pub mod sum;
5
6use std::fmt::Debug;
7use std::marker::PhantomData;
8
9use tract_data::TractResult;
10
11use crate::LADatum;
12
13use super::element_wise_helper::{map_reduce_slice_with_alignment, reduce_slice_with_alignment};
14
15// A reduction kernel from a `run` body. A leading arch ident is for bodies that are inline
16// arch asm or intrinsics, which will not even compile elsewhere: those builds get
17// signature-matched panic stubs instead, so the kernel struct exists everywhere.
18/// Declare a reduction routine: the kernel, its registry descriptor and its accuracy tests, from
19/// one statement. `op` is what the kernel folds, which gives the descriptor's function, the tests'
20/// reference, the identity to start from and how two answers combine; `isa` says which machines
21/// may run it, and therefore which may test it.
22macro_rules! routine_reduce_rust {
23    (arm; $($rest:tt)*) => { routine_reduce_rust!(@ arm, target_arch = "arm"; $($rest)*); };
24    (aarch64; $($rest:tt)*) => {
25        routine_reduce_rust!(@ aarch64, target_arch = "aarch64"; $($rest)*);
26    };
27    (x86_64; $($rest:tt)*) => {
28        routine_reduce_rust!(@ x86_64, target_arch = "x86_64"; $($rest)*);
29    };
30    (riscv64; $($rest:tt)*) => {
31        routine_reduce_rust!(@ riscv64, target_arch = "riscv64"; $($rest)*);
32    };
33    (wasm32; $($rest:tt)*) => {
34        routine_reduce_rust!(@ wasm32,
35            all(target_arch = "wasm32", target_feature = "simd128"); $($rest)*);
36    };
37    (generic; $($rest:tt)*) => { routine_reduce_rust!(@ generic, all(); $($rest)*); };
38
39    // One arm per operation, each naming the identity it starts from and how two answers combine,
40    // then handing the rest on. That is the only place those two facts are written.
41    (@ $arch:ident, $built:meta; $ti:ident, $ker:ident, $nr:expr, $alignment_items:expr,
42     $run:item, op(Max) $(, isa($($isa:ident),+))?) => {
43        routine_reduce_rust!(@@ $arch, $built; $ti, $ker, $nr, $alignment_items, $run, Max,
44            <$ti>::MIN, fn reduce_two(a: $ti, b: $ti) -> $ti { a.max(b) }
45            $(, isa($($isa),+))?);
46    };
47    (@ $arch:ident, $built:meta; $ti:ident, $ker:ident, $nr:expr, $alignment_items:expr,
48     $run:item, op(Min) $(, isa($($isa:ident),+))?) => {
49        routine_reduce_rust!(@@ $arch, $built; $ti, $ker, $nr, $alignment_items, $run, Min,
50            <$ti>::MAX, fn reduce_two(a: $ti, b: $ti) -> $ti { a.min(b) }
51            $(, isa($($isa),+))?);
52    };
53    (@ $arch:ident, $built:meta; $ti:ident, $ker:ident, $nr:expr, $alignment_items:expr,
54     $run:item, op(Sum) $(, isa($($isa:ident),+))?) => {
55        routine_reduce_rust!(@@ $arch, $built; $ti, $ker, $nr, $alignment_items, $run, Sum,
56            <$ti as num_traits::Zero>::zero(), fn reduce_two(a: $ti, b: $ti) -> $ti { a + b }
57            $(, isa($($isa),+))?);
58    };
59
60    (@@ $arch:ident, $built:meta; $ti:ident, $ker:ident, $nr:expr, $alignment_items:expr,
61     $run:item, $op:ident, $neutral:expr, $fold:item $(, isa($($isa:ident),+))?) => {
62        reduce_kernel!(@ $built; $ti, $ker, $nr, $alignment_items, (), $neutral, $run, $fold);
63        paste! {
64            submit_routine!($arch; [<$ti:upper Reduce>], [<Reduce $op>], $ker $(, isa($($isa),+))?);
65            #[cfg(test)]
66            mod [<test_ $ker:snake>] {
67                use super::*;
68                crate::[<$op:snake _frame_tests>]!(
69                    cfg!($built)
70                        && $crate::isa::IsaReq::ANY
71                            $(.needing(&[$($crate::isa::Isa::$isa),+]))?
72                            .satisfied_by($crate::isa::native()),
73                    $ti,
74                    $ker
75                );
76            }
77        }
78    };
79}
80
81macro_rules! reduce_kernel {
82    (arm; $($rest:tt)*) => { reduce_kernel!(@ target_arch = "arm"; $($rest)*); };
83    (aarch64; $($rest:tt)*) => { reduce_kernel!(@ target_arch = "aarch64"; $($rest)*); };
84    (x86_64; $($rest:tt)*) => { reduce_kernel!(@ target_arch = "x86_64"; $($rest)*); };
85    (riscv64; $($rest:tt)*) => { reduce_kernel!(@ target_arch = "riscv64"; $($rest)*); };
86    (wasm32; $($rest:tt)*) => { reduce_kernel!(@ all(target_arch = "wasm32", target_feature = "simd128"); $($rest)*); };
87
88    (@ $built:meta; $ti:ident, $func:ident, $nr:expr, $alignment_items:expr, $params:ty, $neutral:expr, $run:item, $reduce_two:item) => {
89        #[cfg($built)]
90        reduce_kernel!($ti, $func, $nr, $alignment_items, $params, $neutral, $run, $reduce_two);
91        #[cfg(not($built))]
92        reduce_kernel!($ti, $func, $nr, $alignment_items, $params, $neutral,
93            fn run(_vec: &[$ti], _params: $params) -> $ti {
94                panic!(concat!(stringify!($func), ": kernel not built for this target"))
95            },
96            fn reduce_two(_a: $ti, _b: $ti) -> $ti {
97                panic!(concat!(stringify!($func), ": kernel not built for this target"))
98            }
99        );
100    };
101
102    ($ti: ident, $func: ident, $nr: expr, $alignment_items: expr, $params: ty, $neutral: expr, $run: item, $reduce_two: item) => {
103        paste! {
104            #[derive(Copy, Clone, Debug)]
105            #[allow(non_camel_case_types)]
106            pub struct $func;
107
108            impl crate::frame::reduce::ReduceKer<$ti, $params> for $func {
109                #[inline(always)]
110                fn name() -> &'static str {
111                    stringify!($func)
112                }
113                #[inline(always)]
114                fn nr() -> usize {
115                    $nr
116                }
117                #[inline(always)]
118                fn alignment_items() -> usize {
119                    $alignment_items
120                }
121                #[inline(always)]
122                fn alignment_bytes() -> usize {
123                    $alignment_items * std::mem::size_of::<$ti>()
124                }
125                #[inline(always)]
126                fn neutral() -> $ti {
127                    $neutral
128                }
129                $run
130                $reduce_two
131            }
132        }
133    };
134}
135
136pub trait Reduce<T, Params = ()>: Send + Sync + Debug + dyn_clone::DynClone
137where
138    Params: Copy + Send + Sync + Debug + 'static + Default,
139    T: Copy + Debug + PartialEq + Send + Sync,
140{
141    fn name(&self) -> &'static str;
142    fn run(&self, vec: &[T]) -> TractResult<T> {
143        self.run_with_params(vec, Params::default())
144    }
145    fn run_with_params(&self, vec: &[T], params: Params) -> TractResult<T>;
146}
147
148dyn_clone::clone_trait_object!(<T, Params> Reduce<T, Params> where T: Copy, Params: Copy);
149
150#[derive(Debug, Clone, new)]
151pub struct ReduceImpl<K, T, Params = ()>
152where
153    T: LADatum,
154    Params: Copy + Send + Sync + Debug + 'static + Default,
155    K: ReduceKer<T, Params> + Clone,
156{
157    phantom: PhantomData<(K, T, Params)>,
158}
159
160impl<K, T, Params> Reduce<T, Params> for ReduceImpl<K, T, Params>
161where
162    T: LADatum,
163    Params: Copy + Send + Sync + Debug + 'static + Default,
164    K: ReduceKer<T, Params> + Clone,
165{
166    fn name(&self) -> &'static str {
167        K::name()
168    }
169
170    fn run_with_params(&self, vec: &[T], params: Params) -> TractResult<T> {
171        reduce_slice_with_alignment(
172            vec,
173            |data| K::run(data, params),
174            K::nr(),
175            K::alignment_bytes(),
176            K::neutral(),
177            K::reduce_two,
178        )
179    }
180}
181
182pub trait ReduceKer<T, Params = ()>:
183    Send + Sync + Debug + dyn_clone::DynClone + Clone + 'static
184where
185    Params: Copy + Send + Sync + Debug + 'static + Default,
186    T: LADatum,
187{
188    fn name() -> &'static str;
189    fn alignment_bytes() -> usize {
190        Self::alignment_items() * T::datum_type().size_of()
191    }
192    fn alignment_items() -> usize;
193    fn nr() -> usize;
194    fn neutral() -> T;
195    fn reduce_two(a: T, b: T) -> T;
196    fn run(vec: &[T], params: Params) -> T;
197    fn red() -> Box<dyn Reduce<T, Params>> {
198        Box::new(ReduceImpl::<Self, T, Params>::new())
199    }
200}
201
202#[allow(unused_macros)]
203// A map-reduce kernel from a `run` body, arch ident as in `reduce_kernel!`.
204/// Declare a map-reduction routine: the kernel, its registry descriptor and its accuracy tests,
205/// from one statement. One arm per operation, naming the two identities it starts from and how two
206/// answers combine, which is the only place those follow from the operation.
207macro_rules! routine_map_reduce_rust {
208    (arm; $($rest:tt)*) => { routine_map_reduce_rust!(@ arm, target_arch = "arm"; $($rest)*); };
209    (aarch64; $($rest:tt)*) => {
210        routine_map_reduce_rust!(@ aarch64, target_arch = "aarch64"; $($rest)*);
211    };
212    (x86_64; $($rest:tt)*) => {
213        routine_map_reduce_rust!(@ x86_64, target_arch = "x86_64"; $($rest)*);
214    };
215    (riscv64; $($rest:tt)*) => {
216        routine_map_reduce_rust!(@ riscv64, target_arch = "riscv64"; $($rest)*);
217    };
218    (wasm32; $($rest:tt)*) => {
219        routine_map_reduce_rust!(@ wasm32,
220            all(target_arch = "wasm32", target_feature = "simd128"); $($rest)*);
221    };
222    (generic; $($rest:tt)*) => { routine_map_reduce_rust!(@ generic, all(); $($rest)*); };
223
224    (@ $arch:ident, $built:meta; $ti:ident, $ker:ident, $nr:expr, $alignment_items:expr,
225     $run:item, op(Softmax2) $(, isa($($isa:ident),+))?) => {
226        map_reduce_kernel!(@ $built; $ti, $ker, $nr, $alignment_items, $ti,
227            <$ti>::NEG_INFINITY, <$ti as num_traits::Zero>::zero(), $run,
228            fn reduce_two(a: $ti, b: $ti) -> $ti { a + b });
229        paste! {
230            submit_routine!($arch; [<$ti:upper MapReduce>], Softmax2, $ker $(, isa($($isa),+))?);
231            #[cfg(test)]
232            mod [<test_ $ker:snake>] {
233                use super::*;
234                crate::softmax_l2_frame_tests!(
235                    cfg!($built)
236                        && $crate::isa::IsaReq::ANY
237                            $(.needing(&[$($crate::isa::Isa::$isa),+]))?
238                            .satisfied_by($crate::isa::native()),
239                    $ti,
240                    $ker
241                );
242            }
243        }
244    };
245}
246
247macro_rules! map_reduce_kernel {
248    (arm; $($rest:tt)*) => { map_reduce_kernel!(@ target_arch = "arm"; $($rest)*); };
249    (aarch64; $($rest:tt)*) => { map_reduce_kernel!(@ target_arch = "aarch64"; $($rest)*); };
250    (x86_64; $($rest:tt)*) => { map_reduce_kernel!(@ target_arch = "x86_64"; $($rest)*); };
251    (riscv64; $($rest:tt)*) => { map_reduce_kernel!(@ target_arch = "riscv64"; $($rest)*); };
252    (wasm32; $($rest:tt)*) => { map_reduce_kernel!(@ all(target_arch = "wasm32", target_feature = "simd128"); $($rest)*); };
253
254    (@ $built:meta; $ti:ident, $func:ident, $nr:expr, $alignment_items:expr, $params:ty, $map_neutral:expr, $reduce_neutral:expr, $run:item, $reduce_two:item) => {
255        #[cfg($built)]
256        map_reduce_kernel!($ti, $func, $nr, $alignment_items, $params, $map_neutral, $reduce_neutral, $run, $reduce_two);
257        #[cfg(not($built))]
258        map_reduce_kernel!($ti, $func, $nr, $alignment_items, $params, $map_neutral, $reduce_neutral,
259            fn run(_vec: &mut [$ti], _params: $params) -> $ti {
260                panic!(concat!(stringify!($func), ": kernel not built for this target"))
261            },
262            fn reduce_two(_a: $ti, _b: $ti) -> $ti {
263                panic!(concat!(stringify!($func), ": kernel not built for this target"))
264            }
265        );
266    };
267
268    ($ti: ident, $func: ident, $nr: expr, $alignment_items: expr, $params: ty, $map_neutral: expr, $reduce_neutral: expr, $run: item, $reduce_two: item) => {
269        paste! {
270            #[derive(Copy, Clone, Debug)]
271            #[allow(non_camel_case_types)]
272            pub struct $func;
273
274            impl crate::frame::reduce::MapReduceKer<$ti, $params> for $func {
275                #[inline(always)]
276                fn name() -> &'static str {
277                    stringify!($func)
278                }
279                #[inline(always)]
280                fn nr() -> usize {
281                    $nr
282                }
283                #[inline(always)]
284                fn alignment_items() -> usize {
285                    $alignment_items
286                }
287                #[inline(always)]
288                fn alignment_bytes() -> usize {
289                    $alignment_items * std::mem::size_of::<$ti>()
290                }
291                #[inline(always)]
292                fn map_neutral() -> $ti {
293                    $map_neutral
294                }
295                #[inline(always)]
296                fn reduce_neutral() -> $ti {
297                    $reduce_neutral
298                }
299                $run
300                $reduce_two
301            }
302        }
303    };
304}
305
306pub trait MapReduce<T, Params = ()>: Send + Sync + Debug + dyn_clone::DynClone
307where
308    Params: Copy + Send + Sync + Debug + 'static + Default,
309    T: Copy + Debug + PartialEq + Send + Sync,
310{
311    fn name(&self) -> &'static str;
312    fn run(&self, vec: &mut [T]) -> TractResult<T> {
313        self.run_with_params(vec, Params::default())
314    }
315    fn run_with_params(&self, vec: &mut [T], params: Params) -> TractResult<T>;
316}
317
318dyn_clone::clone_trait_object!(<T, Params> MapReduce<T, Params> where T: Copy, Params: Copy);
319
320#[derive(Debug, Clone, new)]
321pub struct MapReduceImpl<K, T, Params = ()>
322where
323    T: LADatum,
324    Params: Copy + Send + Sync + Debug + 'static + Default,
325    K: MapReduceKer<T, Params> + Clone,
326{
327    phantom: PhantomData<(K, T, Params)>,
328}
329
330impl<K, T, Params> MapReduce<T, Params> for MapReduceImpl<K, T, Params>
331where
332    T: LADatum,
333    Params: Copy + Send + Sync + Debug + 'static + Default,
334    K: MapReduceKer<T, Params> + Clone,
335{
336    fn name(&self) -> &'static str {
337        K::name()
338    }
339    fn run_with_params(&self, vec: &mut [T], params: Params) -> TractResult<T> {
340        map_reduce_slice_with_alignment(
341            vec,
342            |data| K::run(data, params),
343            K::nr(),
344            K::alignment_bytes(),
345            K::map_neutral(),
346            K::reduce_neutral(),
347            K::reduce_two,
348        )
349    }
350}
351
352pub trait MapReduceKer<T, Params = ()>:
353    Send + Sync + Debug + dyn_clone::DynClone + Clone + 'static
354where
355    Params: Copy + Send + Sync + Debug + 'static + Default,
356    T: LADatum,
357{
358    fn name() -> &'static str;
359    fn alignment_bytes() -> usize {
360        Self::alignment_items() * T::datum_type().size_of()
361    }
362    fn alignment_items() -> usize;
363    fn nr() -> usize;
364    fn map_neutral() -> T;
365    fn reduce_neutral() -> T;
366    fn reduce_two(a: T, b: T) -> T;
367    fn run(vec: &mut [T], params: Params) -> T;
368    fn red() -> Box<dyn MapReduce<T, Params>> {
369        Box::new(MapReduceImpl::<Self, T, Params>::new())
370    }
371}
372
373#[cfg(test)]
374pub mod test {
375    use super::*;
376    use proptest::test_runner::{TestCaseError, TestCaseResult};
377    use tract_data::internal::*;
378    use tract_data::itertools::Itertools;
379
380    pub fn test_reduce<K: ReduceKer<T, ()>, T: LADatum>(
381        values: &[T],
382        neutral: T,
383        reference_reduce: impl Fn(T, T) -> T,
384    ) -> TestCaseResult {
385        test_reduce_params::<K, T, ()>(values, neutral, reference_reduce, ())
386    }
387
388    pub fn test_reduce_params<K: ReduceKer<T, Params>, T: LADatum, Params>(
389        values: &[T],
390        neutral: T,
391        reference_reducer: impl Fn(T, T) -> T,
392        params: Params,
393    ) -> TestCaseResult
394    where
395        Params: Copy + Send + Sync + Debug + 'static + Default,
396    {
397        crate::setup_test_logger();
398        let op = K::red();
399        let expected = values.iter().fold(neutral, |acc, i| reference_reducer(acc, *i));
400        let found = values;
401        let red = op.run_with_params(found, params).unwrap();
402        tensor0(red)
403            .close_enough(&tensor0(expected), true)
404            .map_err(|e| TestCaseError::fail(e.root_cause().to_string()))?;
405        Ok(())
406    }
407
408    pub fn test_map_reduce<K: MapReduceKer<T, ()>, T: LADatum>(
409        values: &[T],
410        map_neutral: T,
411        neutral: T,
412        reference_map: impl Fn(T) -> T,
413        reference_reduce: impl Fn(T, T) -> T,
414    ) -> TestCaseResult {
415        test_map_reduce_params::<K, T, ()>(
416            values,
417            map_neutral,
418            neutral,
419            reference_map,
420            reference_reduce,
421            (),
422        )
423    }
424
425    pub fn test_map_reduce_params<K: MapReduceKer<T, Params>, T: LADatum, Params>(
426        values: &[T],
427        _neutral: T,
428        map_neutral: T,
429        reference_map: impl Fn(T) -> T,
430        reference_reducer: impl Fn(T, T) -> T,
431        params: Params,
432    ) -> TestCaseResult
433    where
434        Params: Copy + Send + Sync + Debug + 'static + Default,
435    {
436        crate::setup_test_logger();
437        let op = K::red();
438        let mut found = values.to_vec();
439        let expected_values = values.iter().copied().map(reference_map).collect_vec();
440        let expected_reduced =
441            expected_values.iter().fold(map_neutral, |acc, i| reference_reducer(acc, *i));
442        let red = op.run_with_params(&mut found, params).unwrap();
443        tensor1(&found)
444            .close_enough(&tensor1(&expected_values), Approximation::SuperApproximate)
445            .map_err(|e| TestCaseError::fail(e.root_cause().to_string()))?;
446        tensor0(red)
447            .close_enough(&tensor0(expected_reduced), Approximation::SuperApproximate)
448            .map_err(|e| TestCaseError::fail(e.root_cause().to_string()))?;
449        Ok(())
450    }
451}
tract_linalg/frame/reduce/mod.rs

tract_linalg/frame/reduce/
mod.rs
