hermes_simd/dispatch/
gemv.rs

1//! Generic runtime-dispatch register-blocked matrix–vector product (`y += A · x`).
2//!
3//! Plumbs the core register-blocked `gemv` micro-kernel
4//! ([`hermes_simd_core::tiling::TilingStrategy::gemv`]) through runtime backend
5//! selection, mirroring [`super::gemm`]. `A` is row-major `nrows × ncols`; the
6//! result **accumulates** into `y` (`y += A·x`), so callers wanting `y = A·x`
7//! zero `y` first — matching the `axpy`/GEMM accumulate convention.
8//!
9//! # Theorem (operand reuse — why register blocking helps)
10//! GEMV performs `2·nrows·ncols` flops over `nrows·ncols` matrix elements:
11//! arithmetic intensity ≈ 2 flops/element, so it is **memory-bound** and
12//! throughput is governed by operand reuse, not FLOP rate. Blocking `TILE_M`
13//! rows of `A` loads each `x[c..c+lane]` vector **once** and applies it to all
14//! `TILE_M` rows held in independent register accumulators; this cuts `x`
15//! traffic by `TILE_M×` and breaks the per-row FMA dependency chain (one live
16//! accumulator per row). The `nrows mod TILE_M` remainder is handled by a
17//! single-row cleanup, so any shape is supported. `TILE_M` scales with the
18//! register file: wider ISAs block more rows before spilling. ∎
19
20use hermes_simd_core::{
21    align::Unaligned,
22    arch::SimdArch,
23    execution::Unmasked,
24    kernel::SimdKernel,
25    scalar::Scalar,
26    view::{SimdError, SimdView},
27};
28use hermes_simd_macros::runtime_dispatch;
29
30#[runtime_dispatch(avx512f, avx2, neon, scalar)]
31pub(super) fn dispatch_gemv_kernel<T, A>(
32    a: &[T],
33    x: &[T],
34    y: &mut [T],
35    nrows: usize,
36    ncols: usize,
37) -> Result<(), SimdError>
38where
39    T: Scalar,
40    A: SimdArch + SimdKernel<T>,
41{
42    match (
43        SimdView::<T, A, Unaligned, Unmasked, &[T]>::new(a),
44        SimdView::<T, A, Unaligned, Unmasked, &[T]>::new(x),
45    ) {
46        (Some(va), Some(vx)) => {
47            use hermes_simd_core::tiling::{TilingPolicy, TilingStrategy};
48            // GEMV uses only `TILE_M` (row blocking); `TILE_N` is inert here.
49            // Wider register files block more rows to amortize the shared `x`
50            // load — `<1,1>` for scalar avoids any unused-accumulator overhead.
51            if A::LANE_COUNT > 8 {
52                <TilingPolicy<8, 1> as TilingStrategy<T, A, Unaligned>>::gemv(
53                    &va, &vx, y, nrows, ncols,
54                )
55            } else if A::LANE_COUNT > 1 {
56                <TilingPolicy<4, 1> as TilingStrategy<T, A, Unaligned>>::gemv(
57                    &va, &vx, y, nrows, ncols,
58                )
59            } else {
60                <TilingPolicy<1, 1> as TilingStrategy<T, A, Unaligned>>::gemv(
61                    &va, &vx, y, nrows, ncols,
62                )
63            }
64        }
65        // SAFETY: `Unaligned::IS_ALIGNED` is false, so `SimdView::new` skips the
66        // alignment check and is `Some` for every slice (including empty) — this
67        // arm is unreachable. Mirrors `super::gemm::dispatch_tiled_gemm_kernel`.
68        _ => unsafe { core::hint::unreachable_unchecked() },
69    }
70}
71
72#[cfg(test)]
73mod tests {
74    use crate::dispatch::gemv;
75
76    /// Naive reference `y = A·x` (`A` row-major `nrows × ncols`).
77    fn reference(a: &[f64], x: &[f64], nrows: usize, ncols: usize) -> Vec<f64> {
78        (0..nrows)
79            .map(|r| (0..ncols).map(|c| a[r * ncols + c] * x[c]).sum())
80            .collect()
81    }
82
83    fn run_case(nrows: usize, ncols: usize) {
84        // Deterministic, dyadic-exact entries so the SIMD reduction order does
85        // not change the result vs the reference (no rounding divergence).
86        let a: Vec<f64> = (0..nrows * ncols)
87            .map(|i| ((i % 7) as f64 - 3.0) * 0.25)
88            .collect();
89        let x: Vec<f64> = (0..ncols).map(|i| ((i % 5) as f64 - 2.0) * 0.5).collect();
90        let mut y = vec![0.0f64; nrows];
91        gemv::dispatch_gemv::<f64>(&a, &x, &mut y, nrows, ncols).unwrap();
92        let want = reference(&a, &x, nrows, ncols);
93        assert_eq!(y, want, "gemv {nrows}x{ncols} mismatch vs reference");
94    }
95
96    #[test]
97    fn gemv_matches_reference_across_shapes() {
98        // Includes shapes exercising the TILE_M row remainder and the column
99        // SIMD tail (ncols not a multiple of any lane count).
100        for &(m, n) in &[
101            (1, 1),
102            (1, 17),
103            (3, 4),
104            (8, 8),
105            (9, 13),
106            (16, 1),
107            (17, 31),
108            (33, 64),
109            (64, 64),
110        ] {
111            run_case(m, n);
112        }
113    }
114
115    #[test]
116    fn gemv_accumulates_into_y() {
117        let a = vec![1.0f64, 2.0, 3.0, 4.0]; // 2x2
118        let x = vec![1.0f64, 1.0];
119        let mut y = vec![10.0f64, 20.0];
120        gemv::dispatch_gemv::<f64>(&a, &x, &mut y, 2, 2).unwrap();
121        // y += A·x = [10+3, 20+7] = [13, 27]
122        assert_eq!(y, vec![13.0, 27.0]);
123    }
124
125    #[test]
126    fn gemv_rejects_short_operands() {
127        let a = vec![1.0f64; 4];
128        let x = vec![1.0f64; 2];
129        let mut y = vec![0.0f64; 1]; // too short for nrows=2
130        assert!(gemv::dispatch_gemv::<f64>(&a, &x, &mut y, 2, 2).is_err());
131    }
132}
hermes_simd/dispatch/gemv.rs

hermes_simd/dispatch/
gemv.rs
