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

    Line
    Count
    Source
    1
    //! Q4_K quantized matrix-vector tiling implementation.
    2
    3
    use super::config::TilingConfig;
    4
    5
    /// Q4_K superblock constants (per GGML specification)
    6
    pub const Q4K_SUPERBLOCK_SIZE: usize = 256;
    7
    pub const Q4K_SUPERBLOCK_BYTES: usize = 144;
    8
    9
    /// Tiled Q4_K MatVec executor
    10
    ///
    11
    /// Implements TCB-01 pattern: Cache-blocked matvec with 4×1 micro-kernel.
    12
    ///
    13
    /// # Memory Layout
    14
    ///
    15
    /// Weights are stored in Q4_K superblock format (144 bytes per 256 elements):
    16
    /// - d: f16 (2 bytes) - block scale
    17
    /// - dmin: f16 (2 bytes) - block minimum
    18
    /// - scales: 12 bytes - 8 sub-block scales (6-bit packed)
    19
    /// - qs: 128 bytes - 256 quantized values (4-bit packed)
    20
    ///
    21
    /// # Performance Characteristics
    22
    ///
    23
    /// - L2-resident: Process midi_tile.m rows at a time
    24
    /// - Vectorized: 4×1 micro-kernel processes 4 output rows simultaneously
    25
    /// - Aligned: K dimension aligned to Q4_K superblock (256)
    26
    #[derive(Debug, Clone)]
    27
    pub struct TiledQ4KMatvec {
    28
        /// Tiling configuration
    29
        pub config: TilingConfig,
    30
        /// Number of rows (M dimension)
    31
        pub m: usize,
    32
        /// Number of columns (K dimension)
    33
        pub k: usize,
    34
    }
    35
    36
    impl TiledQ4KMatvec {
    37
        /// Create a new tiled Q4K matvec executor
    38
        ///
    39
        /// # Panics
    40
        /// Panics if K is not aligned to Q4_K superblock size (256).
    41
        #[must_use]
    42
    0
        pub fn new(m: usize, k: usize) -> Self {
    43
    0
            assert!(
    44
    0
                k % Q4K_SUPERBLOCK_SIZE == 0,
    45
    0
                "K dimension ({}) must be aligned to Q4_K superblock size ({})",
    46
                k,
    47
                Q4K_SUPERBLOCK_SIZE
    48
            );
    49
    50
    0
            Self {
    51
    0
                config: TilingConfig::cpu_avx2_q4k_matvec(),
    52
    0
                m,
    53
    0
                k,
    54
    0
            }
    55
    0
        }
    56
    57
        /// Get number of superblocks per row
    58
        #[must_use]
    59
    0
        pub fn superblocks_per_row(&self) -> usize {
    60
    0
            self.k / Q4K_SUPERBLOCK_SIZE
    61
    0
        }
    62
    63
        /// Get total number of superblocks
    64
        #[must_use]
    65
    0
        pub fn total_superblocks(&self) -> usize {
    66
    0
            self.m * self.superblocks_per_row()
    67
    0
        }
    68
    69
        /// Get weight bytes offset for a given row
    70
        #[must_use]
    71
        #[inline]
    72
    0
        pub fn weight_row_offset(&self, row: usize) -> usize {
    73
    0
            row * self.superblocks_per_row() * Q4K_SUPERBLOCK_BYTES
    74
    0
        }
    75
    76
        /// Calculate optimal number of parallel rows based on L2 cache
    77
        ///
    78
        /// Goal: Keep working set in L2 (256KB typical)
    79
        /// Working set = midi_tile.m rows × K × sizeof(Q4K) + K × sizeof(f32)
    80
        #[must_use]
    81
    0
        pub fn optimal_parallel_rows(&self, l2_bytes: usize) -> usize {
    82
            // Q4K: 144 bytes per 256 elements = 0.5625 bytes/element
    83
    0
            let row_bytes = (self.k as f32 * 0.5625) as usize;
    84
            // Input vector: K × 4 bytes
    85
    0
            let input_bytes = self.k * 4;
    86
            // Available for rows
    87
    0
            let available = l2_bytes.saturating_sub(input_bytes);
    88
            // Rows that fit (minimum 4 for micro-kernel)
    89
    0
            (available / row_bytes).max(4)
    90
    0
        }
    91
    92
        /// Execute tiled matvec (reference scalar implementation)
    93
        ///
    94
        /// This is the reference implementation for correctness testing.
    95
        /// Actual SIMD implementation would be in the backends.
    96
        ///
    97
        /// For parallel execution, use [`execute_parallel`] when the `parallel` feature is enabled.
    98
    0
        pub fn execute_scalar(&self, weights: &[u8], input: &[f32], output: &mut [f32]) {
    99
    0
            assert_eq!(weights.len(), self.total_superblocks() * Q4K_SUPERBLOCK_BYTES);
    100
    0
            assert_eq!(input.len(), self.k);
    101
    0
            assert_eq!(output.len(), self.m);
    102
    103
    0
            let superblocks_per_row = self.superblocks_per_row();
    104
    105
    0
            for row in 0..self.m {
    106
    0
                let mut sum = 0.0f32;
    107
    0
                let row_offset = row * superblocks_per_row * Q4K_SUPERBLOCK_BYTES;
    108
    109
    0
                for sb in 0..superblocks_per_row {
    110
    0
                    let sb_offset = row_offset + sb * Q4K_SUPERBLOCK_BYTES;
    111
    0
                    let sb_data = &weights[sb_offset..sb_offset + Q4K_SUPERBLOCK_BYTES];
    112
    0
    113
    0
                    // Dequantize and dot product for this superblock
    114
    0
                    let input_offset = sb * Q4K_SUPERBLOCK_SIZE;
    115
    0
                    sum += self.scalar_superblock_dot(sb_data, &input[input_offset..input_offset + Q4K_SUPERBLOCK_SIZE]);
    116
    0
                }
    117
    118
    0
                output[row] = sum;
    119
            }
    120
    0
        }
    121
    122
        /// Execute tiled matvec with parallel row processing
    123
        ///
    124
        /// Uses Rayon to parallelize across rows for multi-core speedup.
    125
        /// Falls back to scalar execution if the `parallel` feature is not enabled.
    126
        ///
    127
        /// # Performance
    128
        ///
    129
        /// Achieves near-linear speedup with core count for large matrices.
    130
        /// For small matrices (< 256 rows), scalar may be faster due to overhead.
    131
        #[cfg(feature = "parallel")]
    132
        pub fn execute_parallel(&self, weights: &[u8], input: &[f32], output: &mut [f32]) {
    133
            use rayon::prelude::*;
    134
    135
            assert_eq!(weights.len(), self.total_superblocks() * Q4K_SUPERBLOCK_BYTES);
    136
            assert_eq!(input.len(), self.k);
    137
            assert_eq!(output.len(), self.m);
    138
    139
            let superblocks_per_row = self.superblocks_per_row();
    140
            let row_stride = superblocks_per_row * Q4K_SUPERBLOCK_BYTES;
    141
    142
            output.par_iter_mut().enumerate().for_each(|(row, out)| {
    143
                let mut sum = 0.0f32;
    144
                let row_offset = row * row_stride;
    145
    146
                for sb in 0..superblocks_per_row {
    147
                    let sb_offset = row_offset + sb * Q4K_SUPERBLOCK_BYTES;
    148
                    let sb_data = &weights[sb_offset..sb_offset + Q4K_SUPERBLOCK_BYTES];
    149
    150
                    let input_offset = sb * Q4K_SUPERBLOCK_SIZE;
    151
                    sum += self.scalar_superblock_dot(sb_data, &input[input_offset..input_offset + Q4K_SUPERBLOCK_SIZE]);
    152
                }
    153
    154
                *out = sum;
    155
            });
    156
        }
    157
    158
        /// Execute tiled matvec with parallel row processing (fallback)
    159
        ///
    160
        /// When `parallel` feature is not enabled, this is equivalent to `execute_scalar`.
    161
        #[cfg(not(feature = "parallel"))]
    162
    0
        pub fn execute_parallel(&self, weights: &[u8], input: &[f32], output: &mut [f32]) {
    163
    0
            self.execute_scalar(weights, input, output);
    164
    0
        }
    165
    166
        /// Scalar dot product for a single Q4_K superblock
    167
        ///
    168
        /// # Performance
    169
        ///
    170
        /// Optimized version with:
    171
        /// - Precomputed scale/min pairs
    172
        /// - Loop unrolling hints
    173
        /// - Minimized branching in inner loop
    174
        #[inline]
    175
    0
        fn scalar_superblock_dot(&self, sb_data: &[u8], input: &[f32]) -> f32 {
    176
            // Read header (hot path optimized)
    177
    0
            let d = f16_to_f32(&sb_data[0..2]);
    178
    0
            let dmin = f16_to_f32(&sb_data[2..4]);
    179
    0
            let scales = &sb_data[4..16];
    180
    0
            let qs = &sb_data[16..144];
    181
    182
            // Precompute all scale/min pairs upfront
    183
    0
            let scale_mins = precompute_scales_mins(scales);
    184
    185
    0
            let mut sum = 0.0f32;
    186
    187
            // Process 256 values in 8 chunks of 32
    188
    0
            for chunk in 0..8 {
    189
    0
                let (sc, m) = scale_mins[chunk];
    190
    0
                let d_scale = d * sc;
    191
    0
                let dm = dmin * m;
    192
    193
    0
                let q_offset = chunk * 16; // 32 nibbles = 16 bytes
    194
    0
                let input_offset = chunk * 32;
    195
    196
                // Process 32 values: low nibbles then high nibbles
    197
                // Manually unroll inner loop for better optimization
    198
    0
                let mut chunk_sum = 0.0f32;
    199
    200
                // Process 16 byte pairs (32 nibbles)
    201
    0
                for i in 0..16 {
    202
    0
                    let byte = qs[q_offset + i];
    203
    0
    204
    0
                    // Extract nibbles
    205
    0
                    let q_lo = (byte & 0x0F) as f32;
    206
    0
                    let q_hi = (byte >> 4) as f32;
    207
    0
    208
    0
                    // Dequantize: val = d * scale * q - dmin * min
    209
    0
                    let val_lo = d_scale * q_lo - dm;
    210
    0
                    let val_hi = d_scale * q_hi - dm;
    211
    0
    212
    0
                    // Accumulate dot product
    213
    0
                    chunk_sum += val_lo * input[input_offset + i];
    214
    0
                    chunk_sum += val_hi * input[input_offset + 16 + i];
    215
    0
                }
    216
    217
    0
                sum += chunk_sum;
    218
            }
    219
    220
    0
            sum
    221
    0
        }
    222
    223
        /// Get tiling statistics for profiling
    224
        #[must_use]
    225
    0
        pub fn stats(&self) -> TilingStats {
    226
    0
            let bytes_per_row = self.superblocks_per_row() * Q4K_SUPERBLOCK_BYTES;
    227
    0
            let total_weight_bytes = self.m * bytes_per_row;
    228
    0
            let input_bytes = self.k * 4;
    229
    0
            let output_bytes = self.m * 4;
    230
    231
    0
            TilingStats {
    232
    0
                total_weight_bytes,
    233
    0
                input_bytes,
    234
    0
                output_bytes,
    235
    0
                superblocks: self.total_superblocks(),
    236
    0
                arithmetic_ops: self.m * self.k * 2, // 2 ops per element (mul + add)
    237
    0
                arithmetic_intensity: (self.m * self.k * 2) as f32 / (total_weight_bytes + input_bytes) as f32,
    238
    0
            }
    239
    0
        }
    240
    }
    241
    242
    /// Statistics for a tiled operation
    243
    #[derive(Debug, Clone)]
    244
    pub struct TilingStats {
    245
        /// Total weight bytes
    246
        pub total_weight_bytes: usize,
    247
        /// Input vector bytes
    248
        pub input_bytes: usize,
    249
        /// Output vector bytes
    250
        pub output_bytes: usize,
    251
        /// Number of superblocks
    252
        pub superblocks: usize,
    253
        /// Total arithmetic operations
    254
        pub arithmetic_ops: usize,
    255
        /// Arithmetic intensity (FLOPS/byte)
    256
        pub arithmetic_intensity: f32,
    257
    }
    258
    259
    /// Convert 2 bytes (f16 IEEE 754) to f32
    260
    ///
    261
    /// Manual implementation to avoid half crate dependency.
    262
    /// Format: 1 sign bit, 5 exponent bits, 10 mantissa bits.
    263
    ///
    264
    /// # Performance
    265
    ///
    266
    /// Optimized for the common case (normal numbers). Special cases (zero,
    267
    /// subnormal, inf, nan) use branches but are rare in practice for model weights.
    268
    #[inline]
    269
    0
    pub fn f16_to_f32(bytes: &[u8]) -> f32 {
    270
    0
        let bits = u16::from_le_bytes([bytes[0], bytes[1]]);
    271
    0
        f16_bits_to_f32(bits)
    272
    0
    }
    273
    274
    /// Fast path f16 to f32 conversion from raw bits
    275
    ///
    276
    /// Optimized version that handles the common case (normal numbers) with
    277
    /// minimal branching. Uses branchless bit manipulation for the hot path.
    278
    #[inline(always)]
    279
    0
    fn f16_bits_to_f32(bits: u16) -> f32 {
    280
    0
        let sign = (bits >> 15) & 0x1;
    281
    0
        let exponent = (bits >> 10) & 0x1F;
    282
    0
        let mantissa = bits & 0x3FF;
    283
    284
        // Fast path: normal numbers (exponent != 0 && exponent != 31)
    285
        // This is the common case for model weights
    286
    0
        if exponent != 0 && exponent != 31 {
    287
            // Branchless conversion for normal numbers
    288
            // f16 bias = 15, f32 bias = 127
    289
    0
            let f32_exp = (exponent as u32 + 112) as u32; // 127 - 15 = 112
    290
    0
            let f32_mant = (mantissa as u32) << 13; // 10 bits -> 23 bits
    291
    0
            let f32_bits = ((sign as u32) << 31) | (f32_exp << 23) | f32_mant;
    292
    0
            return f32::from_bits(f32_bits);
    293
    0
        }
    294
    295
        // Slow path: special cases
    296
    0
        f16_special_to_f32(sign, exponent, mantissa)
    297
    0
    }
    298
    299
    /// Handle f16 special cases (zero, subnormal, inf, nan)
    300
    ///
    301
    /// Cold path - marked to help branch prediction
    302
    #[cold]
    303
    #[inline(never)]
    304
    0
    fn f16_special_to_f32(sign: u16, exponent: u16, mantissa: u16) -> f32 {
    305
    0
        if exponent == 0 {
    306
    0
            if mantissa == 0 {
    307
                // Zero (positive or negative)
    308
    0
                return if sign == 1 { -0.0 } else { 0.0 };
    309
    0
            }
    310
            // Subnormal f16 -> normalized f32
    311
            // 2^-14 as constant to avoid powi() call
    312
            const TWO_POW_NEG_14: f32 = 6.103_515_625e-5; // 2^-14
    313
    0
            let m = mantissa as f32 * (1.0 / 1024.0);
    314
    0
            let result = m * TWO_POW_NEG_14;
    315
    0
            return if sign == 1 { -result } else { result };
    316
    0
        }
    317
    318
        // exponent == 31: Inf or NaN
    319
    0
        if mantissa == 0 {
    320
    0
            if sign == 1 { f32::NEG_INFINITY } else { f32::INFINITY }
    321
        } else {
    322
    0
            f32::NAN
    323
        }
    324
    0
    }
    325
    326
    /// Extract 6-bit scale and min values from packed scales array
    327
    ///
    328
    /// Q4_K uses 6-bit packed scales: 12 bytes encode 8 (scale, min) pairs.
    329
    ///
    330
    /// # Performance
    331
    ///
    332
    /// Uses bitwise operations to avoid branches and bounds checks in the hot path.
    333
    /// The scales array is always 12 bytes, so we use unchecked access after
    334
    /// validating at the entry point.
    335
    #[inline(always)]
    336
    0
    pub fn extract_scale_min_6bit(scales: &[u8], idx: usize) -> (f32, f32) {
    337
    0
        debug_assert!(scales.len() >= 12, "scales array must be at least 12 bytes");
    338
    0
        debug_assert!(idx < 8, "idx must be < 8");
    339
    340
        // Precomputed base offsets: idx * 3 / 2 for idx 0..8
    341
        // [0, 1, 3, 4, 6, 7, 9, 10]
    342
        // Using bitwise: base = idx + (idx >> 1)
    343
    0
        let base = idx + (idx >> 1);
    344
    345
        // Branchless extraction using bitwise selection
    346
        // Even indices: scale = byte[base] & 0x3F
    347
        // Odd indices:  scale = (byte[base] >> 6) | ((byte[base+1] & 0x0F) << 2)
    348
    0
        let is_odd = idx & 1;
    349
    350
        // Safety: base is always < 11 for idx < 8, and scales.len() >= 12
    351
    0
        let b0 = scales[base];
    352
    0
        let b1 = scales[base + 1];
    353
    354
        // Extract scale: branchless using masking
    355
    0
        let scale_even = (b0 & 0x3F) as u32;
    356
    0
        let scale_odd = ((b0 >> 6) | ((b1 & 0x0F) << 2)) as u32;
    357
    0
        let scale = if is_odd == 0 { scale_even } else { scale_odd };
    358
    359
        // Extract min: branchless using masking
    360
    0
        let min_even = ((b0 >> 6) | ((b1 & 0x0F) << 2)) as u32;
    361
        // For odd indices, we need byte at base+2, but use 0 if at boundary
    362
    0
        let b2 = if base + 2 < scales.len() { scales[base + 2] } else { 0 };
    363
    0
        let min_odd = ((b1 >> 4) | ((b2 & 0x03) << 4)) as u32;
    364
    0
        let min = if is_odd == 0 { min_even } else { min_odd };
    365
    366
    0
        (scale as f32, min as f32)
    367
    0
    }
    368
    369
    /// Precompute all 8 scale/min pairs for a Q4_K superblock
    370
    ///
    371
    /// More efficient than calling extract_scale_min_6bit 8 times when
    372
    /// we need all values (which is the common case).
    373
    #[inline]
    374
    0
    fn precompute_scales_mins(scales: &[u8]) -> [(f32, f32); 8] {
    375
    0
        debug_assert!(scales.len() >= 12);
    376
    377
        // Unroll the extraction for all 8 chunks
    378
    0
        [
    379
    0
            extract_scale_min_6bit(scales, 0),
    380
    0
            extract_scale_min_6bit(scales, 1),
    381
    0
            extract_scale_min_6bit(scales, 2),
    382
    0
            extract_scale_min_6bit(scales, 3),
    383
    0
            extract_scale_min_6bit(scales, 4),
    384
    0
            extract_scale_min_6bit(scales, 5),
    385
    0
            extract_scale_min_6bit(scales, 6),
    386
    0
            extract_scale_min_6bit(scales, 7),
    387
    0
        ]
    388
    0
    }