hermes_simd_core/cow/
ops.rs

1//! Operator overload implementations for Clone-on-Write SIMD containers.
2//!
3//! # Safety
4//!
5//! Two obligations recur here. Kernel calls are `#[target_feature]`-gated, and
6//! that precondition holds by construction: a `SimdCow` exists only for an
7//! architecture the host can execute, since its borrowed form comes from
8//! [`SimdView::new`](crate::view::SimdView::new) and its owned constructors
9//! assert the same condition. The second is local — these routines build their
10//! output buffer with `with_capacity` and write it through a raw pointer,
11//! raising the length only once every element is initialized. That avoids both
12//! a zero-fill of a buffer about to be overwritten and any `&mut [T]` spanning
13//! uninitialized elements, so each such site carries a `SAFETY` comment showing
14//! the write coverage. `gather` and `prefix_scan` reserve capacity and fill it
15//! through the view's `*_into_uninit` methods over
16//! [`AlignedVec::spare_capacity_mut`](crate::vec::AlignedVec::spare_capacity_mut),
17//! then raise the length once those report success, so those paths never zero
18//! the buffer either.
19
20use super::SimdCow;
21use crate::align::Alignment;
22use crate::arch::SimdArch;
23use crate::kernel::SimdKernel;
24use crate::ops::ElementOp;
25use crate::scalar::Scalar;
26use crate::vec::AlignedVec;
27
28// ---------------------------------------------------------------------------
29// Operator Overloads with Allocation Reuse
30// ---------------------------------------------------------------------------
31
32fn binary_lhs_inplace<T, Arch, Align, Op>(
33    lhs: &SimdCow<'_, T, Arch, Align>,
34    rhs: &mut AlignedVec<T, Align>,
35    op: Op,
36) where
37    T: Scalar,
38    Arch: SimdArch + SimdKernel<T>,
39    Align: Alignment,
40    Op: ElementOp<T>,
41{
42    let len = lhs.len();
43    assert_eq!(len, rhs.len(), "SIMD length mismatch");
44
45    let lhs_view = lhs.view();
46    let rhs_view = rhs.view_mut::<Arch>();
47
48    let mut chunks_lhs = lhs_view.simd_chunks();
49    let mut chunks_rhs = rhs_view.simd_chunks_mut();
50
51    for (chunk_lhs, mut chunk_rhs) in (&mut chunks_lhs).zip(&mut chunks_rhs) {
52        unsafe {
53            let va = if crate::align::is_aligned_for_arch::<Arch, Align>() {
54                Arch::load_aligned(chunk_lhs.as_ptr())
55            } else {
56                Arch::load_unaligned(chunk_lhs.as_ptr())
57            };
58            let vb = if crate::align::is_aligned_for_arch::<Arch, Align>() {
59                Arch::load_aligned(chunk_rhs.as_ptr())
60            } else {
61                Arch::load_unaligned(chunk_rhs.as_ptr())
62            };
63            let vr = ElementOp::apply::<Arch>(op, va, vb);
64            if crate::align::is_aligned_for_arch::<Arch, Align>() {
65                Arch::store_aligned(chunk_rhs.as_mut_ptr(), vr);
66            } else {
67                Arch::store_unaligned(chunk_rhs.as_mut_ptr(), vr);
68            }
69        }
70    }
71
72    let tail_lhs = chunks_lhs.remainder();
73    let tail_rhs = chunks_rhs.into_remainder();
74    for (&a, b) in tail_lhs.iter().zip(tail_rhs.iter_mut()) {
75        *b = ElementOp::apply_scalar(op, a, *b);
76    }
77}
78
79macro_rules! impl_binary_op {
80    ($op_trait:ident, $op_method:ident, $op_strategy:ty, $op_val:expr, $is_commutative:expr) => {
81        // 1. SimdCow + SimdCow
82        impl<'a, 'b, T, Arch, Align> core::ops::$op_trait<SimdCow<'b, T, Arch, Align>>
83            for SimdCow<'a, T, Arch, Align>
84        where
85            T: Scalar,
86            Arch: SimdArch + SimdKernel<T>,
87            Align: Alignment,
88        {
89            type Output = SimdCow<'static, T, Arch, Align>;
90
91            #[inline]
92            #[allow(unused_mut)]
93            fn $op_method(self, rhs: SimdCow<'b, T, Arch, Align>) -> Self::Output {
94                match (self, rhs) {
95                    (SimdCow::Owned(lhs_vec), rhs) => {
96                        let mut lhs_cow = SimdCow::Owned(lhs_vec);
97                        lhs_cow
98                            .transform_in_place(&rhs, $op_val)
99                            .expect("SIMD length mismatch");
100                        match lhs_cow {
101                            SimdCow::Owned(v) => SimdCow::Owned(v),
102                            _ => unreachable!(),
103                        }
104                    }
105                    (lhs, SimdCow::Owned(mut rhs_vec)) => {
106                        if $is_commutative {
107                            let mut rhs_cow = SimdCow::Owned(rhs_vec);
108                            rhs_cow
109                                .transform_in_place(&lhs, $op_val)
110                                .expect("SIMD length mismatch");
111                            match rhs_cow {
112                                rhs_cow @ SimdCow::Owned(_) => rhs_cow,
113                                _ => unreachable!(),
114                            }
115                        } else {
116                            binary_lhs_inplace::<T, Arch, Align, $op_strategy>(
117                                &lhs,
118                                &mut rhs_vec,
119                                $op_val,
120                            );
121                            SimdCow::Owned(rhs_vec)
122                        }
123                    }
124                    (lhs, rhs) => lhs.zip_cow(&rhs, $op_val).expect("SIMD length mismatch"),
125                }
126            }
127        }
128
129        // 2. SimdCow + &SimdCow
130        impl<'a, 'b, T, Arch, Align> core::ops::$op_trait<&'b SimdCow<'b, T, Arch, Align>>
131            for SimdCow<'a, T, Arch, Align>
132        where
133            T: Scalar,
134            Arch: SimdArch + SimdKernel<T>,
135            Align: Alignment,
136        {
137            type Output = SimdCow<'static, T, Arch, Align>;
138
139            #[inline]
140            #[allow(unused_mut)]
141            fn $op_method(self, rhs: &'b SimdCow<'b, T, Arch, Align>) -> Self::Output {
142                match self {
143                    SimdCow::Owned(lhs_vec) => {
144                        let mut lhs_cow = SimdCow::Owned(lhs_vec);
145                        lhs_cow
146                            .transform_in_place(rhs, $op_val)
147                            .expect("SIMD length mismatch");
148                        match lhs_cow {
149                            SimdCow::Owned(v) => SimdCow::Owned(v),
150                            _ => unreachable!(),
151                        }
152                    }
153                    lhs => lhs.zip_cow(rhs, $op_val).expect("SIMD length mismatch"),
154                }
155            }
156        }
157
158        // 3. &SimdCow + SimdCow
159        impl<'a, 'b, T, Arch, Align> core::ops::$op_trait<SimdCow<'b, T, Arch, Align>>
160            for &'a SimdCow<'a, T, Arch, Align>
161        where
162            T: Scalar,
163            Arch: SimdArch + SimdKernel<T>,
164            Align: Alignment,
165        {
166            type Output = SimdCow<'static, T, Arch, Align>;
167
168            #[inline]
169            fn $op_method(self, rhs: SimdCow<'b, T, Arch, Align>) -> Self::Output {
170                match rhs {
171                    SimdCow::Owned(mut rhs_vec) => {
172                        if $is_commutative {
173                            let mut rhs_cow = SimdCow::Owned(rhs_vec);
174                            rhs_cow
175                                .transform_in_place(self, $op_val)
176                                .expect("SIMD length mismatch");
177                            match rhs_cow {
178                                rhs_cow @ SimdCow::Owned(_) => rhs_cow,
179                                _ => unreachable!(),
180                            }
181                        } else {
182                            binary_lhs_inplace::<T, Arch, Align, $op_strategy>(
183                                self,
184                                &mut rhs_vec,
185                                $op_val,
186                            );
187                            SimdCow::Owned(rhs_vec)
188                        }
189                    }
190                    rhs => self.zip_cow(&rhs, $op_val).expect("SIMD length mismatch"),
191                }
192            }
193        }
194
195        // 4. &SimdCow + &SimdCow
196        impl<'a, 'b, T, Arch, Align> core::ops::$op_trait<&'b SimdCow<'b, T, Arch, Align>>
197            for &'a SimdCow<'a, T, Arch, Align>
198        where
199            T: Scalar,
200            Arch: SimdArch + SimdKernel<T>,
201            Align: Alignment,
202        {
203            type Output = SimdCow<'static, T, Arch, Align>;
204
205            #[inline]
206            fn $op_method(self, rhs: &'b SimdCow<'b, T, Arch, Align>) -> Self::Output {
207                self.zip_cow(rhs, $op_val).expect("SIMD length mismatch")
208            }
209        }
210    };
211}
212
213macro_rules! impl_assign_op {
214    ($op_trait:ident, $op_method:ident, $op_strategy:ty, $op_val:expr) => {
215        // SimdCow += SimdCow
216        impl<'a, 'b, T, Arch, Align> core::ops::$op_trait<SimdCow<'b, T, Arch, Align>>
217            for SimdCow<'a, T, Arch, Align>
218        where
219            'b: 'a,
220            T: Scalar,
221            Arch: SimdArch + SimdKernel<T>,
222            Align: Alignment,
223        {
224            #[inline]
225            fn $op_method(&mut self, rhs: SimdCow<'b, T, Arch, Align>) {
226                self.transform_in_place(&rhs, $op_val)
227                    .expect("SIMD length mismatch");
228            }
229        }
230
231        // SimdCow += &SimdCow
232        impl<'a, 'b, T, Arch, Align> core::ops::$op_trait<&'b SimdCow<'b, T, Arch, Align>>
233            for SimdCow<'a, T, Arch, Align>
234        where
235            'b: 'a,
236            T: Scalar,
237            Arch: SimdArch + SimdKernel<T>,
238            Align: Alignment,
239        {
240            #[inline]
241            fn $op_method(&mut self, rhs: &'b SimdCow<'b, T, Arch, Align>) {
242                self.transform_in_place(rhs, $op_val)
243                    .expect("SIMD length mismatch");
244            }
245        }
246    };
247}
248
249impl_binary_op!(Add, add, crate::ops::Add, crate::ops::Add, true);
250impl_binary_op!(Sub, sub, crate::ops::Sub, crate::ops::Sub, false);
251impl_binary_op!(Mul, mul, crate::ops::Mul, crate::ops::Mul, true);
252impl_binary_op!(Div, div, crate::ops::Div, crate::ops::Div, false);
253impl_binary_op!(BitAnd, bitand, crate::ops::BitAnd, crate::ops::BitAnd, true);
254impl_binary_op!(BitOr, bitor, crate::ops::BitOr, crate::ops::BitOr, true);
255impl_binary_op!(BitXor, bitxor, crate::ops::BitXor, crate::ops::BitXor, true);
256
257impl_assign_op!(AddAssign, add_assign, crate::ops::Add, crate::ops::Add);
258impl_assign_op!(SubAssign, sub_assign, crate::ops::Sub, crate::ops::Sub);
259impl_assign_op!(MulAssign, mul_assign, crate::ops::Mul, crate::ops::Mul);
260impl_assign_op!(DivAssign, div_assign, crate::ops::Div, crate::ops::Div);
261impl_assign_op!(
262    BitAndAssign,
263    bitand_assign,
264    crate::ops::BitAnd,
265    crate::ops::BitAnd
266);
267impl_assign_op!(
268    BitOrAssign,
269    bitor_assign,
270    crate::ops::BitOr,
271    crate::ops::BitOr
272);
273impl_assign_op!(
274    BitXorAssign,
275    bitxor_assign,
276    crate::ops::BitXor,
277    crate::ops::BitXor
278);
hermes_simd_core/cow/ops.rs

hermes_simd_core/cow/
ops.rs
