rten_vecmath/
sum.rs

1use rten_simd::ops::{BitOps, FloatOps, NumOps};
2use rten_simd::{Isa, Simd, SimdIterable, SimdOp, SimdUnaryOp};
3
4use crate::Exp;
5
6/// Computes the sum of a sequence of numbers.
7///
8/// This is more efficient than `slice.iter().sum()` as it computes multiple
9/// partial sums in parallel using SIMD and then sums across the SIMD lanes at
10/// the end. This will produce very slightly different results because the
11/// additions are happening in a different order.
12pub struct Sum<'a> {
13    input: &'a [f32],
14}
15
16impl<'a> Sum<'a> {
17    pub fn new(input: &'a [f32]) -> Self {
18        Sum { input }
19    }
20}
21
22impl SimdOp for Sum<'_> {
23    type Output = f32;
24
25    #[inline(always)]
26    fn eval<I: Isa>(self, isa: I) -> Self::Output {
27        let ops = isa.f32();
28        let vec_sum = self.input.simd_iter(ops).fold_unroll::<4>(
29            ops.zero(),
30            |sum, x| ops.add(sum, x),
31            |sum, x| ops.add(sum, x),
32        );
33        vec_sum.to_array().into_iter().sum()
34    }
35}
36
37/// Computes the sum of squares of a sequence of numbers.
38///
39/// This is conceptually equivalent to `slice.iter().map(|&x| x * x).sum()` but
40/// more efficient as it computes multiple partial sums in parallel using SIMD
41/// and then sums across the SIMD lanes at the end. This will produce very
42/// slightly different results because the additions are happening in a
43/// different order.
44pub struct SumSquare<'a> {
45    input: &'a [f32],
46}
47
48impl<'a> SumSquare<'a> {
49    pub fn new(input: &'a [f32]) -> Self {
50        SumSquare { input }
51    }
52}
53
54impl SimdOp for SumSquare<'_> {
55    type Output = f32;
56
57    #[inline(always)]
58    fn eval<I: Isa>(self, isa: I) -> Self::Output {
59        let ops = isa.f32();
60        let vec_sum = self.input.simd_iter(ops).fold_unroll::<4>(
61            ops.zero(),
62            |sum, x| ops.mul_add(x, x, sum),
63            |sum, x| ops.add(sum, x),
64        );
65        vec_sum.to_array().into_iter().sum()
66    }
67}
68
69/// Computes the sum of absolute values of a sequence of numbers.
70pub struct SumAbs<'a> {
71    input: &'a [f32],
72}
73
74impl<'a> SumAbs<'a> {
75    pub fn new(input: &'a [f32]) -> Self {
76        SumAbs { input }
77    }
78}
79
80impl SimdOp for SumAbs<'_> {
81    type Output = f32;
82
83    #[inline(always)]
84    fn eval<I: Isa>(self, isa: I) -> Self::Output {
85        let ops = isa.f32();
86        let vec_sum = self.input.simd_iter(ops).fold_unroll::<4>(
87            ops.zero(),
88            |sum, x| ops.add(sum, ops.abs(x)),
89            |sum, x| ops.add(sum, x),
90        );
91        vec_sum.to_array().into_iter().sum()
92    }
93}
94
95/// Compute the sum of squares of input with a bias subtracted.
96///
97/// This is a variant of [`SumSquare`] which subtracts a constant value from each
98/// element before squaring it. A typical use case is to compute the variance of
99/// a sequence, which is defined as `mean((X - x_mean)^2)`.
100pub struct SumSquareSub<'a> {
101    input: &'a [f32],
102    offset: f32,
103}
104
105impl<'a> SumSquareSub<'a> {
106    pub fn new(input: &'a [f32], offset: f32) -> Self {
107        SumSquareSub { input, offset }
108    }
109}
110
111impl SimdOp for SumSquareSub<'_> {
112    type Output = f32;
113
114    #[inline(always)]
115    fn eval<I: Isa>(self, isa: I) -> Self::Output {
116        let ops = isa.f32();
117        let offset_vec = ops.splat(self.offset);
118
119        let vec_sum = self.input.simd_iter(ops).fold_unroll::<4>(
120            ops.zero(),
121            |sum, x| {
122                let x_offset = ops.sub(x, offset_vec);
123                ops.mul_add(x_offset, x_offset, sum)
124            },
125            |sum, x| ops.add(sum, x),
126        );
127
128        vec_sum.to_array().into_iter().sum()
129    }
130}
131
132/// Compute the sum of `exp(x - offset)` over the input.
133pub struct SumExpSub<'a> {
134    input: &'a [f32],
135    offset: f32,
136}
137
138impl<'a> SumExpSub<'a> {
139    pub fn new(input: &'a [f32], offset: f32) -> Self {
140        SumExpSub { input, offset }
141    }
142}
143
144impl SimdOp for SumExpSub<'_> {
145    type Output = f32;
146
147    #[inline(always)]
148    fn eval<I: Isa>(self, isa: I) -> Self::Output {
149        let ops = isa.f32();
150        let offset_vec = ops.splat(self.offset);
151
152        let vec_sum = self.input.simd_iter(ops).fold(ops.zero(), |sum, x| {
153            let exp = Exp::apply(isa, ops.sub(x, offset_vec));
154            ops.add(sum, exp)
155        });
156
157        vec_sum.to_array().into_iter().sum()
158    }
159}
160
161#[cfg(test)]
162mod tests {
163    use crate::ulp::assert_ulp_diff_le;
164
165    use super::{Sum, SumAbs, SumExpSub, SumSquare, SumSquareSub};
166    use rten_simd::SimdOp;
167
168    // Chosen to not be a multiple of vector size, so that tail handling is
169    // exercised.
170    const LEN: usize = 100;
171
172    #[test]
173    fn test_sum() {
174        let xs: Vec<f32> = (0..LEN).map(|i| i as f32 * 0.1).collect();
175        let expected_sum: f64 = xs.iter().map(|x| *x as f64).sum();
176        let sum = Sum::new(&xs).dispatch();
177        assert_ulp_diff_le!(sum, expected_sum as f32, 1.0);
178    }
179
180    #[test]
181    fn test_sum_square() {
182        let xs: Vec<f32> = (0..LEN).map(|i| i as f32 * 0.1).collect();
183        let expected_sum: f64 = xs.iter().copied().map(|x| x as f64 * x as f64).sum();
184        let sum = SumSquare::new(&xs).dispatch();
185        assert_ulp_diff_le!(sum, expected_sum as f32, 2.0);
186    }
187
188    #[test]
189    fn test_sum_abs() {
190        let xs: Vec<f32> = (0..LEN).map(|i| (i as f32 * 0.1) - 5.0).collect();
191        let expected_sum: f64 = xs.iter().map(|x| (*x as f64).abs()).sum();
192        let sum = SumAbs::new(&xs).dispatch();
193        assert_ulp_diff_le!(sum, expected_sum as f32, 2.0);
194    }
195
196    #[test]
197    fn test_sum_square_sub() {
198        let xs: Vec<f32> = (0..LEN).map(|i| i as f32 * 0.1).collect();
199        let mean = xs.iter().sum::<f32>() / xs.len() as f32;
200        let expected_sum: f64 = xs
201            .iter()
202            .copied()
203            .map(|x| (x as f64 - mean as f64) * (x as f64 - mean as f64))
204            .sum();
205        let sum = SumSquareSub::new(&xs, mean).dispatch();
206        assert_ulp_diff_le!(sum, expected_sum as f32, 2.0);
207    }
208
209    #[test]
210    fn test_sum_exp_sub() {
211        let xs: Vec<f32> = (0..LEN).map(|i| i as f32 * 0.1).collect();
212        let max = xs.iter().copied().fold(f32::NEG_INFINITY, f32::max);
213        let expected_sum: f64 = xs
214            .iter()
215            .copied()
216            .map(|x| (x as f64 - max as f64).exp())
217            .sum();
218        let sum = SumExpSub::new(&xs, max).dispatch();
219        assert_ulp_diff_le!(sum, expected_sum as f32, 2.0);
220    }
221}
rten_vecmath/sum.rs

rten_vecmath/
sum.rs
