## Coverage Report

#### Created: 2026-01-23 22:55

next uncovered line (L), next uncovered region (R), next uncovered branch (B)

    /home/noah/src/trueno/src/matrix/mod.rs

    Line
    Count
    Source
    1
    //! Matrix operations for Trueno
    2
    //!
    3
    //! Provides 2D matrix operations with SIMD optimization for linear algebra,
    4
    //! machine learning, and scientific computing.
    5
    //!
    6
    //! # Example
    7
    //!
    8
    //! ```
    9
    //! use trueno::Matrix;
    10
    //!
    11
    //! // Create a 2x3 matrix
    12
    //! let m = Matrix::zeros(2, 3);
    13
    //! assert_eq!(m.rows(), 2);
    14
    //! assert_eq!(m.cols(), 3);
    15
    //! ```
    16
    17
    // Allow dead_code for experimental SIMD microkernels kept for future optimization work
    18
    #![allow(dead_code)]
    19
    20
    use crate::{Backend, TruenoError, Vector};
    21
    22
    #[cfg(feature = "tracing")]
    23
    use tracing::instrument;
    24
    25
    /// A 2D matrix with row-major storage
    26
    ///
    27
    /// Data is stored in row-major format (C-style), where consecutive elements
    28
    /// in memory belong to the same row. This is compatible with NumPy's default
    29
    /// layout and optimal for cache locality when accessing rows.
    30
    ///
    31
    /// # Storage Layout
    32
    ///
    33
    /// For a 2x3 matrix:
    34
    /// ```text
    35
    /// [[a, b, c],
    36
    ///  [d, e, f]]
    37
    /// ```
    38
    /// Data is stored as: [a, b, c, d, e, f]
    39
    ///
    40
    /// # Example
    41
    ///
    42
    /// ```
    43
    /// use trueno::Matrix;
    44
    ///
    45
    /// let m = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]).unwrap();
    46
    /// assert_eq!(m.get(0, 0), Some(&1.0));
    47
    /// assert_eq!(m.get(0, 1), Some(&2.0));
    48
    /// assert_eq!(m.get(1, 0), Some(&3.0));
    49
    /// assert_eq!(m.get(1, 1), Some(&4.0));
    50
    /// ```
    51
    #[derive(Debug, Clone, PartialEq)]
    52
    pub struct Matrix<T> {
    53
        rows: usize,
    54
        cols: usize,
    55
        data: Vec<T>,
    56
        backend: Backend,
    57
    }
    58
    59
    impl std::ops::Index<(usize, usize)> for Matrix<f32> {
    60
        type Output = f32;
    61
    62
    0
        fn index(&self, (row, col): (usize, usize)) -> &Self::Output {
    63
    0
            &self.data[row * self.cols + col]
    64
    0
        }
    65
    }
    66
    67
    impl Matrix<f32> {
    68
        /// Creates a new matrix with uninitialized values
    69
        ///
    70
        /// # Arguments
    71
        ///
    72
        /// * `rows` - Number of rows
    73
        /// * `cols` - Number of columns
    74
        ///
    75
        /// # Returns
    76
        ///
    77
        /// A new matrix with dimensions `rows x cols` containing uninitialized values
    78
        ///
    79
        /// # Example
    80
        ///
    81
        /// ```
    82
        /// use trueno::Matrix;
    83
        ///
    84
        /// let m = Matrix::new(3, 4);
    85
        /// assert_eq!(m.rows(), 3);
    86
        /// assert_eq!(m.cols(), 4);
    87
        /// ```
    88
    2
        pub fn new(rows: usize, cols: usize) -> Self {
    89
    2
            let backend = Backend::select_best();
    90
    2
            Matrix {
    91
    2
                rows,
    92
    2
                cols,
    93
    2
                data: vec![0.0; rows * cols],
    94
    2
                backend,
    95
    2
            }
    96
    2
        }
    97
    98
        /// Creates a matrix from a vector of data
    99
        ///
    100
        /// # Arguments
    101
        ///
    102
        /// * `rows` - Number of rows
    103
        /// * `cols` - Number of columns
    104
        /// * `data` - Vector containing matrix elements in row-major order
    105
        ///
    106
        /// # Errors
    107
        ///
    108
        /// Returns `InvalidInput` if `data.len() != rows * cols`
    109
        ///
    110
        /// # Example
    111
        ///
    112
        /// ```
    113
        /// use trueno::Matrix;
    114
        ///
    115
        /// let m = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]).unwrap();
    116
        /// assert_eq!(m.rows(), 2);
    117
        /// assert_eq!(m.cols(), 2);
    118
        /// ```
    119
    0
        pub fn from_vec(rows: usize, cols: usize, data: Vec<f32>) -> Result<Self, TruenoError> {
    120
    0
            if data.len() != rows * cols {
    121
    0
                return Err(TruenoError::InvalidInput(format!(
    122
    0
                    "Data length {} does not match matrix dimensions {}x{} (expected {})",
    123
    0
                    data.len(),
    124
    0
                    rows,
    125
    0
                    cols,
    126
    0
                    rows * cols
    127
    0
                )));
    128
    0
            }
    129
    130
    0
            let backend = Backend::select_best();
    131
    0
            Ok(Matrix {
    132
    0
                rows,
    133
    0
                cols,
    134
    0
                data,
    135
    0
                backend,
    136
    0
            })
    137
    0
        }
    138
    139
        /// Creates a matrix from a vector with a specific backend
    140
        ///
    141
        /// This is useful for testing specific SIMD code paths.
    142
    0
        pub fn from_vec_with_backend(
    143
    0
            rows: usize,
    144
    0
            cols: usize,
    145
    0
            data: Vec<f32>,
    146
    0
            backend: Backend,
    147
    0
        ) -> Self {
    148
    0
            assert_eq!(
    149
    0
                data.len(),
    150
    0
                rows * cols,
    151
    0
                "Data length {} does not match matrix dimensions {}x{}",
    152
    0
                data.len(),
    153
                rows,
    154
                cols
    155
            );
    156
    0
            Matrix {
    157
    0
                rows,
    158
    0
                cols,
    159
    0
                data,
    160
    0
                backend,
    161
    0
            }
    162
    0
        }
    163
    164
        /// Creates a matrix from a slice by copying the data
    165
        ///
    166
        /// This is a convenience method that copies the slice into an owned vector.
    167
        /// For zero-copy scenarios, consider using the data directly with `from_vec`
    168
        /// if you already have an owned `Vec`.
    169
        ///
    170
        /// # Arguments
    171
        ///
    172
        /// * `rows` - Number of rows
    173
        /// * `cols` - Number of columns
    174
        /// * `data` - Slice containing matrix elements in row-major order
    175
        ///
    176
        /// # Errors
    177
        ///
    178
        /// Returns `InvalidInput` if `data.len() != rows * cols`
    179
        ///
    180
        /// # Example
    181
        ///
    182
        /// ```
    183
        /// use trueno::Matrix;
    184
        ///
    185
        /// let data = [1.0, 2.0, 3.0, 4.0];
    186
        /// let m = Matrix::from_slice(2, 2, &data).unwrap();
    187
        /// assert_eq!(m.get(0, 0), Some(&1.0));
    188
        /// ```
    189
    0
        pub fn from_slice(rows: usize, cols: usize, data: &[f32]) -> Result<Self, TruenoError> {
    190
    0
            Self::from_vec(rows, cols, data.to_vec())
    191
    0
        }
    192
    193
        /// Creates a matrix filled with zeros
    194
        ///
    195
        /// # Example
    196
        ///
    197
        /// ```
    198
        /// use trueno::Matrix;
    199
        ///
    200
        /// let m = Matrix::zeros(3, 3);
    201
        /// assert_eq!(m.get(1, 1), Some(&0.0));
    202
        /// ```
    203
    2
        pub fn zeros(rows: usize, cols: usize) -> Self {
    204
    2
            Matrix::new(rows, cols)
    205
    2
        }
    206
    207
        /// Creates a matrix filled with zeros using a specific backend
    208
        /// (Internal use only - reuses backend from parent matrix)
    209
    0
        fn zeros_with_backend(rows: usize, cols: usize, backend: Backend) -> Self {
    210
    0
            Matrix {
    211
    0
                rows,
    212
    0
                cols,
    213
    0
                data: vec![0.0; rows * cols],
    214
    0
                backend,
    215
    0
            }
    216
    0
        }
    217
    218
        /// Creates an identity matrix (square matrix with 1s on diagonal)
    219
        ///
    220
        /// # Example
    221
        ///
    222
        /// ```
    223
        /// use trueno::Matrix;
    224
        ///
    225
        /// let m = Matrix::identity(3);
    226
        /// assert_eq!(m.get(0, 0), Some(&1.0));
    227
        /// assert_eq!(m.get(0, 1), Some(&0.0));
    228
        /// assert_eq!(m.get(1, 1), Some(&1.0));
    229
        /// ```
    230
    0
        pub fn identity(n: usize) -> Self {
    231
    0
            let mut data = vec![0.0; n * n];
    232
    0
            for i in 0..n {
    233
    0
                data[i * n + i] = 1.0;
    234
    0
            }
    235
    0
            let backend = Backend::select_best();
    236
    0
            Matrix {
    237
    0
                rows: n,
    238
    0
                cols: n,
    239
    0
                data,
    240
    0
                backend,
    241
    0
            }
    242
    0
        }
    243
    244
        /// Returns the number of rows
    245
    2
        pub fn rows(&self) -> usize {
    246
    2
            self.rows
    247
    2
        }
    248
    249
        /// Returns the number of columns
    250
    2
        pub fn cols(&self) -> usize {
    251
    2
            self.cols
    252
    2
        }
    253
    254
        /// Returns the shape as (rows, cols)
    255
    0
        pub fn shape(&self) -> (usize, usize) {
    256
    0
            (self.rows, self.cols)
    257
    0
        }
    258
    259
        /// Gets a reference to an element at (row, col)
    260
        ///
    261
        /// Returns `None` if indices are out of bounds
    262
    0
        pub fn get(&self, row: usize, col: usize) -> Option<&f32> {
    263
    0
            if row >= self.rows || col >= self.cols {
    264
    0
                None
    265
            } else {
    266
    0
                self.data.get(row * self.cols + col)
    267
            }
    268
    0
        }
    269
    270
        /// Gets a mutable reference to an element at (row, col)
    271
        ///
    272
        /// Returns `None` if indices are out of bounds
    273
    0
        pub fn get_mut(&mut self, row: usize, col: usize) -> Option<&mut f32> {
    274
    0
            if row >= self.rows || col >= self.cols {
    275
    0
                None
    276
            } else {
    277
    0
                let idx = row * self.cols + col;
    278
    0
                self.data.get_mut(idx)
    279
            }
    280
    0
        }
    281
    282
        /// Returns a reference to the underlying data
    283
    0
        pub fn as_slice(&self) -> &[f32] {
    284
    0
            &self.data
    285
    0
        }
    286
    287
        /// Matrix multiplication (matmul)
    288
        ///
    289
        /// Computes `C = A × B` where A is `m×n`, B is `n×p`, and C is `m×p`.
    290
        ///
    291
        /// # Arguments
    292
        ///
    293
        /// * `other` - The matrix to multiply with (right operand)
    294
        ///
    295
        /// # Returns
    296
        ///
    297
        /// A new matrix containing the result of matrix multiplication
    298
        ///
    299
        /// # Errors
    300
        ///
    301
        /// Returns `InvalidInput` if matrix dimensions are incompatible
    302
        /// (i.e., `self.cols != other.rows`)
    303
        ///
    304
        /// # Example
    305
        ///
    306
        /// ```
    307
        /// use trueno::Matrix;
    308
        ///
    309
        /// let a = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]).unwrap();
    310
        /// let b = Matrix::from_vec(2, 2, vec![5.0, 6.0, 7.0, 8.0]).unwrap();
    311
        /// let c = a.matmul(&b).unwrap();
    312
        ///
    313
        /// // [[1, 2],   [[5, 6],   [[19, 22],
    314
        /// //  [3, 4]] ×  [7, 8]] =  [43, 50]]
    315
        /// assert_eq!(c.get(0, 0), Some(&19.0));
    316
        /// assert_eq!(c.get(0, 1), Some(&22.0));
    317
        /// assert_eq!(c.get(1, 0), Some(&43.0));
    318
        /// assert_eq!(c.get(1, 1), Some(&50.0));
    319
        /// ```
    320
        // =========================================================================
    321
        // HOT PATH - PERFORMANCE CRITICAL
    322
        // =========================================================================
    323
        // Core matrix operation used in neural network forward passes.
    324
        // Changes to inner loops REQUIRE benchmark verification: make bench-check
    325
        // See convolve2d for prohibited patterns in hot loops.
    326
        // =========================================================================
    327
        #[cfg_attr(feature = "tracing", instrument(skip(self, other), fields(dims = %format!("{}x{} @ {}x{}", self.rows, self.cols, other.rows, other.cols))))]
    328
    0
        pub fn matmul(&self, other: &Matrix<f32>) -> Result<Matrix<f32>, TruenoError> {
    329
    0
            if self.cols != other.rows {
    330
    0
                return Err(TruenoError::InvalidInput(format!(
    331
    0
                    "Matrix dimension mismatch for multiplication: {}×{} × {}×{} (inner dimensions {} and {} must match)",
    332
    0
                    self.rows, self.cols, other.rows, other.cols, self.cols, other.rows
    333
    0
                )));
    334
    0
            }
    335
    336
            // Fast path for vector-matrix multiply (rows=1)
    337
            // Common in ML vocab projection: hidden_state @ embedding_transposed
    338
            // 8x faster than general matmul for 1×384 @ 384×51865 pattern
    339
    0
            if self.rows == 1 {
    340
    0
                return self.matmul_vector_matrix(other);
    341
    0
            }
    342
    343
    0
            let mut result = Matrix::zeros_with_backend(self.rows, other.cols, self.backend);
    344
    345
            // Backend selection strategy (empirical - see docs/performance-analysis.md):
    346
            // 1. GPU for large matrices (≥500×500) - 2-10x speedup (measured)
    347
            // 2. SIMD for medium-large matrices (>64×64) - 2-8x speedup
    348
            // 3. Naive for small matrices - lowest overhead
    349
    350
            #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    351
            const GPU_THRESHOLD: usize = 500; // Empirical: 2x at 500×500, 9.6x at 1000×1000
    352
            const SIMD_THRESHOLD: usize = 64;
    353
    354
            // Try GPU first for very large matrices
    355
            #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    356
            {
    357
    0
                if self.rows >= GPU_THRESHOLD
    358
    0
                    && self.cols >= GPU_THRESHOLD
    359
    0
                    && other.cols >= GPU_THRESHOLD
    360
                {
    361
    0
                    if let Ok(gpu_result) = self.matmul_gpu(other) {
    362
    0
                        return Ok(gpu_result);
    363
    0
                    }
    364
                    // GPU failed, fall through to SIMD/naive
    365
    0
                }
    366
            }
    367
    368
            // Use SIMD for medium-large matrices
    369
    0
            if self.rows >= SIMD_THRESHOLD
    370
    0
                || self.cols >= SIMD_THRESHOLD
    371
    0
                || other.cols >= SIMD_THRESHOLD
    372
            {
    373
                // Tiled approach threshold: below this size, tiling beats transpose
    374
                // Based on WASM optimization spec benchmarks
    375
                const TILED_THRESHOLD: usize = 512;
    376
    377
    0
                let max_dim = self.rows.max(self.cols).max(other.cols);
    378
    379
    0
                if max_dim < TILED_THRESHOLD {
    380
                    // Medium matrices: use BLIS on native, tiled on WASM
    381
                    #[cfg(target_arch = "wasm32")]
    382
                    {
    383
                        self.matmul_wasm_tiled(other, &mut result)?;
    384
                    }
    385
                    #[cfg(not(target_arch = "wasm32"))]
    386
                    {
    387
                        // BLIS is faster than tiled for all sizes on native
    388
    0
                        crate::blis::gemm_blis(
    389
    0
                            self.rows,
    390
    0
                            other.cols,
    391
    0
                            self.cols,
    392
    0
                            &self.data,
    393
    0
                            &other.data,
    394
    0
                            &mut result.data,
    395
    0
                            None,
    396
    0
                        )?;
    397
                    }
    398
                } else {
    399
                    // Large matrices: platform-specific optimized paths
    400
                    #[cfg(target_arch = "wasm32")]
    401
                    {
    402
                        // WASM: tiled is always better (no SIMD microkernel advantage)
    403
                        self.matmul_wasm_tiled(other, &mut result)?;
    404
                    }
    405
                    #[cfg(not(target_arch = "wasm32"))]
    406
                    {
    407
                        // Native: use BLIS-style GEMM with register blocking
    408
                        // ~2x faster than old SIMD implementation for large matrices
    409
    0
                        crate::blis::gemm_blis(
    410
    0
                            self.rows,
    411
    0
                            other.cols,
    412
    0
                            self.cols,
    413
    0
                            &self.data,
    414
    0
                            &other.data,
    415
    0
                            &mut result.data,
    416
    0
                            None,
    417
    0
                        )?;
    418
                    }
    419
                }
    420
            } else {
    421
    0
                self.matmul_naive(other, &mut result)?;
    422
            }
    423
    424
    0
            Ok(result)
    425
    0
        }
    426
    427
        /// Batched matrix multiplication for 3D tensors.
    428
        ///
    429
        /// Computes `[batch, m, k] @ [batch, k, n] -> [batch, m, n]` using SIMD for each batch.
    430
        /// This is critical for transformer attention performance.
    431
        ///
    432
        /// # Arguments
    433
        /// * `a_data` - Flattened input A with shape [batch, m, k]
    434
        /// * `b_data` - Flattened input B with shape [batch, k, n]
    435
        /// * `batch` - Batch dimension
    436
        /// * `m` - Rows of A (and output)
    437
        /// * `k` - Columns of A / Rows of B
    438
        /// * `n` - Columns of B (and output)
    439
        ///
    440
        /// # Returns
    441
        /// Flattened output with shape [batch, m, n]
    442
        ///
    443
        /// # Performance
    444
        /// Uses SIMD matmul for each batch slice, achieving ~50 GFLOPS vs ~0.1 GFLOPS naive.
    445
        /// See Williams et al., 2009 (Roofline model) for theoretical analysis.
    446
        #[cfg_attr(
    447
            feature = "tracing",
    448
            instrument(skip(a_data, b_data), fields(batch, m, k, n))
    449
        )]
    450
    0
        pub fn batched_matmul(
    451
    0
            a_data: &[f32],
    452
    0
            b_data: &[f32],
    453
    0
            batch: usize,
    454
    0
            m: usize,
    455
    0
            k: usize,
    456
    0
            n: usize,
    457
    0
        ) -> Result<Vec<f32>, TruenoError> {
    458
    0
            let a_stride = m * k;
    459
    0
            let b_stride = k * n;
    460
    0
            let out_stride = m * n;
    461
    462
            // Validate input sizes
    463
    0
            if a_data.len() != batch * a_stride {
    464
    0
                return Err(TruenoError::InvalidInput(format!(
    465
    0
                    "A data size mismatch: expected {} ({}×{}×{}), got {}",
    466
    0
                    batch * a_stride,
    467
    0
                    batch,
    468
    0
                    m,
    469
    0
                    k,
    470
    0
                    a_data.len()
    471
    0
                )));
    472
    0
            }
    473
    0
            if b_data.len() != batch * b_stride {
    474
    0
                return Err(TruenoError::InvalidInput(format!(
    475
    0
                    "B data size mismatch: expected {} ({}×{}×{}), got {}",
    476
    0
                    batch * b_stride,
    477
    0
                    batch,
    478
    0
                    k,
    479
    0
                    n,
    480
    0
                    b_data.len()
    481
    0
                )));
    482
    0
            }
    483
    484
    0
            let mut output = vec![0.0f32; batch * out_stride];
    485
    486
            // Process each batch using SIMD matmul
    487
    0
            for ba in 0..batch {
    488
    0
                let a_offset = ba * a_stride;
    489
    0
                let b_offset = ba * b_stride;
    490
    0
                let out_offset = ba * out_stride;
    491
    492
                // Create matrix views from slices (no copy - just metadata)
    493
    0
                let a_slice = &a_data[a_offset..a_offset + a_stride];
    494
    0
                let b_slice = &b_data[b_offset..b_offset + b_stride];
    495
    496
                // Use from_slice to avoid copying
    497
    0
                let a_mat = Matrix::from_slice(m, k, a_slice)?;
    498
    0
                let b_mat = Matrix::from_slice(k, n, b_slice)?;
    499
    500
                // SIMD matmul
    501
    0
                let result = a_mat.matmul(&b_mat)?;
    502
    503
                // Copy result to output
    504
    0
                output[out_offset..out_offset + out_stride].copy_from_slice(result.as_slice());
    505
            }
    506
    507
    0
            Ok(output)
    508
    0
        }
    509
    510
        /// Batched matrix multiplication for 4D tensors (attention pattern).
    511
        ///
    512
        /// Computes `[batch, heads, m, k] @ [batch, heads, k, n] -> [batch, heads, m, n]`
    513
        /// This is the exact pattern used in multi-head attention: Q @ K^T and attn @ V.
    514
        ///
    515
        /// # Arguments
    516
        /// * `a_data` - Flattened input A with shape [batch, heads, m, k]
    517
        /// * `b_data` - Flattened input B with shape [batch, heads, k, n]
    518
        /// * `batch` - Batch dimension
    519
        /// * `heads` - Number of attention heads
    520
        /// * `m` - Rows (sequence length for Q)
    521
        /// * `k` - Columns of A / Rows of B (head dimension)
    522
        /// * `n` - Columns of B (sequence length for K^T, or head dim for V)
    523
        ///
    524
        /// # Performance
    525
        /// Processes batch×heads independent matmuls using SIMD.
    526
        /// For Qwen2-0.5B: batch=1, heads=14, m=seq, k=64, n=seq
    527
        #[cfg_attr(
    528
            feature = "tracing",
    529
            instrument(skip(a_data, b_data), fields(batch, heads, m, k, n))
    530
        )]
    531
    0
        pub fn batched_matmul_4d(
    532
    0
            a_data: &[f32],
    533
    0
            b_data: &[f32],
    534
    0
            batch: usize,
    535
    0
            heads: usize,
    536
    0
            m: usize,
    537
    0
            k: usize,
    538
    0
            n: usize,
    539
    0
        ) -> Result<Vec<f32>, TruenoError> {
    540
    0
            let a_head_stride = m * k;
    541
    0
            let b_head_stride = k * n;
    542
    0
            let out_head_stride = m * n;
    543
    0
            let total_heads = batch * heads;
    544
    545
            // Validate input sizes
    546
    0
            let expected_a = total_heads * a_head_stride;
    547
    0
            let expected_b = total_heads * b_head_stride;
    548
    0
            if a_data.len() != expected_a {
    549
    0
                return Err(TruenoError::InvalidInput(format!(
    550
    0
                    "A data size mismatch: expected {} ({}×{}×{}×{}), got {}",
    551
    0
                    expected_a,
    552
    0
                    batch,
    553
    0
                    heads,
    554
    0
                    m,
    555
    0
                    k,
    556
    0
                    a_data.len()
    557
    0
                )));
    558
    0
            }
    559
    0
            if b_data.len() != expected_b {
    560
    0
                return Err(TruenoError::InvalidInput(format!(
    561
    0
                    "B data size mismatch: expected {} ({}×{}×{}×{}), got {}",
    562
    0
                    expected_b,
    563
    0
                    batch,
    564
    0
                    heads,
    565
    0
                    k,
    566
    0
                    n,
    567
    0
                    b_data.len()
    568
    0
                )));
    569
    0
            }
    570
    571
    0
            let mut output = vec![0.0f32; total_heads * out_head_stride];
    572
    573
            // Process each (batch, head) pair using SIMD matmul
    574
    0
            for bh in 0..total_heads {
    575
    0
                let a_offset = bh * a_head_stride;
    576
    0
                let b_offset = bh * b_head_stride;
    577
    0
                let out_offset = bh * out_head_stride;
    578
    579
                // Create matrix views from slices
    580
    0
                let a_slice = &a_data[a_offset..a_offset + a_head_stride];
    581
    0
                let b_slice = &b_data[b_offset..b_offset + b_head_stride];
    582
    583
    0
                let a_mat = Matrix::from_slice(m, k, a_slice)?;
    584
    0
                let b_mat = Matrix::from_slice(k, n, b_slice)?;
    585
    586
                // SIMD matmul
    587
    0
                let result = a_mat.matmul(&b_mat)?;
    588
    589
                // Copy result to output
    590
    0
                output[out_offset..out_offset + out_head_stride].copy_from_slice(result.as_slice());
    591
            }
    592
    593
    0
            Ok(output)
    594
    0
        }
    595
    596
        /// Fast path for vector-matrix multiplication (1×K @ K×N → 1×N)
    597
        ///
    598
        /// This is 8x faster than general matmul for patterns like:
    599
        /// - Vocab projection: hidden_state (1×384) @ embedding_transposed (384×51865)
    600
        /// - Single token decode in Whisper/LLM inference
    601
        ///
    602
        /// Strategy: Outer product accumulation (no transpose needed!)
    603
        /// For result[j] = sum_k(A[0,k] * B[k,j]), we compute:
    604
        ///   result += A[k] * B[k,:]  for each k
    605
        /// This has excellent cache locality since we access entire rows of B.
    606
        #[cfg_attr(feature = "tracing", instrument(skip(self, other), fields(k = self.cols, n = other.cols)))]
    607
    0
        fn matmul_vector_matrix(&self, other: &Matrix<f32>) -> Result<Matrix<f32>, TruenoError> {
    608
    0
            debug_assert_eq!(self.rows, 1);
    609
    610
    0
            let k = self.cols; // Inner dimension
    611
    0
            let n = other.cols; // Output dimension
    612
    613
            // Result is 1×N, initialized to zero
    614
    0
            let mut result = Matrix::zeros_with_backend(1, n, self.backend);
    615
    616
            // Outer product accumulation: result += A[k] * B[k,:]
    617
            // For each k, scale row k of B by A[k] and add to result
    618
            // The compiler will auto-vectorize this inner loop
    619
    0
            for ki in 0..k {
    620
    0
                let a_k = self.data[ki];
    621
    0
                if a_k == 0.0 {
    622
    0
                    continue; // Skip zero multiplications
    623
    0
                }
    624
    625
                // Get row ki of B (contiguous in memory - cache friendly!)
    626
    0
                let b_row_start = ki * n;
    627
    628
                // AXPY: result += a_k * B[ki,:]
    629
                // This loop is auto-vectorized by LLVM with -O2/-O3
    630
    0
                for j in 0..n {
    631
    0
                    result.data[j] += a_k * other.data[b_row_start + j];
    632
    0
                }
    633
            }
    634
    635
    0
            Ok(result)
    636
    0
        }
    637
    638
        /// Naive O(n³) matrix multiplication (baseline for small matrices)
    639
    0
        fn matmul_naive(
    640
    0
            &self,
    641
    0
            other: &Matrix<f32>,
    642
    0
            result: &mut Matrix<f32>,
    643
    0
        ) -> Result<(), TruenoError> {
    644
            // C[i,j] = Σ A[i,k] × B[k,j]
    645
            // SAFETY: Loop bounds are validated by dimension checks in matmul()
    646
    0
            for i in 0..self.rows {
    647
    0
                for j in 0..other.cols {
    648
    0
                    let mut sum = 0.0;
    649
    0
                    for k in 0..self.cols {
    650
    0
                        // Bounds guaranteed: i < self.rows, k < self.cols, j < other.cols
    651
    0
                        sum += self
    652
    0
                            .get(i, k)
    653
    0
                            .expect("matmul_naive: A[i,k] bounds validated by loop")
    654
    0
                            * other
    655
    0
                                .get(k, j)
    656
    0
                                .expect("matmul_naive: B[k,j] bounds validated by loop");
    657
    0
                    }
    658
    0
                    *result
    659
    0
                        .get_mut(i, j)
    660
    0
                        .expect("matmul_naive: C[i,j] bounds validated by loop") = sum;
    661
                }
    662
            }
    663
    0
            Ok(())
    664
    0
        }
    665
    666
        /// AVX2 micro-kernel: Compute 4 rows × 1 column using register blocking (Phase 2)
    667
        ///
    668
        /// This micro-kernel processes 4 rows of matrix A against 1 column of B_transposed
    669
        /// simultaneously, keeping intermediate results in AVX2 registers for efficiency.
    670
        ///
    671
        /// # Performance Benefits
    672
        /// - Loads B-column once, reuses for 4 A-rows (4× reduction in memory bandwidth)
    673
        /// - Uses FMA instructions for fused multiply-add (3× throughput vs separate ops)
    674
        /// - Keeps accumulators in YMM registers (no memory traffic for intermediate results)
    675
        ///
    676
        /// # Safety
    677
        /// - Caller must ensure all slices have the same length
    678
        /// - Must be called on x86_64 with AVX2 support
    679
        #[cfg(target_arch = "x86_64")]
    680
        #[target_feature(enable = "avx2,fma")]
    681
        #[inline]
    682
    0
        unsafe fn matmul_microkernel_4x1_avx2(
    683
    0
            a_rows: [&[f32]; 4],
    684
    0
            b_col: &[f32],
    685
    0
            results: &mut [f32; 4],
    686
    0
        ) {
    687
            use std::arch::x86_64::*;
    688
    689
    0
            let len = b_col.len();
    690
    0
            let chunks = len / 8; // Process 8 f32 elements per iteration (AVX2 = 256 bits)
    691
    692
            // Accumulators for 4 output elements (kept in registers)
    693
    0
            let mut acc0 = _mm256_setzero_ps();
    694
    0
            let mut acc1 = _mm256_setzero_ps();
    695
    0
            let mut acc2 = _mm256_setzero_ps();
    696
    0
            let mut acc3 = _mm256_setzero_ps();
    697
    698
            // Main loop: Process 8 elements at a time
    699
    0
            for i in 0..chunks {
    700
    0
                let offset = i * 8;
    701
    0
    702
    0
                // Load B column (reused for all 4 A rows)
    703
    0
                let b_vec = _mm256_loadu_ps(b_col.as_ptr().add(offset));
    704
    0
    705
    0
                // Load A rows and FMA (Fused Multiply-Add)
    706
    0
                let a0_vec = _mm256_loadu_ps(a_rows[0].as_ptr().add(offset));
    707
    0
                acc0 = _mm256_fmadd_ps(a0_vec, b_vec, acc0);
    708
    0
    709
    0
                let a1_vec = _mm256_loadu_ps(a_rows[1].as_ptr().add(offset));
    710
    0
                acc1 = _mm256_fmadd_ps(a1_vec, b_vec, acc1);
    711
    0
    712
    0
                let a2_vec = _mm256_loadu_ps(a_rows[2].as_ptr().add(offset));
    713
    0
                acc2 = _mm256_fmadd_ps(a2_vec, b_vec, acc2);
    714
    0
    715
    0
                let a3_vec = _mm256_loadu_ps(a_rows[3].as_ptr().add(offset));
    716
    0
                acc3 = _mm256_fmadd_ps(a3_vec, b_vec, acc3);
    717
    0
            }
    718
    719
            // Horizontal sum of each accumulator (reduce 8 elements to 1)
    720
    0
            results[0] = Self::horizontal_sum_avx2(acc0);
    721
    0
            results[1] = Self::horizontal_sum_avx2(acc1);
    722
    0
            results[2] = Self::horizontal_sum_avx2(acc2);
    723
    0
            results[3] = Self::horizontal_sum_avx2(acc3);
    724
    725
            // Handle remainder elements with scalar code
    726
    0
            let remainder_start = chunks * 8;
    727
    0
            if remainder_start < len {
    728
    0
                for i in remainder_start..len {
    729
    0
                    results[0] += a_rows[0][i] * b_col[i];
    730
    0
                    results[1] += a_rows[1][i] * b_col[i];
    731
    0
                    results[2] += a_rows[2][i] * b_col[i];
    732
    0
                    results[3] += a_rows[3][i] * b_col[i];
    733
    0
                }
    734
    0
            }
    735
    0
        }
    736
    737
        /// Helper: Horizontal sum of 8 f32 values in an AVX2 register
    738
        #[cfg(target_arch = "x86_64")]
    739
        #[target_feature(enable = "avx2")]
    740
        #[inline]
    741
    0
        unsafe fn horizontal_sum_avx2(v: std::arch::x86_64::__m256) -> f32 {
    742
            use std::arch::x86_64::*;
    743
    744
            // Sum upper and lower 128-bit lanes
    745
    0
            let sum128 = _mm_add_ps(_mm256_castps256_ps128(v), _mm256_extractf128_ps(v, 1));
    746
    747
            // Horizontal add within 128-bit lane (4 values → 2 values)
    748
    0
            let sum64 = _mm_hadd_ps(sum128, sum128);
    749
    750
            // Horizontal add again (2 values → 1 value)
    751
    0
            let sum32 = _mm_hadd_ps(sum64, sum64);
    752
    753
            // Extract final scalar result
    754
    0
            _mm_cvtss_f32(sum32)
    755
    0
        }
    756
    757
        /// AVX-512 micro-kernel: Compute 8 rows × 1 column using register blocking (Phase 3)
    758
        ///
    759
        /// This micro-kernel processes 8 rows of matrix A against 1 column of B_transposed
    760
        /// simultaneously, keeping intermediate results in AVX-512 registers for efficiency.
    761
        ///
    762
        /// # Performance Benefits
    763
        /// - Processes 16 f32 elements per iteration (vs 8 with AVX2) - 2× throughput
    764
        /// - Loads B-column once, reuses for 8 A-rows (8× reduction in memory bandwidth)
    765
        /// - Uses FMA instructions for fused multiply-add (3× throughput vs separate ops)
    766
        /// - Keeps accumulators in ZMM registers (no memory traffic for intermediate results)
    767
        ///
    768
        /// # Safety
    769
        /// - Caller must ensure all slices have the same length
    770
        /// - Must be called on x86_64 with AVX-512F support
    771
        #[cfg(target_arch = "x86_64")]
    772
        #[target_feature(enable = "avx512f")]
    773
        #[inline]
    774
    0
        unsafe fn matmul_microkernel_8x1_avx512(
    775
    0
            a_rows: [&[f32]; 8],
    776
    0
            b_col: &[f32],
    777
    0
            results: &mut [f32; 8],
    778
    0
        ) {
    779
            use std::arch::x86_64::*;
    780
    781
    0
            let len = b_col.len();
    782
    0
            let chunks = len / 16; // Process 16 f32 elements per iteration (AVX-512 = 512 bits)
    783
    784
            // Accumulators for 8 output elements (kept in ZMM registers)
    785
    0
            let mut acc0 = _mm512_setzero_ps();
    786
    0
            let mut acc1 = _mm512_setzero_ps();
    787
    0
            let mut acc2 = _mm512_setzero_ps();
    788
    0
            let mut acc3 = _mm512_setzero_ps();
    789
    0
            let mut acc4 = _mm512_setzero_ps();
    790
    0
            let mut acc5 = _mm512_setzero_ps();
    791
    0
            let mut acc6 = _mm512_setzero_ps();
    792
    0
            let mut acc7 = _mm512_setzero_ps();
    793
    794
            // Main loop: Process 16 elements at a time
    795
    0
            for i in 0..chunks {
    796
    0
                let offset = i * 16;
    797
    0
    798
    0
                // Load B column (reused for all 8 A rows)
    799
    0
                let b_vec = _mm512_loadu_ps(b_col.as_ptr().add(offset));
    800
    0
    801
    0
                // Load A rows and FMA (Fused Multiply-Add)
    802
    0
                let a0_vec = _mm512_loadu_ps(a_rows[0].as_ptr().add(offset));
    803
    0
                acc0 = _mm512_fmadd_ps(a0_vec, b_vec, acc0);
    804
    0
    805
    0
                let a1_vec = _mm512_loadu_ps(a_rows[1].as_ptr().add(offset));
    806
    0
                acc1 = _mm512_fmadd_ps(a1_vec, b_vec, acc1);
    807
    0
    808
    0
                let a2_vec = _mm512_loadu_ps(a_rows[2].as_ptr().add(offset));
    809
    0
                acc2 = _mm512_fmadd_ps(a2_vec, b_vec, acc2);
    810
    0
    811
    0
                let a3_vec = _mm512_loadu_ps(a_rows[3].as_ptr().add(offset));
    812
    0
                acc3 = _mm512_fmadd_ps(a3_vec, b_vec, acc3);
    813
    0
    814
    0
                let a4_vec = _mm512_loadu_ps(a_rows[4].as_ptr().add(offset));
    815
    0
                acc4 = _mm512_fmadd_ps(a4_vec, b_vec, acc4);
    816
    0
    817
    0
                let a5_vec = _mm512_loadu_ps(a_rows[5].as_ptr().add(offset));
    818
    0
                acc5 = _mm512_fmadd_ps(a5_vec, b_vec, acc5);
    819
    0
    820
    0
                let a6_vec = _mm512_loadu_ps(a_rows[6].as_ptr().add(offset));
    821
    0
                acc6 = _mm512_fmadd_ps(a6_vec, b_vec, acc6);
    822
    0
    823
    0
                let a7_vec = _mm512_loadu_ps(a_rows[7].as_ptr().add(offset));
    824
    0
                acc7 = _mm512_fmadd_ps(a7_vec, b_vec, acc7);
    825
    0
            }
    826
    827
            // Horizontal sum of each accumulator (reduce 16 elements to 1)
    828
    0
            results[0] = _mm512_reduce_add_ps(acc0);
    829
    0
            results[1] = _mm512_reduce_add_ps(acc1);
    830
    0
            results[2] = _mm512_reduce_add_ps(acc2);
    831
    0
            results[3] = _mm512_reduce_add_ps(acc3);
    832
    0
            results[4] = _mm512_reduce_add_ps(acc4);
    833
    0
            results[5] = _mm512_reduce_add_ps(acc5);
    834
    0
            results[6] = _mm512_reduce_add_ps(acc6);
    835
    0
            results[7] = _mm512_reduce_add_ps(acc7);
    836
    837
            // Handle remainder elements with scalar code
    838
    0
            let remainder_start = chunks * 16;
    839
    0
            if remainder_start < len {
    840
    0
                for i in remainder_start..len {
    841
    0
                    results[0] += a_rows[0][i] * b_col[i];
    842
    0
                    results[1] += a_rows[1][i] * b_col[i];
    843
    0
                    results[2] += a_rows[2][i] * b_col[i];
    844
    0
                    results[3] += a_rows[3][i] * b_col[i];
    845
    0
                    results[4] += a_rows[4][i] * b_col[i];
    846
    0
                    results[5] += a_rows[5][i] * b_col[i];
    847
    0
                    results[6] += a_rows[6][i] * b_col[i];
    848
    0
                    results[7] += a_rows[7][i] * b_col[i];
    849
    0
                }
    850
    0
            }
    851
    0
        }
    852
    853
        /// Cache-aware blocked matrix multiplication with SIMD optimization
    854
        ///
    855
        /// Uses 2-level cache blocking (L2/L1) to minimize cache misses:
    856
        /// - L2 blocks: 64×64 (256KB for 3 matrices in f32)
    857
        /// - L1 micro-kernels: 8×8 (768 bytes fits comfortably in L1)
    858
        ///
    859
        /// Performance characteristics:
    860
        /// - Small matrices (<64×64): ~1.2× speedup over naive (overhead dominates)
    861
        /// - Medium matrices (128×128): ~1.5-2× speedup (cache effects visible)
    862
        /// - Large matrices (512×512+): ~2-3× speedup (dramatic cache improvement)
    863
        ///
    864
        /// This is Phase 1 of matmul optimization (Issue #10). Future Phase 2 will
    865
        /// add optional BLAS backend for full NumPy parity on very large matrices.
    866
        /// Helper function to process a single L3 row block for parallel matmul (Phase 4).
    867
        ///
    868
        /// # Safety
    869
        /// When called from parallel code, the caller must ensure that each thread processes
    870
        /// a distinct row range [iii, i3_end) with no overlap. This function is safe because
    871
        /// each thread writes only to its own row range in the result matrix.
    872
        #[cfg(feature = "parallel")]
    873
        #[allow(clippy::too_many_arguments)]
    874
        fn process_l3_row_block_seq(
    875
            iii: usize,
    876
            i3_end: usize,
    877
            a: &Matrix<f32>,
    878
            b_transposed: &Matrix<f32>,
    879
            result: &mut Matrix<f32>,
    880
            l2_block_size: usize,
    881
            l3_block_size: usize,
    882
        ) {
    883
            #[cfg(target_arch = "x86_64")]
    884
            use crate::backends::{avx2::Avx2Backend, sse2::Sse2Backend};
    885
            use crate::backends::{scalar::ScalarBackend, VectorBackend};
    886
    887
            // Process all column blocks for this row block
    888
            for jjj in (0..b_transposed.rows).step_by(l3_block_size) {
    889
                let j3_end = (jjj + l3_block_size).min(b_transposed.rows);
    890
    891
                for kkk in (0..a.cols).step_by(l3_block_size) {
    892
                    let k3_end = (kkk + l3_block_size).min(a.cols);
    893
    894
                    // L2 blocking within L3 blocks
    895
                    for ii in (iii..i3_end).step_by(l2_block_size) {
    896
                        let i_end = (ii + l2_block_size).min(i3_end);
    897
    898
                        for jj in (jjj..j3_end).step_by(l2_block_size) {
    899
                            let j_end = (jj + l2_block_size).min(j3_end);
    900
    901
                            for kk in (kkk..k3_end).step_by(l2_block_size) {
    902
                                let k_end = (kk + l2_block_size).min(k3_end);
    903
                                let block_size = k_end - kk;
    904
    905
                                // Micro-kernel processing
    906
                                #[cfg(target_arch = "x86_64")]
    907
                                let use_microkernel =
    908
                                    matches!(a.backend, Backend::AVX2 | Backend::AVX512);
    909
    910
                                #[cfg(target_arch = "x86_64")]
    911
                                if use_microkernel {
    912
                                    let mut i = ii;
    913
    914
                                    // Process 4 rows at a time with micro-kernel
    915
                                    while i + 4 <= i_end {
    916
                                        let row0_start = i * a.cols + kk;
    917
                                        let row1_start = (i + 1) * a.cols + kk;
    918
                                        let row2_start = (i + 2) * a.cols + kk;
    919
                                        let row3_start = (i + 3) * a.cols + kk;
    920
    921
                                        let a_rows = [
    922
                                            &a.data[row0_start..row0_start + block_size],
    923
                                            &a.data[row1_start..row1_start + block_size],
    924
                                            &a.data[row2_start..row2_start + block_size],
    925
                                            &a.data[row3_start..row3_start + block_size],
    926
                                        ];
    927
    928
                                        for j in jj..j_end {
    929
                                            let col_start = j * b_transposed.cols + kk;
    930
                                            let b_col =
    931
                                                &b_transposed.data[col_start..col_start + block_size];
    932
    933
                                            let mut partial_dots = [0.0f32; 4];
    934
                                            // SAFETY: AVX2 support verified by is_x86_feature_detected!("avx2")
    935
                                            // check in outer scope. Slices a_rows and b_col are bounds-checked
    936
                                            // and properly aligned for SIMD operations.
    937
                                            // SAFETY: CPU feature verified at runtime, slices bounds-checked
    938
                                            unsafe {
    939
                                                Matrix::matmul_microkernel_4x1_avx2(
    940
                                                    a_rows,
    941
                                                    b_col,
    942
                                                    &mut partial_dots,
    943
                                                );
    944
                                            }
    945
    946
                                            result.data[i * result.cols + j] += partial_dots[0];
    947
                                            result.data[(i + 1) * result.cols + j] += partial_dots[1];
    948
                                            result.data[(i + 2) * result.cols + j] += partial_dots[2];
    949
                                            result.data[(i + 3) * result.cols + j] += partial_dots[3];
    950
                                        }
    951
    952
                                        i += 4;
    953
                                    }
    954
    955
                                    // Handle remaining rows (< 4)
    956
                                    for i in i..i_end {
    957
                                        let row_start = i * a.cols + kk;
    958
                                        let a_row = &a.data[row_start..row_start + block_size];
    959
    960
                                        for j in jj..j_end {
    961
                                            let col_start = j * b_transposed.cols + kk;
    962
                                            let b_col =
    963
                                                &b_transposed.data[col_start..col_start + block_size];
    964
    965
                                            // SAFETY: AVX2 verified at runtime, slices bounds-checked
    966
                                            // SAFETY: AVX2 verified at runtime, slices bounds-checked
    967
                                            let partial_dot = unsafe { Avx2Backend::dot(a_row, b_col) };
    968
                                            result.data[i * result.cols + j] += partial_dot;
    969
                                        }
    970
                                    }
    971
                                } else {
    972
                                    // Non-AVX2 path
    973
                                    #[allow(unused_variables)]
    974
                                    for i in ii..i_end {
    975
                                        let row_start = i * a.cols + kk;
    976
                                        let a_row = &a.data[row_start..row_start + block_size];
    977
    978
                                        for j in jj..j_end {
    979
                                            let col_start = j * b_transposed.cols + kk;
    980
                                            let b_col =
    981
                                                &b_transposed.data[col_start..col_start + block_size];
    982
    983
                                            // SAFETY: AVX2 verified at runtime, slices bounds-checked
    984
                                            let partial_dot = unsafe {
    985
                                                match a.backend {
    986
                                                    Backend::Scalar => ScalarBackend::dot(a_row, b_col),
    987
                                                    #[cfg(target_arch = "x86_64")]
    988
                                                    Backend::SSE2 | Backend::AVX => {
    989
                                                        Sse2Backend::dot(a_row, b_col)
    990
                                                    }
    991
                                                    #[cfg(not(target_arch = "x86_64"))]
    992
                                                    Backend::SSE2
    993
                                                    | Backend::AVX
    994
                                                    | Backend::AVX2
    995
                                                    | Backend::AVX512 => {
    996
                                                        ScalarBackend::dot(a_row, b_col)
    997
                                                    }
    998
                                                    #[cfg(any(
    999
                                                        target_arch = "aarch64",
    1000
                                                        target_arch = "arm"
    1001
                                                    ))]
    1002
                                                    Backend::NEON => {
    1003
                                                        use crate::backends::neon::NeonBackend;
    1004
                                                        NeonBackend::dot(a_row, b_col)
    1005
                                                    }
    1006
                                                    #[cfg(not(any(
    1007
                                                        target_arch = "aarch64",
    1008
                                                        target_arch = "arm"
    1009
                                                    )))]
    1010
                                                    Backend::NEON => ScalarBackend::dot(a_row, b_col),
    1011
                                                    #[cfg(target_arch = "wasm32")]
    1012
                                                    Backend::WasmSIMD => {
    1013
                                                        use crate::backends::wasm::WasmBackend;
    1014
                                                        WasmBackend::dot(a_row, b_col)
    1015
                                                    }
    1016
                                                    #[cfg(not(target_arch = "wasm32"))]
    1017
                                                    Backend::WasmSIMD => {
    1018
                                                        ScalarBackend::dot(a_row, b_col)
    1019
                                                    }
    1020
                                                    // Catch-all for GPU, Auto, and any other backends
    1021
                                                    _ => ScalarBackend::dot(a_row, b_col),
    1022
                                                }
    1023
                                            };
    1024
    1025
                                            result.data[i * result.cols + j] += partial_dot;
    1026
                                        }
    1027
                                    }
    1028
                                }
    1029
    1030
                                // Non-x86_64 fallback
    1031
                                #[cfg(not(target_arch = "x86_64"))]
    1032
                                {
    1033
                                    for i in ii..i_end {
    1034
                                        let row_start = i * a.cols + kk;
    1035
                                        let a_row = &a.data[row_start..row_start + block_size];
    1036
    1037
                                        for j in jj..j_end {
    1038
                                            let col_start = j * b_transposed.cols + kk;
    1039
                                            let b_col =
    1040
                                                &b_transposed.data[col_start..col_start + block_size];
    1041
    1042
                                            // SAFETY: AVX2 verified at runtime, slices bounds-checked
    1043
                                            let partial_dot = unsafe {
    1044
                                                match a.backend {
    1045
                                                    Backend::Scalar => ScalarBackend::dot(a_row, b_col),
    1046
                                                    #[cfg(any(
    1047
                                                        target_arch = "aarch64",
    1048
                                                        target_arch = "arm"
    1049
                                                    ))]
    1050
                                                    Backend::NEON => {
    1051
                                                        use crate::backends::neon::NeonBackend;
    1052
                                                        NeonBackend::dot(a_row, b_col)
    1053
                                                    }
    1054
                                                    #[cfg(not(any(
    1055
                                                        target_arch = "aarch64",
    1056
                                                        target_arch = "arm"
    1057
                                                    )))]
    1058
                                                    Backend::NEON => ScalarBackend::dot(a_row, b_col),
    1059
                                                    #[cfg(target_arch = "wasm32")]
    1060
                                                    Backend::WasmSIMD => {
    1061
                                                        use crate::backends::wasm::WasmBackend;
    1062
                                                        WasmBackend::dot(a_row, b_col)
    1063
                                                    }
    1064
                                                    #[cfg(not(target_arch = "wasm32"))]
    1065
                                                    Backend::WasmSIMD => {
    1066
                                                        ScalarBackend::dot(a_row, b_col)
    1067
                                                    }
    1068
                                                    _ => ScalarBackend::dot(a_row, b_col),
    1069
                                                }
    1070
                                            };
    1071
    1072
                                            result.data[i * result.cols + j] += partial_dot;
    1073
                                        }
    1074
                                    }
    1075
                                }
    1076
                            }
    1077
                        }
    1078
                    }
    1079
                }
    1080
            }
    1081
        }
    1082
    1083
    0
        fn matmul_simd(
    1084
    0
            &self,
    1085
    0
            other: &Matrix<f32>,
    1086
    0
            result: &mut Matrix<f32>,
    1087
    0
        ) -> Result<(), TruenoError> {
    1088
            // Cache blocking parameters (tuned for typical x86_64 CPUs)
    1089
            // L2 cache: 256KB typical → 64K f32 elements → 64×64×3 matrices fits
    1090
            const L2_BLOCK_SIZE: usize = 64;
    1091
            // L3 cache: 4-16MB typical → 256×256 blocks for very large matrices (Phase 3)
    1092
            const L3_BLOCK_SIZE: usize = 256;
    1093
            const L3_THRESHOLD: usize = 512; // Use 3-level blocking for matrices ≥512×512
    1094
    1095
            // For small matrices, use simple SIMD approach (blocking overhead too high)
    1096
    0
            if self.rows <= 32 || self.cols <= 32 || other.cols <= 32 {
    1097
    0
                return self.matmul_simd_simple(other, result);
    1098
    0
            }
    1099
    1100
            #[cfg(target_arch = "x86_64")]
    1101
            use crate::backends::{avx2::Avx2Backend, sse2::Sse2Backend};
    1102
            use crate::backends::{scalar::ScalarBackend, VectorBackend};
    1103
    1104
            // Pre-transpose B for better cache locality (columns become rows)
    1105
    0
            let b_transposed = other.transpose();
    1106
    1107
            // Determine if we should use 3-level blocking (Phase 3)
    1108
    0
            let use_l3_blocking =
    1109
    0
                self.rows >= L3_THRESHOLD && self.cols >= L3_THRESHOLD && other.cols >= L3_THRESHOLD;
    1110
    1111
            // Phase 4: Determine if we should use multi-threading (≥1024×1024)
    1112
            #[cfg(feature = "parallel")]
    1113
            const PARALLEL_THRESHOLD: usize = 1024;
    1114
            #[cfg(feature = "parallel")]
    1115
            let use_parallel = self.rows >= PARALLEL_THRESHOLD
    1116
                && self.cols >= PARALLEL_THRESHOLD
    1117
                && other.cols >= PARALLEL_THRESHOLD;
    1118
            #[cfg(not(feature = "parallel"))]
    1119
    0
            let use_parallel = false;
    1120
    1121
    0
            if use_l3_blocking {
    1122
                // ===== Phase 3/4: 3-Level Cache Blocking (L3 → L2 → micro-kernel) =====
    1123
                // For very large matrices (≥512×512), use L3 cache blocking to minimize
    1124
                // cache misses when data doesn't fit in L2 cache
    1125
                //
    1126
                // Hierarchy:
    1127
                // 1. L3 blocks: 256×256 (fits in L3 cache: 4-16MB)
    1128
                // 2. L2 blocks: 64×64 (fits in L2 cache: 256KB)
    1129
                // 3. Micro-kernel: 4×1 for AVX2/AVX512
    1130
                //
    1131
                // Phase 4: For ≥1024×1024, parallelize L3 row blocks with rayon
    1132
    1133
    0
                if use_parallel {
    1134
                    // ===== Phase 4: Parallel 3-Level Cache Blocking (Lock-Free Row Partitioning) =====
    1135
                    #[cfg(feature = "parallel")]
    1136
                    {
    1137
                        use rayon::prelude::*;
    1138
                        use std::sync::atomic::{AtomicPtr, Ordering};
    1139
                        use std::sync::Arc;
    1140
    1141
                        // Lock-free parallelization strategy:
    1142
                        // Each thread processes one L3 row block (256 rows). Since row blocks are
    1143
                        // non-overlapping, threads write to distinct memory regions with no contention.
    1144
                        //
    1145
                        // Safety invariant: Each thread writes to result.data[iii*cols..(i3_end)*cols],
    1146
                        // where iii = block_idx * L3_BLOCK_SIZE. Since L3 blocks don't overlap,
    1147
                        // no two threads write to the same memory location.
    1148
    1149
                        // Store result pointer in Arc<AtomicPtr> for safe sharing
    1150
                        let result_ptr = Arc::new(AtomicPtr::new(result as *mut Matrix<f32>));
    1151
    1152
                        // Calculate number of L3 blocks
    1153
                        let num_blocks = self.rows.div_ceil(L3_BLOCK_SIZE);
    1154
    1155
                        // Process each L3 block in parallel (lock-free)
    1156
                        (0..num_blocks).into_par_iter().for_each(|block_idx| {
    1157
                            let iii = block_idx * L3_BLOCK_SIZE;
    1158
                            let i3_end = (iii + L3_BLOCK_SIZE).min(self.rows);
    1159
    1160
                            // SAFETY: Each thread processes a distinct row range [iii, i3_end).
    1161
                            // No two threads write to overlapping memory locations because:
    1162
                            // 1. L3 blocks partition rows: [0, 256), [256, 512), etc.
    1163
                            // 2. Each thread only modifies result.data[iii*cols..(i3_end)*cols]
    1164
                            // 3. Row ranges are non-overlapping by construction
    1165
                            // 4. All threads complete before function returns (rayon guarantee)
    1166
                            // 5. AtomicPtr ensures proper memory ordering across threads
    1167
                            // SAFETY: CPU feature verified at runtime, slices bounds-checked
    1168
                            unsafe {
    1169
                                let ptr = result_ptr.load(Ordering::Relaxed);
    1170
                                Self::process_l3_row_block_seq(
    1171
                                    iii,
    1172
                                    i3_end,
    1173
                                    self,
    1174
                                    &b_transposed,
    1175
                                    &mut *ptr,
    1176
                                    L2_BLOCK_SIZE,
    1177
                                    L3_BLOCK_SIZE,
    1178
                                );
    1179
                            }
    1180
                        });
    1181
                    }
    1182
    1183
    0
                    return Ok(());
    1184
    0
                }
    1185
    1186
                // ===== Sequential 3-Level Cache Blocking (fallback) =====
    1187
    0
                for iii in (0..self.rows).step_by(L3_BLOCK_SIZE) {
    1188
    0
                    let i3_end = (iii + L3_BLOCK_SIZE).min(self.rows);
    1189
    1190
    0
                    for jjj in (0..other.cols).step_by(L3_BLOCK_SIZE) {
    1191
    0
                        let j3_end = (jjj + L3_BLOCK_SIZE).min(other.cols);
    1192
    1193
    0
                        for kkk in (0..self.cols).step_by(L3_BLOCK_SIZE) {
    1194
    0
                            let k3_end = (kkk + L3_BLOCK_SIZE).min(self.cols);
    1195
    1196
                            // L2 blocking within L3 blocks
    1197
    0
                            for ii in (iii..i3_end).step_by(L2_BLOCK_SIZE) {
    1198
    0
                                let i_end = (ii + L2_BLOCK_SIZE).min(i3_end);
    1199
    1200
    0
                                for jj in (jjj..j3_end).step_by(L2_BLOCK_SIZE) {
    1201
    0
                                    let j_end = (jj + L2_BLOCK_SIZE).min(j3_end);
    1202
    1203
    0
                                    for kk in (kkk..k3_end).step_by(L2_BLOCK_SIZE) {
    1204
    0
                                        let k_end = (kk + L2_BLOCK_SIZE).min(k3_end);
    1205
    0
                                        let block_size = k_end - kk;
    1206
    1207
                                        // Micro-kernel processing
    1208
                                        #[cfg(target_arch = "x86_64")]
    1209
    0
                                        let use_avx512 = matches!(self.backend, Backend::AVX512);
    1210
                                        #[cfg(target_arch = "x86_64")]
    1211
    0
                                        let use_avx2 = matches!(self.backend, Backend::AVX2);
    1212
    1213
                                        #[cfg(target_arch = "x86_64")]
    1214
    0
                                        if use_avx512 {
    1215
                                            // AVX-512 8x1 micro-kernel (Phase 3)
    1216
    0
                                            let mut i = ii;
    1217
    1218
                                            // Process 8 rows at a time with AVX-512 micro-kernel
    1219
    0
                                            while i + 8 <= i_end {
    1220
    0
                                                let a_rows = [
    1221
    0
                                                    &self.data[i * self.cols + kk..(i * self.cols + kk) + block_size],
    1222
    0
                                                    &self.data[(i + 1) * self.cols + kk..((i + 1) * self.cols + kk) + block_size],
    1223
    0
                                                    &self.data[(i + 2) * self.cols + kk..((i + 2) * self.cols + kk) + block_size],
    1224
    0
                                                    &self.data[(i + 3) * self.cols + kk..((i + 3) * self.cols + kk) + block_size],
    1225
    0
                                                    &self.data[(i + 4) * self.cols + kk..((i + 4) * self.cols + kk) + block_size],
    1226
    0
                                                    &self.data[(i + 5) * self.cols + kk..((i + 5) * self.cols + kk) + block_size],
    1227
    0
                                                    &self.data[(i + 6) * self.cols + kk..((i + 6) * self.cols + kk) + block_size],
    1228
    0
                                                    &self.data[(i + 7) * self.cols + kk..((i + 7) * self.cols + kk) + block_size],
    1229
    0
                                                ];
    1230
    1231
    0
                                                for j in jj..j_end {
    1232
    0
                                                    let col_start = j * b_transposed.cols + kk;
    1233
    0
                                                    let b_col = &b_transposed.data
    1234
    0
                                                        [col_start..col_start + block_size];
    1235
    0
    1236
    0
                                                    let mut partial_dots = [0.0f32; 8];
    1237
    0
                                                    // SAFETY: CPU feature verified at runtime, slices bounds-checked
    1238
    0
                                                    unsafe {
    1239
    0
                                                        Self::matmul_microkernel_8x1_avx512(
    1240
    0
                                                            a_rows,
    1241
    0
                                                            b_col,
    1242
    0
                                                            &mut partial_dots,
    1243
    0
                                                        );
    1244
    0
                                                    }
    1245
    0
    1246
    0
                                                    result.data[i * result.cols + j] += partial_dots[0];
    1247
    0
                                                    result.data[(i + 1) * result.cols + j] += partial_dots[1];
    1248
    0
                                                    result.data[(i + 2) * result.cols + j] += partial_dots[2];
    1249
    0
                                                    result.data[(i + 3) * result.cols + j] += partial_dots[3];
    1250
    0
                                                    result.data[(i + 4) * result.cols + j] += partial_dots[4];
    1251
    0
                                                    result.data[(i + 5) * result.cols + j] += partial_dots[5];
    1252
    0
                                                    result.data[(i + 6) * result.cols + j] += partial_dots[6];
    1253
    0
                                                    result.data[(i + 7) * result.cols + j] += partial_dots[7];
    1254
    0
                                                }
    1255
    1256
    0
                                                i += 8;
    1257
                                            }
    1258
    1259
                                            // Handle remaining rows with AVX2 4x1 kernel
    1260
    0
                                            while i + 4 <= i_end {
    1261
    0
                                                let a_rows = [
    1262
    0
                                                    &self.data[i * self.cols + kk..(i * self.cols + kk) + block_size],
    1263
    0
                                                    &self.data[(i + 1) * self.cols + kk..((i + 1) * self.cols + kk) + block_size],
    1264
    0
                                                    &self.data[(i + 2) * self.cols + kk..((i + 2) * self.cols + kk) + block_size],
    1265
    0
                                                    &self.data[(i + 3) * self.cols + kk..((i + 3) * self.cols + kk) + block_size],
    1266
    0
                                                ];
    1267
    1268
    0
                                                for j in jj..j_end {
    1269
    0
                                                    let col_start = j * b_transposed.cols + kk;
    1270
    0
                                                    let b_col = &b_transposed.data[col_start..col_start + block_size];
    1271
    0
    1272
    0
                                                    let mut partial_dots = [0.0f32; 4];
    1273
    0
                                                    // SAFETY: CPU feature verified at runtime, slices bounds-checked
    1274
    0
                                                    unsafe {
    1275
    0
                                                        Self::matmul_microkernel_4x1_avx2(a_rows, b_col, &mut partial_dots);
    1276
    0
                                                    }
    1277
    0
    1278
    0
                                                    result.data[i * result.cols + j] += partial_dots[0];
    1279
    0
                                                    result.data[(i + 1) * result.cols + j] += partial_dots[1];
    1280
    0
                                                    result.data[(i + 2) * result.cols + j] += partial_dots[2];
    1281
    0
                                                    result.data[(i + 3) * result.cols + j] += partial_dots[3];
    1282
    0
                                                }
    1283
    0
                                                i += 4;
    1284
                                            }
    1285
    1286
                                            // Handle remaining rows (< 4)
    1287
    0
                                            for i in i..i_end {
    1288
    0
                                                let row_start = i * self.cols + kk;
    1289
    0
                                                let a_row = &self.data[row_start..row_start + block_size];
    1290
    1291
    0
                                                for j in jj..j_end {
    1292
    0
                                                    let col_start = j * b_transposed.cols + kk;
    1293
    0
                                                    let b_col = &b_transposed.data[col_start..col_start + block_size];
    1294
    0
    1295
    0
                                                    // SAFETY: AVX2 verified at runtime, slices bounds-checked
    1296
    0
                                                    let partial_dot = unsafe { Avx2Backend::dot(a_row, b_col) };
    1297
    0
                                                    result.data[i * result.cols + j] += partial_dot;
    1298
    0
                                                }
    1299
                                            }
    1300
    0
                                        } else if use_avx2 {
    1301
                                            // AVX2 4x1 micro-kernel
    1302
    0
                                            let mut i = ii;
    1303
    1304
                                            // Process 4 rows at a time with micro-kernel
    1305
    0
                                            while i + 4 <= i_end {
    1306
    0
                                                let row0_start = i * self.cols + kk;
    1307
    0
                                                let row1_start = (i + 1) * self.cols + kk;
    1308
    0
                                                let row2_start = (i + 2) * self.cols + kk;
    1309
    0
                                                let row3_start = (i + 3) * self.cols + kk;
    1310
    1311
    0
                                                let a_rows = [
    1312
    0
                                                    &self.data[row0_start..row0_start + block_size],
    1313
    0
                                                    &self.data[row1_start..row1_start + block_size],
    1314
    0
                                                    &self.data[row2_start..row2_start + block_size],
    1315
    0
                                                    &self.data[row3_start..row3_start + block_size],
    1316
    0
                                                ];
    1317
    1318
    0
                                                for j in jj..j_end {
    1319
    0
                                                    let col_start = j * b_transposed.cols + kk;
    1320
    0
                                                    let b_col = &b_transposed.data
    1321
    0
                                                        [col_start..col_start + block_size];
    1322
    0
    1323
    0
                                                    let mut partial_dots = [0.0f32; 4];
    1324
    0
                                                    // SAFETY: CPU feature verified at runtime, slices bounds-checked
    1325
    0
                                                    unsafe {
    1326
    0
                                                        Self::matmul_microkernel_4x1_avx2(
    1327
    0
                                                            a_rows,
    1328
    0
                                                            b_col,
    1329
    0
                                                            &mut partial_dots,
    1330
    0
                                                        );
    1331
    0
                                                    }
    1332
    0
    1333
    0
                                                    result.data[i * result.cols + j] += partial_dots[0];
    1334
    0
                                                    result.data[(i + 1) * result.cols + j] +=
    1335
    0
                                                        partial_dots[1];
    1336
    0
                                                    result.data[(i + 2) * result.cols + j] +=
    1337
    0
                                                        partial_dots[2];
    1338
    0
                                                    result.data[(i + 3) * result.cols + j] +=
    1339
    0
                                                        partial_dots[3];
    1340
    0
                                                }
    1341
    1342
    0
                                                i += 4;
    1343
                                            }
    1344
    1345
                                            // Handle remaining rows (< 4)
    1346
    0
                                            for i in i..i_end {
    1347
    0
                                                let row_start = i * self.cols + kk;
    1348
    0
                                                let a_row =
    1349
    0
                                                    &self.data[row_start..row_start + block_size];
    1350
    1351
    0
                                                for j in jj..j_end {
    1352
    0
                                                    let col_start = j * b_transposed.cols + kk;
    1353
    0
                                                    let b_col = &b_transposed.data
    1354
    0
                                                        [col_start..col_start + block_size];
    1355
    0
    1356
    0
                                                    let partial_dot =
    1357
    0
                                                        // SAFETY: CPU feature verified at runtime, slices bounds-checked
    1358
    0
                                                        unsafe { Avx2Backend::dot(a_row, b_col) };
    1359
    0
                                                    result.data[i * result.cols + j] += partial_dot;
    1360
    0
                                                }
    1361
                                            }
    1362
                                        } else {
    1363
                                            // Non-AVX2 path
    1364
                                            #[allow(unused_variables)]
    1365
    0
                                            for i in ii..i_end {
    1366
    0
                                                let row_start = i * self.cols + kk;
    1367
    0
                                                let a_row =
    1368
    0
                                                    &self.data[row_start..row_start + block_size];
    1369
    1370
    0
                                                for j in jj..j_end {
    1371
    0
                                                    let col_start = j * b_transposed.cols + kk;
    1372
    0
                                                    let b_col = &b_transposed.data
    1373
    0
                                                        [col_start..col_start + block_size];
    1374
    1375
                                                    // SAFETY: AVX2 verified at runtime, slices bounds-checked
    1376
    0
                                                    let partial_dot = unsafe {
    1377
    0
                                                        match self.backend {
    1378
                                                            Backend::Scalar => {
    1379
    0
                                                                ScalarBackend::dot(a_row, b_col)
    1380
                                                            }
    1381
                                                            #[cfg(target_arch = "x86_64")]
    1382
                                                            Backend::SSE2 | Backend::AVX => {
    1383
    0
                                                                Sse2Backend::dot(a_row, b_col)
    1384
                                                            }
    1385
                                                            #[cfg(not(target_arch = "x86_64"))]
    1386
                                                            Backend::SSE2
    1387
                                                            | Backend::AVX
    1388
                                                            | Backend::AVX2
    1389
                                                            | Backend::AVX512 => {
    1390
                                                                ScalarBackend::dot(a_row, b_col)
    1391
                                                            }
    1392
                                                            #[cfg(any(
    1393
                                                                target_arch = "aarch64",
    1394
                                                                target_arch = "arm"
    1395
                                                            ))]
    1396
                                                            Backend::NEON => {
    1397
                                                                use crate::backends::neon::NeonBackend;
    1398
                                                                NeonBackend::dot(a_row, b_col)
    1399
                                                            }
    1400
                                                            #[cfg(not(any(
    1401
                                                                target_arch = "aarch64",
    1402
                                                                target_arch = "arm"
    1403
                                                            )))]
    1404
                                                            Backend::NEON => {
    1405
    0
                                                                ScalarBackend::dot(a_row, b_col)
    1406
                                                            }
    1407
                                                            #[cfg(target_arch = "wasm32")]
    1408
                                                            Backend::WasmSIMD => {
    1409
                                                                use crate::backends::wasm::WasmBackend;
    1410
                                                                WasmBackend::dot(a_row, b_col)
    1411
                                                            }
    1412
                                                            #[cfg(not(target_arch = "wasm32"))]
    1413
                                                            Backend::WasmSIMD => {
    1414
    0
                                                                ScalarBackend::dot(a_row, b_col)
    1415
                                                            }
    1416
                                                            Backend::GPU
    1417
                                                            | Backend::Auto
    1418
                                                            | Backend::AVX2
    1419
                                                            | Backend::AVX512 => {
    1420
    0
                                                                ScalarBackend::dot(a_row, b_col)
    1421
                                                            }
    1422
                                                        }
    1423
                                                    };
    1424
    1425
    0
                                                    result.data[i * result.cols + j] += partial_dot;
    1426
                                                }
    1427
                                            }
    1428
                                        }
    1429
    1430
                                        // Non-x86_64 platforms
    1431
                                        #[cfg(not(target_arch = "x86_64"))]
    1432
                                        for i in ii..i_end {
    1433
                                            let row_start = i * self.cols + kk;
    1434
                                            let a_row = &self.data[row_start..row_start + block_size];
    1435
    1436
                                            for j in jj..j_end {
    1437
                                                let col_start = j * b_transposed.cols + kk;
    1438
                                                let b_col = &b_transposed.data
    1439
                                                    [col_start..col_start + block_size];
    1440
    1441
                                                // SAFETY: AVX2 verified at runtime, slices bounds-checked
    1442
                                                let partial_dot = unsafe {
    1443
                                                    match self.backend {
    1444
                                                        Backend::Scalar => {
    1445
                                                            ScalarBackend::dot(a_row, b_col)
    1446
                                                        }
    1447
                                                        #[cfg(any(
    1448
                                                            target_arch = "aarch64",
    1449
                                                            target_arch = "arm"
    1450
                                                        ))]
    1451
                                                        Backend::NEON => {
    1452
                                                            use crate::backends::neon::NeonBackend;
    1453
                                                            NeonBackend::dot(a_row, b_col)
    1454
                                                        }
    1455
                                                        #[cfg(not(any(
    1456
                                                            target_arch = "aarch64",
    1457
                                                            target_arch = "arm"
    1458
                                                        )))]
    1459
                                                        Backend::NEON => {
    1460
                                                            ScalarBackend::dot(a_row, b_col)
    1461
                                                        }
    1462
                                                        #[cfg(target_arch = "wasm32")]
    1463
                                                        Backend::WasmSIMD => {
    1464
                                                            use crate::backends::wasm::WasmBackend;
    1465
                                                            WasmBackend::dot(a_row, b_col)
    1466
                                                        }
    1467
                                                        #[cfg(not(target_arch = "wasm32"))]
    1468
                                                        Backend::WasmSIMD => {
    1469
                                                            ScalarBackend::dot(a_row, b_col)
    1470
                                                        }
    1471
                                                        _ => ScalarBackend::dot(a_row, b_col),
    1472
                                                    }
    1473
                                                };
    1474
    1475
                                                result.data[i * result.cols + j] += partial_dot;
    1476
                                            }
    1477
                                        }
    1478
                                    }
    1479
                                }
    1480
                            }
    1481
                        }
    1482
                    }
    1483
                }
    1484
            } else {
    1485
                // ===== Phase 1/2: 2-Level Cache Blocking (L2 → micro-kernel) =====
    1486
                // For medium matrices (32-512), use original 2-level blocking
    1487
                //
    1488
                // This path preserves the fast performance for 256×256 and smaller matrices
    1489
                // by avoiding the overhead of 3-level loop nesting
    1490
    1491
    0
                for ii in (0..self.rows).step_by(L2_BLOCK_SIZE) {
    1492
    0
                    let i_end = (ii + L2_BLOCK_SIZE).min(self.rows);
    1493
    1494
    0
                    for jj in (0..other.cols).step_by(L2_BLOCK_SIZE) {
    1495
    0
                        let j_end = (jj + L2_BLOCK_SIZE).min(other.cols);
    1496
    1497
    0
                        for kk in (0..self.cols).step_by(L2_BLOCK_SIZE) {
    1498
    0
                            let k_end = (kk + L2_BLOCK_SIZE).min(self.cols);
    1499
    0
                            let block_size = k_end - kk;
    1500
    1501
                            // Inner loops: Process L2 block with micro-kernel (Phase 2) or SIMD
    1502
                            #[cfg(target_arch = "x86_64")]
    1503
    0
                            let use_microkernel =
    1504
    0
                                matches!(self.backend, Backend::AVX2 | Backend::AVX512);
    1505
    1506
                            #[cfg(target_arch = "x86_64")]
    1507
    0
                            if use_microkernel {
    1508
                                // Phase 2: Use 4×1 micro-kernel for AVX2/AVX512
    1509
    0
                                let mut i = ii;
    1510
    1511
                                // Process 4 rows at a time with micro-kernel
    1512
    0
                                while i + 4 <= i_end {
    1513
                                    // Get 4 consecutive rows of A
    1514
    0
                                    let row0_start = i * self.cols + kk;
    1515
    0
                                    let row1_start = (i + 1) * self.cols + kk;
    1516
    0
                                    let row2_start = (i + 2) * self.cols + kk;
    1517
    0
                                    let row3_start = (i + 3) * self.cols + kk;
    1518
    1519
    0
                                    let a_rows = [
    1520
    0
                                        &self.data[row0_start..row0_start + block_size],
    1521
    0
                                        &self.data[row1_start..row1_start + block_size],
    1522
    0
                                        &self.data[row2_start..row2_start + block_size],
    1523
    0
                                        &self.data[row3_start..row3_start + block_size],
    1524
    0
                                    ];
    1525
    1526
                                    // Process each column of B with the micro-kernel
    1527
    0
                                    for j in jj..j_end {
    1528
    0
                                        let col_start = j * b_transposed.cols + kk;
    1529
    0
                                        let b_col =
    1530
    0
                                            &b_transposed.data[col_start..col_start + block_size];
    1531
    0
    1532
    0
                                        // Compute 4 dot products simultaneously
    1533
    0
                                        let mut partial_dots = [0.0f32; 4];
    1534
    0
                                        // SAFETY: CPU feature verified at runtime, slices bounds-checked
    1535
    0
                                        unsafe {
    1536
    0
                                            Self::matmul_microkernel_4x1_avx2(
    1537
    0
                                                a_rows,
    1538
    0
                                                b_col,
    1539
    0
                                                &mut partial_dots,
    1540
    0
                                            );
    1541
    0
                                        }
    1542
    0
    1543
    0
                                        // Accumulate results
    1544
    0
                                        result.data[i * result.cols + j] += partial_dots[0];
    1545
    0
                                        result.data[(i + 1) * result.cols + j] += partial_dots[1];
    1546
    0
                                        result.data[(i + 2) * result.cols + j] += partial_dots[2];
    1547
    0
                                        result.data[(i + 3) * result.cols + j] += partial_dots[3];
    1548
    0
                                    }
    1549
    1550
    0
                                    i += 4;
    1551
                                }
    1552
    1553
                                // Handle remaining rows (< 4) with standard path
    1554
    0
                                for i in i..i_end {
    1555
    0
                                    let row_start = i * self.cols + kk;
    1556
    0
                                    let a_row = &self.data[row_start..row_start + block_size];
    1557
    1558
    0
                                    for j in jj..j_end {
    1559
    0
                                        let col_start = j * b_transposed.cols + kk;
    1560
    0
                                        let b_col =
    1561
    0
                                            &b_transposed.data[col_start..col_start + block_size];
    1562
    0
    1563
    0
                                        // SAFETY: AVX2 verified at runtime, slices bounds-checked
    1564
    0
                                        let partial_dot = unsafe { Avx2Backend::dot(a_row, b_col) };
    1565
    0
                                        result.data[i * result.cols + j] += partial_dot;
    1566
    0
                                    }
    1567
                                }
    1568
                            } else {
    1569
                                // Phase 1: Standard SIMD path (non-AVX2 backends)
    1570
                                #[allow(unused_variables)]
    1571
    0
                                for i in ii..i_end {
    1572
    0
                                    let row_start = i * self.cols + kk;
    1573
    0
                                    let a_row = &self.data[row_start..row_start + block_size];
    1574
    1575
    0
                                    for j in jj..j_end {
    1576
    0
                                        let col_start = j * b_transposed.cols + kk;
    1577
    0
                                        let b_col =
    1578
    0
                                            &b_transposed.data[col_start..col_start + block_size];
    1579
    1580
                                        // SAFETY: AVX2 verified at runtime, slices bounds-checked
    1581
    0
                                        let partial_dot = unsafe {
    1582
    0
                                            match self.backend {
    1583
    0
                                                Backend::Scalar => ScalarBackend::dot(a_row, b_col),
    1584
                                                #[cfg(target_arch = "x86_64")]
    1585
                                                Backend::SSE2 | Backend::AVX => {
    1586
    0
                                                    Sse2Backend::dot(a_row, b_col)
    1587
                                                }
    1588
                                                #[cfg(not(target_arch = "x86_64"))]
    1589
                                                Backend::SSE2
    1590
                                                | Backend::AVX
    1591
                                                | Backend::AVX2
    1592
                                                | Backend::AVX512 => ScalarBackend::dot(a_row, b_col),
    1593
                                                #[cfg(any(
    1594
                                                    target_arch = "aarch64",
    1595
                                                    target_arch = "arm"
    1596
                                                ))]
    1597
                                                Backend::NEON => {
    1598
                                                    use crate::backends::neon::NeonBackend;
    1599
                                                    NeonBackend::dot(a_row, b_col)
    1600
                                                }
    1601
                                                #[cfg(not(any(
    1602
                                                    target_arch = "aarch64",
    1603
                                                    target_arch = "arm"
    1604
                                                )))]
    1605
    0
                                                Backend::NEON => ScalarBackend::dot(a_row, b_col),
    1606
                                                #[cfg(target_arch = "wasm32")]
    1607
                                                Backend::WasmSIMD => {
    1608
                                                    use crate::backends::wasm::WasmBackend;
    1609
                                                    WasmBackend::dot(a_row, b_col)
    1610
                                                }
    1611
                                                #[cfg(not(target_arch = "wasm32"))]
    1612
    0
                                                Backend::WasmSIMD => ScalarBackend::dot(a_row, b_col),
    1613
                                                Backend::GPU
    1614
                                                | Backend::Auto
    1615
                                                | Backend::AVX2
    1616
    0
                                                | Backend::AVX512 => ScalarBackend::dot(a_row, b_col),
    1617
                                            }
    1618
                                        };
    1619
    1620
    0
                                        result.data[i * result.cols + j] += partial_dot;
    1621
                                    }
    1622
                                }
    1623
                            }
    1624
    1625
                            // Non-x86_64 platforms: Use standard SIMD path
    1626
                            #[cfg(not(target_arch = "x86_64"))]
    1627
                            for i in ii..i_end {
    1628
                                let row_start = i * self.cols + kk;
    1629
                                let a_row = &self.data[row_start..row_start + block_size];
    1630
    1631
                                for j in jj..j_end {
    1632
                                    let col_start = j * b_transposed.cols + kk;
    1633
                                    let b_col = &b_transposed.data[col_start..col_start + block_size];
    1634
    1635
                                    // SAFETY: AVX2 verified at runtime, slices bounds-checked
    1636
                                    let partial_dot = unsafe {
    1637
                                        match self.backend {
    1638
                                            Backend::Scalar => ScalarBackend::dot(a_row, b_col),
    1639
                                            #[cfg(any(target_arch = "aarch64", target_arch = "arm"))]
    1640
                                            Backend::NEON => {
    1641
                                                use crate::backends::neon::NeonBackend;
    1642
                                                NeonBackend::dot(a_row, b_col)
    1643
                                            }
    1644
                                            #[cfg(not(any(
    1645
                                                target_arch = "aarch64",
    1646
                                                target_arch = "arm"
    1647
                                            )))]
    1648
                                            Backend::NEON => ScalarBackend::dot(a_row, b_col),
    1649
                                            #[cfg(target_arch = "wasm32")]
    1650
                                            Backend::WasmSIMD => {
    1651
                                                use crate::backends::wasm::WasmBackend;
    1652
                                                WasmBackend::dot(a_row, b_col)
    1653
                                            }
    1654
                                            #[cfg(not(target_arch = "wasm32"))]
    1655
                                            Backend::WasmSIMD => ScalarBackend::dot(a_row, b_col),
    1656
                                            _ => ScalarBackend::dot(a_row, b_col),
    1657
                                        }
    1658
                                    };
    1659
    1660
                                    result.data[i * result.cols + j] += partial_dot;
    1661
                                }
    1662
                            }
    1663
                        }
    1664
                    }
    1665
                }
    1666
            }
    1667
    1668
    0
            Ok(())
    1669
    0
        }
    1670
    1671
        /// Simple SIMD matrix multiplication without blocking (for small matrices)
    1672
        ///
    1673
        /// This is the pre-blocking implementation that works well for small matrices
    1674
        /// where cache blocking overhead exceeds benefits.
    1675
    0
        fn matmul_simd_simple(
    1676
    0
            &self,
    1677
    0
            other: &Matrix<f32>,
    1678
    0
            result: &mut Matrix<f32>,
    1679
    0
        ) -> Result<(), TruenoError> {
    1680
            #[cfg(target_arch = "x86_64")]
    1681
            use crate::backends::{avx2::Avx2Backend, sse2::Sse2Backend};
    1682
            use crate::backends::{scalar::ScalarBackend, VectorBackend};
    1683
    1684
            // Pre-transpose B for better cache locality
    1685
    0
            let b_transposed = other.transpose();
    1686
    1687
    0
            for i in 0..self.rows {
    1688
    0
                let row_start = i * self.cols;
    1689
    0
                let row_end = row_start + self.cols;
    1690
    0
                let a_row = &self.data[row_start..row_end];
    1691
    1692
    0
                for j in 0..other.cols {
    1693
    0
                    let col_start = j * b_transposed.cols;
    1694
    0
                    let col_end = col_start + b_transposed.cols;
    1695
    0
                    let b_col = &b_transposed.data[col_start..col_end];
    1696
    1697
                    // Compute dot product using SIMD backend directly
    1698
                    // SAFETY: Backend dot() maintains safety invariants
    1699
    0
                    let dot_result = unsafe {
    1700
    0
                        match self.backend {
    1701
    0
                            Backend::Scalar => ScalarBackend::dot(a_row, b_col),
    1702
                            #[cfg(target_arch = "x86_64")]
    1703
    0
                            Backend::SSE2 | Backend::AVX => Sse2Backend::dot(a_row, b_col),
    1704
                            #[cfg(target_arch = "x86_64")]
    1705
    0
                            Backend::AVX2 | Backend::AVX512 => Avx2Backend::dot(a_row, b_col),
    1706
                            #[cfg(not(target_arch = "x86_64"))]
    1707
                            Backend::SSE2 | Backend::AVX | Backend::AVX2 | Backend::AVX512 => {
    1708
                                ScalarBackend::dot(a_row, b_col)
    1709
                            }
    1710
                            #[cfg(any(target_arch = "aarch64", target_arch = "arm"))]
    1711
                            Backend::NEON => {
    1712
                                use crate::backends::neon::NeonBackend;
    1713
                                NeonBackend::dot(a_row, b_col)
    1714
                            }
    1715
                            #[cfg(not(any(target_arch = "aarch64", target_arch = "arm")))]
    1716
    0
                            Backend::NEON => ScalarBackend::dot(a_row, b_col),
    1717
                            #[cfg(target_arch = "wasm32")]
    1718
                            Backend::WasmSIMD => {
    1719
                                use crate::backends::wasm::WasmBackend;
    1720
                                WasmBackend::dot(a_row, b_col)
    1721
                            }
    1722
                            #[cfg(not(target_arch = "wasm32"))]
    1723
    0
                            Backend::WasmSIMD => ScalarBackend::dot(a_row, b_col),
    1724
    0
                            Backend::GPU | Backend::Auto => ScalarBackend::dot(a_row, b_col),
    1725
                        }
    1726
                    };
    1727
    1728
    0
                    result.data[i * result.cols + j] = dot_result;
    1729
                }
    1730
            }
    1731
    1732
    0
            Ok(())
    1733
    0
        }
    1734
    1735
        /// WASM-optimized tiled matrix multiplication with SIMD inner loop
    1736
        ///
    1737
        /// Key optimizations:
    1738
        /// 1. NO transpose - avoids O(n²) memory allocation and copy
    1739
        /// 2. Tiled blocking with SIMD-aligned tile widths
    1740
        /// 3. Inner j-loop uses SIMD (B rows are contiguous in memory)
    1741
        /// 4. Register accumulation to minimize memory traffic
    1742
        ///
    1743
        /// Performance: Targets <30ms for 384×74×384 (Whisper encoder attention)
    1744
    0
        fn matmul_wasm_tiled(
    1745
    0
            &self,
    1746
    0
            other: &Matrix<f32>,
    1747
    0
            result: &mut Matrix<f32>,
    1748
    0
        ) -> Result<(), TruenoError> {
    1749
    0
            let m = self.rows;
    1750
    0
            let k = self.cols;
    1751
    0
            let n = other.cols;
    1752
    1753
            // For each row of A
    1754
    0
            for i in 0..m {
    1755
    0
                let a_row_start = i * k;
    1756
    0
                let result_row_start = i * n;
    1757
    1758
                // For each column of B, compute dot product A[i,:] · B[:,j]
    1759
                // BUT: B[:,j] is not contiguous. Instead, iterate over k and accumulate.
    1760
                //
    1761
                // C[i,j] = Σ_k A[i,k] * B[k,j]
    1762
                //
    1763
                // For efficiency, broadcast A[i,k] and multiply with B[k, j0:j0+width]
    1764
                // This uses SIMD on the contiguous B row segment.
    1765
    1766
                // Process output columns in SIMD-width chunks
    1767
    0
                let simd_width = 8; // AVX2 processes 8 f32s
    1768
    0
                let n_simd = (n / simd_width) * simd_width;
    1769
    1770
                // SIMD portion: columns 0..n_simd
    1771
                // Note: Explicit indexing is intentional for LLVM auto-vectorization.
    1772
                // Iterator patterns prevent the compiler from recognizing the SIMD pattern.
    1773
                #[allow(clippy::needless_range_loop)]
    1774
    0
                for j0 in (0..n_simd).step_by(simd_width) {
    1775
    0
                    let mut acc = [0.0f32; 8];
    1776
    1777
    0
                    for kk in 0..k {
    1778
    0
                        let a_val = self.data[a_row_start + kk];
    1779
    0
                        let b_row_start = kk * n + j0;
    1780
    1781
                        // Multiply a_val with B[kk, j0:j0+8]
    1782
    0
                        for jj in 0..simd_width {
    1783
    0
                            acc[jj] += a_val * other.data[b_row_start + jj];
    1784
    0
                        }
    1785
                    }
    1786
    1787
                    // Write accumulated results
    1788
    0
                    for jj in 0..simd_width {
    1789
    0
                        result.data[result_row_start + j0 + jj] = acc[jj];
    1790
    0
                    }
    1791
                }
    1792
    1793
                // Remainder columns (non-SIMD)
    1794
    0
                for j in n_simd..n {
    1795
    0
                    let mut sum = 0.0f32;
    1796
    0
                    for kk in 0..k {
    1797
    0
                        sum += self.data[a_row_start + kk] * other.data[kk * n + j];
    1798
    0
                    }
    1799
    0
                    result.data[result_row_start + j] = sum;
    1800
                }
    1801
            }
    1802
    1803
    0
            Ok(())
    1804
    0
        }
    1805
    1806
        /// GPU-accelerated matrix multiplication (very large matrices only)
    1807
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    1808
    0
        fn matmul_gpu(&self, other: &Matrix<f32>) -> Result<Matrix<f32>, TruenoError> {
    1809
            use crate::backends::gpu::GpuBackend;
    1810
    1811
            // Check if GPU is available
    1812
    0
            if !GpuBackend::is_available() {
    1813
    0
                return Err(TruenoError::InvalidInput("GPU not available".to_string()));
    1814
    0
            }
    1815
    1816
            // Create GPU backend
    1817
    0
            let mut gpu = GpuBackend::new();
    1818
    1819
            // Execute GPU matmul
    1820
    0
            let result_data = gpu
    1821
    0
                .matmul(&self.data, &other.data, self.rows, self.cols, other.cols)
    1822
    0
                .map_err(|e| TruenoError::InvalidInput(format!("GPU matmul failed: {}", e)))?;
    1823
    1824
            // Create result matrix
    1825
    0
            let mut result = Matrix::zeros(self.rows, other.cols);
    1826
    0
            result.data = result_data;
    1827
    1828
    0
            Ok(result)
    1829
    0
        }
    1830
    1831
        /// Transpose the matrix (swap rows and columns)
    1832
        ///
    1833
        /// Returns a new matrix where element `(i, j)` of the original becomes
    1834
        /// element `(j, i)` in the result.
    1835
        ///
    1836
        /// # Returns
    1837
        ///
    1838
        /// A new matrix with dimensions swapped: if input is `m×n`, output is `n×m`
    1839
        ///
    1840
        /// # Example
    1841
        ///
    1842
        /// ```
    1843
        /// use trueno::Matrix;
    1844
        ///
    1845
        /// let m = Matrix::from_vec(2, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    1846
        /// let t = m.transpose();
    1847
        ///
    1848
        /// // [[1, 2, 3],     [[1, 4],
    1849
        /// //  [4, 5, 6]]  →   [2, 5],
    1850
        /// //                  [3, 6]]
    1851
        /// assert_eq!(t.rows(), 3);
    1852
        /// assert_eq!(t.cols(), 2);
    1853
        /// assert_eq!(t.get(0, 0), Some(&1.0));
    1854
        /// assert_eq!(t.get(0, 1), Some(&4.0));
    1855
        /// assert_eq!(t.get(1, 0), Some(&2.0));
    1856
        /// ```
    1857
        #[cfg_attr(feature = "tracing", instrument(skip(self), fields(dims = %format!("{}x{}", self.rows, self.cols))))]
    1858
    0
        pub fn transpose(&self) -> Matrix<f32> {
    1859
    0
            let mut result = Matrix::zeros_with_backend(self.cols, self.rows, self.backend);
    1860
    1861
            // Use block-wise transpose for better cache locality
    1862
            // Block size of 32 balances cache efficiency for both square and non-square matrices
    1863
            const BLOCK_SIZE: usize = 32;
    1864
    1865
            // For non-square matrices, process output rows sequentially for write coalescing
    1866
            // This ensures writes are sequential in memory regardless of input shape
    1867
            // Fix for issue #65: non-square transpose was slow due to strided writes
    1868
    1869
            // Process in blocks, iterating output rows first for sequential writes
    1870
    0
            for j_block in (0..self.cols).step_by(BLOCK_SIZE) {
    1871
    0
                let j_end = (j_block + BLOCK_SIZE).min(self.cols);
    1872
    1873
    0
                for i_block in (0..self.rows).step_by(BLOCK_SIZE) {
    1874
    0
                    let i_end = (i_block + BLOCK_SIZE).min(self.rows);
    1875
    1876
                    // Within block: iterate output rows (j) in outer loop for sequential writes
    1877
    0
                    for j in j_block..j_end {
    1878
    0
                        let dst_row_start = j * result.cols;
    1879
    0
                        for i in i_block..i_end {
    1880
    0
                            // result[j, i] = self[i, j]
    1881
    0
                            // Sequential write: dst_row_start + i increments by 1
    1882
    0
                            // Strided read: acceptable, CPU prefetch handles this
    1883
    0
                            result.data[dst_row_start + i] = self.data[i * self.cols + j];
    1884
    0
                        }
    1885
                    }
    1886
                }
    1887
            }
    1888
    1889
    0
            result
    1890
    0
        }
    1891
    1892
        /// Matrix-vector multiplication (column vector): A × v
    1893
        ///
    1894
        /// Multiplies this matrix by a column vector, computing `A × v` where the result
    1895
        /// is a column vector with length equal to the number of rows in `A`.
    1896
        ///
    1897
        /// # Mathematical Definition
    1898
        ///
    1899
        /// For an m×n matrix A and an n-dimensional vector v:
    1900
        /// ```text
    1901
        /// result[i] = Σ(j=0 to n-1) A[i,j] × v[j]
    1902
        /// ```
    1903
        ///
    1904
        /// # Arguments
    1905
        ///
    1906
        /// * `v` - Column vector with length equal to `self.cols()`
    1907
        ///
    1908
        /// # Returns
    1909
        ///
    1910
        /// A new vector with length `self.rows()`
    1911
        ///
    1912
        /// # Errors
    1913
        ///
    1914
        /// Returns `InvalidInput` if `v.len() != self.cols()`
    1915
        ///
    1916
        /// # Example
    1917
        ///
    1918
        /// ```
    1919
        /// use trueno::{Matrix, Vector};
    1920
        ///
    1921
        /// let m = Matrix::from_vec(2, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    1922
        /// let v = Vector::from_slice(&[1.0, 2.0, 3.0]);
    1923
        /// let result = m.matvec(&v).unwrap();
    1924
        ///
    1925
        /// // [[1, 2, 3]   [1]   [1×1 + 2×2 + 3×3]   [14]
    1926
        /// //  [4, 5, 6]] × [2] = [4×1 + 5×2 + 6×3] = [32]
    1927
        /// //               [3]
    1928
        /// assert_eq!(result.as_slice(), &[14.0, 32.0]);
    1929
        /// ```
    1930
    0
        pub fn matvec(&self, v: &Vector<f32>) -> Result<Vector<f32>, TruenoError> {
    1931
    0
            if v.len() != self.cols {
    1932
    0
                return Err(TruenoError::InvalidInput(format!(
    1933
    0
                    "Vector length {} does not match matrix columns {} for matrix-vector multiplication",
    1934
    0
                    v.len(),
    1935
    0
                    self.cols
    1936
    0
                )));
    1937
    0
            }
    1938
    1939
            #[cfg(target_arch = "x86_64")]
    1940
            use crate::backends::{avx2::Avx2Backend, sse2::Sse2Backend};
    1941
            use crate::backends::{scalar::ScalarBackend, VectorBackend};
    1942
    1943
    0
            let v_slice = v.as_slice();
    1944
    1945
    0
            let mut result_data = vec![0.0; self.rows];
    1946
    1947
            // Parallel execution for very large matrices (≥4096 rows)
    1948
            // Note: Thread overhead dominates for smaller matrices
    1949
            #[cfg(feature = "parallel")]
    1950
            {
    1951
                const PARALLEL_THRESHOLD: usize = 4096;
    1952
    1953
                if self.rows >= PARALLEL_THRESHOLD {
    1954
                    use rayon::prelude::*;
    1955
                    use std::sync::atomic::{AtomicPtr, Ordering};
    1956
                    use std::sync::Arc;
    1957
    1958
                    let result_ptr = Arc::new(AtomicPtr::new(result_data.as_mut_ptr()));
    1959
    1960
                    // Process rows in parallel - each row computes an independent dot product
    1961
                    (0..self.rows).into_par_iter().for_each(|i| {
    1962
                        let row_start = i * self.cols;
    1963
                        let row = &self.data[row_start..(row_start + self.cols)];
    1964
    1965
                        // SAFETY: CPU feature verified at runtime, slices bounds-checked
    1966
                        let dot_result = unsafe {
    1967
                            #[cfg(target_arch = "x86_64")]
    1968
                            {
    1969
                                match self.backend {
    1970
                                    Backend::AVX2 | Backend::AVX512 => Avx2Backend::dot(row, v_slice),
    1971
                                    Backend::SSE2 | Backend::AVX => Sse2Backend::dot(row, v_slice),
    1972
                                    _ => ScalarBackend::dot(row, v_slice),
    1973
                                }
    1974
                            }
    1975
                            #[cfg(not(target_arch = "x86_64"))]
    1976
                            {
    1977
                                ScalarBackend::dot(row, v_slice)
    1978
                            }
    1979
                        };
    1980
    1981
                        // Write to non-overlapping memory location (thread-safe)
    1982
                        // SAFETY: CPU feature verified at runtime, slices bounds-checked
    1983
                        unsafe {
    1984
                            let ptr = result_ptr.load(Ordering::Relaxed);
    1985
                            *ptr.add(i) = dot_result;
    1986
                        }
    1987
                    });
    1988
    1989
                    return Ok(Vector::from_slice(&result_data));
    1990
                }
    1991
            }
    1992
    1993
            // SIMD-optimized execution: each row-vector product is a dot product
    1994
    0
            for (i, result) in result_data.iter_mut().enumerate() {
    1995
    0
                let row_start = i * self.cols;
    1996
    0
                let row = &self.data[row_start..(row_start + self.cols)];
    1997
    1998
                // Use SIMD dot product for each row
    1999
                // SAFETY: CPU feature verified at runtime, slices bounds-checked
    2000
                *result = unsafe {
    2001
                    #[cfg(target_arch = "x86_64")]
    2002
                    {
    2003
    0
                        match self.backend {
    2004
    0
                            Backend::AVX2 | Backend::AVX512 => Avx2Backend::dot(row, v_slice),
    2005
    0
                            Backend::SSE2 | Backend::AVX => Sse2Backend::dot(row, v_slice),
    2006
    0
                            _ => ScalarBackend::dot(row, v_slice),
    2007
                        }
    2008
                    }
    2009
                    #[cfg(not(target_arch = "x86_64"))]
    2010
                    {
    2011
                        ScalarBackend::dot(row, v_slice)
    2012
                    }
    2013
                };
    2014
            }
    2015
    2016
    0
            Ok(Vector::from_slice(&result_data))
    2017
    0
        }
    2018
    2019
        /// Vector-matrix multiplication (row vector): v^T × A
    2020
        ///
    2021
        /// Multiplies a row vector by this matrix, computing `v^T × A` where the result
    2022
        /// is a row vector with length equal to the number of columns in `A`.
    2023
        ///
    2024
        /// # Mathematical Definition
    2025
        ///
    2026
        /// For an m-dimensional vector v and an m×n matrix A:
    2027
        /// ```text
    2028
        /// result[j] = Σ(i=0 to m-1) v[i] × A[i,j]
    2029
        /// ```
    2030
        ///
    2031
        /// # Arguments
    2032
        ///
    2033
        /// * `v` - Row vector with length equal to `m.rows()`
    2034
        /// * `m` - Matrix to multiply
    2035
        ///
    2036
        /// # Returns
    2037
        ///
    2038
        /// A new vector with length `m.cols()`
    2039
        ///
    2040
        /// # Errors
    2041
        ///
    2042
        /// Returns `InvalidInput` if `v.len() != m.rows()`
    2043
        ///
    2044
        /// # Example
    2045
        ///
    2046
        /// ```
    2047
        /// use trueno::{Matrix, Vector};
    2048
        ///
    2049
        /// let m = Matrix::from_vec(2, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    2050
        /// let v = Vector::from_slice(&[1.0, 2.0]);
    2051
        /// let result = Matrix::vecmat(&v, &m).unwrap();
    2052
        ///
    2053
        /// // [1, 2] × [[1, 2, 3]  = [1×1 + 2×4, 1×2 + 2×5, 1×3 + 2×6]
    2054
        /// //           [4, 5, 6]]
    2055
        /// //         = [9, 12, 15]
    2056
        /// assert_eq!(result.as_slice(), &[9.0, 12.0, 15.0]);
    2057
        /// ```
    2058
    0
        pub fn vecmat(v: &Vector<f32>, m: &Matrix<f32>) -> Result<Vector<f32>, TruenoError> {
    2059
    0
            if v.len() != m.rows {
    2060
    0
                return Err(TruenoError::InvalidInput(format!(
    2061
    0
                    "Vector length {} does not match matrix rows {} for vector-matrix multiplication",
    2062
    0
                    v.len(),
    2063
    0
                    m.rows
    2064
    0
                )));
    2065
    0
            }
    2066
    2067
            // SIMD-optimized implementation using row-wise accumulation
    2068
            // Instead of column-wise access (cache-unfriendly), we compute:
    2069
            // result = Σ(i) v[i] * row_i (cache-friendly, vectorizable)
    2070
            //
    2071
            // This approach:
    2072
            // 1. Sequential row access (cache-friendly vs strided column access)
    2073
            // 2. Uses SIMD scale and add operations
    2074
            // 3. Leverages existing optimized Vector operations
    2075
    2076
    0
            let mut result = Vector::from_slice(&vec![0.0; m.cols]);
    2077
    0
            let v_slice = v.as_slice();
    2078
    2079
            // Accumulate each scaled row into result
    2080
    0
            for (i, &scalar) in v_slice.iter().enumerate().take(m.rows) {
    2081
    0
                let row_start = i * m.cols;
    2082
    0
                let row = &m.data[row_start..(row_start + m.cols)];
    2083
    2084
                // Create vector for this row
    2085
    0
                let row_vec = Vector::from_slice(row);
    2086
    2087
                // result += scalar * row (using SIMD scale and add)
    2088
    0
                let scaled_row = row_vec.scale(scalar)?;
    2089
    0
                result = result.add(&scaled_row)?;
    2090
            }
    2091
    2092
    0
            Ok(result)
    2093
    0
        }
    2094
    2095
        /// Perform 2D convolution with a kernel
    2096
        ///
    2097
        /// Applies a 2D convolution operation using "valid" padding (no padding),
    2098
        /// resulting in an output smaller than the input.
    2099
        ///
    2100
        /// # Arguments
    2101
        ///
    2102
        /// * `kernel` - Convolution kernel (filter) to apply
    2103
        ///
    2104
        /// # Returns
    2105
        ///
    2106
        /// Convolved matrix with dimensions:
    2107
        /// - rows: `input.rows - kernel.rows + 1`
    2108
        /// - cols: `input.cols - kernel.cols + 1`
    2109
        ///
    2110
        /// # Errors
    2111
        ///
    2112
        /// Returns `InvalidInput` if:
    2113
        /// - Kernel is larger than input in any dimension
    2114
        /// - Kernel has even dimensions (center pixel ambiguous)
    2115
        ///
    2116
        /// # Example
    2117
        ///
    2118
        /// ```
    2119
        /// use trueno::Matrix;
    2120
        ///
    2121
        /// // 5x5 input image
    2122
        /// let input = Matrix::from_vec(
    2123
        ///     5, 5,
    2124
        ///     vec![
    2125
        ///         0.0, 0.0, 0.0, 0.0, 0.0,
    2126
        ///         0.0, 0.0, 0.0, 0.0, 0.0,
    2127
        ///         0.0, 0.0, 9.0, 0.0, 0.0,
    2128
        ///         0.0, 0.0, 0.0, 0.0, 0.0,
    2129
        ///         0.0, 0.0, 0.0, 0.0, 0.0,
    2130
        ///     ]
    2131
        /// ).unwrap();
    2132
        ///
    2133
        /// // 3x3 averaging kernel
    2134
        /// let kernel_val = 1.0 / 9.0;
    2135
        /// let kernel = Matrix::from_vec(
    2136
        ///     3, 3,
    2137
        ///     vec![kernel_val; 9]
    2138
        /// ).unwrap();
    2139
        ///
    2140
        /// let result = input.convolve2d(&kernel).unwrap();
    2141
        /// assert_eq!(result.rows(), 3); // 5 - 3 + 1
    2142
        /// assert_eq!(result.cols(), 3);
    2143
        /// ```
    2144
        // =========================================================================
    2145
        // HOT PATH - PERFORMANCE CRITICAL
    2146
        // =========================================================================
    2147
        // This function processes millions of elements for typical image sizes.
    2148
        // Any changes to the inner loop REQUIRE benchmark verification:
    2149
        //   1. Run: make bench-check
    2150
        //   2. Verify no regression >10%
    2151
        //
    2152
        // PROHIBITED in inner loops:
    2153
        //   - .get() / .get_mut() (bounds checking overhead)
    2154
        //   - .expect() / .unwrap() (panic path overhead)
    2155
        //   - Iterator adaptors (closure overhead)
    2156
        //
    2157
        // Use direct indexing with bounds proof documented above the loop.
    2158
        // =========================================================================
    2159
    0
        pub fn convolve2d(&self, kernel: &Matrix<f32>) -> Result<Matrix<f32>, TruenoError> {
    2160
            // Validate kernel size
    2161
    0
            if kernel.rows > self.rows || kernel.cols > self.cols {
    2162
    0
                return Err(TruenoError::InvalidInput(format!(
    2163
    0
                    "Kernel size ({}x{}) larger than input ({}x{})",
    2164
    0
                    kernel.rows, kernel.cols, self.rows, self.cols
    2165
    0
                )));
    2166
    0
            }
    2167
    2168
            // Calculate output dimensions (valid padding)
    2169
    0
            let output_rows = self.rows - kernel.rows + 1;
    2170
    0
            let output_cols = self.cols - kernel.cols + 1;
    2171
    2172
            // Initialize output matrix (reuse parent's backend)
    2173
    0
            let mut result = Matrix::zeros_with_backend(output_rows, output_cols, self.backend);
    2174
    2175
            // Backend selection strategy:
    2176
            // OpComplexity::High - GPU beneficial at >10K elements
    2177
            // GPU for large images (output > 10K elements)
    2178
            // Scalar for smaller images
    2179
    2180
            #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    2181
            const GPU_THRESHOLD: usize = 10_000;
    2182
    2183
            // Try GPU first for large convolutions
    2184
            #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    2185
            {
    2186
    0
                if output_rows * output_cols >= GPU_THRESHOLD {
    2187
                    use crate::backends::gpu::GpuBackend;
    2188
    2189
    0
                    if GpuBackend::is_available() {
    2190
    0
                        if let Ok(gpu_result) =
    2191
    0
                            self.convolve2d_gpu(kernel, &mut result, output_rows, output_cols)
    2192
                        {
    2193
    0
                            return Ok(gpu_result);
    2194
    0
                        }
    2195
                        // Fall through to scalar if GPU fails
    2196
    0
                    }
    2197
    0
                }
    2198
            }
    2199
    2200
            // Scalar baseline implementation - optimized with direct indexing
    2201
            // SAFETY invariant proof:
    2202
            // - output_rows = self.rows - kernel.rows + 1
    2203
            // - output_cols = self.cols - kernel.cols + 1
    2204
            // - For any out_row < output_rows and k_row < kernel.rows:
    2205
            //   in_row = out_row + k_row < (self.rows - kernel.rows + 1) + kernel.rows - 1 = self.rows
    2206
            // - Same logic applies to columns
    2207
            // - All indices are provably within bounds, so we use direct indexing for performance
    2208
    2209
    0
            let input_data = self.as_slice();
    2210
    0
            let kernel_data = kernel.as_slice();
    2211
    0
            let result_data = result.data.as_mut_slice();
    2212
    0
            let input_cols = self.cols;
    2213
    0
            let kernel_cols = kernel.cols;
    2214
    0
            let result_cols = output_cols;
    2215
    2216
    0
            for out_row in 0..output_rows {
    2217
    0
                for out_col in 0..output_cols {
    2218
    0
                    let mut sum = 0.0;
    2219
    2220
                    // Apply kernel - use direct indexing for performance
    2221
    0
                    for k_row in 0..kernel.rows {
    2222
    0
                        let in_row = out_row + k_row;
    2223
    0
                        let input_row_offset = in_row * input_cols;
    2224
    0
                        let kernel_row_offset = k_row * kernel_cols;
    2225
    2226
    0
                        for k_col in 0..kernel.cols {
    2227
    0
                            let in_col = out_col + k_col;
    2228
    0
                            sum += input_data[input_row_offset + in_col]
    2229
    0
                                * kernel_data[kernel_row_offset + k_col];
    2230
    0
                        }
    2231
                    }
    2232
    2233
    0
                    result_data[out_row * result_cols + out_col] = sum;
    2234
                }
    2235
            }
    2236
    2237
    0
            Ok(result)
    2238
    0
        }
    2239
    2240
        /// GPU-accelerated 2D convolution helper
    2241
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    2242
    0
        fn convolve2d_gpu(
    2243
    0
            &self,
    2244
    0
            kernel: &Matrix<f32>,
    2245
    0
            result: &mut Matrix<f32>,
    2246
    0
            _output_rows: usize,
    2247
    0
            _output_cols: usize,
    2248
    0
        ) -> Result<Matrix<f32>, TruenoError> {
    2249
            use crate::backends::gpu::GpuDevice;
    2250
    2251
    0
            let gpu = GpuDevice::new().map_err(TruenoError::InvalidInput)?;
    2252
    2253
    0
            gpu.convolve2d(
    2254
    0
                self.as_slice(),
    2255
    0
                kernel.as_slice(),
    2256
    0
                result.data.as_mut_slice(),
    2257
    0
                self.rows,
    2258
    0
                self.cols,
    2259
    0
                kernel.rows,
    2260
    0
                kernel.cols,
    2261
            )
    2262
    0
            .map_err(TruenoError::InvalidInput)?;
    2263
    2264
    0
            Ok(result.clone())
    2265
    0
        }
    2266
    2267
        /// Lookup embeddings by indices (Issue #61: ML primitives)
    2268
        ///
    2269
        /// Performs embedding lookup where self is the embedding table with shape
    2270
        /// `[vocab_size, embed_dim]` and indices specify which rows to select.
    2271
        ///
    2272
        /// # Arguments
    2273
        ///
    2274
        /// * `indices` - Slice of indices into the embedding table
    2275
        ///
    2276
        /// # Returns
    2277
        ///
    2278
        /// A matrix with shape `[indices.len(), embed_dim]` containing the selected rows
    2279
        ///
    2280
        /// # Errors
    2281
        ///
    2282
        /// Returns `InvalidInput` if any index is out of bounds
    2283
        ///
    2284
        /// # Example
    2285
        ///
    2286
        /// ```
    2287
        /// use trueno::Matrix;
    2288
        ///
    2289
        /// // Create embedding table: 4 words, 3-dimensional embeddings
    2290
        /// let embeddings = Matrix::from_vec(4, 3, vec![
    2291
        ///     1.0, 2.0, 3.0,   // word 0
    2292
        ///     4.0, 5.0, 6.0,   // word 1
    2293
        ///     7.0, 8.0, 9.0,   // word 2
    2294
        ///     10.0, 11.0, 12.0 // word 3
    2295
        /// ]).unwrap();
    2296
        ///
    2297
        /// // Lookup embeddings for indices [1, 3, 0]
    2298
        /// let result = embeddings.embedding_lookup(&[1, 3, 0]).unwrap();
    2299
        ///
    2300
        /// assert_eq!(result.rows(), 3);
    2301
        /// assert_eq!(result.cols(), 3);
    2302
        /// assert_eq!(result.get(0, 0), Some(&4.0)); // word 1
    2303
        /// assert_eq!(result.get(1, 0), Some(&10.0)); // word 3
    2304
        /// assert_eq!(result.get(2, 0), Some(&1.0)); // word 0
    2305
        /// ```
    2306
    0
        pub fn embedding_lookup(&self, indices: &[usize]) -> Result<Matrix<f32>, TruenoError> {
    2307
            // Validate indices
    2308
    0
            for (i, &idx) in indices.iter().enumerate() {
    2309
    0
                if idx >= self.rows {
    2310
    0
                    return Err(TruenoError::InvalidInput(format!(
    2311
    0
                        "Index {} at position {} is out of bounds for embedding table with {} rows",
    2312
    0
                        idx, i, self.rows
    2313
    0
                    )));
    2314
    0
                }
    2315
            }
    2316
    2317
            // Handle empty indices
    2318
    0
            if indices.is_empty() {
    2319
    0
                return Ok(Matrix::zeros_with_backend(0, self.cols, self.backend));
    2320
    0
            }
    2321
    2322
            // Allocate output matrix: [seq_len, embed_dim]
    2323
    0
            let seq_len = indices.len();
    2324
    0
            let embed_dim = self.cols;
    2325
    0
            let mut result = Matrix::zeros_with_backend(seq_len, embed_dim, self.backend);
    2326
    2327
            // Copy rows from embedding table to result
    2328
    0
            for (out_row, &idx) in indices.iter().enumerate() {
    2329
    0
                let src_start = idx * embed_dim;
    2330
    0
                let dst_start = out_row * embed_dim;
    2331
    0
    2332
    0
                // Copy entire row
    2333
    0
                result.data[dst_start..dst_start + embed_dim]
    2334
    0
                    .copy_from_slice(&self.data[src_start..src_start + embed_dim]);
    2335
    0
            }
    2336
    2337
    0
            Ok(result)
    2338
    0
        }
    2339
    2340
        /// Lookup embeddings with gradient tracking support (for training)
    2341
        ///
    2342
        /// Returns both the embeddings and a sparse gradient accumulator.
    2343
        /// This is useful for sparse gradient updates in training.
    2344
        ///
    2345
        /// # Arguments
    2346
        ///
    2347
        /// * `indices` - Slice of indices into the embedding table
    2348
        ///
    2349
        /// # Returns
    2350
        ///
    2351
        /// Tuple of (embeddings, unique_indices) where unique_indices can be used
    2352
        /// for sparse gradient updates
    2353
        ///
    2354
        /// # Errors
    2355
        ///
    2356
        /// Returns `InvalidInput` if any index is out of bounds
    2357
    0
        pub fn embedding_lookup_sparse(
    2358
    0
            &self,
    2359
    0
            indices: &[usize],
    2360
    0
        ) -> Result<(Matrix<f32>, Vec<usize>), TruenoError> {
    2361
    0
            let embeddings = self.embedding_lookup(indices)?;
    2362
    2363
            // Get unique indices for sparse gradient updates
    2364
    0
            let mut unique: Vec<usize> = indices.to_vec();
    2365
    0
            unique.sort_unstable();
    2366
    0
            unique.dedup();
    2367
    2368
    0
            Ok((embeddings, unique))
    2369
    0
        }
    2370
    2371
        /// 2D Max Pooling operation for CNN downsampling
    2372
        ///
    2373
        /// Applies max pooling over a 2D input tensor with specified kernel size and stride.
    2374
        /// Input shape: (height, width), Output shape: ((height - kh) / sh + 1, (width - kw) / sw + 1)
    2375
        ///
    2376
        /// # Arguments
    2377
        /// * `kernel` - (kernel_height, kernel_width) pooling window size
    2378
        /// * `stride` - (stride_height, stride_width) step size
    2379
        ///
    2380
        /// # Examples
    2381
        /// ```
    2382
        /// use trueno::matrix::Matrix;
    2383
        /// let input = Matrix::from_vec(4, 4, vec![
    2384
        ///     1.0, 2.0, 3.0, 4.0,
    2385
        ///     5.0, 6.0, 7.0, 8.0,
    2386
        ///     9.0, 10.0, 11.0, 12.0,
    2387
        ///     13.0, 14.0, 15.0, 16.0,
    2388
        /// ]).unwrap();
    2389
        /// let pooled = input.max_pool2d((2, 2), (2, 2)).unwrap();
    2390
        /// assert_eq!(pooled.shape(), (2, 2));
    2391
        /// assert_eq!(pooled.get(0, 0), Some(&6.0));  // max of [1,2,5,6]
    2392
        /// assert_eq!(pooled.get(1, 1), Some(&16.0)); // max of [11,12,15,16]
    2393
        /// ```
    2394
    0
        pub fn max_pool2d(
    2395
    0
            &self,
    2396
    0
            kernel: (usize, usize),
    2397
    0
            stride: (usize, usize),
    2398
    0
        ) -> Result<Matrix<f32>, TruenoError> {
    2399
    0
            let (kh, kw) = kernel;
    2400
    0
            let (sh, sw) = stride;
    2401
    2402
    0
            if kh == 0 || kw == 0 || sh == 0 || sw == 0 {
    2403
    0
                return Err(TruenoError::InvalidInput(
    2404
    0
                    "Kernel and stride dimensions must be positive".into(),
    2405
    0
                ));
    2406
    0
            }
    2407
    2408
    0
            if kh > self.rows || kw > self.cols {
    2409
    0
                return Err(TruenoError::InvalidInput(format!(
    2410
    0
                    "Kernel size ({}, {}) larger than input ({}, {})",
    2411
    0
                    kh, kw, self.rows, self.cols
    2412
    0
                )));
    2413
    0
            }
    2414
    2415
    0
            let out_h = (self.rows - kh) / sh + 1;
    2416
    0
            let out_w = (self.cols - kw) / sw + 1;
    2417
    0
            let mut result = Matrix::new(out_h, out_w);
    2418
    2419
    0
            for i in 0..out_h {
    2420
    0
                for j in 0..out_w {
    2421
    0
                    let mut max_val = f32::NEG_INFINITY;
    2422
    0
                    for ki in 0..kh {
    2423
    0
                        for kj in 0..kw {
    2424
    0
                            let val = self.data[(i * sh + ki) * self.cols + (j * sw + kj)];
    2425
    0
                            max_val = max_val.max(val);
    2426
    0
                        }
    2427
                    }
    2428
    0
                    result.data[i * out_w + j] = max_val;
    2429
                }
    2430
            }
    2431
    2432
    0
            Ok(result)
    2433
    0
        }
    2434
    2435
        /// 2D Average Pooling operation for CNN downsampling
    2436
        ///
    2437
        /// Applies average pooling over a 2D input tensor with specified kernel size and stride.
    2438
        /// Input shape: (height, width), Output shape: ((height - kh) / sh + 1, (width - kw) / sw + 1)
    2439
        ///
    2440
        /// # Arguments
    2441
        /// * `kernel` - (kernel_height, kernel_width) pooling window size
    2442
        /// * `stride` - (stride_height, stride_width) step size
    2443
        ///
    2444
        /// # Examples
    2445
        /// ```
    2446
        /// use trueno::matrix::Matrix;
    2447
        /// let input = Matrix::from_vec(4, 4, vec![
    2448
        ///     1.0, 2.0, 3.0, 4.0,
    2449
        ///     5.0, 6.0, 7.0, 8.0,
    2450
        ///     9.0, 10.0, 11.0, 12.0,
    2451
        ///     13.0, 14.0, 15.0, 16.0,
    2452
        /// ]).unwrap();
    2453
        /// let pooled = input.avg_pool2d((2, 2), (2, 2)).unwrap();
    2454
        /// assert_eq!(pooled.shape(), (2, 2));
    2455
        /// assert!((pooled.get(0, 0).unwrap() - 3.5).abs() < 1e-5);  // avg of [1,2,5,6]
    2456
        /// ```
    2457
    0
        pub fn avg_pool2d(
    2458
    0
            &self,
    2459
    0
            kernel: (usize, usize),
    2460
    0
            stride: (usize, usize),
    2461
    0
        ) -> Result<Matrix<f32>, TruenoError> {
    2462
    0
            let (kh, kw) = kernel;
    2463
    0
            let (sh, sw) = stride;
    2464
    2465
    0
            if kh == 0 || kw == 0 || sh == 0 || sw == 0 {
    2466
    0
                return Err(TruenoError::InvalidInput(
    2467
    0
                    "Kernel and stride dimensions must be positive".into(),
    2468
    0
                ));
    2469
    0
            }
    2470
    2471
    0
            if kh > self.rows || kw > self.cols {
    2472
    0
                return Err(TruenoError::InvalidInput(format!(
    2473
    0
                    "Kernel size ({}, {}) larger than input ({}, {})",
    2474
    0
                    kh, kw, self.rows, self.cols
    2475
    0
                )));
    2476
    0
            }
    2477
    2478
    0
            let out_h = (self.rows - kh) / sh + 1;
    2479
    0
            let out_w = (self.cols - kw) / sw + 1;
    2480
    0
            let kernel_size = (kh * kw) as f32;
    2481
    0
            let mut result = Matrix::new(out_h, out_w);
    2482
    2483
    0
            for i in 0..out_h {
    2484
    0
                for j in 0..out_w {
    2485
    0
                    let mut sum = 0.0;
    2486
    0
                    for ki in 0..kh {
    2487
    0
                        for kj in 0..kw {
    2488
    0
                            sum += self.data[(i * sh + ki) * self.cols + (j * sw + kj)];
    2489
    0
                        }
    2490
                    }
    2491
    0
                    result.data[i * out_w + j] = sum / kernel_size;
    2492
                }
    2493
            }
    2494
    2495
    0
            Ok(result)
    2496
    0
        }
    2497
    2498
        /// Top-K selection: returns the k largest elements and their indices
    2499
        ///
    2500
        /// Useful for beam search, sampling, and ranking operations.
    2501
        /// Searches row-major order and returns (values, indices) sorted descending.
    2502
        ///
    2503
        /// # Examples
    2504
        /// ```
    2505
        /// use trueno::matrix::Matrix;
    2506
        /// let m = Matrix::from_vec(2, 3, vec![1.0, 5.0, 3.0, 2.0, 6.0, 4.0]).unwrap();
    2507
        /// let (values, indices) = m.topk(2).unwrap();
    2508
        /// assert_eq!(values, vec![6.0, 5.0]);
    2509
        /// assert_eq!(indices, vec![4, 1]);  // flat indices
    2510
        /// ```
    2511
    0
        pub fn topk(&self, k: usize) -> Result<(Vec<f32>, Vec<usize>), TruenoError> {
    2512
    0
            if k == 0 {
    2513
    0
                return Ok((vec![], vec![]));
    2514
    0
            }
    2515
    2516
    0
            let k = k.min(self.data.len());
    2517
    0
            let mut indexed: Vec<(usize, f32)> = self.data.iter().copied().enumerate().collect();
    2518
    2519
            // Partial sort - only sort k elements
    2520
    0
            indexed.select_nth_unstable_by(k.saturating_sub(1), |a, b| {
    2521
    0
                b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal)
    2522
    0
            });
    2523
    2524
    0
            indexed.truncate(k);
    2525
    0
            indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
    2526
    2527
    0
            let values: Vec<f32> = indexed.iter().map(|(_, v)| *v).collect();
    2528
    0
            let indices: Vec<usize> = indexed.iter().map(|(i, _)| *i).collect();
    2529
    2530
    0
            Ok((values, indices))
    2531
    0
        }
    2532
    2533
        /// Gather elements along axis using indices
    2534
        ///
    2535
        /// For 2D matrix with axis=0: output[i] = self[indices[i], :]
    2536
        /// For 2D matrix with axis=1: output[:, i] = self[:, indices[i]]
    2537
        ///
    2538
        /// # Examples
    2539
        /// ```
    2540
        /// use trueno::matrix::Matrix;
    2541
        /// let m = Matrix::from_vec(3, 2, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    2542
        /// let gathered = m.gather(&[2, 0], 0).unwrap();  // Select rows 2 and 0
    2543
        /// assert_eq!(gathered.shape(), (2, 2));
    2544
        /// assert_eq!(gathered.get(0, 0), Some(&5.0));  // Row 2
    2545
        /// assert_eq!(gathered.get(1, 0), Some(&1.0));  // Row 0
    2546
        /// ```
    2547
    0
        pub fn gather(&self, indices: &[usize], axis: usize) -> Result<Matrix<f32>, TruenoError> {
    2548
    0
            match axis {
    2549
                0 => {
    2550
                    // Gather rows
    2551
    0
                    let mut result = Matrix::new(indices.len(), self.cols);
    2552
    0
                    for (out_i, &idx) in indices.iter().enumerate() {
    2553
    0
                        if idx >= self.rows {
    2554
    0
                            return Err(TruenoError::InvalidInput(format!(
    2555
    0
                                "Index {} out of bounds for axis 0 with size {}",
    2556
    0
                                idx, self.rows
    2557
    0
                            )));
    2558
    0
                        }
    2559
    0
                        for j in 0..self.cols {
    2560
    0
                            result.data[out_i * self.cols + j] = self.data[idx * self.cols + j];
    2561
    0
                        }
    2562
                    }
    2563
    0
                    Ok(result)
    2564
                }
    2565
                1 => {
    2566
                    // Gather columns
    2567
    0
                    let mut result = Matrix::new(self.rows, indices.len());
    2568
    0
                    for i in 0..self.rows {
    2569
    0
                        for (out_j, &idx) in indices.iter().enumerate() {
    2570
    0
                            if idx >= self.cols {
    2571
    0
                                return Err(TruenoError::InvalidInput(format!(
    2572
    0
                                    "Index {} out of bounds for axis 1 with size {}",
    2573
    0
                                    idx, self.cols
    2574
    0
                                )));
    2575
    0
                            }
    2576
    0
                            result.data[i * indices.len() + out_j] = self.data[i * self.cols + idx];
    2577
                        }
    2578
                    }
    2579
    0
                    Ok(result)
    2580
                }
    2581
    0
                _ => Err(TruenoError::InvalidInput(format!(
    2582
    0
                    "Axis {} not supported for 2D matrix (use 0 or 1)",
    2583
    0
                    axis
    2584
    0
                ))),
    2585
            }
    2586
    0
        }
    2587
    2588
        /// Pad matrix with a constant value
    2589
        ///
    2590
        /// # Arguments
    2591
        /// * `padding` - ((top, bottom), (left, right)) padding amounts
    2592
        /// * `value` - constant value to pad with (usually 0.0)
    2593
        ///
    2594
        /// # Examples
    2595
        /// ```
    2596
        /// use trueno::matrix::Matrix;
    2597
        /// let m = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]).unwrap();
    2598
        /// let padded = m.pad(((1, 1), (1, 1)), 0.0).unwrap();
    2599
        /// assert_eq!(padded.shape(), (4, 4));
    2600
        /// assert_eq!(padded.get(0, 0), Some(&0.0));  // top-left padding
    2601
        /// assert_eq!(padded.get(1, 1), Some(&1.0));  // original (0,0)
    2602
        /// ```
    2603
    0
        pub fn pad(
    2604
    0
            &self,
    2605
    0
            padding: ((usize, usize), (usize, usize)),
    2606
    0
            value: f32,
    2607
    0
        ) -> Result<Matrix<f32>, TruenoError> {
    2608
    0
            let ((top, bottom), (left, right)) = padding;
    2609
    0
            let new_rows = self.rows + top + bottom;
    2610
    0
            let new_cols = self.cols + left + right;
    2611
    2612
    0
            let mut result = Matrix::from_vec(new_rows, new_cols, vec![value; new_rows * new_cols])?;
    2613
    2614
            // Copy original data
    2615
    0
            for i in 0..self.rows {
    2616
    0
                for j in 0..self.cols {
    2617
    0
                    result.data[(i + top) * new_cols + (j + left)] = self.data[i * self.cols + j];
    2618
    0
                }
    2619
            }
    2620
    2621
    0
            Ok(result)
    2622
    0
        }
    2623
    }
    2624
    2625
    2626
    // Tests (~2.6K lines extracted for TDG compliance)
    2627
    #[cfg(test)]
    2628
    mod tests;