## 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/mod.rs

    Line
    Count
    Source
    1
    //! BLIS-Style Matrix Multiplication
    2
    //!
    3
    //! High-performance GEMM implementation based on the BLIS framework.
    4
    //!
    5
    //! # References
    6
    //!
    7
    //! - Goto, K., & Van de Geijn, R. A. (2008). Anatomy of High-Performance Matrix Multiplication.
    8
    //!   ACM TOMS, 34(3). <https://doi.org/10.1145/1356052.1356053>
    9
    //! - Van Zee, F. G., & Van de Geijn, R. A. (2015). BLIS: A Framework for Rapidly Instantiating
    10
    //!   BLAS Functionality. ACM TOMS, 41(3). <https://doi.org/10.1145/2764454>
    11
    //! - Low, T. M., et al. (2016). Analytical Modeling Is Enough for High-Performance BLIS.
    12
    //!   ACM TOMS, 43(2). <https://doi.org/10.1145/2925987>
    13
    //!
    14
    //! # Toyota Production System Integration
    15
    //!
    16
    //! - **Jidoka**: Runtime guards that stop on numerical errors (see [`jidoka`] module)
    17
    //! - **Poka-Yoke**: Compile-time type safety for panel dimensions
    18
    //! - **Heijunka**: Load-balanced parallel execution
    19
    //! - **Kaizen**: Performance tracking for continuous improvement (see [`profiler`] module)
    20
    //!
    21
    //! # Module Structure
    22
    //!
    23
    //! - [`jidoka`]: Runtime validation guards (stop-on-defect)
    24
    //! - [`profiler`]: Performance tracking at all BLIS hierarchy levels
    25
    //! - [`microkernels`]: High-performance SIMD compute kernels
    26
    //! - [`backend_selection`]: Automatic CPU/GPU backend selection
    27
    28
    pub mod backend_selection;
    29
    pub mod jidoka;
    30
    pub mod microkernels;
    31
    pub mod profiler;
    32
    33
    // Re-export jidoka types for backwards compatibility
    34
    pub use jidoka::{JidokaError, JidokaGuard};
    35
    36
    // Re-export profiler types for backwards compatibility
    37
    pub use profiler::{BlisLevelStats, BlisProfileLevel, BlisProfiler, KaizenMetrics};
    38
    39
    // Re-export microkernel functions
    40
    pub use microkernels::{microkernel_scalar, microkernel_8x6_avx2, microkernel_8x6_avx2_asm, microkernel_8x6_true_asm};
    41
    #[cfg(target_arch = "aarch64")]
    42
    pub use microkernels::microkernel_8x8_neon;
    43
    44
    // Re-export backend selection types
    45
    pub use backend_selection::{
    46
        BackendCostModel, BrickLevel, ComputeBackend, PtxMicrokernelSpec, RooflineResult,
    47
        UnifiedBrickProfiler, WgslMicrokernelSpec, gemm_auto,
    48
    };
    49
    50
    use std::time::Instant;
    51
    52
    use crate::error::TruenoError;
    53
    54
    // ============================================================================
    55
    // BLIS Configuration Constants
    56
    // ============================================================================
    57
    58
    /// Microkernel row dimension (AVX2: 8 f32 per ymm register)
    59
    pub const MR: usize = 8;
    60
    61
    /// Microkernel column dimension (6 columns fit in remaining registers)
    62
    pub const NR: usize = 6;
    63
    64
    /// K-dimension blocking for L1 cache (256 elements = 1KB)
    65
    pub const KC: usize = 256;
    66
    67
    /// M-dimension blocking for L2 cache
    68
    pub const MC: usize = 72;
    69
    70
    /// N-dimension blocking for L3 cache
    71
    pub const NC: usize = 4096;
    72
    73
    // ============================================================================
    74
    // Phase 1: Scalar Reference Implementation
    75
    // ============================================================================
    76
    77
    /// Scalar reference GEMM for Jidoka validation
    78
    ///
    79
    /// Computes C += A * B where:
    80
    /// - A is M x K (row-major)
    81
    /// - B is K x N (row-major)
    82
    /// - C is M x N (row-major)
    83
    ///
    84
    /// This is the "gold standard" implementation used to validate optimized versions.
    85
    ///
    86
    /// # References
    87
    ///
    88
    /// This implements the naive O(MNK) algorithm as described in
    89
    /// Golub & Van Loan (2013), Matrix Computations, 4th ed., Algorithm 1.1.1.
    90
    0
    pub fn gemm_reference(
    91
    0
        m: usize,
    92
    0
        n: usize,
    93
    0
        k: usize,
    94
    0
        a: &[f32],
    95
    0
        b: &[f32],
    96
    0
        c: &mut [f32],
    97
    0
    ) -> Result<(), TruenoError> {
    98
        // Poka-yoke: dimension validation
    99
    0
        if a.len() != m * k {
    100
    0
            return Err(TruenoError::InvalidInput(format!(
    101
    0
                "A size mismatch: expected {}x{}={}, got {}",
    102
    0
                m,
    103
    0
                k,
    104
    0
                m * k,
    105
    0
                a.len()
    106
    0
            )));
    107
    0
        }
    108
    0
        if b.len() != k * n {
    109
    0
            return Err(TruenoError::InvalidInput(format!(
    110
    0
                "B size mismatch: expected {}x{}={}, got {}",
    111
    0
                k,
    112
    0
                n,
    113
    0
                k * n,
    114
    0
                b.len()
    115
    0
            )));
    116
    0
        }
    117
    0
        if c.len() != m * n {
    118
    0
            return Err(TruenoError::InvalidInput(format!(
    119
    0
                "C size mismatch: expected {}x{}={}, got {}",
    120
    0
                m,
    121
    0
                n,
    122
    0
                m * n,
    123
    0
                c.len()
    124
    0
            )));
    125
    0
        }
    126
    127
        // Scalar triple-nested loop
    128
    0
        for i in 0..m {
    129
    0
            for j in 0..n {
    130
    0
                let mut sum = 0.0f32;
    131
    0
                for p in 0..k {
    132
    0
                    sum += a[i * k + p] * b[p * n + j];
    133
    0
                }
    134
    0
                c[i * n + j] += sum;
    135
            }
    136
        }
    137
    138
    0
        Ok(())
    139
    0
    }
    140
    141
    /// Scalar reference GEMM with Jidoka validation
    142
    ///
    143
    /// Same as `gemm_reference` but validates outputs against known-good computation.
    144
    0
    pub fn gemm_reference_with_jidoka(
    145
    0
        m: usize,
    146
    0
        n: usize,
    147
    0
        k: usize,
    148
    0
        a: &[f32],
    149
    0
        b: &[f32],
    150
    0
        c: &mut [f32],
    151
    0
        guard: &JidokaGuard,
    152
    0
    ) -> Result<(), JidokaError> {
    153
        // Check inputs for NaN/Inf
    154
    0
        for (idx, &val) in a.iter().enumerate() {
    155
    0
            if idx % guard.sample_rate == 0 {
    156
    0
                guard.check_input(val, "matrix A")?;
    157
    0
            }
    158
        }
    159
    0
        for (idx, &val) in b.iter().enumerate() {
    160
    0
            if idx % guard.sample_rate == 0 {
    161
    0
                guard.check_input(val, "matrix B")?;
    162
    0
            }
    163
        }
    164
    165
        // Compute with validation
    166
    0
        for i in 0..m {
    167
    0
            for j in 0..n {
    168
    0
                let mut sum = 0.0f32;
    169
    0
                for p in 0..k {
    170
    0
                    sum += a[i * k + p] * b[p * n + j];
    171
    0
                }
    172
    0
                let output = c[i * n + j] + sum;
    173
    174
                // Jidoka: check output
    175
    0
                if (i * n + j) % guard.sample_rate == 0 {
    176
    0
                    if output.is_nan() {
    177
    0
                        return Err(JidokaError::NaNDetected { location: "output" });
    178
    0
                    }
    179
    0
                    if output.is_infinite() {
    180
    0
                        return Err(JidokaError::InfDetected { location: "output" });
    181
    0
                    }
    182
    0
                }
    183
    184
    0
                c[i * n + j] = output;
    185
            }
    186
        }
    187
    188
    0
        Ok(())
    189
    0
    }
    190
    191
    // ============================================================================
    192
    // Phase 2: Microkernel (MR=8, NR=6)
    193
    // ============================================================================
    194
    195
    // Phase 3: Cache-Optimized Packing
    196
    // ============================================================================
    197
    198
    /// Pack A into MC x KC panel with MR-aligned micro-panels
    199
    ///
    200
    /// Memory layout (Van Zee & Van de Geijn, 2015, Fig. 4):
    201
    /// Original A (row-major):     Packed A (column-major micro-panels):
    202
    /// [a00 a01 a02 ...]           [a00 a10 a20 ... a(MR-1)0 | a01 a11 ...]
    203
    /// [a10 a11 a12 ...]            \____ MR elements ____/
    204
    ///
    205
    /// This layout ensures:
    206
    /// 1. Sequential access in the microkernel
    207
    /// 2. Optimal cache line utilization
    208
    /// 3. Aligned loads for SIMD
    209
    0
    pub fn pack_a(
    210
    0
        a: &[f32],
    211
    0
        lda: usize,  // Leading dimension of A (number of columns in original)
    212
    0
        mc: usize,   // Number of rows to pack
    213
    0
        kc: usize,   // Number of columns to pack
    214
    0
        packed: &mut [f32],
    215
    0
    ) {
    216
    0
        let mut pack_idx = 0;
    217
    218
        // Process MR rows at a time
    219
    0
        let full_panels = mc / MR;
    220
    0
        let remainder = mc % MR;
    221
    222
    0
        for panel in 0..full_panels {
    223
    0
            let row_start = panel * MR;
    224
    225
    0
            for col in 0..kc {
    226
    0
                for row in 0..MR {
    227
    0
                    packed[pack_idx] = a[(row_start + row) * lda + col];
    228
    0
                    pack_idx += 1;
    229
    0
                }
    230
            }
    231
        }
    232
    233
        // Handle remainder rows (pad with zeros)
    234
    0
        if remainder > 0 {
    235
    0
            let row_start = full_panels * MR;
    236
    237
    0
            for col in 0..kc {
    238
    0
                for row in 0..MR {
    239
    0
                    if row < remainder {
    240
    0
                        packed[pack_idx] = a[(row_start + row) * lda + col];
    241
    0
                    } else {
    242
    0
                        packed[pack_idx] = 0.0; // Zero padding
    243
    0
                    }
    244
    0
                    pack_idx += 1;
    245
                }
    246
            }
    247
    0
        }
    248
    0
    }
    249
    250
    /// Pack B into KC x NC panel with NR-aligned micro-panels
    251
    ///
    252
    /// Memory layout:
    253
    /// Original B (row-major):     Packed B (row-major micro-panels):
    254
    /// [b00 b01 b02 ...]           [b00 b01 ... b(NR-1) | b10 b11 ...]
    255
    /// [b10 b11 b12 ...]            \____ NR elements ____/
    256
    0
    pub fn pack_b(
    257
    0
        b: &[f32],
    258
    0
        ldb: usize,  // Leading dimension of B (number of columns in original)
    259
    0
        kc: usize,   // Number of rows to pack
    260
    0
        nc: usize,   // Number of columns to pack
    261
    0
        packed: &mut [f32],
    262
    0
    ) {
    263
    0
        let mut pack_idx = 0;
    264
    265
    0
        let full_panels = nc / NR;
    266
    0
        let remainder = nc % NR;
    267
    268
    0
        for panel in 0..full_panels {
    269
    0
            let col_start = panel * NR;
    270
    271
    0
            for row in 0..kc {
    272
    0
                for col in 0..NR {
    273
    0
                    packed[pack_idx] = b[row * ldb + col_start + col];
    274
    0
                    pack_idx += 1;
    275
    0
                }
    276
            }
    277
        }
    278
    279
        // Handle remainder columns (pad with zeros)
    280
    0
        if remainder > 0 {
    281
    0
            let col_start = full_panels * NR;
    282
    283
    0
            for row in 0..kc {
    284
    0
                for col in 0..NR {
    285
    0
                    if col < remainder {
    286
    0
                        packed[pack_idx] = b[row * ldb + col_start + col];
    287
    0
                    } else {
    288
    0
                        packed[pack_idx] = 0.0;
    289
    0
                    }
    290
    0
                    pack_idx += 1;
    291
                }
    292
            }
    293
    0
        }
    294
    0
    }
    295
    296
    /// Compute required packed A buffer size
    297
    #[inline]
    298
    0
    pub fn packed_a_size(mc: usize, kc: usize) -> usize {
    299
    0
        let panels = (mc + MR - 1) / MR;
    300
    0
        panels * MR * kc
    301
    0
    }
    302
    303
    /// Compute required packed B buffer size
    304
    #[inline]
    305
    0
    pub fn packed_b_size(kc: usize, nc: usize) -> usize {
    306
    0
        let panels = (nc + NR - 1) / NR;
    307
    0
        panels * NR * kc
    308
    0
    }
    309
    310
    // ============================================================================
    311
    // Phase 4: Cache-Blocked GEMM
    312
    // ============================================================================
    313
    314
    /// BLIS-style blocked GEMM
    315
    ///
    316
    /// Implements the 5-loop BLIS algorithm (Van Zee & Van de Geijn, 2015):
    317
    /// Loop 5 (jc): N dimension, L3 blocking
    318
    /// Loop 4 (pc): K dimension, L2 blocking
    319
    /// Loop 3 (ic): M dimension, L1 blocking
    320
    /// Loop 2 (jr): Microkernel columns
    321
    /// Loop 1 (ir): Microkernel rows
    322
    0
    pub fn gemm_blis(
    323
    0
        m: usize,
    324
    0
        n: usize,
    325
    0
        k: usize,
    326
    0
        a: &[f32],
    327
    0
        b: &[f32],
    328
    0
        c: &mut [f32],
    329
    0
        mut profiler: Option<&mut BlisProfiler>,
    330
    0
    ) -> Result<(), TruenoError> {
    331
        // Dimension validation (Poka-yoke)
    332
    0
        if a.len() != m * k {
    333
    0
            return Err(TruenoError::InvalidInput(format!(
    334
    0
                "A size mismatch: expected {}, got {}",
    335
    0
                m * k,
    336
    0
                a.len()
    337
    0
            )));
    338
    0
        }
    339
    0
        if b.len() != k * n {
    340
    0
            return Err(TruenoError::InvalidInput(format!(
    341
    0
                "B size mismatch: expected {}, got {}",
    342
    0
                k * n,
    343
    0
                b.len()
    344
    0
            )));
    345
    0
        }
    346
    0
        if c.len() != m * n {
    347
    0
            return Err(TruenoError::InvalidInput(format!(
    348
    0
                "C size mismatch: expected {}, got {}",
    349
    0
                m * n,
    350
    0
                c.len()
    351
    0
            )));
    352
    0
        }
    353
    354
        // Handle edge cases
    355
    0
        if m == 0 || n == 0 || k == 0 {
    356
    0
            return Ok(());
    357
    0
        }
    358
    359
        // Small matrix: use reference implementation
    360
    0
        if m * n * k < 4096 {
    361
    0
            return gemm_reference(m, n, k, a, b, c);
    362
    0
        }
    363
    364
    0
        let start = Instant::now();
    365
    366
        // Allocate packing buffers
    367
    0
        let mc = MC.min(m);
    368
    0
        let nc = NC.min(n);
    369
    0
        let kc = KC.min(k);
    370
    371
    0
        let mut packed_a = vec![0.0f32; packed_a_size(mc, kc)];
    372
    0
        let mut packed_b = vec![0.0f32; packed_b_size(kc, nc)];
    373
    374
        // Workspace for microkernel output (column-major)
    375
    0
        let mut c_micro = vec![0.0f32; MR * NR];
    376
    377
        // Loop 5: jc (N dimension, L3 blocking)
    378
    0
        for jc in (0..n).step_by(NC) {
    379
    0
            let nc_block = NC.min(n - jc);
    380
    381
            // Loop 4: pc (K dimension, L2 blocking)
    382
    0
            for pc in (0..k).step_by(KC) {
    383
    0
                let kc_block = KC.min(k - pc);
    384
    385
                // Pack B panel: B[pc:pc+kc, jc:jc+nc] -> packed_b
    386
    0
                let pack_start = Instant::now();
    387
    0
                pack_b_block(b, n, pc, jc, kc_block, nc_block, &mut packed_b);
    388
    0
                if let Some(ref mut prof) = profiler.as_deref_mut() {
    389
    0
                    prof.record(BlisProfileLevel::Pack, pack_start.elapsed().as_nanos() as u64, 0);
    390
    0
                }
    391
    392
                // Loop 3: ic (M dimension, L1 blocking)
    393
    0
                for ic in (0..m).step_by(MC) {
    394
    0
                    let mc_block = MC.min(m - ic);
    395
    396
                    // Pack A panel: A[ic:ic+mc, pc:pc+kc] -> packed_a
    397
    0
                    let pack_start = Instant::now();
    398
    0
                    pack_a_block(a, k, ic, pc, mc_block, kc_block, &mut packed_a);
    399
    0
                    if let Some(ref mut prof) = profiler.as_deref_mut() {
    400
    0
                        prof.record(BlisProfileLevel::Pack, pack_start.elapsed().as_nanos() as u64, 0);
    401
    0
                    }
    402
    403
                    // Midi profiling
    404
    0
                    let midi_start = Instant::now();
    405
    406
                    // Loop 2: jr (microkernel columns)
    407
    0
                    for jr in (0..nc_block).step_by(NR) {
    408
    0
                        let nr_block = NR.min(nc_block - jr);
    409
    410
                        // Loop 1: ir (microkernel rows)
    411
    0
                        for ir in (0..mc_block).step_by(MR) {
    412
    0
                            let mr_block = MR.min(mc_block - ir);
    413
    414
                            // Compute microkernel
    415
    0
                            let micro_start = Instant::now();
    416
    417
                            // Get packed panel pointers
    418
    0
                            let a_panel = &packed_a[(ir / MR) * MR * kc_block..];
    419
    0
                            let b_panel = &packed_b[(jr / NR) * NR * kc_block..];
    420
    421
                            // Load existing C values into micro workspace for accumulation
    422
                            // GEMM computes C += A*B, so we always load C first
    423
    0
                            c_micro.fill(0.0); // Zero padding area
    424
    0
                            for jj in 0..nr_block {
    425
    0
                                for ii in 0..mr_block {
    426
    0
                                    c_micro[jj * MR + ii] = c[(ic + ir + ii) * n + (jc + jr + jj)];
    427
    0
                                }
    428
                            }
    429
    430
                            // Call microkernel (use Phase 2c true ASM for 70%+ FMA utilization)
    431
                            #[cfg(target_arch = "x86_64")]
    432
                            {
    433
    0
                                if is_x86_feature_detected!("avx2") && is_x86_feature_detected!("fma") {
    434
    0
                                    if mr_block == MR && nr_block == NR {
    435
    0
                                        unsafe {
    436
    0
                                            // Use true inline ASM for 70%+ FMA utilization
    437
    0
                                            microkernel_8x6_true_asm(
    438
    0
                                                kc_block,
    439
    0
                                                a_panel.as_ptr(),
    440
    0
                                                b_panel.as_ptr(),
    441
    0
                                                c_micro.as_mut_ptr(),
    442
    0
                                                MR,
    443
    0
                                            );
    444
    0
                                        }
    445
    0
                                    } else {
    446
    0
                                        microkernel_scalar(kc_block, a_panel, b_panel, &mut c_micro, MR);
    447
    0
                                    }
    448
    0
                                } else {
    449
    0
                                    microkernel_scalar(kc_block, a_panel, b_panel, &mut c_micro, MR);
    450
    0
                                }
    451
                            }
    452
    453
                            #[cfg(target_arch = "aarch64")]
    454
                            {
    455
                                // Use scalar for now; NEON kernel has different dimensions
    456
                                microkernel_scalar(kc_block, a_panel, b_panel, &mut c_micro, MR);
    457
                            }
    458
    459
                            #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
    460
                            {
    461
                                microkernel_scalar(kc_block, a_panel, b_panel, &mut c_micro, MR);
    462
                            }
    463
    464
                            // Store results back to C
    465
    0
                            for jj in 0..nr_block {
    466
    0
                                for ii in 0..mr_block {
    467
    0
                                    c[(ic + ir + ii) * n + (jc + jr + jj)] = c_micro[jj * MR + ii];
    468
    0
                                }
    469
                            }
    470
    471
    0
                            if let Some(ref mut prof) = profiler.as_deref_mut() {
    472
    0
                                let flops = 2 * mr_block * nr_block * kc_block;
    473
    0
                                prof.record(
    474
    0
                                    BlisProfileLevel::Micro,
    475
    0
                                    micro_start.elapsed().as_nanos() as u64,
    476
    0
                                    flops as u64,
    477
    0
                                );
    478
    0
                            }
    479
                        }
    480
                    }
    481
    482
    0
                    if let Some(ref mut prof) = profiler.as_deref_mut() {
    483
    0
                        let flops = 2 * mc_block * nc_block * kc_block;
    484
    0
                        prof.record(
    485
    0
                            BlisProfileLevel::Midi,
    486
    0
                            midi_start.elapsed().as_nanos() as u64,
    487
    0
                            flops as u64,
    488
    0
                        );
    489
    0
                    }
    490
                }
    491
            }
    492
        }
    493
    494
    0
        if let Some(prof) = profiler {
    495
    0
            let flops = 2 * m * n * k;
    496
    0
            prof.record(
    497
    0
                BlisProfileLevel::Macro,
    498
    0
                start.elapsed().as_nanos() as u64,
    499
    0
                flops as u64,
    500
    0
            );
    501
    0
        }
    502
    503
    0
        Ok(())
    504
    0
    }
    505
    506
    /// Pack A block from row-major source
    507
    0
    fn pack_a_block(
    508
    0
        a: &[f32],
    509
    0
        lda: usize,
    510
    0
        row_start: usize,
    511
    0
        col_start: usize,
    512
    0
        rows: usize,
    513
    0
        cols: usize,
    514
    0
        packed: &mut [f32],
    515
    0
    ) {
    516
    0
        let mut pack_idx = 0;
    517
    0
        let panels = (rows + MR - 1) / MR;
    518
    519
    0
        for panel in 0..panels {
    520
    0
            let ir = panel * MR;
    521
    0
            let mr_actual = MR.min(rows - ir);
    522
    523
    0
            for col in 0..cols {
    524
    0
                for row in 0..MR {
    525
    0
                    if row < mr_actual {
    526
    0
                        packed[pack_idx] = a[(row_start + ir + row) * lda + col_start + col];
    527
    0
                    } else {
    528
    0
                        packed[pack_idx] = 0.0;
    529
    0
                    }
    530
    0
                    pack_idx += 1;
    531
                }
    532
            }
    533
        }
    534
    0
    }
    535
    536
    /// Pack B block from row-major source
    537
    0
    fn pack_b_block(
    538
    0
        b: &[f32],
    539
    0
        ldb: usize,
    540
    0
        row_start: usize,
    541
    0
        col_start: usize,
    542
    0
        rows: usize,
    543
    0
        cols: usize,
    544
    0
        packed: &mut [f32],
    545
    0
    ) {
    546
    0
        let mut pack_idx = 0;
    547
    0
        let panels = (cols + NR - 1) / NR;
    548
    549
    0
        for panel in 0..panels {
    550
    0
            let jr = panel * NR;
    551
    0
            let nr_actual = NR.min(cols - jr);
    552
    553
    0
            for row in 0..rows {
    554
    0
                for col in 0..NR {
    555
    0
                    if col < nr_actual {
    556
    0
                        packed[pack_idx] = b[(row_start + row) * ldb + col_start + jr + col];
    557
    0
                    } else {
    558
    0
                        packed[pack_idx] = 0.0;
    559
    0
                    }
    560
    0
                    pack_idx += 1;
    561
                }
    562
            }
    563
        }
    564
    0
    }
    565
    566
    // ============================================================================
    567
    // Phase 5: Parallel GEMM with Heijunka
    568
    // ============================================================================
    569
    570
    /// Heijunka (load-leveling) scheduler for parallel GEMM
    571
    #[derive(Debug, Clone)]
    572
    pub struct HeijunkaScheduler {
    573
        /// Number of threads
    574
        pub num_threads: usize,
    575
        /// Target load variance threshold
    576
        pub variance_threshold: f32,
    577
    }
    578
    579
    impl Default for HeijunkaScheduler {
    580
    0
        fn default() -> Self {
    581
            #[cfg(feature = "parallel")]
    582
            let threads = rayon::current_num_threads();
    583
            #[cfg(not(feature = "parallel"))]
    584
    0
            let threads = 1;
    585
    586
    0
            Self {
    587
    0
                num_threads: threads,
    588
    0
                variance_threshold: 0.05, // 5% variance target
    589
    0
            }
    590
    0
        }
    591
    }
    592
    593
    impl HeijunkaScheduler {
    594
        /// Partition M dimension into balanced chunks
    595
    0
        pub fn partition_m(&self, m: usize, mc: usize) -> Vec<std::ops::Range<usize>> {
    596
    0
            let num_blocks = (m + mc - 1) / mc;
    597
    0
            let blocks_per_thread = num_blocks / self.num_threads;
    598
    0
            let remainder = num_blocks % self.num_threads;
    599
    600
    0
            let mut partitions = Vec::with_capacity(self.num_threads);
    601
    0
            let mut start_block = 0;
    602
    603
    0
            for t in 0..self.num_threads {
    604
    0
                let extra = if t < remainder { 1 } else { 0 };
    605
    0
                let thread_blocks = blocks_per_thread + extra;
    606
    607
    0
                let start_row = start_block * mc;
    608
    0
                let end_row = ((start_block + thread_blocks) * mc).min(m);
    609
    610
    0
                if start_row < end_row {
    611
    0
                    partitions.push(start_row..end_row);
    612
    0
                }
    613
    614
    0
                start_block += thread_blocks;
    615
            }
    616
    617
    0
            partitions
    618
    0
        }
    619
    }
    620
    621
    /// Parallel BLIS GEMM using Rayon
    622
    #[cfg(feature = "parallel")]
    623
    pub fn gemm_blis_parallel(
    624
        m: usize,
    625
        n: usize,
    626
        k: usize,
    627
        a: &[f32],
    628
        b: &[f32],
    629
        c: &mut [f32],
    630
    ) -> Result<(), TruenoError> {
    631
        use rayon::prelude::*;
    632
    633
        // Dimension validation
    634
        if a.len() != m * k || b.len() != k * n || c.len() != m * n {
    635
            return Err(TruenoError::InvalidInput("Dimension mismatch".to_string()));
    636
        }
    637
    638
        // Small matrices: single-threaded
    639
        if m * n * k < 1_000_000 {
    640
            return gemm_blis(m, n, k, a, b, c, None);
    641
        }
    642
    643
        let scheduler = HeijunkaScheduler::default();
    644
        let partitions = scheduler.partition_m(m, MC);
    645
    646
        // Pack B once (shared across threads)
    647
        let nc = NC.min(n);
    648
        let kc = KC.min(k);
    649
        let packed_b_total_size = ((n + NR - 1) / NR) * ((k + KC - 1) / KC) * packed_b_size(kc, nc);
    650
        let _packed_b = std::sync::Arc::new(std::sync::RwLock::new(vec![0.0f32; packed_b_total_size]));
    651
    652
        // Parallel over M partitions
    653
        let c_ptr = c.as_mut_ptr() as usize;
    654
        let _c_len = c.len();
    655
    656
        partitions.into_par_iter().for_each(|m_range| {
    657
            let m_local = m_range.len();
    658
            let m_start = m_range.start;
    659
    660
            // Local A slice
    661
            let a_local = &a[m_start * k..(m_start + m_local) * k];
    662
    663
            // Local C slice (unsafe but safe due to non-overlapping partitions)
    664
            let c_local = unsafe {
    665
                let ptr = c_ptr as *mut f32;
    666
                std::slice::from_raw_parts_mut(ptr.add(m_start * n), m_local * n)
    667
            };
    668
    669
            // Run local GEMM
    670
            let _ = gemm_blis(m_local, n, k, a_local, b, c_local, None);
    671
        });
    672
    673
        Ok(())
    674
    }
    675
    676
    /// Non-parallel fallback
    677
    #[cfg(not(feature = "parallel"))]
    678
    0
    pub fn gemm_blis_parallel(
    679
    0
        m: usize,
    680
    0
        n: usize,
    681
    0
        k: usize,
    682
    0
        a: &[f32],
    683
    0
        b: &[f32],
    684
    0
        c: &mut [f32],
    685
    0
    ) -> Result<(), TruenoError> {
    686
    0
        gemm_blis(m, n, k, a, b, c, None)
    687
    0
    }
    688
    689
    // Public API
    690
    // ============================================================================
    691
    692
    /// High-performance GEMM using BLIS algorithm
    693
    ///
    694
    /// Computes C += A * B where:
    695
    /// - A is M x K (row-major)
    696
    /// - B is K x N (row-major)
    697
    /// - C is M x N (row-major)
    698
    ///
    699
    /// Automatically selects single-threaded or parallel execution based on matrix size.
    700
    0
    pub fn gemm(
    701
    0
        m: usize,
    702
    0
        n: usize,
    703
    0
        k: usize,
    704
    0
        a: &[f32],
    705
    0
        b: &[f32],
    706
    0
        c: &mut [f32],
    707
    0
    ) -> Result<(), TruenoError> {
    708
        #[cfg(feature = "parallel")]
    709
        {
    710
            gemm_blis_parallel(m, n, k, a, b, c)
    711
        }
    712
        #[cfg(not(feature = "parallel"))]
    713
        {
    714
    0
            gemm_blis(m, n, k, a, b, c, None)
    715
        }
    716
    0
    }
    717
    718
    /// GEMM with profiling enabled
    719
    0
    pub fn gemm_profiled(
    720
    0
        m: usize,
    721
    0
        n: usize,
    722
    0
        k: usize,
    723
    0
        a: &[f32],
    724
    0
        b: &[f32],
    725
    0
        c: &mut [f32],
    726
    0
        profiler: &mut BlisProfiler,
    727
    0
    ) -> Result<(), TruenoError> {
    728
    0
        gemm_blis(m, n, k, a, b, c, Some(profiler))
    729
    0
    }
    730
    731
    // ============================================================================
    732
    // Matrix Transpose (SIMD-optimized)
    733
    // ============================================================================
    734
    735
    /// Transpose a matrix: B = A^T
    736
    ///
    737
    /// SIMD-optimized for large matrices (>=64 elements).
    738
    /// Uses cache-efficient 8x8 blocking with manual unrolling.
    739
    ///
    740
    /// # Arguments
    741
    ///
    742
    /// * `rows` - Number of rows in A (cols in B)
    743
    /// * `cols` - Number of cols in A (rows in B)
    744
    /// * `a` - Input matrix A (rows x cols, row-major)
    745
    /// * `b` - Output matrix B (cols x rows, row-major)
    746
    ///
    747
    /// # Returns
    748
    ///
    749
    /// `Ok(())` on success, `Err` if dimensions mismatch
    750
    0
    pub fn transpose(rows: usize, cols: usize, a: &[f32], b: &mut [f32]) -> Result<(), TruenoError> {
    751
    0
        let expected = rows * cols;
    752
    0
        if a.len() != expected || b.len() != expected {
    753
    0
            return Err(TruenoError::InvalidInput(format!(
    754
    0
                "transpose size mismatch: a[{}], b[{}], expected {}",
    755
    0
                a.len(),
    756
    0
                b.len(),
    757
    0
                expected
    758
    0
            )));
    759
    0
        }
    760
    761
        // For small matrices, use simple scalar transpose
    762
    0
        if expected < 64 {
    763
    0
            for r in 0..rows {
    764
    0
                for c in 0..cols {
    765
    0
                    b[c * rows + r] = a[r * cols + c];
    766
    0
                }
    767
            }
    768
    0
            return Ok(());
    769
    0
        }
    770
    771
        // Cache-efficient blocked transpose for larger matrices
    772
        // 8x8 blocks to maximize cache line utilization
    773
        const BLOCK: usize = 8;
    774
    775
        // Process full blocks
    776
    0
        let row_blocks = rows / BLOCK;
    777
    0
        let col_blocks = cols / BLOCK;
    778
    779
    0
        for rb in 0..row_blocks {
    780
    0
            for cb in 0..col_blocks {
    781
    0
                let row_start = rb * BLOCK;
    782
    0
                let col_start = cb * BLOCK;
    783
    784
                // Transpose 8x8 block with manual unrolling
    785
    0
                for i in 0..BLOCK {
    786
    0
                    for j in 0..BLOCK {
    787
    0
                        let src = (row_start + i) * cols + (col_start + j);
    788
    0
                        let dst = (col_start + j) * rows + (row_start + i);
    789
    0
                        b[dst] = a[src];
    790
    0
                    }
    791
                }
    792
            }
    793
        }
    794
    795
        // Handle remaining columns (right edge)
    796
    0
        let col_remainder_start = col_blocks * BLOCK;
    797
    0
        if col_remainder_start < cols {
    798
    0
            for r in 0..(row_blocks * BLOCK) {
    799
    0
                for c in col_remainder_start..cols {
    800
    0
                    b[c * rows + r] = a[r * cols + c];
    801
    0
                }
    802
            }
    803
    0
        }
    804
    805
        // Handle remaining rows (bottom edge)
    806
    0
        let row_remainder_start = row_blocks * BLOCK;
    807
    0
        if row_remainder_start < rows {
    808
    0
            for r in row_remainder_start..rows {
    809
    0
                for c in 0..cols {
    810
    0
                    b[c * rows + r] = a[r * cols + c];
    811
    0
                }
    812
            }
    813
    0
        }
    814
    815
    0
        Ok(())
    816
    0
    }
    817
    818
    #[cfg(test)]
    819
    mod tests;