tract_linalg/generic/
reduce.rs

1// Reduce<max> generic implementation
2pub mod max {
3    pub use tract_data::internal::f16;
4
5    reduce_impl_wrap!(
6        f32,
7        SMax4,
8        4,
9        4,
10        (),
11        f32::MIN,
12        fn run(x: &[f32], _: ()) -> f32 {
13            debug_assert!(x.len() % Self::nr() == 0);
14            debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
15            *x.iter().max_by(|a, b| a.total_cmp(b)).unwrap()
16        },
17        fn reduce_two(a: f32, b: f32) -> f32 {
18            a.max(b)
19        }
20    );
21
22    reduce_impl_wrap!(
23        f16,
24        HMax8,
25        8,
26        8,
27        (),
28        f16::MIN,
29        fn run(x: &[f16], _: ()) -> f16 {
30            debug_assert!(x.len() % Self::nr() == 0);
31            debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
32            *x.iter().max_by(|a, b| a.total_cmp(b)).unwrap()
33        },
34        fn reduce_two(a: f16, b: f16) -> f16 {
35            a.max(b)
36        }
37    );
38
39    #[cfg(test)]
40    #[macro_use]
41    pub mod s {
42        crate::max_frame_tests!(true, f32, crate::generic::reduce::max::SMax4);
43    }
44
45    #[cfg(test)]
46    #[macro_use]
47    pub mod h {
48        use super::*;
49        crate::max_frame_tests!(true, f16, crate::generic::reduce::max::HMax8);
50    }
51}
52
53// Reduce<min> generic implementation
54pub mod min {
55    pub use tract_data::internal::f16;
56
57    reduce_impl_wrap!(
58        f32,
59        SMin4,
60        4,
61        4,
62        (),
63        f32::MAX,
64        fn run(x: &[f32], _: ()) -> f32 {
65            debug_assert!(x.len() % Self::nr() == 0);
66            debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
67            *x.iter().min_by(|a, b| a.total_cmp(b)).unwrap()
68        },
69        fn reduce_two(a: f32, b: f32) -> f32 {
70            a.min(b)
71        }
72    );
73
74    #[cfg(test)]
75    #[macro_use]
76    pub mod s {
77        crate::min_frame_tests!(true, f32, crate::generic::reduce::min::SMin4);
78    }
79}
80
81// Reduce<sum> generic implementation
82pub mod sum {
83    use crate::num_traits::Zero;
84    pub use tract_data::internal::f16;
85
86    reduce_impl_wrap!(
87        f32,
88        SSum4,
89        4,
90        4,
91        (),
92        0.0,
93        fn run(x: &[f32], _: ()) -> f32 {
94            debug_assert!(x.len() % Self::nr() == 0);
95            debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
96            x.iter().sum::<f32>()
97        },
98        fn reduce_two(a: f32, b: f32) -> f32 {
99            a + b
100        }
101    );
102
103    reduce_impl_wrap!(
104        f16,
105        HSum8,
106        8,
107        8,
108        (),
109        f16::zero(),
110        fn run(x: &[f16], _: ()) -> f16 {
111            debug_assert!(x.len() % Self::nr() == 0);
112            debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
113            x.iter().sum::<f16>()
114        },
115        fn reduce_two(a: f16, b: f16) -> f16 {
116            a + b
117        }
118    );
119
120    #[cfg(test)]
121    #[macro_use]
122    pub mod s {
123        crate::sum_frame_tests!(true, f32, crate::generic::reduce::sum::SSum4);
124    }
125
126    #[cfg(test)]
127    #[macro_use]
128    pub mod h {
129        use super::*;
130        crate::sum_frame_tests!(true, f16, crate::generic::reduce::sum::HSum8);
131    }
132}
133
134// Softmax generic implementation
135pub mod softmax_l2 {
136    use crate::num_traits::Zero;
137    use tract_data::internal::f16;
138
139    map_reduce_impl_wrap!(
140        f32,
141        SSoftMaxL2,
142        4,
143        4,
144        f32,
145        f32::MIN,
146        0.0,
147        fn run(x: &mut [f32], max: f32) -> f32 {
148            debug_assert!(x.len() % Self::nr() == 0);
149            debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
150            let mut sum = 0.;
151            for v in x.iter_mut() {
152                let y = *v - max;
153                let y = fast_compact_exp_f32(y);
154                *v = y;
155                sum += y;
156            }
157            sum
158        },
159        fn reduce_two(a: f32, b: f32) -> f32 {
160            a + b
161        }
162    );
163
164    map_reduce_impl_wrap!(
165        f16,
166        HSoftMaxL2,
167        8,
168        8,
169        f16,
170        f16::MIN,
171        f16::zero(),
172        fn run(x: &mut [f16], max: f16) -> f16 {
173            debug_assert!(x.len() % Self::nr() == 0);
174            debug_assert!(x.as_ptr() as usize % Self::alignment_bytes() == 0);
175            let mut sum = f16::zero();
176            for v in x.iter_mut() {
177                let y = *v - max;
178                let y = f16::from_f32(fast_compact_exp_f32(y.to_f32()));
179                *v = y;
180                sum += y;
181            }
182            sum
183        },
184        fn reduce_two(a: f16, b: f16) -> f16 {
185            a + b
186        }
187    );
188
189    // ported from https://github.com/gnuradio/volk/blob/master/kernels/volk/volk_32f_expfast_32f.h
190    // probably inspired from https://nic.schraudolph.org/pubs/Schraudolph99.pdf
191    // not that the cast to u32 deals with negative right, while implem in volk code are wrong in some
192    // corner cases (need a max(0,x) before the u32 conversion)
193    pub fn fast_compact_exp_f32(v: f32) -> f32 {
194        const MLN2: f32 = 0.6931471805f32;
195        const A: f32 = 8388608.0f32;
196        const B: f32 = 1065353216.0f32;
197        const C: f32 = 60801.0f32;
198        const SLOPE: f32 = A / MLN2;
199        const OFFSET: f32 = B - C;
200        f32::from_bits(((SLOPE * v) + OFFSET) as u32)
201    }
202
203    #[cfg(test)]
204    #[macro_use]
205    pub mod s {
206        crate::softmax_l2_frame_tests!(true, f32, super::SSoftMaxL2);
207    }
208
209    #[cfg(test)]
210    #[macro_use]
211    pub mod h {
212        use super::*;
213        crate::softmax_l2_frame_tests!(true, f16, HSoftMaxL2);
214    }
215}
tract_linalg/generic/reduce.rs

tract_linalg/generic/
reduce.rs
