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
15macro_rules! reduce_impl_wrap {
16 ($ti: ident, $func: ident, $nr: expr, $alignment_items: expr, $params: ty, $neutral: expr, $run: item, $reduce_two: item) => {
17 paste! {
18 #[derive(Copy, Clone, Debug)]
19 #[allow(non_camel_case_types)]
20 pub struct $func;
21
22 impl crate::frame::reduce::ReduceKer<$ti, $params> for $func {
23 #[inline(always)]
24 fn name() -> &'static str {
25 stringify!($func)
26 }
27 #[inline(always)]
28 fn nr() -> usize {
29 $nr
30 }
31 #[inline(always)]
32 fn alignment_items() -> usize {
33 $alignment_items
34 }
35 #[inline(always)]
36 fn alignment_bytes() -> usize {
37 $alignment_items * std::mem::size_of::<$ti>()
38 }
39 #[inline(always)]
40 fn neutral() -> $ti {
41 $neutral
42 }
43 $run
44 $reduce_two
45 }
46 }
47 };
48}
49
50pub trait Reduce<T, Params = ()>: Send + Sync + Debug + dyn_clone::DynClone
51where
52 Params: Copy + Send + Sync + Debug + 'static + Default,
53 T: Copy + Debug + PartialEq + Send + Sync,
54{
55 fn name(&self) -> &'static str;
56 fn run(&self, vec: &[T]) -> TractResult<T> {
57 self.run_with_params(vec, Params::default())
58 }
59 fn run_with_params(&self, vec: &[T], params: Params) -> TractResult<T>;
60}
61
62dyn_clone::clone_trait_object!(<T, Params> Reduce<T, Params> where T: Copy, Params: Copy);
63
64#[derive(Debug, Clone, new)]
65pub struct ReduceImpl<K, T, Params = ()>
66where
67 T: LADatum,
68 Params: Copy + Send + Sync + Debug + 'static + Default,
69 K: ReduceKer<T, Params> + Clone,
70{
71 phantom: PhantomData<(K, T, Params)>,
72}
73
74impl<K, T, Params> Reduce<T, Params> for ReduceImpl<K, T, Params>
75where
76 T: LADatum,
77 Params: Copy + Send + Sync + Debug + 'static + Default,
78 K: ReduceKer<T, Params> + Clone,
79{
80 fn name(&self) -> &'static str {
81 K::name()
82 }
83
84 fn run_with_params(&self, vec: &[T], params: Params) -> TractResult<T> {
85 reduce_slice_with_alignment(
86 vec,
87 |data| K::run(data, params),
88 K::nr(),
89 K::alignment_bytes(),
90 K::neutral(),
91 K::reduce_two,
92 )
93 }
94}
95
96pub trait ReduceKer<T, Params = ()>:
97 Send + Sync + Debug + dyn_clone::DynClone + Clone + 'static
98where
99 Params: Copy + Send + Sync + Debug + 'static + Default,
100 T: LADatum,
101{
102 fn name() -> &'static str;
103 fn alignment_bytes() -> usize {
104 Self::alignment_items() * T::datum_type().size_of()
105 }
106 fn alignment_items() -> usize;
107 fn nr() -> usize;
108 fn neutral() -> T;
109 fn reduce_two(a: T, b: T) -> T;
110 fn run(vec: &[T], params: Params) -> T;
111 fn red() -> Box<dyn Reduce<T, Params>> {
112 Box::new(ReduceImpl::<Self, T, Params>::new())
113 }
114}
115
116#[allow(unused_macros)]
117macro_rules! map_reduce_impl_wrap {
118 ($ti: ident, $func: ident, $nr: expr, $alignment_items: expr, $params: ty, $map_neutral: expr, $reduce_neutral: expr, $run: item, $reduce_two: item) => {
119 paste! {
120 #[derive(Copy, Clone, Debug)]
121 #[allow(non_camel_case_types)]
122 pub struct $func;
123
124 impl crate::frame::reduce::MapReduceKer<$ti, $params> for $func {
125 #[inline(always)]
126 fn name() -> &'static str {
127 stringify!($func)
128 }
129 #[inline(always)]
130 fn nr() -> usize {
131 $nr
132 }
133 #[inline(always)]
134 fn alignment_items() -> usize {
135 $alignment_items
136 }
137 #[inline(always)]
138 fn alignment_bytes() -> usize {
139 $alignment_items * std::mem::size_of::<$ti>()
140 }
141 #[inline(always)]
142 fn map_neutral() -> $ti {
143 $map_neutral
144 }
145 #[inline(always)]
146 fn reduce_neutral() -> $ti {
147 $reduce_neutral
148 }
149 $run
150 $reduce_two
151 }
152 }
153 };
154}
155
156pub trait MapReduce<T, Params = ()>: Send + Sync + Debug + dyn_clone::DynClone
157where
158 Params: Copy + Send + Sync + Debug + 'static + Default,
159 T: Copy + Debug + PartialEq + Send + Sync,
160{
161 fn name(&self) -> &'static str;
162 fn run(&self, vec: &mut [T]) -> TractResult<T> {
163 self.run_with_params(vec, Params::default())
164 }
165 fn run_with_params(&self, vec: &mut [T], params: Params) -> TractResult<T>;
166}
167
168dyn_clone::clone_trait_object!(<T, Params> MapReduce<T, Params> where T: Copy, Params: Copy);
169
170#[derive(Debug, Clone, new)]
171pub struct MapReduceImpl<K, T, Params = ()>
172where
173 T: LADatum,
174 Params: Copy + Send + Sync + Debug + 'static + Default,
175 K: MapReduceKer<T, Params> + Clone,
176{
177 phantom: PhantomData<(K, T, Params)>,
178}
179
180impl<K, T, Params> MapReduce<T, Params> for MapReduceImpl<K, T, Params>
181where
182 T: LADatum,
183 Params: Copy + Send + Sync + Debug + 'static + Default,
184 K: MapReduceKer<T, Params> + Clone,
185{
186 fn name(&self) -> &'static str {
187 K::name()
188 }
189 fn run_with_params(&self, vec: &mut [T], params: Params) -> TractResult<T> {
190 map_reduce_slice_with_alignment(
191 vec,
192 |data| K::run(data, params),
193 K::nr(),
194 K::alignment_bytes(),
195 K::map_neutral(),
196 K::reduce_neutral(),
197 K::reduce_two,
198 )
199 }
200}
201
202pub trait MapReduceKer<T, Params = ()>:
203 Send + Sync + Debug + dyn_clone::DynClone + Clone + 'static
204where
205 Params: Copy + Send + Sync + Debug + 'static + Default,
206 T: LADatum,
207{
208 fn name() -> &'static str;
209 fn alignment_bytes() -> usize {
210 Self::alignment_items() * T::datum_type().size_of()
211 }
212 fn alignment_items() -> usize;
213 fn nr() -> usize;
214 fn map_neutral() -> T;
215 fn reduce_neutral() -> T;
216 fn reduce_two(a: T, b: T) -> T;
217 fn run(vec: &mut [T], params: Params) -> T;
218 fn red() -> Box<dyn MapReduce<T, Params>> {
219 Box::new(MapReduceImpl::<Self, T, Params>::new())
220 }
221}
222
223#[cfg(test)]
224pub mod test {
225 use super::*;
226 use proptest::test_runner::{TestCaseError, TestCaseResult};
227 use tract_data::internal::*;
228 use tract_data::itertools::Itertools;
229
230 pub fn test_reduce<K: ReduceKer<T, ()>, T: LADatum>(
231 values: &[T],
232 neutral: T,
233 reference_reduce: impl Fn(T, T) -> T,
234 ) -> TestCaseResult {
235 test_reduce_params::<K, T, ()>(values, neutral, reference_reduce, ())
236 }
237
238 pub fn test_reduce_params<K: ReduceKer<T, Params>, T: LADatum, Params>(
239 values: &[T],
240 neutral: T,
241 reference_reducer: impl Fn(T, T) -> T,
242 params: Params,
243 ) -> TestCaseResult
244 where
245 Params: Copy + Send + Sync + Debug + 'static + Default,
246 {
247 crate::setup_test_logger();
248 let op = K::red();
249 let expected = values.iter().fold(neutral, |acc, i| reference_reducer(acc, *i));
250 let found = values;
251 let red = op.run_with_params(found, params).unwrap();
252 tensor0(red)
253 .close_enough(&tensor0(expected), true)
254 .map_err(|e| TestCaseError::fail(e.root_cause().to_string()))?;
255 Ok(())
256 }
257
258 pub fn test_map_reduce<K: MapReduceKer<T, ()>, T: LADatum>(
259 values: &[T],
260 map_neutral: T,
261 neutral: T,
262 reference_map: impl Fn(T) -> T,
263 reference_reduce: impl Fn(T, T) -> T,
264 ) -> TestCaseResult {
265 test_map_reduce_params::<K, T, ()>(
266 values,
267 map_neutral,
268 neutral,
269 reference_map,
270 reference_reduce,
271 (),
272 )
273 }
274
275 pub fn test_map_reduce_params<K: MapReduceKer<T, Params>, T: LADatum, Params>(
276 values: &[T],
277 _neutral: T,
278 map_neutral: T,
279 reference_map: impl Fn(T) -> T,
280 reference_reducer: impl Fn(T, T) -> T,
281 params: Params,
282 ) -> TestCaseResult
283 where
284 Params: Copy + Send + Sync + Debug + 'static + Default,
285 {
286 crate::setup_test_logger();
287 let op = K::red();
288 let mut found = values.to_vec();
289 let expected_values = values.iter().copied().map(reference_map).collect_vec();
290 let expected_reduced =
291 expected_values.iter().fold(map_neutral, |acc, i| reference_reducer(acc, *i));
292 let red = op.run_with_params(&mut found, params).unwrap();
293 tensor1(&found)
294 .close_enough(&tensor1(&expected_values), Approximation::SuperApproximate)
295 .map_err(|e| TestCaseError::fail(e.root_cause().to_string()))?;
296 tensor0(red)
297 .close_enough(&tensor0(expected_reduced), Approximation::SuperApproximate)
298 .map_err(|e| TestCaseError::fail(e.root_cause().to_string()))?;
299 Ok(())
300 }
301}