Skip to main content

hermes_simd/dispatch/
ops.rs

1use super::{complex, modular, simd_ops::SimdOps};
2use hermes_simd_core::scalar::Scalar as ScalarTrait;
3use hermes_simd_core::sparse::{
4    BlockedCooData, CsrData, DenseWithMaskData, SellPData, ValidatedData,
5};
6use hermes_simd_core::view::SimdError;
7
8/// Computes the sum of elements in the slice using runtime-dispatched SIMD.
9#[inline(always)]
10pub fn sum<T: SimdOps>(data: &[T]) -> T {
11    T::sum(data)
12}
13
14/// Computes the minimum element of the slice using runtime-dispatched SIMD.
15///
16/// Returns `T::MAX_VALUE` for empty slices.
17#[inline(always)]
18pub fn min<T: SimdOps>(data: &[T]) -> T {
19    T::min(data)
20}
21
22/// Computes the maximum element of the slice using runtime-dispatched SIMD.
23///
24/// Returns `T::MIN_VALUE` for empty slices.
25#[inline(always)]
26pub fn max<T: SimdOps>(data: &[T]) -> T {
27    T::max(data)
28}
29
30/// Reduces the slice to `Σ |x|` (L1-norm accumulator); `T::ZERO` for empty.
31#[inline(always)]
32pub fn abs_sum<T: SimdOps>(data: &[T]) -> T {
33    T::abs_sum(data)
34}
35
36/// Reduces the slice to `max |x|` (∞-norm accumulator); `T::ZERO` for empty.
37#[inline(always)]
38pub fn abs_max<T: SimdOps>(data: &[T]) -> T {
39    T::abs_max(data)
40}
41
42/// Multiplies every element of `data` by `scalar` in-place.
43#[inline(always)]
44pub fn scale<T: SimdOps>(data: &mut [T], scalar: T) {
45    T::scale(data, scalar)
46}
47
48/// Returns the first minimum, or `None` for empty or NaN-containing data.
49#[inline(always)]
50pub fn argmin<T: SimdOps>(data: &[T]) -> Option<(usize, T)> {
51    T::argmin(data)
52}
53
54/// Returns the first maximum, or `None` for empty or NaN-containing data.
55#[inline(always)]
56pub fn argmax<T: SimdOps>(data: &[T]) -> Option<(usize, T)> {
57    T::argmax(data)
58}
59
60/// Computes the dot product of two slices using runtime-dispatched SIMD.
61#[inline(always)]
62pub fn dot<T: SimdOps>(a: &[T], b: &[T]) -> Result<T, SimdError> {
63    T::dot(a, b)
64}
65
66/// Fused row update `out[i] += alpha * x[i]` (AXPY) via runtime-dispatched
67/// SIMD with no temporary allocation. Errors on length mismatch.
68#[inline(always)]
69pub fn axpy<T: SimdOps>(alpha: T, x: &[T], out: &mut [T]) -> Result<(), SimdError> {
70    T::axpy(alpha, x, out)
71}
72
73/// Fused ternary update `out[i] += alpha * a[i] * b[i]` without a temporary.
74///
75/// # Errors
76/// Returns [`SimdError::LengthMismatch`] when `a`, `b`, and `out` do not have
77/// equal lengths.
78#[inline(always)]
79pub fn axpy_mul<T: SimdOps>(alpha: T, a: &[T], b: &[T], out: &mut [T]) -> Result<(), SimdError> {
80    T::axpy_mul(alpha, a, b, out)
81}
82
83/// Fused multi-row update `out[row, i] += alphas[row] * x[i]` via one
84/// runtime-dispatched SIMD kernel. `out` is a row-major strided window.
85#[inline(always)]
86pub fn axpy_rows<T: SimdOps>(
87    alphas: &[T],
88    x: &[T],
89    out: &mut [T],
90    row_stride: usize,
91    rows: usize,
92    cols: usize,
93) -> Result<(), SimdError> {
94    T::axpy_rows(alphas, x, out, row_stride, rows, cols)
95}
96
97/// Fused batched multi-row update:
98/// `out[row, i] += sum_k alphas[k, row] * x_panel[k, i]` via one
99/// runtime-dispatched SIMD kernel. `alphas` is depth-major with `rows`
100/// elements per depth, `x_panel` is depth-major with `cols` elements per
101/// depth, and `out` is a row-major strided window.
102#[inline(always)]
103pub fn axpy_rows_batch<T: SimdOps>(
104    alphas: &[T],
105    x_panel: &[T],
106    out: &mut [T],
107    row_stride: usize,
108    rows: usize,
109    depth: usize,
110    cols: usize,
111) -> Result<(), SimdError> {
112    T::axpy_rows_batch(alphas, x_panel, out, row_stride, rows, depth, cols)
113}
114
115/// Computes the elementwise multiplication of two slices and writes to `out`.
116#[inline(always)]
117pub fn elementwise_mul<T: SimdOps>(a: &[T], b: &[T], out: &mut [T]) -> Result<(), SimdError> {
118    T::elementwise_mul(a, b, out)
119}
120
121/// Computes the elementwise sum of two slices and writes to `out`.
122#[inline(always)]
123pub fn elementwise_add<T: SimdOps>(a: &[T], b: &[T], out: &mut [T]) -> Result<(), SimdError> {
124    T::elementwise_add(a, b, out)
125}
126
127/// Computes the elementwise difference of two slices and writes to `out`.
128#[inline(always)]
129pub fn elementwise_sub<T: SimdOps>(a: &[T], b: &[T], out: &mut [T]) -> Result<(), SimdError> {
130    T::elementwise_sub(a, b, out)
131}
132
133/// Computes the elementwise quotient of two slices and writes to `out`.
134#[inline(always)]
135pub fn elementwise_div<T: SimdOps>(a: &[T], b: &[T], out: &mut [T]) -> Result<(), SimdError> {
136    T::elementwise_div(a, b, out)
137}
138
139/// Executes one exact modular radix-2 NTT butterfly stage over `u64` residues.
140#[inline]
141pub fn ntt_butterfly_stage_u64(
142    data: &mut [u64],
143    stage_len: usize,
144    twiddles: &[u64],
145    modulus: u64,
146) -> Result<(), SimdError> {
147    modular::ntt_butterfly_stage_u64(data, stage_len, twiddles, modulus)
148}
149
150/// Computes the sum of elements matching a boolean mask.
151#[inline(always)]
152pub fn masked_sum<T: SimdOps>(data: &[T], mask: &[bool]) -> T {
153    T::masked_sum(data, mask)
154}
155
156/// Computes the dot product of elements matching a boolean mask.
157#[inline(always)]
158pub fn masked_dot<T: SimdOps>(a: &[T], b: &[T], mask: &[bool]) -> Result<T, SimdError> {
159    T::masked_dot(a, b, mask)
160}
161
162/// Computes the elementwise sum of elements matching a boolean mask.
163#[inline(always)]
164pub fn masked_add<T: SimdOps>(
165    a: &[T],
166    b: &[T],
167    mask: &[bool],
168    out: &mut [T],
169) -> Result<(), SimdError> {
170    T::masked_add(a, b, mask, out)
171}
172
173/// Computes sparse SpMV using CSR: `y += A · x`.
174///
175/// # Panics
176/// Panics if `x.len() < ncols` or `y.len() < nrows`. Structural CSR validation
177/// is performed by [`ValidatedData::new`] before this function can be called.
178#[inline(always)]
179pub fn spmv_csr<T: SimdOps>(data: ValidatedData<CsrData<'_, T>>, x: &[T], y: &mut [T]) {
180    T::spmv_csr(data, x, y)
181}
182
183/// Computes sparse SpMV using const-generic Blocked-COO tiles.
184///
185/// # Panics
186/// Panics if `x.len() < ncols` or `y.len() < nrows`. Structural Blocked-COO
187/// validation is performed by [`ValidatedData::new`] before this function can be
188/// called.
189#[inline(always)]
190pub fn spmv_bcoo<T: SimdOps, const BM: usize, const BN: usize>(
191    data: ValidatedData<BlockedCooData<'_, T, BM, BN>>,
192    x: &[T],
193    y: &mut [T],
194) {
195    T::spmv_bcoo::<BM, BN>(data, x, y)
196}
197
198/// Computes sparse SpMV using Dense-with-Mask.
199#[inline(always)]
200pub fn spmv_dense_masked<T: SimdOps>(data: DenseWithMaskData<'_, T>, x: &[T], y: &mut [T]) {
201    T::spmv_dense_masked(data, x, y)
202}
203
204/// Computes sparse SpMV using const-generic Sliced ELLPACK (SELL-p).
205///
206/// # Panics
207/// Panics if `x.len() < ncols` or `y.len() < nrows`. Structural SELL-p
208/// validation is performed by [`ValidatedData::new`] before this function can be
209/// called.
210#[inline(always)]
211pub fn spmv_sellp<T: SimdOps, const C: usize>(
212    data: ValidatedData<SellPData<'_, T, C>>,
213    x: &[T],
214    y: &mut [T],
215) {
216    T::spmv_sellp::<C>(data, x, y)
217}
218
219/// Computes register-blocked tiled GEMM: `c += A * B`.
220#[inline(always)]
221pub fn tiled_gemm<T: SimdOps>(
222    a: &[T],
223    b: &[T],
224    c: &mut [T],
225    m: usize,
226    n: usize,
227    k: usize,
228) -> Result<(), SimdError> {
229    T::tiled_gemm(a, b, c, m, n, k)
230}
231
232/// Computes register-blocked GEMV `y += A · x` with runtime backend selection.
233///
234/// `a` is row-major `nrows × ncols`; the product **accumulates** into `y`
235/// (zero `y` first for `y = A·x`). See [`gemv()`] for the
236/// operand-reuse theorem.
237///
238/// # Errors
239/// [`SimdError::LengthMismatch`] if `a.len() < nrows·ncols`, `x.len() < ncols`,
240/// or `y.len() < nrows`.
241#[inline(always)]
242pub fn gemv<T: SimdOps>(
243    a: &[T],
244    x: &[T],
245    y: &mut [T],
246    nrows: usize,
247    ncols: usize,
248) -> Result<(), SimdError> {
249    T::gemv(a, x, y, nrows, ncols)
250}
251
252/// Computes register-blocked transposed GEMV `y += Aᵀ · x` with runtime backend
253/// selection — the complement of [`gemv()`].
254///
255/// `a` is row-major `nrows × ncols`, `x` length `nrows`, `y` length `ncols`; the
256/// product **accumulates** into `y` (zero `y` first for `y = Aᵀ·x`). See
257/// [`gemv_transpose()`] for the operand-reuse theorem.
258///
259/// # Errors
260/// [`SimdError::LengthMismatch`] if `a.len() < nrows·ncols`, `x.len() < nrows`,
261/// or `y.len() < ncols`.
262#[inline(always)]
263pub fn gemv_transpose<T: SimdOps>(
264    a: &[T],
265    x: &[T],
266    y: &mut [T],
267    nrows: usize,
268    ncols: usize,
269) -> Result<(), SimdError> {
270    T::gemv_transpose(a, x, y, nrows, ncols)
271}
272
273/// Computes register-blocked sub-matrix GEMV `y += A · x` with row stride `lda`,
274/// runtime backend selection. `A` is a row-major `nrows × ncols` block with
275/// leading dimension `lda ≥ ncols`; `lda = ncols` is the packed [`gemv()`].
276/// Accumulates into `y`.
277///
278/// # Errors
279/// [`SimdError::LengthMismatch`] if `lda < ncols`, `a.len() < (nrows−1)·lda +
280/// ncols`, `x.len() < ncols`, or `y.len() < nrows`.
281#[inline(always)]
282pub fn gemv_strided<T: SimdOps>(
283    a: &[T],
284    x: &[T],
285    y: &mut [T],
286    nrows: usize,
287    ncols: usize,
288    lda: usize,
289) -> Result<(), SimdError> {
290    T::gemv_strided(a, x, y, nrows, ncols, lda)
291}
292
293/// Computes register-blocked transposed sub-matrix GEMV `y += Aᵀ · x` with row
294/// stride `lda`, runtime backend selection. `lda = ncols` is the packed
295/// [`gemv_transpose()`]. Accumulates into `y`.
296///
297/// # Errors
298/// [`SimdError::LengthMismatch`] if `lda < ncols`, `a.len() < (nrows−1)·lda +
299/// ncols`, `x.len() < nrows`, or `y.len() < ncols`.
300#[inline(always)]
301pub fn gemv_transpose_strided<T: SimdOps>(
302    a: &[T],
303    x: &[T],
304    y: &mut [T],
305    nrows: usize,
306    ncols: usize,
307    lda: usize,
308) -> Result<(), SimdError> {
309    T::gemv_transpose_strided(a, x, y, nrows, ncols, lda)
310}
311
312/// Multiplies interleaved complex values in-place using a monomorphized SIMD architecture.
313///
314/// Inputs are primitive lane slices in `[re0, im0, re1, im1, ...]` order. `a`
315/// is updated with `a[i] * b[i]`; when `CONJ_B` is true, the operation is
316/// `a[i] * conj(b[i])`.
317#[inline]
318pub fn interleaved_complex_mul_assign<T, A, const CONJ_B: bool>(
319    a: &mut [T],
320    b: &[T],
321) -> Result<(), SimdError>
322where
323    T: ScalarTrait + core::ops::Neg<Output = T>,
324    A: hermes_simd_core::arch::SimdArch + hermes_simd_core::kernel::SimdKernel<T>,
325{
326    complex::interleaved_complex_mul_assign::<T, A, CONJ_B>(a, b)
327}
328
329/// Computes an interleaved complex dot product using a monomorphized SIMD architecture.
330///
331/// Inputs are primitive lane slices in `[re0, im0, re1, im1, ...]` order. The
332/// result is `(re, im)` for `sum(a[i] * b[i])`; when `CONJ_B` is true, the
333/// operation is `sum(a[i] * conj(b[i]))`.
334#[inline]
335pub fn interleaved_complex_dot<T, A, const CONJ_B: bool>(
336    a: &[T],
337    b: &[T],
338) -> Result<(T, T), SimdError>
339where
340    T: ScalarTrait + core::ops::Neg<Output = T>,
341    A: hermes_simd_core::arch::SimdArch + hermes_simd_core::kernel::SimdKernel<T>,
342{
343    complex::interleaved_complex_dot::<T, A, CONJ_B>(a, b)
344}
345
346/// Multiplies interleaved complex values in-place using Hermes runtime provider selection.
347#[inline]
348pub fn interleaved_complex_mul_assign_runtime<T, const CONJ_B: bool>(
349    a: &mut [T],
350    b: &[T],
351) -> Result<(), SimdError>
352where
353    T: SimdOps + core::ops::Neg<Output = T>,
354{
355    T::interleaved_complex_mul_assign::<CONJ_B>(a, b)
356}
357
358/// Computes an interleaved complex dot product using Hermes runtime provider selection.
359#[inline]
360pub fn interleaved_complex_dot_runtime<T, const CONJ_B: bool>(
361    a: &[T],
362    b: &[T],
363) -> Result<(T, T), SimdError>
364where
365    T: SimdOps + core::ops::Neg<Output = T>,
366{
367    T::interleaved_complex_dot::<CONJ_B>(a, b)
368}
369
370/// Computes the horizontal sum of population counts of all elements using runtime-dispatched SIMD.
371#[inline(always)]
372pub fn reduce_popcount<T: SimdOps>(data: &[T]) -> usize {
373    T::reduce_popcount(data)
374}
375
376/// Computes the horizontal sum of population counts of `a[i] & b[i]` using runtime-dispatched SIMD.
377#[inline(always)]
378pub fn reduce_popcount_and<T: SimdOps>(a: &[T], b: &[T]) -> Result<usize, SimdError> {
379    T::reduce_popcount_and(a, b)
380}
381
382/// Computes the horizontal sum of population counts of `a[i] | b[i]` using runtime-dispatched SIMD.
383#[inline(always)]
384pub fn reduce_popcount_or<T: SimdOps>(a: &[T], b: &[T]) -> Result<usize, SimdError> {
385    T::reduce_popcount_or(a, b)
386}
387
388/// Computes the horizontal sum of population counts of `a[i] ^ b[i]` (Hamming distance) using runtime-dispatched SIMD.
389#[inline(always)]
390pub fn reduce_popcount_xor<T: SimdOps>(a: &[T], b: &[T]) -> Result<usize, SimdError> {
391    T::reduce_popcount_xor(a, b)
392}