hermes_simd/dispatch/
simd_ops.rs

1use super::popcount::{
2    dispatch_reduce_popcount, dispatch_reduce_popcount_and, dispatch_reduce_popcount_or,
3    dispatch_reduce_popcount_xor,
4};
5use super::{
6    abs_reduce, argmax, argmin, axpy, binary, complex, dot, gemm, gemv, gemv_strided,
7    gemv_transpose, gemv_transpose_strided, masked, max, min, scale, sparse, sum,
8};
9use hermes_simd_core::scalar::Scalar as ScalarTrait;
10use hermes_simd_core::sparse::{
11    BlockedCooData, CsrData, DenseWithMaskData, SellPData, ValidatedData,
12};
13use hermes_simd_core::view::SimdError;
14use hermes_simd_core::{Add, Div, Mul, Sub};
15#[cfg(not(any(target_arch = "x86", target_arch = "x86_64", target_arch = "aarch64")))]
16use hermes_simd_intrinsics::Scalar as ScalarArch;
17#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
18#[allow(unused_imports)]
19use hermes_simd_intrinsics::{Avx2, Avx512, Neon, Scalar as ScalarArch};
20#[cfg(target_arch = "aarch64")]
21use hermes_simd_intrinsics::{Neon, Scalar as ScalarArch};
22
23mod private {
24    pub trait Sealed {}
25}
26
27impl private::Sealed for f32 {}
28impl private::Sealed for f64 {}
29impl private::Sealed for i8 {}
30impl private::Sealed for i16 {}
31impl private::Sealed for i32 {}
32
33impl private::Sealed for eunomia::F16 {}
34impl private::Sealed for eunomia::F32 {}
35impl private::Sealed for eunomia::F64 {}
36impl private::Sealed for eunomia::Bf16 {}
37impl private::Sealed for eunomia::Bf8 {}
38impl private::Sealed for eunomia::Bf4 {}
39impl private::Sealed for eunomia::F8 {}
40impl private::Sealed for eunomia::F4 {}
41impl private::Sealed for eunomia::I8 {}
42impl private::Sealed for eunomia::I16 {}
43impl private::Sealed for eunomia::I32 {}
44
45/// Sealed extension trait implementing dynamic runtime SIMD dispatch for any `T: Scalar`.
46pub trait SimdOps: ScalarTrait + private::Sealed {
47    /// Reduces the slice to its sum.
48    fn sum(data: &[Self]) -> Self;
49    /// Reduces the slice to `Σ |x|` (L1-norm accumulator); `T::ZERO` for empty.
50    fn abs_sum(data: &[Self]) -> Self;
51    /// Reduces the slice to `max |x|` (∞-norm accumulator); `T::ZERO` for empty.
52    fn abs_max(data: &[Self]) -> Self;
53    /// Reduces the slice to its minimum element.
54    ///
55    /// Returns `T::MAX_VALUE` for empty slices (the identity element for min).
56    fn min(data: &[Self]) -> Self;
57    /// Reduces the slice to its maximum element.
58    ///
59    /// Returns `T::MIN_VALUE` for empty slices (the identity element for max).
60    fn max(data: &[Self]) -> Self;
61    /// Multiplies every element by `scalar` in-place.
62    fn scale(data: &mut [Self], scalar: Self);
63    /// Returns `Some((index, value))` of the minimum element, or `None` for empty.
64    fn argmin(data: &[Self]) -> Option<(usize, Self)>;
65    /// Returns `Some((index, value))` of the maximum element, or `None` for empty.
66    fn argmax(data: &[Self]) -> Option<(usize, Self)>;
67    /// Computes the dot product of two slices.
68    fn dot(a: &[Self], b: &[Self]) -> Result<Self, SimdError>;
69    /// Fused row update `out[i] += alpha * x[i]` (AXPY) with no temporaries.
70    fn axpy(alpha: Self, x: &[Self], out: &mut [Self]) -> Result<(), SimdError>;
71    /// Fused ternary update `out[i] += alpha * a[i] * b[i]` with no temporary.
72    fn axpy_mul(alpha: Self, a: &[Self], b: &[Self], out: &mut [Self]) -> Result<(), SimdError>;
73    /// Fused multi-row update `out[row, i] += alphas[row] * x[i]`.
74    fn axpy_rows(
75        alphas: &[Self],
76        x: &[Self],
77        out: &mut [Self],
78        row_stride: usize,
79        rows: usize,
80        cols: usize,
81    ) -> Result<(), SimdError>;
82    /// Fused batched multi-row update:
83    /// `out[row, i] += sum_k alphas[k, row] * x_panel[k, i]`.
84    fn axpy_rows_batch(
85        alphas: &[Self],
86        x_panel: &[Self],
87        out: &mut [Self],
88        row_stride: usize,
89        rows: usize,
90        depth: usize,
91        cols: usize,
92    ) -> Result<(), SimdError>;
93    /// Computes the elementwise product and writes to `out`.
94    fn elementwise_mul(a: &[Self], b: &[Self], out: &mut [Self]) -> Result<(), SimdError>;
95    /// Computes the elementwise sum `a[i] + b[i]` and writes to `out`.
96    fn elementwise_add(a: &[Self], b: &[Self], out: &mut [Self]) -> Result<(), SimdError>;
97    /// Computes the elementwise difference `a[i] - b[i]` and writes to `out`.
98    fn elementwise_sub(a: &[Self], b: &[Self], out: &mut [Self]) -> Result<(), SimdError>;
99    /// Computes the elementwise quotient `a[i] / b[i]` and writes to `out`.
100    fn elementwise_div(a: &[Self], b: &[Self], out: &mut [Self]) -> Result<(), SimdError>;
101    /// Computes the sum of elements matching a boolean mask.
102    fn masked_sum(data: &[Self], mask: &[bool]) -> Self;
103    /// Computes the dot product of elements matching a boolean mask.
104    fn masked_dot(a: &[Self], b: &[Self], mask: &[bool]) -> Result<Self, SimdError>;
105    /// Computes the elementwise sum of elements matching a boolean mask.
106    fn masked_add(a: &[Self], b: &[Self], mask: &[bool], out: &mut [Self])
107        -> Result<(), SimdError>;
108    /// Computes sparse SpMV using CSR.
109    fn spmv_csr(data: ValidatedData<CsrData<'_, Self>>, x: &[Self], y: &mut [Self]);
110    /// Computes sparse SpMV using const-generic Blocked-COO tiles.
111    fn spmv_bcoo<const BM: usize, const BN: usize>(
112        data: ValidatedData<BlockedCooData<'_, Self, BM, BN>>,
113        x: &[Self],
114        y: &mut [Self],
115    );
116    /// Computes sparse SpMV using Dense-with-Mask.
117    fn spmv_dense_masked(data: DenseWithMaskData<'_, Self>, x: &[Self], y: &mut [Self]);
118    /// Computes sparse SpMV using const-generic Sliced ELLPACK (SELL-p).
119    fn spmv_sellp<const C: usize>(
120        data: ValidatedData<SellPData<'_, Self, C>>,
121        x: &[Self],
122        y: &mut [Self],
123    );
124    /// Computes register-blocked tiled GEMM: `c += A * B`.
125    fn tiled_gemm(
126        a: &[Self],
127        b: &[Self],
128        c: &mut [Self],
129        m: usize,
130        n: usize,
131        k: usize,
132    ) -> Result<(), SimdError>;
133    /// Computes register-blocked GEMV: `y += A * x` (`A` row-major `nrows × ncols`).
134    fn gemv(
135        a: &[Self],
136        x: &[Self],
137        y: &mut [Self],
138        nrows: usize,
139        ncols: usize,
140    ) -> Result<(), SimdError>;
141    /// Computes register-blocked transposed GEMV: `y += Aᵀ * x`
142    /// (`A` row-major `nrows × ncols`, `x` length `nrows`, `y` length `ncols`).
143    fn gemv_transpose(
144        a: &[Self],
145        x: &[Self],
146        y: &mut [Self],
147        nrows: usize,
148        ncols: usize,
149    ) -> Result<(), SimdError>;
150    /// Computes register-blocked sub-matrix GEMV: `y += A * x` with row stride
151    /// `lda ≥ ncols` (`lda = ncols` is the packed [`Self::gemv`]).
152    fn gemv_strided(
153        a: &[Self],
154        x: &[Self],
155        y: &mut [Self],
156        nrows: usize,
157        ncols: usize,
158        lda: usize,
159    ) -> Result<(), SimdError>;
160    /// Computes register-blocked transposed sub-matrix GEMV: `y += Aᵀ * x` with
161    /// row stride `lda ≥ ncols` (`lda = ncols` is the packed [`Self::gemv_transpose`]).
162    fn gemv_transpose_strided(
163        a: &[Self],
164        x: &[Self],
165        y: &mut [Self],
166        nrows: usize,
167        ncols: usize,
168        lda: usize,
169    ) -> Result<(), SimdError>;
170    /// Multiplies interleaved complex lanes in-place: `a[k] *= b[k]`
171    /// (`a[k] *= conj(b[k])` when `CONJ_B`).
172    fn interleaved_complex_mul_assign<const CONJ_B: bool>(
173        a: &mut [Self],
174        b: &[Self],
175    ) -> Result<(), SimdError>
176    where
177        Self: core::ops::Neg<Output = Self>;
178    /// Computes the interleaved complex dot product `(re, im)` of `sum(a[k] * b[k])`
179    /// (`sum(a[k] * conj(b[k]))` when `CONJ_B`).
180    fn interleaved_complex_dot<const CONJ_B: bool>(
181        a: &[Self],
182        b: &[Self],
183    ) -> Result<(Self, Self), SimdError>
184    where
185        Self: core::ops::Neg<Output = Self>;
186    /// Computes the horizontal sum of population counts of all elements.
187    fn reduce_popcount(data: &[Self]) -> usize;
188    /// Computes the horizontal sum of population counts of `a[i] & b[i]`.
189    fn reduce_popcount_and(a: &[Self], b: &[Self]) -> Result<usize, SimdError>;
190    /// Computes the horizontal sum of population counts of `a[i] | b[i]`.
191    fn reduce_popcount_or(a: &[Self], b: &[Self]) -> Result<usize, SimdError>;
192    /// Computes the horizontal sum of population counts of `a[i] ^ b[i]` (Hamming distance).
193    fn reduce_popcount_xor(a: &[Self], b: &[Self]) -> Result<usize, SimdError>;
194}
195
196/// Method bodies shared verbatim by the three target-gated `SimdOps`
197/// blanket impls below, which differ only in the architecture-kernel
198/// bound each `where` clause requires. Defining them once keeps the
199/// dispatch facade DRY and behavior identical across targets.
200macro_rules! impl_simd_ops_methods {
201    () => {
202        #[inline(always)]
203        fn sum(data: &[Self]) -> Self {
204            sum::dispatch_sum::<Self>(data)
205        }
206        #[inline(always)]
207        fn abs_sum(data: &[Self]) -> Self {
208            abs_reduce::dispatch_abs_sum::<Self>(data)
209        }
210        #[inline(always)]
211        fn abs_max(data: &[Self]) -> Self {
212            abs_reduce::dispatch_abs_max::<Self>(data)
213        }
214        #[inline(always)]
215        fn min(data: &[Self]) -> Self {
216            min::dispatch_min::<Self>(data)
217        }
218        #[inline(always)]
219        fn max(data: &[Self]) -> Self {
220            max::dispatch_max::<Self>(data)
221        }
222        #[inline(always)]
223        fn scale(data: &mut [Self], scalar: Self) {
224            scale::dispatch_scale::<Self>(data, scalar)
225        }
226        #[inline(always)]
227        fn argmin(data: &[Self]) -> Option<(usize, Self)> {
228            argmin::dispatch_argmin::<Self>(data)
229        }
230        #[inline(always)]
231        fn argmax(data: &[Self]) -> Option<(usize, Self)> {
232            argmax::dispatch_argmax::<Self>(data)
233        }
234        #[inline(always)]
235        fn dot(a: &[Self], b: &[Self]) -> Result<Self, SimdError> {
236            dot::dispatch_dot::<Self>(a, b)
237        }
238        #[inline(always)]
239        fn axpy(alpha: Self, x: &[Self], out: &mut [Self]) -> Result<(), SimdError> {
240            axpy::dispatch_axpy::<Self>(alpha, x, out)
241        }
242        #[inline(always)]
243        fn axpy_mul(
244            alpha: Self,
245            a: &[Self],
246            b: &[Self],
247            out: &mut [Self],
248        ) -> Result<(), SimdError> {
249            axpy::dispatch_axpy_mul::<Self>(alpha, a, b, out)
250        }
251        #[inline(always)]
252        fn axpy_rows(
253            alphas: &[Self],
254            x: &[Self],
255            out: &mut [Self],
256            row_stride: usize,
257            rows: usize,
258            cols: usize,
259        ) -> Result<(), SimdError> {
260            axpy::dispatch_axpy_rows::<Self>(alphas, x, out, row_stride, rows, cols)
261        }
262        #[inline(always)]
263        fn axpy_rows_batch(
264            alphas: &[Self],
265            x_panel: &[Self],
266            out: &mut [Self],
267            row_stride: usize,
268            rows: usize,
269            depth: usize,
270            cols: usize,
271        ) -> Result<(), SimdError> {
272            axpy::dispatch_axpy_rows_batch::<Self>(
273                alphas, x_panel, out, row_stride, rows, depth, cols,
274            )
275        }
276        #[inline(always)]
277        fn elementwise_mul(a: &[Self], b: &[Self], out: &mut [Self]) -> Result<(), SimdError> {
278            binary::dispatch_elementwise_binary::<Self, Mul>(a, b, out, Mul)
279        }
280        #[inline(always)]
281        fn elementwise_add(a: &[Self], b: &[Self], out: &mut [Self]) -> Result<(), SimdError> {
282            binary::dispatch_elementwise_binary::<Self, Add>(a, b, out, Add)
283        }
284        #[inline(always)]
285        fn elementwise_sub(a: &[Self], b: &[Self], out: &mut [Self]) -> Result<(), SimdError> {
286            binary::dispatch_elementwise_binary::<Self, Sub>(a, b, out, Sub)
287        }
288        #[inline(always)]
289        fn elementwise_div(a: &[Self], b: &[Self], out: &mut [Self]) -> Result<(), SimdError> {
290            binary::dispatch_elementwise_binary::<Self, Div>(a, b, out, Div)
291        }
292        #[inline(always)]
293        fn masked_sum(data: &[Self], mask: &[bool]) -> Self {
294            masked::dispatch_masked_sum::<Self>(data, mask)
295        }
296        #[inline(always)]
297        fn masked_dot(a: &[Self], b: &[Self], mask: &[bool]) -> Result<Self, SimdError> {
298            masked::dispatch_masked_dot::<Self>(a, b, mask)
299        }
300        #[inline(always)]
301        fn masked_add(
302            a: &[Self],
303            b: &[Self],
304            mask: &[bool],
305            out: &mut [Self],
306        ) -> Result<(), SimdError> {
307            masked::dispatch_masked_add::<Self>(a, b, mask, out)
308        }
309        #[inline(always)]
310        fn spmv_csr(data: ValidatedData<CsrData<'_, Self>>, x: &[Self], y: &mut [Self]) {
311            sparse::dispatch_spmv_csr::<Self>(data, x, y)
312        }
313        #[inline(always)]
314        fn spmv_bcoo<const BM: usize, const BN: usize>(
315            data: ValidatedData<BlockedCooData<'_, Self, BM, BN>>,
316            x: &[Self],
317            y: &mut [Self],
318        ) {
319            // Runtime-dispatched like the other sparse kernels (was hardcoded to
320            // ScalarArch, which left the SIMD BlockedCoo paths dead at runtime).
321            sparse::dispatch_spmv_bcoo::<Self, BM, BN>(data, x, y)
322        }
323        #[inline(always)]
324        fn spmv_dense_masked(data: DenseWithMaskData<'_, Self>, x: &[Self], y: &mut [Self]) {
325            sparse::dispatch_spmv_dense_masked::<Self>(data, x, y)
326        }
327        #[inline(always)]
328        fn spmv_sellp<const C: usize>(
329            data: ValidatedData<SellPData<'_, Self, C>>,
330            x: &[Self],
331            y: &mut [Self],
332        ) {
333            sparse::dispatch_spmv_sellp::<Self, C>(data, x, y)
334        }
335        #[inline(always)]
336        fn tiled_gemm(
337            a: &[Self],
338            b: &[Self],
339            c: &mut [Self],
340            m: usize,
341            n: usize,
342            k: usize,
343        ) -> Result<(), SimdError> {
344            gemm::dispatch_tiled_gemm::<Self>(a, b, c, m, n, k)
345        }
346        #[inline(always)]
347        fn gemv(
348            a: &[Self],
349            x: &[Self],
350            y: &mut [Self],
351            nrows: usize,
352            ncols: usize,
353        ) -> Result<(), SimdError> {
354            gemv::dispatch_gemv::<Self>(a, x, y, nrows, ncols)
355        }
356        #[inline(always)]
357        fn gemv_transpose(
358            a: &[Self],
359            x: &[Self],
360            y: &mut [Self],
361            nrows: usize,
362            ncols: usize,
363        ) -> Result<(), SimdError> {
364            gemv_transpose::dispatch_gemv_transpose::<Self>(a, x, y, nrows, ncols)
365        }
366        #[inline(always)]
367        fn gemv_strided(
368            a: &[Self],
369            x: &[Self],
370            y: &mut [Self],
371            nrows: usize,
372            ncols: usize,
373            lda: usize,
374        ) -> Result<(), SimdError> {
375            gemv_strided::dispatch_gemv_strided::<Self>(a, x, y, nrows, ncols, lda)
376        }
377        #[inline(always)]
378        fn gemv_transpose_strided(
379            a: &[Self],
380            x: &[Self],
381            y: &mut [Self],
382            nrows: usize,
383            ncols: usize,
384            lda: usize,
385        ) -> Result<(), SimdError> {
386            gemv_transpose_strided::dispatch_gemv_transpose_strided::<Self>(
387                a, x, y, nrows, ncols, lda,
388            )
389        }
390        #[inline(always)]
391        fn interleaved_complex_mul_assign<const CONJ_B: bool>(
392            a: &mut [Self],
393            b: &[Self],
394        ) -> Result<(), SimdError>
395        where
396            Self: core::ops::Neg<Output = Self>,
397        {
398            complex::dispatch_interleaved_complex_mul_assign::<Self, CONJ_B>(a, b)
399        }
400        #[inline(always)]
401        fn interleaved_complex_dot<const CONJ_B: bool>(
402            a: &[Self],
403            b: &[Self],
404        ) -> Result<(Self, Self), SimdError>
405        where
406            Self: core::ops::Neg<Output = Self>,
407        {
408            complex::dispatch_interleaved_complex_dot::<Self, CONJ_B>(a, b)
409        }
410        #[inline(always)]
411        fn reduce_popcount(data: &[Self]) -> usize {
412            dispatch_reduce_popcount::<Self>(data)
413        }
414        #[inline(always)]
415        fn reduce_popcount_and(a: &[Self], b: &[Self]) -> Result<usize, SimdError> {
416            dispatch_reduce_popcount_and::<Self>(a, b)
417        }
418        #[inline(always)]
419        fn reduce_popcount_or(a: &[Self], b: &[Self]) -> Result<usize, SimdError> {
420            dispatch_reduce_popcount_or::<Self>(a, b)
421        }
422        #[inline(always)]
423        fn reduce_popcount_xor(a: &[Self], b: &[Self]) -> Result<usize, SimdError> {
424            dispatch_reduce_popcount_xor::<Self>(a, b)
425        }
426    };
427}
428
429/// x86/x86_64 specialized generic implementation of SimdOps.
430#[cfg(any(target_arch = "x86", target_arch = "x86_64"))]
431impl<T> SimdOps for T
432where
433    T: ScalarTrait + private::Sealed,
434    ScalarArch: hermes_simd_core::kernel::SimdKernel<T>,
435    Avx2: hermes_simd_core::kernel::SimdKernel<T>,
436    Avx512: hermes_simd_core::kernel::SimdKernel<T>,
437{
438    impl_simd_ops_methods!();
439}
440
441/// AArch64 specialized generic implementation of SimdOps.
442#[cfg(target_arch = "aarch64")]
443impl<T> SimdOps for T
444where
445    T: ScalarTrait + private::Sealed,
446    ScalarArch: hermes_simd_core::kernel::SimdKernel<T>,
447    Neon: hermes_simd_core::kernel::SimdKernel<T>,
448{
449    impl_simd_ops_methods!();
450}
451
452/// Fallback generic implementation of SimdOps.
453#[cfg(not(any(target_arch = "x86", target_arch = "x86_64", target_arch = "aarch64")))]
454impl<T> SimdOps for T
455where
456    T: ScalarTrait + private::Sealed,
457    ScalarArch: hermes_simd_core::kernel::SimdKernel<T>,
458{
459    impl_simd_ops_methods!();
460}
hermes_simd/dispatch/simd_ops.rs

hermes_simd/dispatch/
simd_ops.rs
