hermes_simd_core/ops/
elementwise.rs

1//! Pairwise lane-wise elementwise operation strategies.
2//!
3//! `ElementOp<T>` is a sealed ZST trait used by `zip_reduce`, `zip_cow`, and
4//! `transform_in_place` to parameterize binary vector operations without code
5//! duplication. All `apply` impls are `#[inline(always)]` — DCE removes unused strategies.
6
7use crate::kernel::SimdKernel;
8use crate::scalar::Scalar;
9
10// ---------------------------------------------------------------------------
11// ElementOp — pairwise lane-wise operation
12// ---------------------------------------------------------------------------
13
14/// Sealed ZST trait for pairwise SIMD elementwise operations.
15///
16/// Used by `zip_reduce`, `zip_cow`, and `transform_in_place` to parameterize binary
17/// vector operations without code duplication.
18///
19/// # Scalar Tail
20///
21/// `apply_scalar(a, b)` handles elements that do not fill a complete SIMD vector.
22/// Implementations use direct `T: Scalar` arithmetic so no vector load/store
23/// boundary conditions apply. The default does NOT exist — every impl must provide
24/// both `apply` (vector) and `apply_scalar` (scalar element).
25pub trait ElementOp<T: Scalar>: crate::private::Sealed + Copy + 'static {
26    /// Apply the operation to two vectors lane-wise.
27    ///
28    /// Takes `self` by value — for ZSTs this is free; for `Clamp` it captures the bounds.
29    ///
30    /// # Safety
31    /// Processor must support the target feature of `Arch`.
32    unsafe fn apply<Arch: SimdKernel<T>>(self, a: Arch::Vector, b: Arch::Vector) -> Arch::Vector;
33
34    /// Apply the operation to two individual scalar elements.
35    ///
36    /// Used for the SIMD tail (elements that do not fill a complete vector).
37    /// Takes `self` by value — for ZSTs this is free; for `Clamp` it captures the bounds.
38    fn apply_scalar(self, a: T, b: T) -> T;
39}
40
41// ---------------------------------------------------------------------------
42// Concrete elementwise ZSTs
43// ---------------------------------------------------------------------------
44
45/// Elementwise multiplication: `a[i] * b[i]`.
46#[derive(Debug, Clone, Copy, PartialEq, Eq)]
47pub struct Mul;
48
49/// Elementwise addition: `a[i] + b[i]`.
50#[derive(Debug, Clone, Copy, PartialEq, Eq)]
51pub struct Add;
52
53/// Elementwise subtraction: `a[i] - b[i]`.
54#[derive(Debug, Clone, Copy, PartialEq, Eq)]
55pub struct Sub;
56
57/// Elementwise division: `a[i] / b[i]`.
58#[derive(Debug, Clone, Copy, PartialEq, Eq)]
59pub struct Div;
60
61/// Elementwise bitwise AND: `a[i] & b[i]`.
62#[derive(Debug, Clone, Copy, PartialEq, Eq)]
63pub struct BitAnd;
64
65/// Elementwise bitwise OR: `a[i] | b[i]`.
66#[derive(Debug, Clone, Copy, PartialEq, Eq)]
67pub struct BitOr;
68
69/// Elementwise bitwise XOR: `a[i] ^ b[i]`.
70#[derive(Debug, Clone, Copy, PartialEq, Eq)]
71pub struct BitXor;
72
73/// Fused-multiply-add elementwise operation: `out[i] = a[i] * b[i] + a[i]` (binary form).
74///
75/// As an `ElementOp`, this interprets the two operand vectors as `a` and `b`, and computes
76/// `fmadd(a, b, zero)` — i.e. `a * b` with hardware FMA precision, accumulating into zero.
77/// To use as a ternary `a*b + c` accumulation, call `Arch::fmadd` directly.
78///
79/// # Zero-Cost Guarantee
80///
81/// `size_of::<FmaAdd>() == 0`. Monomorphization over `Arch` eliminates the ZST entirely;
82/// the call site reduces to a direct `Arch::fmadd` instruction.
83#[derive(Debug, Clone, Copy, PartialEq, Eq)]
84pub struct FmaAdd;
85
86/// Elementwise clamp: `min(max(a[i], lo), hi)`.
87///
88/// Unlike the other strategies, `Clamp` carries its bounds as value fields — the
89/// bounds are `T`-typed and cannot be encoded as const generics for a generic `T`.
90/// The struct is `Copy` and 2×`size_of::<T>()` bytes; the compiler monomorphizes per
91/// `(T, Arch)` pair.
92///
93/// # Usage
94///
95/// Elementwise clamp: `a[i].clamp(lo, hi)`.
96///
97/// For binary elementwise use via `zip_cow`, the second operand is ignored —
98/// bounds are carried in the struct and broadcast-splat at each SIMD iteration.
99/// For unary use via `map_unary`, pass `Clamp { lo, hi }` directly.
100///
101/// Inlined by the optimizer at each monomorphization site: `Clamp<T>` bounds
102/// become register-constant splat values with no indirect reads.
103#[derive(Debug, Clone, Copy, PartialEq)]
104pub struct Clamp<T: Copy> {
105    /// Lower bound (inclusive).
106    pub lo: T,
107    /// Upper bound (inclusive).
108    pub hi: T,
109}
110
111impl<T: Copy> Clamp<T> {
112    /// Construct a new `Clamp` strategy with the given bounds.
113    #[inline(always)]
114    pub fn new(lo: T, hi: T) -> Self {
115        Self { lo, hi }
116    }
117}
118
119// ---------------------------------------------------------------------------
120// Sealing impls
121// ---------------------------------------------------------------------------
122
123impl crate::private::Sealed for Mul {}
124impl crate::private::Sealed for Add {}
125impl crate::private::Sealed for Sub {}
126impl crate::private::Sealed for Div {}
127impl crate::private::Sealed for BitAnd {}
128impl crate::private::Sealed for BitOr {}
129impl crate::private::Sealed for BitXor {}
130impl crate::private::Sealed for FmaAdd {}
131impl<T: Copy + 'static> crate::private::Sealed for Clamp<T> {}
132
133// ---------------------------------------------------------------------------
134// ElementOp impls
135// ---------------------------------------------------------------------------
136
137impl<T: Scalar> ElementOp<T> for Mul {
138    #[inline(always)]
139    unsafe fn apply<Arch: SimdKernel<T>>(self, a: Arch::Vector, b: Arch::Vector) -> Arch::Vector {
140        Arch::mul(a, b)
141    }
142    #[inline(always)]
143    fn apply_scalar(self, a: T, b: T) -> T {
144        a * b
145    }
146}
147
148impl<T: Scalar> ElementOp<T> for Add {
149    #[inline(always)]
150    unsafe fn apply<Arch: SimdKernel<T>>(self, a: Arch::Vector, b: Arch::Vector) -> Arch::Vector {
151        Arch::add(a, b)
152    }
153    #[inline(always)]
154    fn apply_scalar(self, a: T, b: T) -> T {
155        a + b
156    }
157}
158
159impl<T: Scalar> ElementOp<T> for Sub {
160    #[inline(always)]
161    unsafe fn apply<Arch: SimdKernel<T>>(self, a: Arch::Vector, b: Arch::Vector) -> Arch::Vector {
162        Arch::sub(a, b)
163    }
164    #[inline(always)]
165    fn apply_scalar(self, a: T, b: T) -> T {
166        a - b
167    }
168}
169
170impl<T: Scalar> ElementOp<T> for Div {
171    #[inline(always)]
172    unsafe fn apply<Arch: SimdKernel<T>>(self, a: Arch::Vector, b: Arch::Vector) -> Arch::Vector {
173        Arch::div(a, b)
174    }
175    #[inline(always)]
176    fn apply_scalar(self, a: T, b: T) -> T {
177        a / b
178    }
179}
180
181impl<T: Scalar> ElementOp<T> for BitAnd {
182    #[inline(always)]
183    unsafe fn apply<Arch: SimdKernel<T>>(self, a: Arch::Vector, b: Arch::Vector) -> Arch::Vector {
184        Arch::bitand(a, b)
185    }
186    #[inline(always)]
187    fn apply_scalar(self, a: T, b: T) -> T {
188        a.bitand(b)
189    }
190}
191
192impl<T: Scalar> ElementOp<T> for BitOr {
193    #[inline(always)]
194    unsafe fn apply<Arch: SimdKernel<T>>(self, a: Arch::Vector, b: Arch::Vector) -> Arch::Vector {
195        Arch::bitor(a, b)
196    }
197    #[inline(always)]
198    fn apply_scalar(self, a: T, b: T) -> T {
199        a.bitor(b)
200    }
201}
202
203impl<T: Scalar> ElementOp<T> for BitXor {
204    #[inline(always)]
205    unsafe fn apply<Arch: SimdKernel<T>>(self, a: Arch::Vector, b: Arch::Vector) -> Arch::Vector {
206        Arch::bitxor(a, b)
207    }
208    #[inline(always)]
209    fn apply_scalar(self, a: T, b: T) -> T {
210        a.bitxor(b)
211    }
212}
213
214impl<T: Scalar> ElementOp<T> for FmaAdd {
215    /// `fmadd(a, b, zero)` — uses hardware FMA where available.
216    ///
217    /// # Safety
218    /// Processor must support the target feature of `Arch`.
219    #[inline(always)]
220    unsafe fn apply<Arch: SimdKernel<T>>(self, a: Arch::Vector, b: Arch::Vector) -> Arch::Vector {
221        // Accumulate into a zero register: a[i] * b[i] + 0.
222        let zero = Arch::zero();
223        Arch::fmadd(a, b, zero)
224    }
225
226    /// Scalar tail: `a * b` (scalar `mul`; the addend is the implicit zero).
227    #[inline(always)]
228    fn apply_scalar(self, a: T, b: T) -> T {
229        // scalar_fmadd(a, b, 0) — uses T's scalar FMA implementation.
230        a.scalar_fmadd(b, T::ZERO)
231    }
232}
233
234impl<T: Scalar + Copy> ElementOp<T> for Clamp<T> {
235    /// Clamp lanes: `min(max(a_lane, lo_splat), hi_splat)`.
236    ///
237    /// The second operand `_b` is unused — `Clamp` is a unary operation whose
238    /// bounds are carried in the struct. Use `clamp_cow` or `transform_with_clamp`
239    /// to apply this as a true single-operand transform.
240    ///
241    /// The compiler hoists the `splat(lo)` and `splat(hi)` outside the vectorized
242    /// loop because `self` is captured by value in each `zip_cow` iteration.
243    #[inline(always)]
244    unsafe fn apply<Arch: SimdKernel<T>>(self, a: Arch::Vector, _b: Arch::Vector) -> Arch::Vector {
245        let lo_vec = Arch::splat(self.lo);
246        let hi_vec = Arch::splat(self.hi);
247        Arch::min(Arch::max(a, lo_vec), hi_vec)
248    }
249
250    #[inline(always)]
251    fn apply_scalar(self, a: T, _b: T) -> T {
252        // min(max(a, lo), hi) — scalar path for the vector tail.
253        a.max_scalar(self.lo).min_scalar(self.hi)
254    }
255}
hermes_simd_core/ops/elementwise.rs

hermes_simd_core/ops/
elementwise.rs
