Skip to main content

hermes_simd/dispatch/
gemv_transpose.rs

1//! Generic runtime-dispatch transposed matrix–vector product (`y += Aᵀ · x`).
2//!
3//! Plumbs the core register-blocked transposed-GEMV micro-kernel
4//! ([`hermes_simd_core::tiling::TilingStrategy::gemv_transpose`]) through runtime
5//! backend selection, the complement of [`super::gemv()`]. `A` is row-major
6//! `nrows × ncols`, `x` length `nrows`, `y` length `ncols`; the result
7//! **accumulates** into `y` (`y += Aᵀ·x`), so callers wanting `y = Aᵀ·x` zero
8//! `y` first.
9//!
10//! # Theorem (output reuse, reduction-free)
11//! `Aᵀx = Σᵢ xᵢ·A[i,:]` — a sum of the rows of `A` scaled by `x`. Each row is
12//! contiguous, so the update vectorizes across the `ncols` output lanes with **no
13//! horizontal reduction** (unlike `A·x`). Blocking `TILE_N` output lane-chunks of
14//! `y` in registers reuses each accumulator across all `nrows` rows and breaks the
15//! per-chunk FMA dependency chain; `TILE_N` scales with the register file. ∎
16
17use hermes_simd_core::{
18    align::Unaligned,
19    arch::SimdArch,
20    execution::Unmasked,
21    kernel::SimdKernel,
22    scalar::Scalar,
23    view::{SimdError, SimdView},
24};
25use hermes_simd_macros::runtime_dispatch;
26
27#[runtime_dispatch(avx512f, avx2, neon, scalar)]
28pub(super) fn dispatch_gemv_transpose_kernel<T, A>(
29    a: &[T],
30    x: &[T],
31    y: &mut [T],
32    nrows: usize,
33    ncols: usize,
34) -> Result<(), SimdError>
35where
36    T: Scalar,
37    A: SimdArch + SimdKernel<T>,
38{
39    match (
40        SimdView::<T, A, Unaligned, Unmasked, &[T]>::new(a),
41        SimdView::<T, A, Unaligned, Unmasked, &[T]>::new(x),
42    ) {
43        (Some(va), Some(vx)) => {
44            use hermes_simd_core::tiling::{TilingPolicy, TilingStrategy};
45            // `gemv_transpose` blocks `TILE_N` output lane-chunks (TILE_M inert);
46            // wider register files block more chunks before spilling.
47            if A::LANE_COUNT > 8 {
48                <TilingPolicy<1, 8> as TilingStrategy<T, A, Unaligned>>::gemv_transpose(
49                    &va, &vx, y, nrows, ncols,
50                )
51            } else if A::LANE_COUNT > 1 {
52                <TilingPolicy<1, 4> as TilingStrategy<T, A, Unaligned>>::gemv_transpose(
53                    &va, &vx, y, nrows, ncols,
54                )
55            } else {
56                <TilingPolicy<1, 1> as TilingStrategy<T, A, Unaligned>>::gemv_transpose(
57                    &va, &vx, y, nrows, ncols,
58                )
59            }
60        }
61        // SAFETY: `Unaligned` skips the alignment check, so `SimdView::new` is
62        // `Some` for every slice — this arm is unreachable (mirrors `super::gemv`).
63        _ => unsafe { core::hint::unreachable_unchecked() },
64    }
65}
66
67#[cfg(test)]
68mod tests {
69    use crate::dispatch::gemv_transpose;
70
71    /// Naive reference `y = Aᵀ·x` (`A` row-major `nrows × ncols`).
72    fn reference(a: &[f64], x: &[f64], nrows: usize, ncols: usize) -> Vec<f64> {
73        let mut y = vec![0.0f64; ncols];
74        for (i, &xi) in x.iter().enumerate().take(nrows) {
75            for (j, yj) in y.iter_mut().enumerate() {
76                *yj += a[i * ncols + j] * xi;
77            }
78        }
79        y
80    }
81
82    fn run_case(nrows: usize, ncols: usize) {
83        // Dyadic-exact entries: the column-accumulation order matches the
84        // reference, so no rounding divergence (assert exact equality).
85        let a: Vec<f64> = (0..nrows * ncols)
86            .map(|i| ((i % 7) as f64 - 3.0) * 0.25)
87            .collect();
88        let x: Vec<f64> = (0..nrows).map(|i| ((i % 5) as f64 - 2.0) * 0.5).collect();
89        let mut y = vec![0.0f64; ncols];
90        gemv_transpose::dispatch_gemv_transpose::<f64>(&a, &x, &mut y, nrows, ncols).unwrap();
91        let want = reference(&a, &x, nrows, ncols);
92        assert_eq!(
93            y, want,
94            "gemv_transpose {nrows}x{ncols} mismatch vs reference"
95        );
96    }
97
98    #[test]
99    fn gemv_transpose_matches_reference_across_shapes() {
100        // Shapes exercising the TILE_N output-chunk remainder and the column tail.
101        for &(m, n) in &[
102            (1, 1),
103            (1, 33),
104            (4, 3),
105            (8, 8),
106            (13, 9),
107            (1, 64),
108            (31, 17),
109            (64, 33),
110            (64, 64),
111        ] {
112            run_case(m, n);
113        }
114    }
115
116    #[test]
117    fn gemv_transpose_accumulates_into_y() {
118        let a = vec![1.0f64, 2.0, 3.0, 4.0]; // 2x2: rows [1,2], [3,4]
119        let x = vec![1.0f64, 1.0];
120        let mut y = vec![10.0f64, 20.0];
121        gemv_transpose::dispatch_gemv_transpose::<f64>(&a, &x, &mut y, 2, 2).unwrap();
122        // y += Aᵀ·x = [10 + (1+3), 20 + (2+4)] = [14, 26]
123        assert_eq!(y, vec![14.0, 26.0]);
124    }
125
126    #[test]
127    fn gemv_transpose_rejects_short_operands() {
128        let a = vec![1.0f64; 4];
129        let x = vec![1.0f64; 2];
130        let mut y = vec![0.0f64; 1]; // too short for ncols=2
131        assert!(gemv_transpose::dispatch_gemv_transpose::<f64>(&a, &x, &mut y, 2, 2).is_err());
132    }
133}