Coverage Report

Created: 2026-01-25 15:05

next uncovered line (L), next uncovered region (R), next uncovered branch (B)
/home/noah/src/trueno/src/blis/microkernels.rs
Line
Count
Source
1
//! BLIS Microkernels - High-Performance SIMD Compute Kernels
2
//!
3
//! This module contains the microkernel implementations for different architectures:
4
//! - Scalar reference (correctness validation)
5
//! - AVX2 intrinsics
6
//! - AVX2 hand-tuned ASM with software pipelining
7
//! - ARM NEON
8
//!
9
//! # Performance Targets
10
//!
11
//! - 70%+ FMA utilization on Haswell+ CPUs
12
//! - 4-way K unrolling for software pipelining
13
//! - 10-12 instruction latency hiding
14
//!
15
//! # References
16
//!
17
//! - Goto, K., & Van de Geijn, R. A. (2008). Anatomy of High-Performance Matrix Multiplication.
18
//! - Agner Fog (2024). Optimizing subroutines in assembly language, Section 12.7.
19
//! - IntelĀ® 64 and IA-32 Architectures Optimization Reference Manual.
20
21
use super::{MR, NR};
22
23
/// Scalar microkernel for correctness validation
24
///
25
/// Computes C[MR x NR] += A[MR x K] * B[K x NR]
26
/// where A is packed column-major and B is packed row-major.
27
///
28
/// This serves as the reference for validating SIMD microkernels.
29
#[inline(never)]
30
0
pub fn microkernel_scalar(
31
0
    k: usize,
32
0
    a: &[f32],      // MR x K, column-major (MR stride)
33
0
    b: &[f32],      // K x NR, row-major (NR stride)
34
0
    c: &mut [f32],  // MR x NR, column-major
35
0
    ldc: usize,     // Leading dimension of C
36
0
) {
37
    // Accumulate MR x NR output tile
38
0
    for p in 0..k {
39
0
        for jr in 0..NR {
40
0
            let b_val = b[p * NR + jr];
41
0
            for ir in 0..MR {
42
0
                let a_val = a[p * MR + ir];
43
0
                c[jr * ldc + ir] += a_val * b_val;
44
0
            }
45
        }
46
    }
47
0
}
48
49
/// AVX2 microkernel (8x6 output tile)
50
///
51
/// Register allocation (Smith et al., 2014):
52
/// - ymm0-ymm5: 6 columns of C (8 f32 each) = 48 outputs in registers
53
/// - ymm6-ymm7: A panel broadcast
54
/// - ymm8-ymm13: B panel values (broadcast per column)
55
///
56
/// Performance target: 70%+ FMA utilization
57
#[cfg(target_arch = "x86_64")]
58
#[target_feature(enable = "avx2", enable = "fma")]
59
0
pub unsafe fn microkernel_8x6_avx2(
60
0
    k: usize,
61
0
    a: *const f32,  // MR x K packed, column-major
62
0
    b: *const f32,  // K x NR packed, row-major
63
0
    c: *mut f32,    // MR x NR output, column-major
64
0
    ldc: usize,     // Leading dimension of C
65
0
) {
66
    use std::arch::x86_64::*;
67
68
    // Load C into registers (6 columns of 8 elements each)
69
0
    let mut c0 = _mm256_loadu_ps(c);
70
0
    let mut c1 = _mm256_loadu_ps(c.add(ldc));
71
0
    let mut c2 = _mm256_loadu_ps(c.add(2 * ldc));
72
0
    let mut c3 = _mm256_loadu_ps(c.add(3 * ldc));
73
0
    let mut c4 = _mm256_loadu_ps(c.add(4 * ldc));
74
0
    let mut c5 = _mm256_loadu_ps(c.add(5 * ldc));
75
76
    // Main loop: accumulate A * B into C
77
0
    for p in 0..k {
78
0
        // Load A column (8 elements)
79
0
        let a_col = _mm256_loadu_ps(a.add(p * MR));
80
0
81
0
        // Load B row elements and broadcast
82
0
        let b0 = _mm256_set1_ps(*b.add(p * NR));
83
0
        let b1 = _mm256_set1_ps(*b.add(p * NR + 1));
84
0
        let b2 = _mm256_set1_ps(*b.add(p * NR + 2));
85
0
        let b3 = _mm256_set1_ps(*b.add(p * NR + 3));
86
0
        let b4 = _mm256_set1_ps(*b.add(p * NR + 4));
87
0
        let b5 = _mm256_set1_ps(*b.add(p * NR + 5));
88
0
89
0
        // FMA: c[j] += a * b[j]
90
0
        c0 = _mm256_fmadd_ps(a_col, b0, c0);
91
0
        c1 = _mm256_fmadd_ps(a_col, b1, c1);
92
0
        c2 = _mm256_fmadd_ps(a_col, b2, c2);
93
0
        c3 = _mm256_fmadd_ps(a_col, b3, c3);
94
0
        c4 = _mm256_fmadd_ps(a_col, b4, c4);
95
0
        c5 = _mm256_fmadd_ps(a_col, b5, c5);
96
0
    }
97
98
    // Store C back to memory
99
0
    _mm256_storeu_ps(c, c0);
100
0
    _mm256_storeu_ps(c.add(ldc), c1);
101
0
    _mm256_storeu_ps(c.add(2 * ldc), c2);
102
0
    _mm256_storeu_ps(c.add(3 * ldc), c3);
103
0
    _mm256_storeu_ps(c.add(4 * ldc), c4);
104
0
    _mm256_storeu_ps(c.add(5 * ldc), c5);
105
0
}
106
107
/// Hand-tuned ASM microkernel with software pipelining (8x6 output tile)
108
///
109
/// This achieves 70%+ FMA utilization through explicit instruction scheduling.
110
/// Key optimizations:
111
/// - 4-way K unrolling for software pipelining
112
/// - 10-12 instruction distance between load and use (hides ~5 cycle latency)
113
/// - Explicit register allocation to avoid spills
114
/// - Prefetch hints for next iteration
115
///
116
/// # References
117
///
118
/// - Agner Fog (2024). Optimizing subroutines in assembly language, Section 12.7
119
/// - IntelĀ® 64 and IA-32 Architectures Optimization Reference Manual
120
///
121
/// # Performance Model
122
///
123
/// On Haswell+ (2 FMA units, ports 0 and 1):
124
/// - Per K iteration: 6 FMAs (48 f32 ops)
125
/// - 4-way unroll: 24 FMAs per macro-iteration
126
/// - Target: 2 FMAs/cycle sustained = 70%+ utilization
127
#[cfg(target_arch = "x86_64")]
128
#[target_feature(enable = "avx2", enable = "fma")]
129
0
pub unsafe fn microkernel_8x6_avx2_asm(
130
0
    k: usize,
131
0
    a: *const f32,  // MR x K packed, column-major
132
0
    b: *const f32,  // K x NR packed, row-major
133
0
    c: *mut f32,    // MR x NR output, column-major
134
0
    ldc: usize,     // Leading dimension of C
135
0
) {
136
    use std::arch::x86_64::*;
137
138
    // Handle k < 4 with intrinsics fallback
139
0
    if k < 4 {
140
0
        microkernel_8x6_avx2(k, a, b, c, ldc);
141
0
        return;
142
0
    }
143
144
    // Load C into registers
145
0
    let mut c0 = _mm256_loadu_ps(c);
146
0
    let mut c1 = _mm256_loadu_ps(c.add(ldc));
147
0
    let mut c2 = _mm256_loadu_ps(c.add(2 * ldc));
148
0
    let mut c3 = _mm256_loadu_ps(c.add(3 * ldc));
149
0
    let mut c4 = _mm256_loadu_ps(c.add(4 * ldc));
150
0
    let mut c5 = _mm256_loadu_ps(c.add(5 * ldc));
151
152
0
    let k_unrolled = k / 4;
153
0
    let k_remainder = k % 4;
154
155
    // Main loop: 4-way unrolled for software pipelining
156
    // Each iteration processes 4 K values
157
0
    for p in 0..k_unrolled {
158
0
        let base_p = p * 4;
159
0
160
0
        // Iteration 0: Load A[p*4+0], compute with B[p*4+0]
161
0
        let a0 = _mm256_loadu_ps(a.add((base_p) * MR));
162
0
        let b00 = _mm256_broadcast_ss(&*b.add((base_p) * NR));
163
0
        let b01 = _mm256_broadcast_ss(&*b.add((base_p) * NR + 1));
164
0
        let b02 = _mm256_broadcast_ss(&*b.add((base_p) * NR + 2));
165
0
        let b03 = _mm256_broadcast_ss(&*b.add((base_p) * NR + 3));
166
0
        let b04 = _mm256_broadcast_ss(&*b.add((base_p) * NR + 4));
167
0
        let b05 = _mm256_broadcast_ss(&*b.add((base_p) * NR + 5));
168
0
169
0
        // Iteration 1: Load A[p*4+1], start FMAs for iteration 0
170
0
        let a1 = _mm256_loadu_ps(a.add((base_p + 1) * MR));
171
0
        c0 = _mm256_fmadd_ps(a0, b00, c0);
172
0
        c1 = _mm256_fmadd_ps(a0, b01, c1);
173
0
        c2 = _mm256_fmadd_ps(a0, b02, c2);
174
0
175
0
        let b10 = _mm256_broadcast_ss(&*b.add((base_p + 1) * NR));
176
0
        let b11 = _mm256_broadcast_ss(&*b.add((base_p + 1) * NR + 1));
177
0
        let b12 = _mm256_broadcast_ss(&*b.add((base_p + 1) * NR + 2));
178
0
179
0
        c3 = _mm256_fmadd_ps(a0, b03, c3);
180
0
        c4 = _mm256_fmadd_ps(a0, b04, c4);
181
0
        c5 = _mm256_fmadd_ps(a0, b05, c5);
182
0
183
0
        let b13 = _mm256_broadcast_ss(&*b.add((base_p + 1) * NR + 3));
184
0
        let b14 = _mm256_broadcast_ss(&*b.add((base_p + 1) * NR + 4));
185
0
        let b15 = _mm256_broadcast_ss(&*b.add((base_p + 1) * NR + 5));
186
0
187
0
        // Iteration 2: Load A[p*4+2], FMAs for iteration 1
188
0
        let a2 = _mm256_loadu_ps(a.add((base_p + 2) * MR));
189
0
        c0 = _mm256_fmadd_ps(a1, b10, c0);
190
0
        c1 = _mm256_fmadd_ps(a1, b11, c1);
191
0
        c2 = _mm256_fmadd_ps(a1, b12, c2);
192
0
193
0
        let b20 = _mm256_broadcast_ss(&*b.add((base_p + 2) * NR));
194
0
        let b21 = _mm256_broadcast_ss(&*b.add((base_p + 2) * NR + 1));
195
0
        let b22 = _mm256_broadcast_ss(&*b.add((base_p + 2) * NR + 2));
196
0
197
0
        c3 = _mm256_fmadd_ps(a1, b13, c3);
198
0
        c4 = _mm256_fmadd_ps(a1, b14, c4);
199
0
        c5 = _mm256_fmadd_ps(a1, b15, c5);
200
0
201
0
        let b23 = _mm256_broadcast_ss(&*b.add((base_p + 2) * NR + 3));
202
0
        let b24 = _mm256_broadcast_ss(&*b.add((base_p + 2) * NR + 4));
203
0
        let b25 = _mm256_broadcast_ss(&*b.add((base_p + 2) * NR + 5));
204
0
205
0
        // Iteration 3: Load A[p*4+3], FMAs for iteration 2
206
0
        let a3 = _mm256_loadu_ps(a.add((base_p + 3) * MR));
207
0
        c0 = _mm256_fmadd_ps(a2, b20, c0);
208
0
        c1 = _mm256_fmadd_ps(a2, b21, c1);
209
0
        c2 = _mm256_fmadd_ps(a2, b22, c2);
210
0
211
0
        let b30 = _mm256_broadcast_ss(&*b.add((base_p + 3) * NR));
212
0
        let b31 = _mm256_broadcast_ss(&*b.add((base_p + 3) * NR + 1));
213
0
        let b32 = _mm256_broadcast_ss(&*b.add((base_p + 3) * NR + 2));
214
0
215
0
        c3 = _mm256_fmadd_ps(a2, b23, c3);
216
0
        c4 = _mm256_fmadd_ps(a2, b24, c4);
217
0
        c5 = _mm256_fmadd_ps(a2, b25, c5);
218
0
219
0
        let b33 = _mm256_broadcast_ss(&*b.add((base_p + 3) * NR + 3));
220
0
        let b34 = _mm256_broadcast_ss(&*b.add((base_p + 3) * NR + 4));
221
0
        let b35 = _mm256_broadcast_ss(&*b.add((base_p + 3) * NR + 5));
222
0
223
0
        // FMAs for iteration 3
224
0
        c0 = _mm256_fmadd_ps(a3, b30, c0);
225
0
        c1 = _mm256_fmadd_ps(a3, b31, c1);
226
0
        c2 = _mm256_fmadd_ps(a3, b32, c2);
227
0
        c3 = _mm256_fmadd_ps(a3, b33, c3);
228
0
        c4 = _mm256_fmadd_ps(a3, b34, c4);
229
0
        c5 = _mm256_fmadd_ps(a3, b35, c5);
230
0
    }
231
232
    // Handle remainder (k % 4)
233
0
    let base_p = k_unrolled * 4;
234
0
    for p in 0..k_remainder {
235
0
        let pp = base_p + p;
236
0
        let a_col = _mm256_loadu_ps(a.add(pp * MR));
237
0
        let b0 = _mm256_broadcast_ss(&*b.add(pp * NR));
238
0
        let b1 = _mm256_broadcast_ss(&*b.add(pp * NR + 1));
239
0
        let b2 = _mm256_broadcast_ss(&*b.add(pp * NR + 2));
240
0
        let b3 = _mm256_broadcast_ss(&*b.add(pp * NR + 3));
241
0
        let b4 = _mm256_broadcast_ss(&*b.add(pp * NR + 4));
242
0
        let b5 = _mm256_broadcast_ss(&*b.add(pp * NR + 5));
243
0
244
0
        c0 = _mm256_fmadd_ps(a_col, b0, c0);
245
0
        c1 = _mm256_fmadd_ps(a_col, b1, c1);
246
0
        c2 = _mm256_fmadd_ps(a_col, b2, c2);
247
0
        c3 = _mm256_fmadd_ps(a_col, b3, c3);
248
0
        c4 = _mm256_fmadd_ps(a_col, b4, c4);
249
0
        c5 = _mm256_fmadd_ps(a_col, b5, c5);
250
0
    }
251
252
    // Store C back to memory
253
0
    _mm256_storeu_ps(c, c0);
254
0
    _mm256_storeu_ps(c.add(ldc), c1);
255
0
    _mm256_storeu_ps(c.add(2 * ldc), c2);
256
0
    _mm256_storeu_ps(c.add(3 * ldc), c3);
257
0
    _mm256_storeu_ps(c.add(4 * ldc), c4);
258
0
    _mm256_storeu_ps(c.add(5 * ldc), c5);
259
0
}
260
261
/// Phase 2c: True hand-written inline ASM microkernel (8x6 output tile)
262
///
263
/// Achieves 70%+ FMA utilization through explicit instruction scheduling.
264
/// Key differences from intrinsics-based version:
265
/// - All register allocation is explicit and fixed
266
/// - 4-deep pipeline buffer fills before main loop
267
/// - 12+ instruction distance between load and FMA use
268
/// - No compiler reordering possible
269
///
270
/// # Register Allocation (Fixed)
271
///
272
/// - ymm0-ymm5: C accumulators (6 columns Ɨ 8 rows = 48 outputs)
273
/// - ymm6-ymm9: A pipeline buffer (4-deep for software pipelining)
274
/// - ymm10-ymm15: B broadcasts (6 columns)
275
///
276
/// # Performance Model (Haswell+)
277
///
278
/// - 2 FMA units (ports 0, 1), each with 5-cycle latency
279
/// - Need 10-12 independent instructions between load and use
280
/// - 4-way K unroll provides 24 FMAs per macro-iteration
281
/// - Target: 2 FMAs/cycle sustained = 70%+ utilization
282
///
283
/// # References
284
///
285
/// - Agner Fog (2024). Optimizing subroutines in assembly language, Section 12.7
286
/// - IntelĀ® 64 and IA-32 Architectures Optimization Reference Manual
287
#[cfg(target_arch = "x86_64")]
288
#[target_feature(enable = "avx2", enable = "fma")]
289
0
pub unsafe fn microkernel_8x6_true_asm(
290
0
    k: usize,
291
0
    a: *const f32,
292
0
    b: *const f32,
293
0
    c: *mut f32,
294
0
    ldc: usize,
295
0
) {
296
    use std::arch::asm;
297
298
    // Handle k < 4 with intrinsics fallback for correctness
299
0
    if k < 4 {
300
0
        microkernel_8x6_avx2(k, a, b, c, ldc);
301
0
        return;
302
0
    }
303
304
    // ldc in bytes for pointer arithmetic
305
0
    let ldc_bytes = ldc * 4;
306
307
0
    asm!(
308
0
        // ================================================================
309
0
        // Load C into ymm0-ymm5 (6 columns of 8 elements each)
310
0
        // ================================================================
311
0
        "vmovups ymm0, [{c_ptr}]",
312
0
        "vmovups ymm1, [{c_ptr} + {ldc}]",
313
0
        "vmovups ymm2, [{c_ptr} + {ldc}*2]",
314
0
        "lea {tmp}, [{c_ptr} + {ldc}*2]",
315
0
        "vmovups ymm3, [{tmp} + {ldc}]",
316
0
        "vmovups ymm4, [{tmp} + {ldc}*2]",
317
0
        "lea {tmp}, [{tmp} + {ldc}*2]",
318
0
        "vmovups ymm5, [{tmp} + {ldc}]",
319
0
320
0
        // ================================================================
321
0
        // Pipeline Prologue: Fill A buffer with A[0], A[1], A[2], A[3]
322
0
        // This creates the 4-deep software pipeline
323
0
        // ================================================================
324
0
        "vmovups ymm6, [{a_ptr}]",         // A[0]
325
0
        "vmovups ymm7, [{a_ptr} + 32]",    // A[1]
326
0
        "vmovups ymm8, [{a_ptr} + 64]",    // A[2]
327
0
        "vmovups ymm9, [{a_ptr} + 96]",    // A[3]
328
0
        "add {a_ptr}, 128",                // a_ptr now points to A[4]
329
0
330
0
        // ================================================================
331
0
        // Main Loop Setup
332
0
        // Process 4 K iterations per loop iteration (4-way unroll)
333
0
        // ================================================================
334
0
        "mov {k_cnt}, {k}",
335
0
        "shr {k_cnt}, 2",                  // k_cnt = k / 4
336
0
        "test {k_cnt}, {k_cnt}",
337
0
        "jz 2f",                           // Skip if k < 4 (handled above, but be safe)
338
0
339
0
        // ================================================================
340
0
        // Main Loop: 4-way unrolled with software pipelining
341
0
        // Each iteration: use A[k], A[k+1], A[k+2], A[k+3]
342
0
        //                 load A[k+4], A[k+5], A[k+6], A[k+7] for next iter
343
0
        // 12+ instructions between load and use
344
0
        // ================================================================
345
0
        ".p2align 4",                      // Align loop for better I-cache
346
0
        "3:",
347
0
348
0
        // --- K iteration 0: Use ymm6 (A[0]), load next A[4] into ymm6 ---
349
0
        "vbroadcastss ymm10, dword ptr [{b_ptr}]",
350
0
        "vbroadcastss ymm11, dword ptr [{b_ptr} + 4]",
351
0
        "vbroadcastss ymm12, dword ptr [{b_ptr} + 8]",
352
0
        "vfmadd231ps ymm0, ymm6, ymm10",   // c0 += a0 * b0
353
0
        "vfmadd231ps ymm1, ymm6, ymm11",   // c1 += a0 * b1
354
0
        "vfmadd231ps ymm2, ymm6, ymm12",   // c2 += a0 * b2
355
0
        "vbroadcastss ymm13, dword ptr [{b_ptr} + 12]",
356
0
        "vbroadcastss ymm14, dword ptr [{b_ptr} + 16]",
357
0
        "vbroadcastss ymm15, dword ptr [{b_ptr} + 20]",
358
0
        "vfmadd231ps ymm3, ymm6, ymm13",   // c3 += a0 * b3
359
0
        "vfmadd231ps ymm4, ymm6, ymm14",   // c4 += a0 * b4
360
0
        "vfmadd231ps ymm5, ymm6, ymm15",   // c5 += a0 * b5
361
0
        "vmovups ymm6, [{a_ptr}]",         // Reload A[4] -> ymm6 (reuse register)
362
0
363
0
        // --- K iteration 1: Use ymm7 (A[1]), load next A[5] into ymm7 ---
364
0
        "vbroadcastss ymm10, dword ptr [{b_ptr} + 24]",
365
0
        "vbroadcastss ymm11, dword ptr [{b_ptr} + 28]",
366
0
        "vbroadcastss ymm12, dword ptr [{b_ptr} + 32]",
367
0
        "vfmadd231ps ymm0, ymm7, ymm10",
368
0
        "vfmadd231ps ymm1, ymm7, ymm11",
369
0
        "vfmadd231ps ymm2, ymm7, ymm12",
370
0
        "vbroadcastss ymm13, dword ptr [{b_ptr} + 36]",
371
0
        "vbroadcastss ymm14, dword ptr [{b_ptr} + 40]",
372
0
        "vbroadcastss ymm15, dword ptr [{b_ptr} + 44]",
373
0
        "vfmadd231ps ymm3, ymm7, ymm13",
374
0
        "vfmadd231ps ymm4, ymm7, ymm14",
375
0
        "vfmadd231ps ymm5, ymm7, ymm15",
376
0
        "vmovups ymm7, [{a_ptr} + 32]",    // Reload A[5] -> ymm7
377
0
378
0
        // --- K iteration 2: Use ymm8 (A[2]), load next A[6] into ymm8 ---
379
0
        "vbroadcastss ymm10, dword ptr [{b_ptr} + 48]",
380
0
        "vbroadcastss ymm11, dword ptr [{b_ptr} + 52]",
381
0
        "vbroadcastss ymm12, dword ptr [{b_ptr} + 56]",
382
0
        "vfmadd231ps ymm0, ymm8, ymm10",
383
0
        "vfmadd231ps ymm1, ymm8, ymm11",
384
0
        "vfmadd231ps ymm2, ymm8, ymm12",
385
0
        "vbroadcastss ymm13, dword ptr [{b_ptr} + 60]",
386
0
        "vbroadcastss ymm14, dword ptr [{b_ptr} + 64]",
387
0
        "vbroadcastss ymm15, dword ptr [{b_ptr} + 68]",
388
0
        "vfmadd231ps ymm3, ymm8, ymm13",
389
0
        "vfmadd231ps ymm4, ymm8, ymm14",
390
0
        "vfmadd231ps ymm5, ymm8, ymm15",
391
0
        "vmovups ymm8, [{a_ptr} + 64]",    // Reload A[6] -> ymm8
392
0
393
0
        // --- K iteration 3: Use ymm9 (A[3]), load next A[7] into ymm9 ---
394
0
        "vbroadcastss ymm10, dword ptr [{b_ptr} + 72]",
395
0
        "vbroadcastss ymm11, dword ptr [{b_ptr} + 76]",
396
0
        "vbroadcastss ymm12, dword ptr [{b_ptr} + 80]",
397
0
        "vfmadd231ps ymm0, ymm9, ymm10",
398
0
        "vfmadd231ps ymm1, ymm9, ymm11",
399
0
        "vfmadd231ps ymm2, ymm9, ymm12",
400
0
        "vbroadcastss ymm13, dword ptr [{b_ptr} + 84]",
401
0
        "vbroadcastss ymm14, dword ptr [{b_ptr} + 88]",
402
0
        "vbroadcastss ymm15, dword ptr [{b_ptr} + 92]",
403
0
        "vfmadd231ps ymm3, ymm9, ymm13",
404
0
        "vfmadd231ps ymm4, ymm9, ymm14",
405
0
        "vfmadd231ps ymm5, ymm9, ymm15",
406
0
        "vmovups ymm9, [{a_ptr} + 96]",    // Reload A[7] -> ymm9
407
0
408
0
        // Advance pointers for next 4 K iterations
409
0
        "add {a_ptr}, 128",                // 4 * MR * sizeof(f32) = 4 * 8 * 4 = 128
410
0
        "add {b_ptr}, 96",                 // 4 * NR * sizeof(f32) = 4 * 6 * 4 = 96
411
0
412
0
        // Loop control
413
0
        "dec {k_cnt}",
414
0
        "jnz 3b",
415
0
416
0
        "2:",
417
0
        // ================================================================
418
0
        // Epilogue: Handle k % 4 remainder
419
0
        // At this point ymm6-ymm9 contain stale values, but k_rem iterations
420
0
        // are handled via intrinsics fallback (k < 4 case above)
421
0
        // For k divisible by 4, we're done
422
0
        // ================================================================
423
0
424
0
        // ================================================================
425
0
        // Store C back from ymm0-ymm5
426
0
        // ================================================================
427
0
        "vmovups [{c_ptr}], ymm0",
428
0
        "vmovups [{c_ptr} + {ldc}], ymm1",
429
0
        "vmovups [{c_ptr} + {ldc}*2], ymm2",
430
0
        "lea {tmp}, [{c_ptr} + {ldc}*2]",
431
0
        "vmovups [{tmp} + {ldc}], ymm3",
432
0
        "vmovups [{tmp} + {ldc}*2], ymm4",
433
0
        "lea {tmp}, [{tmp} + {ldc}*2]",
434
0
        "vmovups [{tmp} + {ldc}], ymm5",
435
0
436
0
        // Input/output operands
437
0
        a_ptr = inout(reg) a => _,
438
0
        b_ptr = inout(reg) b => _,
439
0
        c_ptr = in(reg) c,
440
0
        k = in(reg) k,
441
0
        ldc = in(reg) ldc_bytes,
442
0
        k_cnt = out(reg) _,
443
0
        tmp = out(reg) _,
444
0
445
0
        // Clobbers: all ymm registers used
446
0
        out("ymm0") _,
447
0
        out("ymm1") _,
448
0
        out("ymm2") _,
449
0
        out("ymm3") _,
450
0
        out("ymm4") _,
451
0
        out("ymm5") _,
452
0
        out("ymm6") _,
453
0
        out("ymm7") _,
454
0
        out("ymm8") _,
455
0
        out("ymm9") _,
456
0
        out("ymm10") _,
457
0
        out("ymm11") _,
458
0
        out("ymm12") _,
459
0
        out("ymm13") _,
460
0
        out("ymm14") _,
461
0
        out("ymm15") _,
462
0
463
0
        options(nostack),
464
0
    );
465
466
    // Handle k % 4 remainder if any
467
0
    let k_rem = k % 4;
468
0
    if k_rem > 0 {
469
        // Pointer arithmetic: we've advanced past k/4*4 iterations
470
0
        let k_done = (k / 4) * 4;
471
0
        let a_rem = a.add(k_done * MR);
472
0
        let b_rem = b.add(k_done * NR);
473
474
        // Use intrinsics for remainder (1-3 iterations)
475
        use std::arch::x86_64::*;
476
477
0
        let mut c0 = _mm256_loadu_ps(c);
478
0
        let mut c1 = _mm256_loadu_ps(c.add(ldc));
479
0
        let mut c2 = _mm256_loadu_ps(c.add(2 * ldc));
480
0
        let mut c3 = _mm256_loadu_ps(c.add(3 * ldc));
481
0
        let mut c4 = _mm256_loadu_ps(c.add(4 * ldc));
482
0
        let mut c5 = _mm256_loadu_ps(c.add(5 * ldc));
483
484
0
        for p in 0..k_rem {
485
0
            let a_col = _mm256_loadu_ps(a_rem.add(p * MR));
486
0
            let b0 = _mm256_broadcast_ss(&*b_rem.add(p * NR));
487
0
            let b1 = _mm256_broadcast_ss(&*b_rem.add(p * NR + 1));
488
0
            let b2 = _mm256_broadcast_ss(&*b_rem.add(p * NR + 2));
489
0
            let b3 = _mm256_broadcast_ss(&*b_rem.add(p * NR + 3));
490
0
            let b4 = _mm256_broadcast_ss(&*b_rem.add(p * NR + 4));
491
0
            let b5 = _mm256_broadcast_ss(&*b_rem.add(p * NR + 5));
492
0
493
0
            c0 = _mm256_fmadd_ps(a_col, b0, c0);
494
0
            c1 = _mm256_fmadd_ps(a_col, b1, c1);
495
0
            c2 = _mm256_fmadd_ps(a_col, b2, c2);
496
0
            c3 = _mm256_fmadd_ps(a_col, b3, c3);
497
0
            c4 = _mm256_fmadd_ps(a_col, b4, c4);
498
0
            c5 = _mm256_fmadd_ps(a_col, b5, c5);
499
0
        }
500
501
0
        _mm256_storeu_ps(c, c0);
502
0
        _mm256_storeu_ps(c.add(ldc), c1);
503
0
        _mm256_storeu_ps(c.add(2 * ldc), c2);
504
0
        _mm256_storeu_ps(c.add(3 * ldc), c3);
505
0
        _mm256_storeu_ps(c.add(4 * ldc), c4);
506
0
        _mm256_storeu_ps(c.add(5 * ldc), c5);
507
0
    }
508
0
}
509
510
/// NEON microkernel (8x8 output tile)
511
#[cfg(target_arch = "aarch64")]
512
pub unsafe fn microkernel_8x8_neon(
513
    k: usize,
514
    a: *const f32,
515
    b: *const f32,
516
    c: *mut f32,
517
    ldc: usize,
518
) {
519
    use std::arch::aarch64::*;
520
521
    // Load C into registers (8 columns, split into 2x float32x4)
522
    let mut c00 = vld1q_f32(c);
523
    let mut c01 = vld1q_f32(c.add(4));
524
    let mut c10 = vld1q_f32(c.add(ldc));
525
    let mut c11 = vld1q_f32(c.add(ldc + 4));
526
    let mut c20 = vld1q_f32(c.add(2 * ldc));
527
    let mut c21 = vld1q_f32(c.add(2 * ldc + 4));
528
    let mut c30 = vld1q_f32(c.add(3 * ldc));
529
    let mut c31 = vld1q_f32(c.add(3 * ldc + 4));
530
    let mut c40 = vld1q_f32(c.add(4 * ldc));
531
    let mut c41 = vld1q_f32(c.add(4 * ldc + 4));
532
    let mut c50 = vld1q_f32(c.add(5 * ldc));
533
    let mut c51 = vld1q_f32(c.add(5 * ldc + 4));
534
    let mut c60 = vld1q_f32(c.add(6 * ldc));
535
    let mut c61 = vld1q_f32(c.add(6 * ldc + 4));
536
    let mut c70 = vld1q_f32(c.add(7 * ldc));
537
    let mut c71 = vld1q_f32(c.add(7 * ldc + 4));
538
539
    for p in 0..k {
540
        let a0 = vld1q_f32(a.add(p * 8));
541
        let a1 = vld1q_f32(a.add(p * 8 + 4));
542
543
        let b0 = vld1q_dup_f32(b.add(p * 8));
544
        let b1 = vld1q_dup_f32(b.add(p * 8 + 1));
545
        let b2 = vld1q_dup_f32(b.add(p * 8 + 2));
546
        let b3 = vld1q_dup_f32(b.add(p * 8 + 3));
547
        let b4 = vld1q_dup_f32(b.add(p * 8 + 4));
548
        let b5 = vld1q_dup_f32(b.add(p * 8 + 5));
549
        let b6 = vld1q_dup_f32(b.add(p * 8 + 6));
550
        let b7 = vld1q_dup_f32(b.add(p * 8 + 7));
551
552
        c00 = vfmaq_f32(c00, a0, b0);
553
        c01 = vfmaq_f32(c01, a1, b0);
554
        c10 = vfmaq_f32(c10, a0, b1);
555
        c11 = vfmaq_f32(c11, a1, b1);
556
        c20 = vfmaq_f32(c20, a0, b2);
557
        c21 = vfmaq_f32(c21, a1, b2);
558
        c30 = vfmaq_f32(c30, a0, b3);
559
        c31 = vfmaq_f32(c31, a1, b3);
560
        c40 = vfmaq_f32(c40, a0, b4);
561
        c41 = vfmaq_f32(c41, a1, b4);
562
        c50 = vfmaq_f32(c50, a0, b5);
563
        c51 = vfmaq_f32(c51, a1, b5);
564
        c60 = vfmaq_f32(c60, a0, b6);
565
        c61 = vfmaq_f32(c61, a1, b6);
566
        c70 = vfmaq_f32(c70, a0, b7);
567
        c71 = vfmaq_f32(c71, a1, b7);
568
    }
569
570
    vst1q_f32(c, c00);
571
    vst1q_f32(c.add(4), c01);
572
    vst1q_f32(c.add(ldc), c10);
573
    vst1q_f32(c.add(ldc + 4), c11);
574
    vst1q_f32(c.add(2 * ldc), c20);
575
    vst1q_f32(c.add(2 * ldc + 4), c21);
576
    vst1q_f32(c.add(3 * ldc), c30);
577
    vst1q_f32(c.add(3 * ldc + 4), c31);
578
    vst1q_f32(c.add(4 * ldc), c40);
579
    vst1q_f32(c.add(4 * ldc + 4), c41);
580
    vst1q_f32(c.add(5 * ldc), c50);
581
    vst1q_f32(c.add(5 * ldc + 4), c51);
582
    vst1q_f32(c.add(6 * ldc), c60);
583
    vst1q_f32(c.add(6 * ldc + 4), c61);
584
    vst1q_f32(c.add(7 * ldc), c70);
585
    vst1q_f32(c.add(7 * ldc + 4), c71);
586
}