hermes_simd/dispatch/
gemv_strided.rs

1//! Generic runtime-dispatch register-blocked sub-matrix GEMV (`y += A · x`,
2//! row stride `lda`).
3//!
4//! Generalizes [`super::gemv()`] to a row-major **sub-matrix**: `nrows × ncols`
5//! with leading dimension `lda ≥ ncols` (rows contiguous over `ncols`, spaced
6//! `lda` apart). `lda = ncols` recovers the packed `gemv`. This admits matvec
7//! over a trailing/leading block of a larger buffer — e.g. the column-major
8//! trailing block of a reflector apply, whose columns are contiguous but spaced
9//! by the buffer's row count — without copying the block out. Result
10//! **accumulates** into `y`.
11
12use hermes_simd_core::{
13    align::Unaligned,
14    arch::SimdArch,
15    execution::Unmasked,
16    kernel::SimdKernel,
17    scalar::Scalar,
18    view::{SimdError, SimdView},
19};
20use hermes_simd_macros::runtime_dispatch;
21
22#[runtime_dispatch(avx512f, avx2, neon, scalar)]
23pub(super) fn dispatch_gemv_strided_kernel<T, A>(
24    a: &[T],
25    x: &[T],
26    y: &mut [T],
27    nrows: usize,
28    ncols: usize,
29    lda: usize,
30) -> Result<(), SimdError>
31where
32    T: Scalar,
33    A: SimdArch + SimdKernel<T>,
34{
35    match (
36        SimdView::<T, A, Unaligned, Unmasked, &[T]>::new(a),
37        SimdView::<T, A, Unaligned, Unmasked, &[T]>::new(x),
38    ) {
39        (Some(va), Some(vx)) => {
40            use hermes_simd_core::tiling::{TilingPolicy, TilingStrategy};
41            if A::LANE_COUNT > 8 {
42                <TilingPolicy<8, 1> as TilingStrategy<T, A, Unaligned>>::gemv_strided(
43                    &va, &vx, y, nrows, ncols, lda,
44                )
45            } else if A::LANE_COUNT > 1 {
46                <TilingPolicy<4, 1> as TilingStrategy<T, A, Unaligned>>::gemv_strided(
47                    &va, &vx, y, nrows, ncols, lda,
48                )
49            } else {
50                <TilingPolicy<1, 1> as TilingStrategy<T, A, Unaligned>>::gemv_strided(
51                    &va, &vx, y, nrows, ncols, lda,
52                )
53            }
54        }
55        // SAFETY: `Unaligned` skips the alignment check, so `SimdView::new` is
56        // `Some` for every slice — this arm is unreachable (mirrors `super::gemv`).
57        _ => unsafe { core::hint::unreachable_unchecked() },
58    }
59}
60
61#[cfg(test)]
62mod tests {
63    use crate::dispatch::{gemv, gemv_strided};
64
65    /// Naive reference `y = A·x` over a sub-matrix with row stride `lda`.
66    fn reference(a: &[f64], x: &[f64], nrows: usize, ncols: usize, lda: usize) -> Vec<f64> {
67        (0..nrows)
68            .map(|r| (0..ncols).map(|c| a[r * lda + c] * x[c]).sum())
69            .collect()
70    }
71
72    #[test]
73    fn gemv_strided_matches_reference_over_submatrix() {
74        // A 6x10 backing buffer; operate on the 4x6 sub-matrix (lda=10, ncols=6).
75        let lda = 10usize;
76        let (nrows, ncols) = (4usize, 6usize);
77        let a: Vec<f64> = (0..nrows * lda)
78            .map(|i| ((i % 9) as f64 - 4.0) * 0.25)
79            .collect();
80        let x: Vec<f64> = (0..ncols).map(|i| ((i % 5) as f64 - 2.0) * 0.5).collect();
81        let mut y = vec![0.0f64; nrows];
82        gemv_strided::dispatch_gemv_strided::<f64>(&a, &x, &mut y, nrows, ncols, lda).unwrap();
83        assert_eq!(y, reference(&a, &x, nrows, ncols, lda));
84    }
85
86    #[test]
87    fn gemv_strided_packed_equals_gemv() {
88        // lda == ncols must agree bit-for-bit with the packed gemv.
89        let (nrows, ncols) = (9usize, 13usize);
90        let a: Vec<f64> = (0..nrows * ncols)
91            .map(|i| ((i % 7) as f64 - 3.0) * 0.5)
92            .collect();
93        let x: Vec<f64> = (0..ncols).map(|i| ((i % 4) as f64 - 1.0) * 0.25).collect();
94        let mut y_strided = vec![0.0f64; nrows];
95        let mut y_packed = vec![0.0f64; nrows];
96        gemv_strided::dispatch_gemv_strided::<f64>(&a, &x, &mut y_strided, nrows, ncols, ncols)
97            .unwrap();
98        gemv::dispatch_gemv::<f64>(&a, &x, &mut y_packed, nrows, ncols).unwrap();
99        assert_eq!(y_strided, y_packed);
100    }
101
102    #[test]
103    fn gemv_strided_rejects_lda_below_ncols_and_short_spans() {
104        let a = vec![1.0f64; 40];
105        let x = vec![1.0f64; 6];
106        let mut y = vec![0.0f64; 4];
107        // lda < ncols is invalid.
108        assert!(gemv_strided::dispatch_gemv_strided::<f64>(&a, &x, &mut y, 4, 6, 5).is_err());
109        // a too short for the requested sub-matrix span.
110        let short = vec![1.0f64; 10];
111        assert!(gemv_strided::dispatch_gemv_strided::<f64>(&short, &x, &mut y, 4, 6, 10).is_err());
112    }
113
114    #[test]
115    fn gemv_strided_rejects_dimension_overflow() {
116        use hermes_simd_core::view::SimdError;
117        // Adversarial stride: `(nrows-1)·lda + ncols` overflows `usize`. Unchecked
118        // (release `overflow-checks = false`), the product wraps to a small value,
119        // the `a_len < a_needed` guard passes, and the kernel issues an OOB SIMD
120        // load. `x`/`y` are sized so the only failing condition is the A-span. The
121        // checked span arithmetic must reject with the exact variant.
122        let a = vec![1.0f64; 40];
123        let x = vec![1.0f64; 6];
124        let mut y = vec![0.0f64; 2];
125        let r = gemv_strided::dispatch_gemv_strided::<f64>(&a, &x, &mut y, 2, 6, usize::MAX);
126        assert_eq!(r, Err(SimdError::LengthMismatch));
127    }
128}
hermes_simd/dispatch/gemv_strided.rs

hermes_simd/dispatch/
gemv_strided.rs
