numeric_array/
simd.rs

1//! Context-sensitive SIMD operations
2//!
3//! These aren't exactly numeric-specific, but they can make use of SIMD instructions.
4
5use core::mem::ManuallyDrop;
6
7use super::*;
8
9use generic_array::internals::ArrayConsumer;
10
11/// Selects elements from one array or another using `self` as a mask.
12pub unsafe trait Select<T, N: ArrayLength> {
13    /// Selects elements from one array or another using `self` as a mask.
14    ///
15    /// Example:
16    ///
17    /// ```ignore
18    /// use numeric_array::simd::Select;
19    ///
20    /// let mask = narr![bool; true, false, false, true];
21    ///
22    /// let a = narr![i32; 1, 2, 3, 4];
23    /// let b = narr![i32; 5, 6, 7, 8];
24    ///
25    /// // Compiles to vblendvps on my machine
26    /// let selected = mask.select(a, b);
27    ///
28    /// assert_eq!(selected, narr![i32; 1, 6, 7, 4]);
29    /// ```
30    ///
31    /// NOTE: The default implementation of this will clamp the index values to the length of the array.
32    fn select(self, true_values: NumericArray<T, N>, false_values: NumericArray<T, N>) -> NumericArray<T, N>;
33}
34
35/// Rearranges one numeric array into another using the supplied indices
36pub unsafe trait Permute<T, N: ArrayLength> {
37    /// Performs the permutation
38    fn permute(self, values: &NumericArray<T, N>) -> NumericArray<T, N>;
39}
40
41unsafe impl<T, I, N: ArrayLength> Permute<T, N> for NumericArray<I, N>
42where
43    T: Clone,
44    I: Copy + Into<usize>,
45{
46    #[inline]
47    fn permute(self, values: &NumericArray<T, N>) -> NumericArray<T, N> {
48        NumericArray::from_iter(self.iter().map(|index| unsafe {
49            let index: usize = (*index).into();
50
51            values.get_unchecked(index.min(N::to_usize())).clone()
52        }))
53    }
54}
55
56unsafe impl<T, N: ArrayLength> Select<T, N> for NumericArray<bool, N> {
57    #[inline]
58    fn select(self, true_values: NumericArray<T, N>, false_values: NumericArray<T, N>) -> NumericArray<T, N> {
59        if core::mem::needs_drop::<T>() {
60            unsafe {
61                let mut true_values = ArrayConsumer::new(true_values.0);
62                let mut false_values = ArrayConsumer::new(false_values.0);
63
64                let (true_values_iter, true_values_position) = true_values.iter_position();
65                let (false_values_iter, false_values_position) = false_values.iter_position();
66
67                NumericArray::from_iter(self.0.iter().zip(true_values_iter.zip(false_values_iter)).map(|(m, (t, f))| {
68                    let t = ptr::read(t);
69                    let f = ptr::read(f);
70
71                    *true_values_position += 1;
72                    *false_values_position = *true_values_position;
73
74                    match *m {
75                        true => t,
76                        false => f,
77                    }
78                }))
79            }
80        } else {
81            let true_values = ManuallyDrop::new(true_values);
82            let mut values = ManuallyDrop::new(false_values);
83
84            for (mask, (v, t)) in self.iter().zip(values.iter_mut().zip(true_values.iter())) {
85                if *mask {
86                    unsafe { ptr::copy_nonoverlapping(t, v, 1) };
87                }
88            }
89
90            ManuallyDrop::into_inner(values)
91        }
92    }
93}
numeric_array/simd.rs

numeric_array/
simd.rs
