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

    Line
    Count
    Source
    1
    //! Backend Selection and Cost Model
    2
    //!
    3
    //! Automatic selection between CPU (SIMD), CUDA (PTX), and wgpu (WGSL) backends
    4
    //! based on the 5× PCIe rule and roofline analysis.
    5
    //!
    6
    //! # Philosophy
    7
    //!
    8
    //! Uses Gregg & Hazelwood (2011) "5× PCIe rule": GPU worthwhile when
    9
    //! compute time exceeds 5× data transfer time.
    10
    //!
    11
    //! # References
    12
    //!
    13
    //! - Gregg, C., & Hazelwood, K. (2011). Where is the Data? Why You Cannot
    14
    //!   Debate CPU vs. GPU Performance Without the Answer. IEEE ISPASS.
    15
    //! - Volkov, V. (2010). Better Performance at Lower Occupancy.
    16
    17
    #[cfg(target_arch = "x86_64")]
    18
    use std::arch::is_x86_feature_detected;
    19
    20
    use super::profiler::BlisProfiler;
    21
    use super::{gemm_blis, TruenoError};
    22
    23
    ///
    24
    /// Maps to different ISA targets:
    25
    /// - Cpu: x86 asm (AVX2/AVX-512), ARM asm (NEON)
    26
    /// - Gpu: PTX (CUDA), wgpu compute shaders
    27
    /// - Wgpu: WGSL for cross-platform GPU (Vulkan/Metal/DX12/WebGPU)
    28
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
    29
    pub enum ComputeBackend {
    30
        /// CPU SIMD backend (AVX2, AVX-512, NEON, SSE2)
    31
        Cpu,
    32
        /// NVIDIA GPU backend (PTX)
    33
        #[allow(dead_code)]
    34
        Gpu,
    35
        /// Cross-platform GPU backend (wgpu/WGSL)
    36
        #[allow(dead_code)]
    37
        Wgpu,
    38
        /// Scalar fallback (no SIMD)
    39
        Scalar,
    40
    }
    41
    42
    /// ComputeBrick hierarchy level
    43
    ///
    44
    /// Maps BLIS loop structure to brick abstraction:
    45
    /// - Nano: Microkernel (MR×NR×K) - register file
    46
    /// - Micro: Midi loop (MC×NC×KC) - L1/L2 cache
    47
    /// - Meso: Macro loop (full M×N×K) - L3/DRAM
    48
    #[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
    49
    pub enum BrickLevel {
    50
        /// Register-level compute (MR×NR tile)
    51
        Nano,
    52
        /// Cache-level compute (MC×NC block)
    53
        Micro,
    54
        /// Memory-level compute (full matrix)
    55
        Meso,
    56
    }
    57
    58
    /// Cost model for backend selection
    59
    ///
    60
    /// Based on Gregg & Hazelwood (2011): GPU worthwhile when compute > 5× transfer
    61
    #[derive(Debug, Clone)]
    62
    pub struct BackendCostModel {
    63
        /// PCIe bandwidth in GB/s (e.g., 15.75 for PCIe 3.0 x16)
    64
        pub pcie_bandwidth_gbps: f64,
    65
        /// GPU peak TFLOP/s
    66
        pub gpu_peak_tflops: f64,
    67
        /// CPU peak GFLOP/s
    68
        pub cpu_peak_gflops: f64,
    69
        /// Minimum problem size for GPU (elements)
    70
        pub gpu_min_elements: usize,
    71
    }
    72
    73
    impl Default for BackendCostModel {
    74
    0
        fn default() -> Self {
    75
    0
            Self {
    76
    0
                pcie_bandwidth_gbps: 15.75,  // PCIe 3.0 x16
    77
    0
                gpu_peak_tflops: 10.0,        // Mid-range GPU
    78
    0
                cpu_peak_gflops: 400.0,       // Modern AVX2 CPU
    79
    0
                gpu_min_elements: 1_000_000,  // ~1M elements
    80
    0
            }
    81
    0
        }
    82
    }
    83
    84
    impl BackendCostModel {
    85
        /// Select optimal backend based on 5× PCIe rule
    86
        ///
    87
        /// # References
    88
        ///
    89
        /// Gregg, C., & Hazelwood, K. (2011). Where is the Data? Why You Cannot
    90
        /// Debate CPU vs. GPU Performance Without the Answer. IEEE ISPASS.
    91
    0
        pub fn select_backend(&self, m: usize, n: usize, k: usize) -> ComputeBackend {
    92
    0
            let flops = 2 * m * n * k;
    93
    0
            let bytes = 4 * (m * k + k * n + m * n); // f32 = 4 bytes
    94
    0
            let arithmetic_intensity = flops as f64 / bytes as f64;
    95
    96
            // Ridge point: where compute = memory bandwidth
    97
    0
            let ridge_point = self.gpu_peak_tflops * 1000.0 / self.pcie_bandwidth_gbps;
    98
    99
            // GPU worthwhile if:
    100
            // 1. High arithmetic intensity (compute-bound)
    101
            // 2. Problem size exceeds minimum threshold
    102
            // 3. Transfer time is amortized (5× rule)
    103
    0
            let elements = m * n * k;
    104
    0
            if arithmetic_intensity > ridge_point && elements > self.gpu_min_elements {
    105
                // Check if wgpu available at runtime
    106
                #[cfg(feature = "wgpu")]
    107
    0
                return ComputeBackend::Wgpu;
    108
    109
                #[cfg(all(not(feature = "wgpu"), feature = "cuda"))]
    110
                return ComputeBackend::Gpu;
    111
    112
                #[allow(unreachable_code)]
    113
                ComputeBackend::Cpu
    114
            } else {
    115
                // CPU is better for small problems or memory-bound workloads
    116
                #[cfg(target_arch = "x86_64")]
    117
                {
    118
    0
                    if is_x86_feature_detected!("avx2") {
    119
    0
                        return ComputeBackend::Cpu;
    120
    0
                    }
    121
                }
    122
                #[cfg(target_arch = "aarch64")]
    123
                {
    124
                    return ComputeBackend::Cpu;
    125
                }
    126
    0
                ComputeBackend::Scalar
    127
            }
    128
    0
        }
    129
    130
        /// Estimate execution time in microseconds
    131
    0
        pub fn estimate_time_us(&self, m: usize, n: usize, k: usize, backend: ComputeBackend) -> f64 {
    132
    0
            let flops = 2.0 * m as f64 * n as f64 * k as f64;
    133
    0
            let bytes = 4.0 * (m * k + k * n + m * n) as f64;
    134
    135
    0
            match backend {
    136
                ComputeBackend::Gpu | ComputeBackend::Wgpu => {
    137
                    // Transfer time + compute time
    138
    0
                    let transfer_us = bytes / (self.pcie_bandwidth_gbps * 1e3);
    139
    0
                    let compute_us = flops / (self.gpu_peak_tflops * 1e6);
    140
    0
                    transfer_us + compute_us
    141
                }
    142
                ComputeBackend::Cpu => {
    143
    0
                    flops / (self.cpu_peak_gflops * 1e3)
    144
                }
    145
                ComputeBackend::Scalar => {
    146
                    // Assume 1 GFLOP/s for scalar
    147
    0
                    flops / 1e3
    148
                }
    149
            }
    150
    0
        }
    151
    }
    152
    153
    /// Unified profiler for all backends
    154
    ///
    155
    /// Collects metrics across CPU (RDTSC), GPU (CUDA events), and wgpu (timestamp queries)
    156
    #[derive(Debug, Clone, Default)]
    157
    pub struct UnifiedBrickProfiler {
    158
        /// CPU profiling stats
    159
        pub cpu_stats: BlisProfiler,
    160
        /// Selected backend for this run
    161
        pub backend: Option<ComputeBackend>,
    162
        /// Total elements processed
    163
        pub total_elements: u64,
    164
        /// Backend selection decisions
    165
        pub selection_history: Vec<(usize, usize, usize, ComputeBackend)>,
    166
    }
    167
    168
    impl UnifiedBrickProfiler {
    169
        /// Create a new unified profiler
    170
    0
        pub fn new() -> Self {
    171
    0
            Self {
    172
    0
                cpu_stats: BlisProfiler::enabled(),
    173
    0
                backend: None,
    174
    0
                total_elements: 0,
    175
    0
                selection_history: Vec::new(),
    176
    0
            }
    177
    0
        }
    178
    179
        /// Record backend selection
    180
    0
        pub fn record_selection(&mut self, m: usize, n: usize, k: usize, backend: ComputeBackend) {
    181
    0
            self.backend = Some(backend);
    182
    0
            self.total_elements += (m * n) as u64;
    183
    0
            self.selection_history.push((m, n, k, backend));
    184
    0
        }
    185
    186
        /// Get roofline analysis for current backend
    187
    0
        pub fn roofline_analysis(&self, m: usize, n: usize, k: usize) -> RooflineResult {
    188
    0
            let cost = BackendCostModel::default();
    189
    0
            let flops = 2.0 * m as f64 * n as f64 * k as f64;
    190
    0
            let bytes = 4.0 * (m * k + k * n + m * n) as f64;
    191
    0
            let ai = flops / bytes;
    192
    193
    0
            let ridge_point = match self.backend.unwrap_or(ComputeBackend::Cpu) {
    194
                ComputeBackend::Gpu | ComputeBackend::Wgpu => {
    195
    0
                    cost.gpu_peak_tflops * 1000.0 / cost.pcie_bandwidth_gbps
    196
                }
    197
                ComputeBackend::Cpu | ComputeBackend::Scalar => {
    198
    0
                    cost.cpu_peak_gflops / 50.0 // ~50 GB/s memory bandwidth
    199
                }
    200
            };
    201
    202
    0
            if ai < ridge_point {
    203
    0
                RooflineResult::MemoryBound { ai, ridge_point }
    204
            } else {
    205
    0
                RooflineResult::ComputeBound { ai, ridge_point }
    206
            }
    207
    0
        }
    208
    209
        /// Generate summary report
    210
    0
        pub fn summary(&self) -> String {
    211
    0
            let mut s = String::new();
    212
    0
            s.push_str("Unified Brick Profiler Summary\n");
    213
    0
            s.push_str("==============================\n");
    214
    0
            s.push_str(&format!(
    215
    0
                "Backend: {:?}\n",
    216
    0
                self.backend.unwrap_or(ComputeBackend::Scalar)
    217
    0
            ));
    218
    0
            s.push_str(&format!("Total elements: {}\n", self.total_elements));
    219
    0
            s.push_str(&format!(
    220
    0
                "Selections: {} decisions\n",
    221
    0
                self.selection_history.len()
    222
    0
            ));
    223
    0
            s.push_str("\nCPU Stats:\n");
    224
    0
            s.push_str(&self.cpu_stats.summary());
    225
    0
            s
    226
    0
        }
    227
    }
    228
    229
    /// Roofline model result
    230
    #[derive(Debug, Clone, Copy)]
    231
    pub enum RooflineResult {
    232
        /// Workload is memory-bound (AI < ridge point)
    233
        MemoryBound {
    234
            /// Arithmetic intensity (FLOP/byte)
    235
            ai: f64,
    236
            /// Ridge point where compute = memory
    237
            ridge_point: f64,
    238
        },
    239
        /// Workload is compute-bound (AI > ridge point)
    240
        ComputeBound {
    241
            /// Arithmetic intensity (FLOP/byte)
    242
            ai: f64,
    243
            /// Ridge point where compute = memory
    244
            ridge_point: f64,
    245
        },
    246
    }
    247
    248
    impl RooflineResult {
    249
        /// Get arithmetic intensity
    250
    0
        pub fn arithmetic_intensity(&self) -> f64 {
    251
    0
            match self {
    252
    0
                RooflineResult::MemoryBound { ai, .. } => *ai,
    253
    0
                RooflineResult::ComputeBound { ai, .. } => *ai,
    254
            }
    255
    0
        }
    256
    257
        /// Check if compute-bound
    258
    0
        pub fn is_compute_bound(&self) -> bool {
    259
    0
            matches!(self, RooflineResult::ComputeBound { .. })
    260
    0
        }
    261
    }
    262
    263
    /// PTX microkernel definition (for documentation and future CUDA support)
    264
    ///
    265
    /// This is a specification for the GPU microkernel. Actual PTX code generation
    266
    /// would be done by the trueno-ptx crate.
    267
    ///
    268
    /// # References
    269
    ///
    270
    /// - NVIDIA PTX ISA Reference Manual
    271
    /// - Volkov, V. (2010). Better Performance at Lower Occupancy.
    272
    #[derive(Debug, Clone)]
    273
    pub struct PtxMicrokernelSpec {
    274
        /// PTX version (e.g., "8.0")
    275
        pub ptx_version: &'static str,
    276
        /// Target SM architecture (e.g., "sm_80")
    277
        pub sm_target: &'static str,
    278
        /// Register count per thread
    279
        pub registers_per_thread: u32,
    280
        /// Shared memory bytes per block
    281
        pub smem_bytes: usize,
    282
        /// Thread block dimensions
    283
        pub block_dim: (u32, u32, u32),
    284
        /// Tile dimensions (MR, NR)
    285
        pub tile_dim: (usize, usize),
    286
    }
    287
    288
    impl Default for PtxMicrokernelSpec {
    289
    0
        fn default() -> Self {
    290
    0
            Self {
    291
    0
                ptx_version: "8.0",
    292
    0
                sm_target: "sm_80",
    293
    0
                registers_per_thread: 64,
    294
    0
                smem_bytes: 48 * 1024, // 48KB shared memory
    295
    0
                block_dim: (16, 16, 1),
    296
    0
                tile_dim: (16, 16), // 16x16 output tile per warp
    297
    0
            }
    298
    0
        }
    299
    }
    300
    301
    /// WGSL microkernel specification (for wgpu backend)
    302
    ///
    303
    /// Defines the compute shader for matrix multiplication.
    304
    #[derive(Debug, Clone)]
    305
    pub struct WgslMicrokernelSpec {
    306
        /// Workgroup size (x, y, z)
    307
        pub workgroup_size: (u32, u32, u32),
    308
        /// Tile dimensions (MR, NR)
    309
        pub tile_dim: (usize, usize),
    310
        /// Use shared memory for tiling
    311
        pub use_shared_memory: bool,
    312
    }
    313
    314
    impl Default for WgslMicrokernelSpec {
    315
    0
        fn default() -> Self {
    316
    0
            Self {
    317
    0
                workgroup_size: (8, 8, 1),
    318
    0
                tile_dim: (8, 8),
    319
    0
                use_shared_memory: true,
    320
    0
            }
    321
    0
        }
    322
    }
    323
    324
    impl WgslMicrokernelSpec {
    325
        /// Generate WGSL shader source
    326
        ///
    327
        /// This generates a basic tiled GEMM shader. For production use,
    328
        /// this would be optimized with coalesced memory access and bank conflict avoidance.
    329
    0
        pub fn generate_wgsl(&self) -> String {
    330
    0
            format!(
    331
    0
                r#"// WGSL GEMM Microkernel
    332
    0
    // Generated by trueno BLIS module
    333
    0
    // Tile: {}x{}, Workgroup: {}x{}x{}
    334
    0
    335
    0
    struct GemmParams {{
    336
    0
        m: u32,
    337
    0
        n: u32,
    338
    0
        k: u32,
    339
    0
        alpha: f32,
    340
    0
        beta: f32,
    341
    0
    }}
    342
    0
    343
    0
    @group(0) @binding(0) var<uniform> params: GemmParams;
    344
    0
    @group(0) @binding(1) var<storage, read> a: array<f32>;
    345
    0
    @group(0) @binding(2) var<storage, read> b: array<f32>;
    346
    0
    @group(0) @binding(3) var<storage, read_write> c: array<f32>;
    347
    0
    348
    0
    var<workgroup> tile_a: array<f32, {tile_a_size}>;
    349
    0
    var<workgroup> tile_b: array<f32, {tile_b_size}>;
    350
    0
    351
    0
    @compute @workgroup_size({wx}, {wy}, {wz})
    352
    0
    fn main(
    353
    0
        @builtin(global_invocation_id) global_id: vec3<u32>,
    354
    0
        @builtin(local_invocation_id) local_id: vec3<u32>,
    355
    0
        @builtin(workgroup_id) group_id: vec3<u32>,
    356
    0
    ) {{
    357
    0
        let row = global_id.y;
    358
    0
        let col = global_id.x;
    359
    0
    360
    0
        if (row >= params.m || col >= params.n) {{
    361
    0
            return;
    362
    0
        }}
    363
    0
    364
    0
        var sum: f32 = 0.0;
    365
    0
    366
    0
        // Tile over K dimension
    367
    0
        let num_tiles = (params.k + {tile_k}u - 1u) / {tile_k}u;
    368
    0
    369
    0
        for (var t: u32 = 0u; t < num_tiles; t++) {{
    370
    0
            let k_base = t * {tile_k}u;
    371
    0
    372
    0
            // Load tile_a and tile_b into shared memory
    373
    0
            // (simplified - production code would have proper coalescing)
    374
    0
            let k_idx = k_base + local_id.x;
    375
    0
            if (row < params.m && k_idx < params.k) {{
    376
    0
                tile_a[local_id.y * {tile_k}u + local_id.x] = a[row * params.k + k_idx];
    377
    0
            }}
    378
    0
            if (k_idx < params.k && col < params.n) {{
    379
    0
                tile_b[local_id.y * {tile_k}u + local_id.x] = b[k_idx * params.n + col];
    380
    0
            }}
    381
    0
    382
    0
            workgroupBarrier();
    383
    0
    384
    0
            // Compute partial sum
    385
    0
            for (var kk: u32 = 0u; kk < {tile_k}u; kk++) {{
    386
    0
                if (k_base + kk < params.k) {{
    387
    0
                    sum += tile_a[local_id.y * {tile_k}u + kk] * tile_b[kk * {tile_k}u + local_id.x];
    388
    0
                }}
    389
    0
            }}
    390
    0
    391
    0
            workgroupBarrier();
    392
    0
        }}
    393
    0
    394
    0
        // Store result
    395
    0
        let c_idx = row * params.n + col;
    396
    0
        c[c_idx] = params.alpha * sum + params.beta * c[c_idx];
    397
    0
    }}
    398
    0
    "#,
    399
                self.tile_dim.0,
    400
                self.tile_dim.1,
    401
                self.workgroup_size.0,
    402
                self.workgroup_size.1,
    403
                self.workgroup_size.2,
    404
    0
                tile_a_size = self.tile_dim.0 * self.tile_dim.0,
    405
    0
                tile_b_size = self.tile_dim.0 * self.tile_dim.1,
    406
                wx = self.workgroup_size.0,
    407
                wy = self.workgroup_size.1,
    408
                wz = self.workgroup_size.2,
    409
                tile_k = self.tile_dim.0,
    410
            )
    411
    0
        }
    412
    }
    413
    414
    /// GEMM with automatic backend selection
    415
    ///
    416
    /// Uses the 5× PCIe rule to select between CPU (asm) and GPU (PTX/WGSL) backends.
    417
    0
    pub fn gemm_auto(
    418
    0
        m: usize,
    419
    0
        n: usize,
    420
    0
        k: usize,
    421
    0
        a: &[f32],
    422
    0
        b: &[f32],
    423
    0
        c: &mut [f32],
    424
    0
        profiler: Option<&mut UnifiedBrickProfiler>,
    425
    0
    ) -> Result<(), TruenoError> {
    426
    0
        let cost_model = BackendCostModel::default();
    427
    0
        let backend = cost_model.select_backend(m, n, k);
    428
    429
    0
        if let Some(prof) = profiler {
    430
    0
            prof.record_selection(m, n, k, backend);
    431
    0
        }
    432
    433
    0
        match backend {
    434
            ComputeBackend::Cpu | ComputeBackend::Scalar => {
    435
                // Use BLIS CPU implementation
    436
    0
                gemm_blis(m, n, k, a, b, c, None)
    437
            }
    438
            ComputeBackend::Gpu => {
    439
                // PTX backend (stub - requires CUDA support)
    440
                // For now, fall back to CPU
    441
    0
                gemm_blis(m, n, k, a, b, c, None)
    442
            }
    443
            ComputeBackend::Wgpu => {
    444
                // WGSL backend (stub - requires wgpu support)
    445
                // For now, fall back to CPU
    446
    0
                gemm_blis(m, n, k, a, b, c, None)
    447
            }
    448
        }
    449
    0
    }