## 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/matrix/ops/arithmetic.rs

    Line
    Count
    Source
    1
    //! Matrix arithmetic operations
    2
    //!
    3
    //! This module provides matrix multiplication and related operations:
    4
    //! - `matmul()` - Standard matrix multiplication with SIMD optimization
    5
    //! - `batched_matmul()` - Batched 3D tensor multiplication
    6
    //! - `batched_matmul_4d()` - 4D tensor multiplication for attention
    7
    //!
    8
    //! ## Domain Separation (PMAT-018)
    9
    //!
    10
    //! Arithmetic operations (multiplication, addition) are separate from storage
    11
    //! operations (allocation, indexing). This allows optimizing compute kernels
    12
    //! independently of memory layout decisions.
    13
    //!
    14
    //! ## Performance Hierarchy
    15
    //!
    16
    //! 1. GPU for large matrices (≥500×500) - 2-10x speedup
    17
    //! 2. BLIS/SIMD for medium-large matrices (>64×64) - 2-8x speedup
    18
    //! 3. Naive for small matrices - lowest overhead
    19
    20
    use crate::TruenoError;
    21
    22
    #[cfg(feature = "tracing")]
    23
    use tracing::instrument;
    24
    25
    use super::super::Matrix;
    26
    27
    impl Matrix<f32> {
    28
        /// Matrix multiplication (matmul)
    29
        ///
    30
        /// Computes `C = A × B` where A is `m×n`, B is `n×p`, and C is `m×p`.
    31
        ///
    32
        /// # Arguments
    33
        ///
    34
        /// * `other` - The matrix to multiply with (right operand)
    35
        ///
    36
        /// # Returns
    37
        ///
    38
        /// A new matrix containing the result of matrix multiplication
    39
        ///
    40
        /// # Errors
    41
        ///
    42
        /// Returns `InvalidInput` if matrix dimensions are incompatible
    43
        /// (i.e., `self.cols != other.rows`)
    44
        ///
    45
        /// # Example
    46
        ///
    47
        /// ```
    48
        /// use trueno::Matrix;
    49
        ///
    50
        /// let a = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]).unwrap();
    51
        /// let b = Matrix::from_vec(2, 2, vec![5.0, 6.0, 7.0, 8.0]).unwrap();
    52
        /// let c = a.matmul(&b).unwrap();
    53
        ///
    54
        /// // [[1, 2],   [[5, 6],   [[19, 22],
    55
        /// //  [3, 4]] ×  [7, 8]] =  [43, 50]]
    56
        /// assert_eq!(c.get(0, 0), Some(&19.0));
    57
        /// assert_eq!(c.get(0, 1), Some(&22.0));
    58
        /// assert_eq!(c.get(1, 0), Some(&43.0));
    59
        /// assert_eq!(c.get(1, 1), Some(&50.0));
    60
        /// ```
    61
        // =========================================================================
    62
        // HOT PATH - PERFORMANCE CRITICAL
    63
        // =========================================================================
    64
        // Core matrix operation used in neural network forward passes.
    65
        // Changes to inner loops REQUIRE benchmark verification: make bench-check
    66
        // =========================================================================
    67
        #[cfg_attr(feature = "tracing", instrument(skip(self, other), fields(dims = %format!("{}x{} @ {}x{}", self.rows, self.cols, other.rows, other.cols))))]
    68
    0
        pub fn matmul(&self, other: &Matrix<f32>) -> Result<Matrix<f32>, TruenoError> {
    69
    0
            if self.cols != other.rows {
    70
    0
                return Err(TruenoError::InvalidInput(format!(
    71
    0
                    "Matrix dimension mismatch for multiplication: {}×{} × {}×{} (inner dimensions {} and {} must match)",
    72
    0
                    self.rows, self.cols, other.rows, other.cols, self.cols, other.rows
    73
    0
                )));
    74
    0
            }
    75
    76
            // Fast path for vector-matrix multiply (rows=1)
    77
    0
            if self.rows == 1 {
    78
    0
                return self.matmul_vector_matrix(other);
    79
    0
            }
    80
    81
    0
            let mut result = Matrix::zeros_with_backend(self.rows, other.cols, self.backend);
    82
    83
            #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    84
            const GPU_THRESHOLD: usize = 500;
    85
            const SIMD_THRESHOLD: usize = 64;
    86
    87
            // Try GPU first for very large matrices
    88
            #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    89
            {
    90
    0
                if self.rows >= GPU_THRESHOLD
    91
    0
                    && self.cols >= GPU_THRESHOLD
    92
    0
                    && other.cols >= GPU_THRESHOLD
    93
                {
    94
    0
                    if let Ok(gpu_result) = self.matmul_gpu(other) {
    95
    0
                        return Ok(gpu_result);
    96
    0
                    }
    97
    0
                }
    98
            }
    99
    100
            // Use SIMD for medium-large matrices
    101
    0
            if self.rows >= SIMD_THRESHOLD
    102
    0
                || self.cols >= SIMD_THRESHOLD
    103
    0
                || other.cols >= SIMD_THRESHOLD
    104
            {
    105
                #[cfg(target_arch = "wasm32")]
    106
                {
    107
                    self.matmul_wasm_tiled(other, &mut result)?;
    108
                }
    109
                #[cfg(not(target_arch = "wasm32"))]
    110
                {
    111
    0
                    crate::blis::gemm_blis(
    112
    0
                        self.rows,
    113
    0
                        other.cols,
    114
    0
                        self.cols,
    115
    0
                        &self.data,
    116
    0
                        &other.data,
    117
    0
                        &mut result.data,
    118
    0
                        None,
    119
    0
                    )?;
    120
                }
    121
            } else {
    122
    0
                self.matmul_naive(other, &mut result)?;
    123
            }
    124
    125
    0
            Ok(result)
    126
    0
        }
    127
    128
        /// Batched matrix multiplication for 3D tensors.
    129
        ///
    130
        /// Computes `[batch, m, k] @ [batch, k, n] -> [batch, m, n]` using SIMD for each batch.
    131
        #[cfg_attr(
    132
            feature = "tracing",
    133
            instrument(skip(a_data, b_data), fields(batch, m, k, n))
    134
        )]
    135
    0
        pub fn batched_matmul(
    136
    0
            a_data: &[f32],
    137
    0
            b_data: &[f32],
    138
    0
            batch: usize,
    139
    0
            m: usize,
    140
    0
            k: usize,
    141
    0
            n: usize,
    142
    0
        ) -> Result<Vec<f32>, TruenoError> {
    143
    0
            let a_stride = m * k;
    144
    0
            let b_stride = k * n;
    145
    0
            let out_stride = m * n;
    146
    147
    0
            if a_data.len() != batch * a_stride {
    148
    0
                return Err(TruenoError::InvalidInput(format!(
    149
    0
                    "A data size mismatch: expected {} ({}×{}×{}), got {}",
    150
    0
                    batch * a_stride, batch, m, k, a_data.len()
    151
    0
                )));
    152
    0
            }
    153
    0
            if b_data.len() != batch * b_stride {
    154
    0
                return Err(TruenoError::InvalidInput(format!(
    155
    0
                    "B data size mismatch: expected {} ({}×{}×{}), got {}",
    156
    0
                    batch * b_stride, batch, k, n, b_data.len()
    157
    0
                )));
    158
    0
            }
    159
    160
    0
            let mut output = vec![0.0f32; batch * out_stride];
    161
    162
    0
            for ba in 0..batch {
    163
    0
                let a_offset = ba * a_stride;
    164
    0
                let b_offset = ba * b_stride;
    165
    0
                let out_offset = ba * out_stride;
    166
    167
    0
                let a_slice = &a_data[a_offset..a_offset + a_stride];
    168
    0
                let b_slice = &b_data[b_offset..b_offset + b_stride];
    169
    170
    0
                let a_mat = Matrix::from_slice(m, k, a_slice)?;
    171
    0
                let b_mat = Matrix::from_slice(k, n, b_slice)?;
    172
    173
    0
                let result = a_mat.matmul(&b_mat)?;
    174
    0
                output[out_offset..out_offset + out_stride].copy_from_slice(result.as_slice());
    175
            }
    176
    177
    0
            Ok(output)
    178
    0
        }
    179
    180
        /// Batched matrix multiplication for 4D tensors (attention pattern).
    181
        ///
    182
        /// Computes `[batch, heads, m, k] @ [batch, heads, k, n] -> [batch, heads, m, n]`
    183
        #[cfg_attr(
    184
            feature = "tracing",
    185
            instrument(skip(a_data, b_data), fields(batch, heads, m, k, n))
    186
        )]
    187
    0
        pub fn batched_matmul_4d(
    188
    0
            a_data: &[f32],
    189
    0
            b_data: &[f32],
    190
    0
            batch: usize,
    191
    0
            heads: usize,
    192
    0
            m: usize,
    193
    0
            k: usize,
    194
    0
            n: usize,
    195
    0
        ) -> Result<Vec<f32>, TruenoError> {
    196
    0
            let a_head_stride = m * k;
    197
    0
            let b_head_stride = k * n;
    198
    0
            let out_head_stride = m * n;
    199
    0
            let total_heads = batch * heads;
    200
    201
    0
            let expected_a = total_heads * a_head_stride;
    202
    0
            let expected_b = total_heads * b_head_stride;
    203
    0
            if a_data.len() != expected_a {
    204
    0
                return Err(TruenoError::InvalidInput(format!(
    205
    0
                    "A data size mismatch: expected {} ({}×{}×{}×{}), got {}",
    206
    0
                    expected_a, batch, heads, m, k, a_data.len()
    207
    0
                )));
    208
    0
            }
    209
    0
            if b_data.len() != expected_b {
    210
    0
                return Err(TruenoError::InvalidInput(format!(
    211
    0
                    "B data size mismatch: expected {} ({}×{}×{}×{}), got {}",
    212
    0
                    expected_b, batch, heads, k, n, b_data.len()
    213
    0
                )));
    214
    0
            }
    215
    216
    0
            let mut output = vec![0.0f32; total_heads * out_head_stride];
    217
    218
    0
            for bh in 0..total_heads {
    219
    0
                let a_offset = bh * a_head_stride;
    220
    0
                let b_offset = bh * b_head_stride;
    221
    0
                let out_offset = bh * out_head_stride;
    222
    223
    0
                let a_slice = &a_data[a_offset..a_offset + a_head_stride];
    224
    0
                let b_slice = &b_data[b_offset..b_offset + b_head_stride];
    225
    226
    0
                let a_mat = Matrix::from_slice(m, k, a_slice)?;
    227
    0
                let b_mat = Matrix::from_slice(k, n, b_slice)?;
    228
    229
    0
                let result = a_mat.matmul(&b_mat)?;
    230
    0
                output[out_offset..out_offset + out_head_stride].copy_from_slice(result.as_slice());
    231
            }
    232
    233
    0
            Ok(output)
    234
    0
        }
    235
    236
        /// Fast path for vector-matrix multiplication (1×K @ K×N → 1×N)
    237
        #[cfg_attr(feature = "tracing", instrument(skip(self, other), fields(k = self.cols, n = other.cols)))]
    238
    0
        fn matmul_vector_matrix(&self, other: &Matrix<f32>) -> Result<Matrix<f32>, TruenoError> {
    239
    0
            debug_assert_eq!(self.rows, 1);
    240
    241
    0
            let k = self.cols;
    242
    0
            let n = other.cols;
    243
    0
            let mut result = Matrix::zeros_with_backend(1, n, self.backend);
    244
    245
    0
            for ki in 0..k {
    246
    0
                let a_k = self.data[ki];
    247
    0
                if a_k == 0.0 {
    248
    0
                    continue;
    249
    0
                }
    250
    251
    0
                let b_row_start = ki * n;
    252
    0
                for j in 0..n {
    253
    0
                    result.data[j] += a_k * other.data[b_row_start + j];
    254
    0
                }
    255
            }
    256
    257
    0
            Ok(result)
    258
    0
        }
    259
    260
        /// Naive O(n³) matrix multiplication (baseline for small matrices)
    261
    0
        fn matmul_naive(
    262
    0
            &self,
    263
    0
            other: &Matrix<f32>,
    264
    0
            result: &mut Matrix<f32>,
    265
    0
        ) -> Result<(), TruenoError> {
    266
    0
            for i in 0..self.rows {
    267
    0
                for j in 0..other.cols {
    268
    0
                    let mut sum = 0.0;
    269
    0
                    for k in 0..self.cols {
    270
    0
                        sum += self.get(i, k).expect("bounds validated")
    271
    0
                            * other.get(k, j).expect("bounds validated");
    272
    0
                    }
    273
    0
                    *result.get_mut(i, j).expect("bounds validated") = sum;
    274
                }
    275
            }
    276
    0
            Ok(())
    277
    0
        }
    278
    279
        /// WASM-optimized tiled matrix multiplication
    280
        #[allow(dead_code)]
    281
    0
        fn matmul_wasm_tiled(
    282
    0
            &self,
    283
    0
            other: &Matrix<f32>,
    284
    0
            result: &mut Matrix<f32>,
    285
    0
        ) -> Result<(), TruenoError> {
    286
    0
            let m = self.rows;
    287
    0
            let k = self.cols;
    288
    0
            let n = other.cols;
    289
    290
    0
            for i in 0..m {
    291
    0
                let a_row_start = i * k;
    292
    0
                let result_row_start = i * n;
    293
    294
    0
                let simd_width = 8;
    295
    0
                let n_simd = (n / simd_width) * simd_width;
    296
    297
                #[allow(clippy::needless_range_loop)]
    298
    0
                for j0 in (0..n_simd).step_by(simd_width) {
    299
    0
                    let mut acc = [0.0f32; 8];
    300
    301
    0
                    for kk in 0..k {
    302
    0
                        let a_val = self.data[a_row_start + kk];
    303
    0
                        let b_row_start = kk * n + j0;
    304
    305
    0
                        for jj in 0..simd_width {
    306
    0
                            acc[jj] += a_val * other.data[b_row_start + jj];
    307
    0
                        }
    308
                    }
    309
    310
    0
                    for jj in 0..simd_width {
    311
    0
                        result.data[result_row_start + j0 + jj] = acc[jj];
    312
    0
                    }
    313
                }
    314
    315
    0
                for j in n_simd..n {
    316
    0
                    let mut sum = 0.0f32;
    317
    0
                    for kk in 0..k {
    318
    0
                        sum += self.data[a_row_start + kk] * other.data[kk * n + j];
    319
    0
                    }
    320
    0
                    result.data[result_row_start + j] = sum;
    321
                }
    322
            }
    323
    324
    0
            Ok(())
    325
    0
        }
    326
    327
        /// GPU-accelerated matrix multiplication
    328
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    329
    0
        fn matmul_gpu(&self, other: &Matrix<f32>) -> Result<Matrix<f32>, TruenoError> {
    330
            use crate::backends::gpu::GpuBackend;
    331
    332
    0
            if !GpuBackend::is_available() {
    333
    0
                return Err(TruenoError::InvalidInput("GPU not available".to_string()));
    334
    0
            }
    335
    336
    0
            let mut gpu = GpuBackend::new();
    337
    0
            let result_data = gpu
    338
    0
                .matmul(&self.data, &other.data, self.rows, self.cols, other.cols)
    339
    0
                .map_err(|e| TruenoError::InvalidInput(format!("GPU matmul failed: {}", e)))?;
    340
    341
    0
            let mut result = Matrix::zeros(self.rows, other.cols);
    342
    0
            result.data = result_data;
    343
    344
    0
            Ok(result)
    345
    0
        }
    346
    }
    347
    348
    #[cfg(test)]
    349
    mod tests {
    350
        use super::*;
    351
    352
        #[test]
    353
        fn test_matmul_basic() {
    354
            let a = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]).unwrap();
    355
            let b = Matrix::from_vec(2, 2, vec![5.0, 6.0, 7.0, 8.0]).unwrap();
    356
            let c = a.matmul(&b).unwrap();
    357
    358
            assert_eq!(c.get(0, 0), Some(&19.0));
    359
            assert_eq!(c.get(0, 1), Some(&22.0));
    360
            assert_eq!(c.get(1, 0), Some(&43.0));
    361
            assert_eq!(c.get(1, 1), Some(&50.0));
    362
        }
    363
    364
        #[test]
    365
        fn test_matmul_dimension_mismatch() {
    366
            let a = Matrix::from_vec(2, 3, vec![1.0; 6]).unwrap();
    367
            let b = Matrix::from_vec(2, 2, vec![1.0; 4]).unwrap();
    368
            assert!(a.matmul(&b).is_err());
    369
        }
    370
    371
        #[test]
    372
        fn test_matmul_identity() {
    373
            let a = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]).unwrap();
    374
            let i = Matrix::identity(2);
    375
            let result = a.matmul(&i).unwrap();
    376
            assert_eq!(result.as_slice(), a.as_slice());
    377
        }
    378
    379
        #[test]
    380
        fn test_batched_matmul() {
    381
            let a = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]; // 2 batches of 2×2
    382
            let b = vec![1.0, 0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0]; // 2 identity matrices
    383
            let result = Matrix::batched_matmul(&a, &b, 2, 2, 2, 2).unwrap();
    384
            assert_eq!(result, a); // A × I = A
    385
        }
    386
    }