hermes_simd_core/view/
select.rs

1//! Conditional SIMD blend/select operations on [`SimdView`].
2//!
3//! `select` performs a lane-wise conditional merge: for each index `i`,
4//! the output is `self[i]` if `mask[i]` is true, else `other[i]`.
5//!
6//! # Architecture mapping
7//!
8//! | Method | Instruction family |
9//! |---|---|
10//! | `blend` (float mask) | AVX-512 `_mm512_mask_blend_ps`, AVX2 `_mm256_blendv_ps`, NEON `vbslq_f32` |
11//! | Scalar fallback | `if mask[i] { self[i] } else { other[i] }` |
12//!
13//! # Zero-Cost Contract
14//!
15//! Selection is monomorphized per `(T, Arch, Align)`. The `Align` ZST governs
16//! which load instruction is emitted; `Arch` is a ZST erased after codegen.
17
18use crate::align::Alignment;
19use crate::arch::SimdArch;
20use crate::execution::ExecutionMode;
21use crate::kernel::SimdKernel;
22use crate::scalar::Scalar;
23use crate::vec::AlignedVec;
24use crate::view::{SimdError, SimdView};
25
26impl<'a, T: 'a, Arch: SimdArch + SimdKernel<T>, Align: Alignment, Mode: ExecutionMode, Ref: 'a>
27    SimdView<'a, T, Arch, Align, Mode, Ref>
28where
29    T: Scalar,
30{
31    /// Lane-wise conditional select: `out[i] = if mask[i] { self[i] } else { other[i] }`.
32    ///
33    /// Allocates one `AlignedVec<T, Align>` of `self.len()` elements.
34    ///
35    /// # Errors
36    ///
37    /// Returns [`SimdError::LengthMismatch`] if `other.len() != self.len()`, or
38    /// [`SimdError::InsufficientOutputLength`] if `mask.len() < self.len()`.
39    ///
40    /// # Implementation
41    ///
42    /// The scalar fallback loop is the authoritative path. Hardware backends override
43    /// `SimdKernel::blend` with the matching intrinsic; the compiler selects the
44    /// appropriate specialization at monomorphization.
45    pub fn select<ORef>(
46        &self,
47        mask: &[bool],
48        other: &SimdView<'_, T, Arch, Align, Mode, ORef>,
49    ) -> Result<AlignedVec<T, Align>, SimdError>
50    where
51        ORef: 'a,
52    {
53        super::check_lengths_equal(self.len(), other.len())?;
54        if mask.len() < self.len() {
55            return Err(SimdError::InsufficientOutputLength);
56        }
57
58        let a = self.as_slice();
59        let b = other.as_slice();
60        let len = a.len();
61
62        let mut out: AlignedVec<T, Align> = AlignedVec::with_capacity(len);
63        // SAFETY: every element is written below.
64        unsafe {
65            out.set_len(len);
66        }
67        let out_slice = out.as_mut_slice();
68
69        let lane_count = Arch::LANE_COUNT;
70        let mut i = 0;
71        unsafe {
72            while i + lane_count <= len {
73                let m = Arch::mask_from_bools(&mask[i..i + lane_count]);
74                let vb = if crate::align::is_aligned_for_arch::<Arch, Align>() {
75                    Arch::load_aligned(b.as_ptr().add(i))
76                } else {
77                    Arch::load_unaligned(b.as_ptr().add(i))
78                };
79                let v_res = Arch::masked_load_unaligned(a.as_ptr().add(i), m, vb);
80                if crate::align::is_aligned_for_arch::<Arch, Align>() {
81                    Arch::store_aligned(out_slice.as_mut_ptr().add(i), v_res);
82                } else {
83                    Arch::store_unaligned(out_slice.as_mut_ptr().add(i), v_res);
84                }
85                i += lane_count;
86            }
87        }
88        for j in i..len {
89            out_slice[j] = if mask[j] { a[j] } else { b[j] };
90        }
91
92        Ok(out)
93    }
94
95    /// Lane-wise conditional negate: `out[i] = if mask[i] { -self[i] } else { self[i] }`.
96    ///
97    /// Allocates one `AlignedVec<T, Align>`. Mask length must be `≥ self.len()`.
98    ///
99    /// # Errors
100    ///
101    /// Returns [`SimdError::InsufficientOutputLength`] if `mask.len() < self.len()`.
102    pub fn masked_negate(&self, mask: &[bool]) -> Result<AlignedVec<T, Align>, SimdError>
103    where
104        T: core::ops::Neg<Output = T>,
105    {
106        if mask.len() < self.len() {
107            return Err(SimdError::InsufficientOutputLength);
108        }
109
110        let data = self.as_slice();
111        let len = data.len();
112        let mut out: AlignedVec<T, Align> = AlignedVec::with_capacity(len);
113        // SAFETY: every element is written in the loop below.
114        unsafe {
115            out.set_len(len);
116        }
117        let out_slice = out.as_mut_slice();
118
119        let lane_count = Arch::LANE_COUNT;
120        let mut i = 0;
121        unsafe {
122            while i + lane_count <= len {
123                let v = if crate::align::is_aligned_for_arch::<Arch, Align>() {
124                    Arch::load_aligned(data.as_ptr().add(i))
125                } else {
126                    Arch::load_unaligned(data.as_ptr().add(i))
127                };
128                let m = Arch::mask_from_bools(&mask[i..i + lane_count]);
129                let vmask = Arch::mask_to_vector(m);
130                let neg_v = Arch::neg(v);
131                let v_res = Arch::blend(vmask, neg_v, v);
132                if crate::align::is_aligned_for_arch::<Arch, Align>() {
133                    Arch::store_aligned(out_slice.as_mut_ptr().add(i), v_res);
134                } else {
135                    Arch::store_unaligned(out_slice.as_mut_ptr().add(i), v_res);
136                }
137                i += lane_count;
138            }
139        }
140        for j in i..len {
141            out_slice[j] = if mask[j] { -data[j] } else { data[j] };
142        }
143
144        Ok(out)
145    }
146}
hermes_simd_core/view/select.rs

hermes_simd_core/view/
select.rs
