rten_vecmath/
min_max.rs

1use rten_simd::ops::{BitOps, NumOps};
2use rten_simd::{Isa, Simd, SimdIterable, SimdOp};
3
4/// Compute the minimum and maximum values in a slice of floats.
5///
6/// For an empty slice, returns `(+infinity, -infinity)` — the identity values
7/// for min and max on the extended reals.
8pub struct MinMax<'a> {
9    input: &'a [f32],
10}
11
12impl<'a> MinMax<'a> {
13    pub fn new(input: &'a [f32]) -> Self {
14        MinMax { input }
15    }
16}
17
18impl SimdOp for MinMax<'_> {
19    type Output = (f32, f32);
20
21    #[inline(always)]
22    fn eval<I: Isa>(self, isa: I) -> Self::Output {
23        let ops = isa.f32();
24        let [vec_min, vec_max] = self.input.simd_iter(ops).fold_n_unroll::<2, 4>(
25            [ops.splat(f32::INFINITY), ops.splat(f32::NEG_INFINITY)],
26            #[inline(always)]
27            |[min, max], x| [ops.min(x, min), ops.max(x, max)],
28            #[inline(always)]
29            |[min_a, max_a], [min_b, max_b]| [ops.min(min_a, min_b), ops.max(max_a, max_b)],
30        );
31        let min = vec_min
32            .to_array()
33            .as_ref()
34            .iter()
35            .fold(f32::INFINITY, |min, x| x.min(min));
36        let max = vec_max
37            .to_array()
38            .as_ref()
39            .iter()
40            .fold(f32::NEG_INFINITY, |max, x| x.max(max));
41        (min, max)
42    }
43}
44
45/// Compute the maximum value in a slice, propagating NaNs.
46///
47/// For an empty slice, returns `-infinity` — the identity value for max on
48/// the extended reals.
49pub struct MaxNum<'a, T> {
50    input: &'a [T],
51}
52
53impl<'a, T> MaxNum<'a, T> {
54    pub fn new(input: &'a [T]) -> Self {
55        MaxNum { input }
56    }
57}
58
59impl<'a> SimdOp for MaxNum<'a, f32> {
60    type Output = f32;
61
62    #[inline(always)]
63    fn eval<I: Isa>(self, isa: I) -> Self::Output {
64        let ops = isa.f32();
65
66        let max_num = |max, x| {
67            let not_nan = ops.eq(x, x);
68            let new_max = ops.max(max, x);
69            ops.select(new_max, x, not_nan)
70        };
71
72        let vec_max = self.input.simd_iter(ops).fold_unroll::<2>(
73            ops.splat(f32::NEG_INFINITY),
74            max_num,
75            max_num,
76        );
77
78        vec_max
79            .to_array()
80            .as_ref()
81            .iter()
82            .copied()
83            .fold(f32::NEG_INFINITY, |max, x| {
84                if x.is_nan() {
85                    x
86                } else if max.is_nan() {
87                    max
88                } else {
89                    x.max(max)
90                }
91            })
92    }
93}
94
95/// Compute the minimum value in a slice, propagating NaNs.
96///
97/// For an empty slice, returns `+infinity` — the identity value for min on
98/// the extended reals.
99pub struct MinNum<'a, T> {
100    input: &'a [T],
101}
102
103impl<'a, T> MinNum<'a, T> {
104    pub fn new(input: &'a [T]) -> Self {
105        MinNum { input }
106    }
107}
108
109impl<'a> SimdOp for MinNum<'a, f32> {
110    type Output = f32;
111
112    #[inline(always)]
113    fn eval<I: Isa>(self, isa: I) -> Self::Output {
114        let ops = isa.f32();
115
116        let min_num = |min, x| {
117            let not_nan = ops.eq(x, x);
118            let new_min = ops.min(min, x);
119            ops.select(new_min, x, not_nan)
120        };
121
122        let vec_min = self
123            .input
124            .simd_iter(ops)
125            .fold(ops.splat(f32::INFINITY), min_num);
126
127        vec_min
128            .to_array()
129            .as_ref()
130            .iter()
131            .copied()
132            .fold(f32::INFINITY, |min, x| {
133                if x.is_nan() {
134                    x
135                } else if min.is_nan() {
136                    min
137                } else {
138                    x.min(min)
139                }
140            })
141    }
142}
143
144#[cfg(test)]
145mod tests {
146    use super::{MaxNum, MinMax, MinNum};
147    use rten_simd::SimdOp;
148
149    // Chosen to not be a multiple of vector size, so that tail handling is
150    // exercised.
151    const LEN: usize = 100;
152
153    fn reference_min_max(xs: &[f32]) -> (f32, f32) {
154        let min = xs.iter().fold(f32::MAX, |min, x| x.min(min));
155        let max = xs.iter().fold(f32::MIN, |max, x| x.max(max));
156        (min, max)
157    }
158
159    #[test]
160    fn test_min_max() {
161        let xs: Vec<f32> = (0..LEN).map(|i| i as f32 * 0.1).collect();
162        let expected = reference_min_max(&xs);
163        let min_max = MinMax::new(&xs).dispatch();
164        assert_eq!(min_max, expected);
165    }
166
167    #[test]
168    fn test_max_num() {
169        let xs: Vec<f32> = (0..LEN).map(|i| i as f32 * 0.1).collect();
170        let (_, expected_max) = reference_min_max(&xs);
171        let max = MaxNum::new(&xs).dispatch();
172        assert_eq!(max, expected_max);
173
174        let xs = [0.1, 1.0, 0.2, f32::NAN, 0.4, 0.5, 0.6];
175        let max = MaxNum::new(&xs).dispatch();
176        assert!(max.is_nan());
177    }
178
179    #[test]
180    fn test_min_num() {
181        let xs: Vec<f32> = (0..LEN).map(|i| i as f32 * 0.1).collect();
182        let (expected_min, _) = reference_min_max(&xs);
183        let min = MinNum::new(&xs).dispatch();
184        assert_eq!(min, expected_min);
185
186        let xs = [0.1, 1.0, 0.2, f32::NAN, 0.4, 0.5, 0.6];
187        let min = MinNum::new(&xs).dispatch();
188        assert!(min.is_nan());
189    }
190
191    // For an empty slice, min and max return their identity values on the
192    // extended reals (+/-infinity). This matches the ONNX spec for ReduceMin
193    // and ReduceMax.
194    #[test]
195    fn test_empty() {
196        let xs: [f32; 0] = [];
197        assert_eq!(MinNum::new(&xs).dispatch(), f32::INFINITY);
198        assert_eq!(MaxNum::new(&xs).dispatch(), f32::NEG_INFINITY);
199        assert_eq!(
200            MinMax::new(&xs).dispatch(),
201            (f32::INFINITY, f32::NEG_INFINITY)
202        );
203    }
204}
rten_vecmath/min_max.rs

rten_vecmath/
min_max.rs
