1use 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 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 _ => unsafe { core::hint::unreachable_unchecked() },
64 }
65}
66
67#[cfg(test)]
68mod tests {
69 use crate::dispatch::gemv_transpose;
70
71 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 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 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]; 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 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]; assert!(gemv_transpose::dispatch_gemv_transpose::<f64>(&a, &x, &mut y, 2, 2).is_err());
132 }
133}