moirai_utils/simd/
mod.rs

1//! Generic SIMD-aware vector operations.
2//!
3//! Public operations are expressed once over sealed scalar traits. The concrete
4//! scalar type is selected at the call site, and monomorphization removes the
5//! trait layer before execution. Native ISA kernels are private backend details.
6
7mod arch;
8mod scalar;
9
10#[cfg(test)]
11mod tests;
12
13pub use arch::{has_avx2_support, has_neon_support};
14pub use scalar::{SimdReal, SimdScalar};
15
16#[inline]
17fn record_dispatch(vectorized: bool, elements: usize) {
18    if vectorized {
19        crate::global_simd_counter().record_vectorized_op(elements);
20    } else {
21        crate::global_simd_counter().record_scalar_op(elements);
22    }
23}
24
25#[inline]
26fn assert_same_len<T>(left: &[T], right: &[T], result: &[T]) {
27    assert_eq!(left.len(), right.len(), "input slices must match");
28    assert_eq!(left.len(), result.len(), "output slice must match inputs");
29}
30
31/// Returns whether `T` has a native vector backend on the current CPU.
32#[inline]
33pub fn has_native_vector_path<T: SimdScalar>() -> bool {
34    T::native_vector_available()
35}
36
37/// Adds two slices into `result`.
38///
39/// The operation executes in the native precision of `T`. For scalar types with
40/// a private native vector backend, the backend is selected with static dispatch
41/// after runtime ISA detection; all other scalar types use the monomorphized
42/// scalar loop.
43#[inline]
44pub fn add<T: SimdScalar>(left: &[T], right: &[T], result: &mut [T]) {
45    assert_same_len(left, right, result);
46    let vectorized = T::uses_native_vector_path(left.len());
47    T::add_slices(left, right, result);
48    record_dispatch(vectorized, left.len());
49}
50
51/// Multiplies two slices into `result`.
52#[inline]
53pub fn mul<T: SimdScalar>(left: &[T], right: &[T], result: &mut [T]) {
54    assert_same_len(left, right, result);
55    let vectorized = T::uses_native_vector_path(left.len());
56    T::mul_slices(left, right, result);
57    record_dispatch(vectorized, left.len());
58}
59
60/// Computes a native-precision dot product.
61#[inline]
62pub fn dot<T: SimdScalar>(left: &[T], right: &[T]) -> T {
63    assert_eq!(left.len(), right.len(), "input slices must match");
64    let vectorized = T::uses_native_vector_path(left.len());
65    let result = T::dot_slice(left, right);
66    record_dispatch(vectorized, left.len());
67    result
68}
69
70/// Computes a native-precision sum.
71#[inline]
72pub fn sum<T: SimdScalar>(data: &[T]) -> T {
73    let vectorized = T::uses_native_vector_path(data.len());
74    let result = T::sum_slice(data);
75    record_dispatch(vectorized, data.len());
76    result
77}
78
79/// Computes a native-precision arithmetic mean.
80#[inline]
81pub fn mean<T: SimdReal>(data: &[T]) -> T {
82    assert!(!data.is_empty(), "mean requires at least one value");
83    let vectorized = T::uses_native_vector_path(data.len());
84    let result = T::mean_slice(data);
85    record_dispatch(vectorized, data.len());
86    result
87}
88
89/// Computes a native-precision population variance.
90#[inline]
91pub fn variance<T: SimdReal>(data: &[T]) -> T {
92    assert!(!data.is_empty(), "variance requires at least one value");
93    let vectorized = T::uses_native_vector_path(data.len());
94    let result = T::variance_slice(data);
95    record_dispatch(vectorized, data.len());
96    result
97}
98
99/// Multiplies two row-major square matrices.
100///
101/// `N` is the matrix dimension. Each slice must contain exactly `N * N`
102/// elements. The const parameter makes the structural arity visible to the
103/// optimizer without encoding the dimension in the function name.
104#[inline]
105pub fn matrix_mul_square<T: SimdScalar, const N: usize>(left: &[T], right: &[T], result: &mut [T]) {
106    T::matrix_mul_square::<N>(left, right, result);
107    record_dispatch(T::matrix_vector_path_available::<N>(), result.len());
108}
moirai_utils/simd/mod.rs

moirai_utils/simd/
mod.rs
