nam_rs/math/gemm/gemv/
kernel_macro.rs

1// SPDX-License-Identifier: Apache-2.0
2// Copyright (c) 2026 Fábio Henrique de Lima Silva (fhl.bsb@gmail.com) All rights reserved.
3
4/// GEMV kernel macro — generates platform-specific FMADD accumulate loops.
5///
6/// Parameterized by SIMD width (4 = AVX2/256-bit, 8 = AVX-512/512-bit)
7/// and all relevant SIMD operations as inline closures.
8/// Both variants use 8 independent accumulators to maximize FMA throughput.
9///
10/// # Safety
11///
12/// Caller must ensure valid pointer arithmetic and slice bounds.
13#[macro_export]
14macro_rules! gemv_kernel {
15    (
16        4,
17        $is_fused:expr,
18        $out_c:expr,
19        $out_len:expr,
20        $in_frame:expr,
21        $weights:expr,
22        $bias:expr,
23        $out_frame:expr,
24        $do_bias:expr,
25        $setzero:expr,
26        $load_out:expr,
27        $load_bias:expr,
28        $add_ps:expr,
29        $load_weight:expr,
30        $fmadd_ps:expr,
31        $store_ps:expr
32    ) => {
33        let mut acc0 = if $is_fused {
34            let mut acc = $load_out($out_c);
35            if $do_bias {
36                acc = $add_ps(acc, $load_bias($out_c));
37            }
38            acc
39        } else {
40            if $do_bias {
41                $load_bias($out_c)
42            } else {
43                $setzero()
44            }
45        };
46        let mut acc1 = $setzero();
47        let mut acc2 = $setzero();
48        let mut acc3 = $setzero();
49        let mut acc4 = $setzero();
50        let mut acc5 = $setzero();
51        let mut acc6 = $setzero();
52        let mut acc7 = $setzero();
53
54        let mut in_c = 0;
55        let in_len = $in_frame.len();
56        while in_c + 8 <= in_len {
57            _mm_prefetch::<_MM_HINT_T0>($in_frame.as_ptr().wrapping_add(in_c + 64) as *const i8);
58
59            let vs0 = _mm256_set1_ps(*$in_frame.get_unchecked(in_c));
60            let vs1 = _mm256_set1_ps(*$in_frame.get_unchecked(in_c + 1));
61            let vs2 = _mm256_set1_ps(*$in_frame.get_unchecked(in_c + 2));
62            let vs3 = _mm256_set1_ps(*$in_frame.get_unchecked(in_c + 3));
63            let vs4 = _mm256_set1_ps(*$in_frame.get_unchecked(in_c + 4));
64            let vs5 = _mm256_set1_ps(*$in_frame.get_unchecked(in_c + 5));
65            let vs6 = _mm256_set1_ps(*$in_frame.get_unchecked(in_c + 6));
66            let vs7 = _mm256_set1_ps(*$in_frame.get_unchecked(in_c + 7));
67
68            let w_ptr = $weights.as_ptr().add(in_c * $out_len + $out_c);
69            let w0 = $load_weight(w_ptr);
70            acc0 = $fmadd_ps(vs0, w0, acc0);
71
72            let w1 = $load_weight(w_ptr.add($out_len));
73            acc1 = $fmadd_ps(vs1, w1, acc1);
74
75            let w2 = $load_weight(w_ptr.add(2 * $out_len));
76            acc2 = $fmadd_ps(vs2, w2, acc2);
77
78            let w3 = $load_weight(w_ptr.add(3 * $out_len));
79            acc3 = $fmadd_ps(vs3, w3, acc3);
80
81            let w4 = $load_weight(w_ptr.add(4 * $out_len));
82            acc4 = $fmadd_ps(vs4, w4, acc4);
83
84            let w5 = $load_weight(w_ptr.add(5 * $out_len));
85            acc5 = $fmadd_ps(vs5, w5, acc5);
86
87            let w6 = $load_weight(w_ptr.add(6 * $out_len));
88            acc6 = $fmadd_ps(vs6, w6, acc6);
89
90            let w7 = $load_weight(w_ptr.add(7 * $out_len));
91            acc7 = $fmadd_ps(vs7, w7, acc7);
92
93            in_c += 8;
94        }
95
96        acc0 = $add_ps(acc0, acc1);
97        acc2 = $add_ps(acc2, acc3);
98        acc4 = $add_ps(acc4, acc5);
99        acc6 = $add_ps(acc6, acc7);
100        acc0 = $add_ps(acc0, acc2);
101        acc4 = $add_ps(acc4, acc6);
102        acc0 = $add_ps(acc0, acc4);
103
104        while in_c < in_len {
105            let vs = _mm256_set1_ps(*$in_frame.get_unchecked(in_c));
106            let weight_ptr = $weights.as_ptr().add(in_c * $out_len + $out_c);
107            let vw = $load_weight(weight_ptr);
108            acc0 = $fmadd_ps(vs, vw, acc0);
109            in_c += 1;
110        }
111
112        $store_ps($out_c, acc0);
113    };
114
115    (
116        8,
117        $is_fused:expr,
118        $out_c:expr,
119        $out_len:expr,
120        $in_frame:expr,
121        $weights:expr,
122        $bias:expr,
123        $out_frame:expr,
124        $do_bias:expr,
125        $setzero:expr,
126        $load_out:expr,
127        $load_bias:expr,
128        $add_ps:expr,
129        $load_weight:expr,
130        $fmadd_ps:expr,
131        $store_ps:expr
132    ) => {
133        let mut acc0 = if $is_fused {
134            let mut acc = $load_out($out_c);
135            if $do_bias {
136                acc = $add_ps(acc, $load_bias($out_c));
137            }
138            acc
139        } else {
140            if $do_bias {
141                $load_bias($out_c)
142            } else {
143                $setzero()
144            }
145        };
146        let mut acc1 = $setzero();
147        let mut acc2 = $setzero();
148        let mut acc3 = $setzero();
149        let mut acc4 = $setzero();
150        let mut acc5 = $setzero();
151        let mut acc6 = $setzero();
152        let mut acc7 = $setzero();
153
154        let mut in_c = 0;
155        let in_len = $in_frame.len();
156        while in_c + 8 <= in_len {
157            _mm_prefetch::<_MM_HINT_T0>($in_frame.as_ptr().wrapping_add(in_c + 64) as *const i8);
158
159            let vs0 = _mm512_set1_ps(*$in_frame.get_unchecked(in_c));
160            let vs1 = _mm512_set1_ps(*$in_frame.get_unchecked(in_c + 1));
161            let vs2 = _mm512_set1_ps(*$in_frame.get_unchecked(in_c + 2));
162            let vs3 = _mm512_set1_ps(*$in_frame.get_unchecked(in_c + 3));
163            let vs4 = _mm512_set1_ps(*$in_frame.get_unchecked(in_c + 4));
164            let vs5 = _mm512_set1_ps(*$in_frame.get_unchecked(in_c + 5));
165            let vs6 = _mm512_set1_ps(*$in_frame.get_unchecked(in_c + 6));
166            let vs7 = _mm512_set1_ps(*$in_frame.get_unchecked(in_c + 7));
167
168            let w_ptr = $weights.as_ptr().add(in_c * $out_len + $out_c);
169            let w0 = $load_weight(w_ptr);
170            acc0 = $fmadd_ps(vs0, w0, acc0);
171
172            let w1 = $load_weight(w_ptr.add($out_len));
173            acc1 = $fmadd_ps(vs1, w1, acc1);
174
175            let w2 = $load_weight(w_ptr.add(2 * $out_len));
176            acc2 = $fmadd_ps(vs2, w2, acc2);
177
178            let w3 = $load_weight(w_ptr.add(3 * $out_len));
179            acc3 = $fmadd_ps(vs3, w3, acc3);
180
181            let w4 = $load_weight(w_ptr.add(4 * $out_len));
182            acc4 = $fmadd_ps(vs4, w4, acc4);
183
184            let w5 = $load_weight(w_ptr.add(5 * $out_len));
185            acc5 = $fmadd_ps(vs5, w5, acc5);
186
187            let w6 = $load_weight(w_ptr.add(6 * $out_len));
188            acc6 = $fmadd_ps(vs6, w6, acc6);
189
190            let w7 = $load_weight(w_ptr.add(7 * $out_len));
191            acc7 = $fmadd_ps(vs7, w7, acc7);
192
193            in_c += 8;
194        }
195
196        acc0 = $add_ps(acc0, acc1);
197        acc2 = $add_ps(acc2, acc3);
198        acc4 = $add_ps(acc4, acc5);
199        acc6 = $add_ps(acc6, acc7);
200        acc0 = $add_ps(acc0, acc2);
201        acc4 = $add_ps(acc4, acc6);
202        acc0 = $add_ps(acc0, acc4);
203
204        while in_c < in_len {
205            let vs = _mm512_set1_ps(*$in_frame.get_unchecked(in_c));
206            let weight_ptr = $weights.as_ptr().add(in_c * $out_len + $out_c);
207            let vw = $load_weight(weight_ptr);
208            acc0 = $fmadd_ps(vs, vw, acc0);
209            in_c += 1;
210        }
211
212        $store_ps($out_c, acc0);
213    };
214}
nam_rs/math/gemm/gemv/kernel_macro.rs

nam_rs/math/gemm/gemv/
kernel_macro.rs
