trueno/

matrix.rs

//! Matrix operations for Trueno
//!
//! Provides 2D matrix operations with SIMD optimization for linear algebra,
//! machine learning, and scientific computing.
//!
//! # Example
//!
//! ```
//! use trueno::Matrix;
//!
//! // Create a 2x3 matrix
//! let m = Matrix::zeros(2, 3);
//! assert_eq!(m.rows(), 2);
//! assert_eq!(m.cols(), 3);
//! ```

use crate::{Backend, TruenoError, Vector};

#[cfg(feature = "tracing")]
use tracing::instrument;

/// A 2D matrix with row-major storage
///
/// Data is stored in row-major format (C-style), where consecutive elements
/// in memory belong to the same row. This is compatible with NumPy's default
/// layout and optimal for cache locality when accessing rows.
///
/// # Storage Layout
///
/// For a 2x3 matrix:
/// ```text
/// [[a, b, c],
///  [d, e, f]]
/// ```
/// Data is stored as: [a, b, c, d, e, f]
///
/// # Example
///
/// ```
/// use trueno::Matrix;
///
/// let m = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]).unwrap();
/// assert_eq!(m.get(0, 0), Some(&1.0));
/// assert_eq!(m.get(0, 1), Some(&2.0));
/// assert_eq!(m.get(1, 0), Some(&3.0));
/// assert_eq!(m.get(1, 1), Some(&4.0));
/// ```
#[derive(Debug, Clone, PartialEq)]
pub struct Matrix<T> {
    rows: usize,
    cols: usize,
    data: Vec<T>,
    backend: Backend,
}

impl Matrix<f32> {
    /// Creates a new matrix with uninitialized values
    ///
    /// # Arguments
    ///
    /// * `rows` - Number of rows
    /// * `cols` - Number of columns
    ///
    /// # Returns
    ///
    /// A new matrix with dimensions `rows x cols` containing uninitialized values
    ///
    /// # Example
    ///
    /// ```
    /// use trueno::Matrix;
    ///
    /// let m = Matrix::new(3, 4);
    /// assert_eq!(m.rows(), 3);
    /// assert_eq!(m.cols(), 4);
    /// ```
    pub fn new(rows: usize, cols: usize) -> Self {
        let backend = Backend::select_best();
        Matrix {
            rows,
            cols,
            data: vec![0.0; rows * cols],
            backend,
        }
    }

    /// Creates a matrix from a vector of data
    ///
    /// # Arguments
    ///
    /// * `rows` - Number of rows
    /// * `cols` - Number of columns
    /// * `data` - Vector containing matrix elements in row-major order
    ///
    /// # Errors
    ///
    /// Returns `InvalidInput` if `data.len() != rows * cols`
    ///
    /// # Example
    ///
    /// ```
    /// use trueno::Matrix;
    ///
    /// let m = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]).unwrap();
    /// assert_eq!(m.rows(), 2);
    /// assert_eq!(m.cols(), 2);
    /// ```
    pub fn from_vec(rows: usize, cols: usize, data: Vec<f32>) -> Result<Self, TruenoError> {
        if data.len() != rows * cols {
            return Err(TruenoError::InvalidInput(format!(
                "Data length {} does not match matrix dimensions {}x{} (expected {})",
                data.len(),
                rows,
                cols,
                rows * cols
            )));
        }

        let backend = Backend::select_best();
        Ok(Matrix {
            rows,
            cols,
            data,
            backend,
        })
    }

    /// Creates a matrix from a slice by copying the data
    ///
    /// This is a convenience method that copies the slice into an owned vector.
    /// For zero-copy scenarios, consider using the data directly with `from_vec`
    /// if you already have an owned `Vec`.
    ///
    /// # Arguments
    ///
    /// * `rows` - Number of rows
    /// * `cols` - Number of columns
    /// * `data` - Slice containing matrix elements in row-major order
    ///
    /// # Errors
    ///
    /// Returns `InvalidInput` if `data.len() != rows * cols`
    ///
    /// # Example
    ///
    /// ```
    /// use trueno::Matrix;
    ///
    /// let data = [1.0, 2.0, 3.0, 4.0];
    /// let m = Matrix::from_slice(2, 2, &data).unwrap();
    /// assert_eq!(m.get(0, 0), Some(&1.0));
    /// ```
    pub fn from_slice(rows: usize, cols: usize, data: &[f32]) -> Result<Self, TruenoError> {
        Self::from_vec(rows, cols, data.to_vec())
    }

    /// Creates a matrix filled with zeros
    ///
    /// # Example
    ///
    /// ```
    /// use trueno::Matrix;
    ///
    /// let m = Matrix::zeros(3, 3);
    /// assert_eq!(m.get(1, 1), Some(&0.0));
    /// ```
    pub fn zeros(rows: usize, cols: usize) -> Self {
        Matrix::new(rows, cols)
    }

    /// Creates a matrix filled with zeros using a specific backend
    /// (Internal use only - reuses backend from parent matrix)
    fn zeros_with_backend(rows: usize, cols: usize, backend: Backend) -> Self {
        Matrix {
            rows,
            cols,
            data: vec![0.0; rows * cols],
            backend,
        }
    }

    /// Creates an identity matrix (square matrix with 1s on diagonal)
    ///
    /// # Example
    ///
    /// ```
    /// use trueno::Matrix;
    ///
    /// let m = Matrix::identity(3);
    /// assert_eq!(m.get(0, 0), Some(&1.0));
    /// assert_eq!(m.get(0, 1), Some(&0.0));
    /// assert_eq!(m.get(1, 1), Some(&1.0));
    /// ```
    pub fn identity(n: usize) -> Self {
        let mut data = vec![0.0; n * n];
        for i in 0..n {
            data[i * n + i] = 1.0;
        }
        let backend = Backend::select_best();
        Matrix {
            rows: n,
            cols: n,
            data,
            backend,
        }
    }

    /// Returns the number of rows
    pub fn rows(&self) -> usize {
        self.rows
    }

    /// Returns the number of columns
    pub fn cols(&self) -> usize {
        self.cols
    }

    /// Returns the shape as (rows, cols)
    pub fn shape(&self) -> (usize, usize) {
        (self.rows, self.cols)
    }

    /// Gets a reference to an element at (row, col)
    ///
    /// Returns `None` if indices are out of bounds
    pub fn get(&self, row: usize, col: usize) -> Option<&f32> {
        if row >= self.rows || col >= self.cols {
            None
        } else {
            self.data.get(row * self.cols + col)
        }
    }

    /// Gets a mutable reference to an element at (row, col)
    ///
    /// Returns `None` if indices are out of bounds
    pub fn get_mut(&mut self, row: usize, col: usize) -> Option<&mut f32> {
        if row >= self.rows || col >= self.cols {
            None
        } else {
            let idx = row * self.cols + col;
            self.data.get_mut(idx)
        }
    }

    /// Returns a reference to the underlying data
    pub fn as_slice(&self) -> &[f32] {
        &self.data
    }

    /// Matrix multiplication (matmul)
    ///
    /// Computes `C = A × B` where A is `m×n`, B is `n×p`, and C is `m×p`.
    ///
    /// # Arguments
    ///
    /// * `other` - The matrix to multiply with (right operand)
    ///
    /// # Returns
    ///
    /// A new matrix containing the result of matrix multiplication
    ///
    /// # Errors
    ///
    /// Returns `InvalidInput` if matrix dimensions are incompatible
    /// (i.e., `self.cols != other.rows`)
    ///
    /// # Example
    ///
    /// ```
    /// use trueno::Matrix;
    ///
    /// let a = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]).unwrap();
    /// let b = Matrix::from_vec(2, 2, vec![5.0, 6.0, 7.0, 8.0]).unwrap();
    /// let c = a.matmul(&b).unwrap();
    ///
    /// // [[1, 2],   [[5, 6],   [[19, 22],
    /// //  [3, 4]] ×  [7, 8]] =  [43, 50]]
    /// assert_eq!(c.get(0, 0), Some(&19.0));
    /// assert_eq!(c.get(0, 1), Some(&22.0));
    /// assert_eq!(c.get(1, 0), Some(&43.0));
    /// assert_eq!(c.get(1, 1), Some(&50.0));
    /// ```
    #[cfg_attr(feature = "tracing", instrument(skip(self, other), fields(dims = %format!("{}x{} @ {}x{}", self.rows, self.cols, other.rows, other.cols))))]
    pub fn matmul(&self, other: &Matrix<f32>) -> Result<Matrix<f32>, TruenoError> {
        if self.cols != other.rows {
            return Err(TruenoError::InvalidInput(format!(
                "Matrix dimension mismatch for multiplication: {}×{} × {}×{} (inner dimensions {} and {} must match)",
                self.rows, self.cols, other.rows, other.cols, self.cols, other.rows
            )));
        }

        // Fast path for vector-matrix multiply (rows=1)
        // Common in ML vocab projection: hidden_state @ embedding_transposed
        // 8x faster than general matmul for 1×384 @ 384×51865 pattern
        if self.rows == 1 {
            return self.matmul_vector_matrix(other);
        }

        let mut result = Matrix::zeros_with_backend(self.rows, other.cols, self.backend);

        // Backend selection strategy (empirical - see docs/performance-analysis.md):
        // 1. GPU for large matrices (≥500×500) - 2-10x speedup (measured)
        // 2. SIMD for medium-large matrices (>64×64) - 2-8x speedup
        // 3. Naive for small matrices - lowest overhead

        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
        const GPU_THRESHOLD: usize = 500; // Empirical: 2x at 500×500, 9.6x at 1000×1000
        const SIMD_THRESHOLD: usize = 64;

        // Try GPU first for very large matrices
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
        {
            if self.rows >= GPU_THRESHOLD
                && self.cols >= GPU_THRESHOLD
                && other.cols >= GPU_THRESHOLD
            {
                if let Ok(gpu_result) = self.matmul_gpu(other) {
                    return Ok(gpu_result);
                }
                // GPU failed, fall through to SIMD/naive
            }
        }

        // Use SIMD for medium-large matrices
        if self.rows >= SIMD_THRESHOLD
            || self.cols >= SIMD_THRESHOLD
            || other.cols >= SIMD_THRESHOLD
        {
            // Tiled approach threshold: below this size, tiling beats transpose
            // Based on WASM optimization spec benchmarks
            const TILED_THRESHOLD: usize = 512;

            let max_dim = self.rows.max(self.cols).max(other.cols);

            if max_dim < TILED_THRESHOLD {
                // Medium matrices: use tiled approach (no transpose overhead)
                // Works well for both WASM and native for matrices up to ~512
                self.matmul_wasm_tiled(other, &mut result)?;
            } else {
                // Large matrices: platform-specific optimized paths
                #[cfg(target_arch = "wasm32")]
                {
                    // WASM: tiled is always better (no SIMD microkernel advantage)
                    self.matmul_wasm_tiled(other, &mut result)?;
                }
                #[cfg(not(target_arch = "wasm32"))]
                {
                    // Native: use AVX2/NEON SIMD with cache blocking
                    self.matmul_simd(other, &mut result)?;
                }
            }
        } else {
            self.matmul_naive(other, &mut result)?;
        }

        Ok(result)
    }

    /// Batched matrix multiplication for 3D tensors.
    ///
    /// Computes `[batch, m, k] @ [batch, k, n] -> [batch, m, n]` using SIMD for each batch.
    /// This is critical for transformer attention performance.
    ///
    /// # Arguments
    /// * `a_data` - Flattened input A with shape [batch, m, k]
    /// * `b_data` - Flattened input B with shape [batch, k, n]
    /// * `batch` - Batch dimension
    /// * `m` - Rows of A (and output)
    /// * `k` - Columns of A / Rows of B
    /// * `n` - Columns of B (and output)
    ///
    /// # Returns
    /// Flattened output with shape [batch, m, n]
    ///
    /// # Performance
    /// Uses SIMD matmul for each batch slice, achieving ~50 GFLOPS vs ~0.1 GFLOPS naive.
    /// See Williams et al., 2009 (Roofline model) for theoretical analysis.
    #[cfg_attr(feature = "tracing", instrument(skip(a_data, b_data), fields(batch, m, k, n)))]
    pub fn batched_matmul(
        a_data: &[f32],
        b_data: &[f32],
        batch: usize,
        m: usize,
        k: usize,
        n: usize,
    ) -> Result<Vec<f32>, TruenoError> {
        let a_stride = m * k;
        let b_stride = k * n;
        let out_stride = m * n;

        // Validate input sizes
        if a_data.len() != batch * a_stride {
            return Err(TruenoError::InvalidInput(format!(
                "A data size mismatch: expected {} ({}×{}×{}), got {}",
                batch * a_stride, batch, m, k, a_data.len()
            )));
        }
        if b_data.len() != batch * b_stride {
            return Err(TruenoError::InvalidInput(format!(
                "B data size mismatch: expected {} ({}×{}×{}), got {}",
                batch * b_stride, batch, k, n, b_data.len()
            )));
        }

        let mut output = vec![0.0f32; batch * out_stride];

        // Process each batch using SIMD matmul
        for ba in 0..batch {
            let a_offset = ba * a_stride;
            let b_offset = ba * b_stride;
            let out_offset = ba * out_stride;

            // Create matrix views from slices (no copy - just metadata)
            let a_slice = &a_data[a_offset..a_offset + a_stride];
            let b_slice = &b_data[b_offset..b_offset + b_stride];

            // Use from_slice to avoid copying
            let a_mat = Matrix::from_slice(m, k, a_slice)?;
            let b_mat = Matrix::from_slice(k, n, b_slice)?;

            // SIMD matmul
            let result = a_mat.matmul(&b_mat)?;

            // Copy result to output
            output[out_offset..out_offset + out_stride].copy_from_slice(result.as_slice());
        }

        Ok(output)
    }

    /// Batched matrix multiplication for 4D tensors (attention pattern).
    ///
    /// Computes `[batch, heads, m, k] @ [batch, heads, k, n] -> [batch, heads, m, n]`
    /// This is the exact pattern used in multi-head attention: Q @ K^T and attn @ V.
    ///
    /// # Arguments
    /// * `a_data` - Flattened input A with shape [batch, heads, m, k]
    /// * `b_data` - Flattened input B with shape [batch, heads, k, n]
    /// * `batch` - Batch dimension
    /// * `heads` - Number of attention heads
    /// * `m` - Rows (sequence length for Q)
    /// * `k` - Columns of A / Rows of B (head dimension)
    /// * `n` - Columns of B (sequence length for K^T, or head dim for V)
    ///
    /// # Performance
    /// Processes batch×heads independent matmuls using SIMD.
    /// For Qwen2-0.5B: batch=1, heads=14, m=seq, k=64, n=seq
    #[cfg_attr(feature = "tracing", instrument(skip(a_data, b_data), fields(batch, heads, m, k, n)))]
    pub fn batched_matmul_4d(
        a_data: &[f32],
        b_data: &[f32],
        batch: usize,
        heads: usize,
        m: usize,
        k: usize,
        n: usize,
    ) -> Result<Vec<f32>, TruenoError> {
        let a_head_stride = m * k;
        let b_head_stride = k * n;
        let out_head_stride = m * n;
        let total_heads = batch * heads;

        // Validate input sizes
        let expected_a = total_heads * a_head_stride;
        let expected_b = total_heads * b_head_stride;
        if a_data.len() != expected_a {
            return Err(TruenoError::InvalidInput(format!(
                "A data size mismatch: expected {} ({}×{}×{}×{}), got {}",
                expected_a, batch, heads, m, k, a_data.len()
            )));
        }
        if b_data.len() != expected_b {
            return Err(TruenoError::InvalidInput(format!(
                "B data size mismatch: expected {} ({}×{}×{}×{}), got {}",
                expected_b, batch, heads, k, n, b_data.len()
            )));
        }

        let mut output = vec![0.0f32; total_heads * out_head_stride];

        // Process each (batch, head) pair using SIMD matmul
        for bh in 0..total_heads {
            let a_offset = bh * a_head_stride;
            let b_offset = bh * b_head_stride;
            let out_offset = bh * out_head_stride;

            // Create matrix views from slices
            let a_slice = &a_data[a_offset..a_offset + a_head_stride];
            let b_slice = &b_data[b_offset..b_offset + b_head_stride];

            let a_mat = Matrix::from_slice(m, k, a_slice)?;
            let b_mat = Matrix::from_slice(k, n, b_slice)?;

            // SIMD matmul
            let result = a_mat.matmul(&b_mat)?;

            // Copy result to output
            output[out_offset..out_offset + out_head_stride].copy_from_slice(result.as_slice());
        }

        Ok(output)
    }

    /// Fast path for vector-matrix multiplication (1×K @ K×N → 1×N)
    ///
    /// This is 8x faster than general matmul for patterns like:
    /// - Vocab projection: hidden_state (1×384) @ embedding_transposed (384×51865)
    /// - Single token decode in Whisper/LLM inference
    ///
    /// Strategy: Outer product accumulation (no transpose needed!)
    /// For result[j] = sum_k(A[0,k] * B[k,j]), we compute:
    ///   result += A[k] * B[k,:]  for each k
    /// This has excellent cache locality since we access entire rows of B.
    #[cfg_attr(feature = "tracing", instrument(skip(self, other), fields(k = self.cols, n = other.cols)))]
    fn matmul_vector_matrix(&self, other: &Matrix<f32>) -> Result<Matrix<f32>, TruenoError> {
        debug_assert_eq!(self.rows, 1);

        let k = self.cols; // Inner dimension
        let n = other.cols; // Output dimension

        // Result is 1×N, initialized to zero
        let mut result = Matrix::zeros_with_backend(1, n, self.backend);

        // Outer product accumulation: result += A[k] * B[k,:]
        // For each k, scale row k of B by A[k] and add to result
        // The compiler will auto-vectorize this inner loop
        for ki in 0..k {
            let a_k = self.data[ki];
            if a_k == 0.0 {
                continue; // Skip zero multiplications
            }

            // Get row ki of B (contiguous in memory - cache friendly!)
            let b_row_start = ki * n;

            // AXPY: result += a_k * B[ki,:]
            // This loop is auto-vectorized by LLVM with -O2/-O3
            for j in 0..n {
                result.data[j] += a_k * other.data[b_row_start + j];
            }
        }

        Ok(result)
    }

    /// Naive O(n³) matrix multiplication (baseline for small matrices)
    fn matmul_naive(
        &self,
        other: &Matrix<f32>,
        result: &mut Matrix<f32>,
    ) -> Result<(), TruenoError> {
        // C[i,j] = Σ A[i,k] × B[k,j]
        // SAFETY: Loop bounds are validated by dimension checks in matmul()
        for i in 0..self.rows {
            for j in 0..other.cols {
                let mut sum = 0.0;
                for k in 0..self.cols {
                    // Bounds guaranteed: i < self.rows, k < self.cols, j < other.cols
                    sum += self
                        .get(i, k)
                        .expect("matmul_naive: A[i,k] bounds validated by loop")
                        * other
                            .get(k, j)
                            .expect("matmul_naive: B[k,j] bounds validated by loop");
                }
                *result
                    .get_mut(i, j)
                    .expect("matmul_naive: C[i,j] bounds validated by loop") = sum;
            }
        }
        Ok(())
    }

    /// AVX2 micro-kernel: Compute 4 rows × 1 column using register blocking (Phase 2)
    ///
    /// This micro-kernel processes 4 rows of matrix A against 1 column of B_transposed
    /// simultaneously, keeping intermediate results in AVX2 registers for efficiency.
    ///
    /// # Performance Benefits
    /// - Loads B-column once, reuses for 4 A-rows (4× reduction in memory bandwidth)
    /// - Uses FMA instructions for fused multiply-add (3× throughput vs separate ops)
    /// - Keeps accumulators in YMM registers (no memory traffic for intermediate results)
    ///
    /// # Safety
    /// - Caller must ensure all slices have the same length
    /// - Must be called on x86_64 with AVX2 support
    #[cfg(target_arch = "x86_64")]
    #[target_feature(enable = "avx2,fma")]
    #[inline]
    unsafe fn matmul_microkernel_4x1_avx2(
        a_rows: [&[f32]; 4],
        b_col: &[f32],
        results: &mut [f32; 4],
    ) {
        use std::arch::x86_64::*;

        let len = b_col.len();
        let chunks = len / 8; // Process 8 f32 elements per iteration (AVX2 = 256 bits)

        // Accumulators for 4 output elements (kept in registers)
        let mut acc0 = _mm256_setzero_ps();
        let mut acc1 = _mm256_setzero_ps();
        let mut acc2 = _mm256_setzero_ps();
        let mut acc3 = _mm256_setzero_ps();

        // Main loop: Process 8 elements at a time
        for i in 0..chunks {
            let offset = i * 8;

            // Load B column (reused for all 4 A rows)
            let b_vec = _mm256_loadu_ps(b_col.as_ptr().add(offset));

            // Load A rows and FMA (Fused Multiply-Add)
            let a0_vec = _mm256_loadu_ps(a_rows[0].as_ptr().add(offset));
            acc0 = _mm256_fmadd_ps(a0_vec, b_vec, acc0);

            let a1_vec = _mm256_loadu_ps(a_rows[1].as_ptr().add(offset));
            acc1 = _mm256_fmadd_ps(a1_vec, b_vec, acc1);

            let a2_vec = _mm256_loadu_ps(a_rows[2].as_ptr().add(offset));
            acc2 = _mm256_fmadd_ps(a2_vec, b_vec, acc2);

            let a3_vec = _mm256_loadu_ps(a_rows[3].as_ptr().add(offset));
            acc3 = _mm256_fmadd_ps(a3_vec, b_vec, acc3);
        }

        // Horizontal sum of each accumulator (reduce 8 elements to 1)
        results[0] = Self::horizontal_sum_avx2(acc0);
        results[1] = Self::horizontal_sum_avx2(acc1);
        results[2] = Self::horizontal_sum_avx2(acc2);
        results[3] = Self::horizontal_sum_avx2(acc3);

        // Handle remainder elements with scalar code
        let remainder_start = chunks * 8;
        if remainder_start < len {
            for i in remainder_start..len {
                results[0] += a_rows[0][i] * b_col[i];
                results[1] += a_rows[1][i] * b_col[i];
                results[2] += a_rows[2][i] * b_col[i];
                results[3] += a_rows[3][i] * b_col[i];
            }
        }
    }

    /// Helper: Horizontal sum of 8 f32 values in an AVX2 register
    #[cfg(target_arch = "x86_64")]
    #[target_feature(enable = "avx2")]
    #[inline]
    unsafe fn horizontal_sum_avx2(v: std::arch::x86_64::__m256) -> f32 {
        use std::arch::x86_64::*;

        // Sum upper and lower 128-bit lanes
        let sum128 = _mm_add_ps(_mm256_castps256_ps128(v), _mm256_extractf128_ps(v, 1));

        // Horizontal add within 128-bit lane (4 values → 2 values)
        let sum64 = _mm_hadd_ps(sum128, sum128);

        // Horizontal add again (2 values → 1 value)
        let sum32 = _mm_hadd_ps(sum64, sum64);

        // Extract final scalar result
        _mm_cvtss_f32(sum32)
    }

    /// Cache-aware blocked matrix multiplication with SIMD optimization
    ///
    /// Uses 2-level cache blocking (L2/L1) to minimize cache misses:
    /// - L2 blocks: 64×64 (256KB for 3 matrices in f32)
    /// - L1 micro-kernels: 8×8 (768 bytes fits comfortably in L1)
    ///
    /// Performance characteristics:
    /// - Small matrices (<64×64): ~1.2× speedup over naive (overhead dominates)
    /// - Medium matrices (128×128): ~1.5-2× speedup (cache effects visible)
    /// - Large matrices (512×512+): ~2-3× speedup (dramatic cache improvement)
    ///
    /// This is Phase 1 of matmul optimization (Issue #10). Future Phase 2 will
    /// add optional BLAS backend for full NumPy parity on very large matrices.
    /// Helper function to process a single L3 row block for parallel matmul (Phase 4).
    ///
    /// # Safety
    /// When called from parallel code, the caller must ensure that each thread processes
    /// a distinct row range [iii, i3_end) with no overlap. This function is safe because
    /// each thread writes only to its own row range in the result matrix.
    #[cfg(feature = "parallel")]
    #[allow(clippy::too_many_arguments)]
    fn process_l3_row_block_seq(
        iii: usize,
        i3_end: usize,
        a: &Matrix<f32>,
        b_transposed: &Matrix<f32>,
        result: &mut Matrix<f32>,
        l2_block_size: usize,
        l3_block_size: usize,
    ) {
        #[cfg(target_arch = "x86_64")]
        use crate::backends::{avx2::Avx2Backend, sse2::Sse2Backend};
        use crate::backends::{scalar::ScalarBackend, VectorBackend};

        // Process all column blocks for this row block
        for jjj in (0..b_transposed.rows).step_by(l3_block_size) {
            let j3_end = (jjj + l3_block_size).min(b_transposed.rows);

            for kkk in (0..a.cols).step_by(l3_block_size) {
                let k3_end = (kkk + l3_block_size).min(a.cols);

                // L2 blocking within L3 blocks
                for ii in (iii..i3_end).step_by(l2_block_size) {
                    let i_end = (ii + l2_block_size).min(i3_end);

                    for jj in (jjj..j3_end).step_by(l2_block_size) {
                        let j_end = (jj + l2_block_size).min(j3_end);

                        for kk in (kkk..k3_end).step_by(l2_block_size) {
                            let k_end = (kk + l2_block_size).min(k3_end);
                            let block_size = k_end - kk;

                            // Micro-kernel processing
                            #[cfg(target_arch = "x86_64")]
                            let use_microkernel =
                                matches!(a.backend, Backend::AVX2 | Backend::AVX512);

                            #[cfg(target_arch = "x86_64")]
                            if use_microkernel {
                                let mut i = ii;

                                // Process 4 rows at a time with micro-kernel
                                while i + 4 <= i_end {
                                    let row0_start = i * a.cols + kk;
                                    let row1_start = (i + 1) * a.cols + kk;
                                    let row2_start = (i + 2) * a.cols + kk;
                                    let row3_start = (i + 3) * a.cols + kk;

                                    let a_rows = [
                                        &a.data[row0_start..row0_start + block_size],
                                        &a.data[row1_start..row1_start + block_size],
                                        &a.data[row2_start..row2_start + block_size],
                                        &a.data[row3_start..row3_start + block_size],
                                    ];

                                    for j in jj..j_end {
                                        let col_start = j * b_transposed.cols + kk;
                                        let b_col =
                                            &b_transposed.data[col_start..col_start + block_size];

                                        let mut partial_dots = [0.0f32; 4];
                                        unsafe {
                                            Matrix::matmul_microkernel_4x1_avx2(
                                                a_rows,
                                                b_col,
                                                &mut partial_dots,
                                            );
                                        }

                                        result.data[i * result.cols + j] += partial_dots[0];
                                        result.data[(i + 1) * result.cols + j] += partial_dots[1];
                                        result.data[(i + 2) * result.cols + j] += partial_dots[2];
                                        result.data[(i + 3) * result.cols + j] += partial_dots[3];
                                    }

                                    i += 4;
                                }

                                // Handle remaining rows (< 4)
                                for i in i..i_end {
                                    let row_start = i * a.cols + kk;
                                    let a_row = &a.data[row_start..row_start + block_size];

                                    for j in jj..j_end {
                                        let col_start = j * b_transposed.cols + kk;
                                        let b_col =
                                            &b_transposed.data[col_start..col_start + block_size];

                                        let partial_dot = unsafe { Avx2Backend::dot(a_row, b_col) };
                                        result.data[i * result.cols + j] += partial_dot;
                                    }
                                }
                            } else {
                                // Non-AVX2 path
                                #[allow(unused_variables)]
                                for i in ii..i_end {
                                    let row_start = i * a.cols + kk;
                                    let a_row = &a.data[row_start..row_start + block_size];

                                    for j in jj..j_end {
                                        let col_start = j * b_transposed.cols + kk;
                                        let b_col =
                                            &b_transposed.data[col_start..col_start + block_size];

                                        let partial_dot = unsafe {
                                            match a.backend {
                                                Backend::Scalar => ScalarBackend::dot(a_row, b_col),
                                                #[cfg(target_arch = "x86_64")]
                                                Backend::SSE2 | Backend::AVX => {
                                                    Sse2Backend::dot(a_row, b_col)
                                                }
                                                #[cfg(not(target_arch = "x86_64"))]
                                                Backend::SSE2
                                                | Backend::AVX
                                                | Backend::AVX2
                                                | Backend::AVX512 => {
                                                    ScalarBackend::dot(a_row, b_col)
                                                }
                                                #[cfg(any(
                                                    target_arch = "aarch64",
                                                    target_arch = "arm"
                                                ))]
                                                Backend::NEON => {
                                                    use crate::backends::neon::NeonBackend;
                                                    NeonBackend::dot(a_row, b_col)
                                                }
                                                #[cfg(not(any(
                                                    target_arch = "aarch64",
                                                    target_arch = "arm"
                                                )))]
                                                Backend::NEON => ScalarBackend::dot(a_row, b_col),
                                                #[cfg(target_arch = "wasm32")]
                                                Backend::WasmSIMD => {
                                                    use crate::backends::wasm::WasmBackend;
                                                    WasmBackend::dot(a_row, b_col)
                                                }
                                                #[cfg(not(target_arch = "wasm32"))]
                                                Backend::WasmSIMD => {
                                                    ScalarBackend::dot(a_row, b_col)
                                                }
                                                // Catch-all for GPU, Auto, and any other backends
                                                _ => ScalarBackend::dot(a_row, b_col),
                                            }
                                        };

                                        result.data[i * result.cols + j] += partial_dot;
                                    }
                                }
                            }

                            // Non-x86_64 fallback
                            #[cfg(not(target_arch = "x86_64"))]
                            {
                                for i in ii..i_end {
                                    let row_start = i * a.cols + kk;
                                    let a_row = &a.data[row_start..row_start + block_size];

                                    for j in jj..j_end {
                                        let col_start = j * b_transposed.cols + kk;
                                        let b_col =
                                            &b_transposed.data[col_start..col_start + block_size];

                                        let partial_dot = unsafe {
                                            match a.backend {
                                                Backend::Scalar => ScalarBackend::dot(a_row, b_col),
                                                #[cfg(any(
                                                    target_arch = "aarch64",
                                                    target_arch = "arm"
                                                ))]
                                                Backend::NEON => {
                                                    use crate::backends::neon::NeonBackend;
                                                    NeonBackend::dot(a_row, b_col)
                                                }
                                                #[cfg(not(any(
                                                    target_arch = "aarch64",
                                                    target_arch = "arm"
                                                )))]
                                                Backend::NEON => ScalarBackend::dot(a_row, b_col),
                                                #[cfg(target_arch = "wasm32")]
                                                Backend::WasmSIMD => {
                                                    use crate::backends::wasm::WasmBackend;
                                                    WasmBackend::dot(a_row, b_col)
                                                }
                                                #[cfg(not(target_arch = "wasm32"))]
                                                Backend::WasmSIMD => {
                                                    ScalarBackend::dot(a_row, b_col)
                                                }
                                                _ => ScalarBackend::dot(a_row, b_col),
                                            }
                                        };

                                        result.data[i * result.cols + j] += partial_dot;
                                    }
                                }
                            }
                        }
                    }
                }
            }
        }
    }

    fn matmul_simd(
        &self,
        other: &Matrix<f32>,
        result: &mut Matrix<f32>,
    ) -> Result<(), TruenoError> {
        // Cache blocking parameters (tuned for typical x86_64 CPUs)
        // L2 cache: 256KB typical → 64K f32 elements → 64×64×3 matrices fits
        const L2_BLOCK_SIZE: usize = 64;
        // L3 cache: 4-16MB typical → 256×256 blocks for very large matrices (Phase 3)
        const L3_BLOCK_SIZE: usize = 256;
        const L3_THRESHOLD: usize = 512; // Use 3-level blocking for matrices ≥512×512

        // For small matrices, use simple SIMD approach (blocking overhead too high)
        if self.rows <= 32 || self.cols <= 32 || other.cols <= 32 {
            return self.matmul_simd_simple(other, result);
        }

        #[cfg(target_arch = "x86_64")]
        use crate::backends::{avx2::Avx2Backend, sse2::Sse2Backend};
        use crate::backends::{scalar::ScalarBackend, VectorBackend};

        // Pre-transpose B for better cache locality (columns become rows)
        let b_transposed = other.transpose();

        // Determine if we should use 3-level blocking (Phase 3)
        let use_l3_blocking =
            self.rows >= L3_THRESHOLD && self.cols >= L3_THRESHOLD && other.cols >= L3_THRESHOLD;

        // Phase 4: Determine if we should use multi-threading (≥1024×1024)
        #[cfg(feature = "parallel")]
        const PARALLEL_THRESHOLD: usize = 1024;
        #[cfg(feature = "parallel")]
        let use_parallel = self.rows >= PARALLEL_THRESHOLD
            && self.cols >= PARALLEL_THRESHOLD
            && other.cols >= PARALLEL_THRESHOLD;
        #[cfg(not(feature = "parallel"))]
        let use_parallel = false;

        if use_l3_blocking {
            // ===== Phase 3/4: 3-Level Cache Blocking (L3 → L2 → micro-kernel) =====
            // For very large matrices (≥512×512), use L3 cache blocking to minimize
            // cache misses when data doesn't fit in L2 cache
            //
            // Hierarchy:
            // 1. L3 blocks: 256×256 (fits in L3 cache: 4-16MB)
            // 2. L2 blocks: 64×64 (fits in L2 cache: 256KB)
            // 3. Micro-kernel: 4×1 for AVX2/AVX512
            //
            // Phase 4: For ≥1024×1024, parallelize L3 row blocks with rayon

            if use_parallel {
                // ===== Phase 4: Parallel 3-Level Cache Blocking (Lock-Free Row Partitioning) =====
                #[cfg(feature = "parallel")]
                {
                    use rayon::prelude::*;
                    use std::sync::atomic::{AtomicPtr, Ordering};
                    use std::sync::Arc;

                    // Lock-free parallelization strategy:
                    // Each thread processes one L3 row block (256 rows). Since row blocks are
                    // non-overlapping, threads write to distinct memory regions with no contention.
                    //
                    // Safety invariant: Each thread writes to result.data[iii*cols..(i3_end)*cols],
                    // where iii = block_idx * L3_BLOCK_SIZE. Since L3 blocks don't overlap,
                    // no two threads write to the same memory location.

                    // Store result pointer in Arc<AtomicPtr> for safe sharing
                    let result_ptr = Arc::new(AtomicPtr::new(result as *mut Matrix<f32>));

                    // Calculate number of L3 blocks
                    let num_blocks = self.rows.div_ceil(L3_BLOCK_SIZE);

                    // Process each L3 block in parallel (lock-free)
                    (0..num_blocks).into_par_iter().for_each(|block_idx| {
                        let iii = block_idx * L3_BLOCK_SIZE;
                        let i3_end = (iii + L3_BLOCK_SIZE).min(self.rows);

                        // SAFETY: Each thread processes a distinct row range [iii, i3_end).
                        // No two threads write to overlapping memory locations because:
                        // 1. L3 blocks partition rows: [0, 256), [256, 512), etc.
                        // 2. Each thread only modifies result.data[iii*cols..(i3_end)*cols]
                        // 3. Row ranges are non-overlapping by construction
                        // 4. All threads complete before function returns (rayon guarantee)
                        // 5. AtomicPtr ensures proper memory ordering across threads
                        unsafe {
                            let ptr = result_ptr.load(Ordering::Relaxed);
                            Self::process_l3_row_block_seq(
                                iii,
                                i3_end,
                                self,
                                &b_transposed,
                                &mut *ptr,
                                L2_BLOCK_SIZE,
                                L3_BLOCK_SIZE,
                            );
                        }
                    });
                }

                return Ok(());
            }

            // ===== Sequential 3-Level Cache Blocking (fallback) =====
            for iii in (0..self.rows).step_by(L3_BLOCK_SIZE) {
                let i3_end = (iii + L3_BLOCK_SIZE).min(self.rows);

                for jjj in (0..other.cols).step_by(L3_BLOCK_SIZE) {
                    let j3_end = (jjj + L3_BLOCK_SIZE).min(other.cols);

                    for kkk in (0..self.cols).step_by(L3_BLOCK_SIZE) {
                        let k3_end = (kkk + L3_BLOCK_SIZE).min(self.cols);

                        // L2 blocking within L3 blocks
                        for ii in (iii..i3_end).step_by(L2_BLOCK_SIZE) {
                            let i_end = (ii + L2_BLOCK_SIZE).min(i3_end);

                            for jj in (jjj..j3_end).step_by(L2_BLOCK_SIZE) {
                                let j_end = (jj + L2_BLOCK_SIZE).min(j3_end);

                                for kk in (kkk..k3_end).step_by(L2_BLOCK_SIZE) {
                                    let k_end = (kk + L2_BLOCK_SIZE).min(k3_end);
                                    let block_size = k_end - kk;

                                    // Micro-kernel processing
                                    #[cfg(target_arch = "x86_64")]
                                    let use_microkernel =
                                        matches!(self.backend, Backend::AVX2 | Backend::AVX512);

                                    #[cfg(target_arch = "x86_64")]
                                    if use_microkernel {
                                        let mut i = ii;

                                        // Process 4 rows at a time with micro-kernel
                                        while i + 4 <= i_end {
                                            let row0_start = i * self.cols + kk;
                                            let row1_start = (i + 1) * self.cols + kk;
                                            let row2_start = (i + 2) * self.cols + kk;
                                            let row3_start = (i + 3) * self.cols + kk;

                                            let a_rows = [
                                                &self.data[row0_start..row0_start + block_size],
                                                &self.data[row1_start..row1_start + block_size],
                                                &self.data[row2_start..row2_start + block_size],
                                                &self.data[row3_start..row3_start + block_size],
                                            ];

                                            for j in jj..j_end {
                                                let col_start = j * b_transposed.cols + kk;
                                                let b_col = &b_transposed.data
                                                    [col_start..col_start + block_size];

                                                let mut partial_dots = [0.0f32; 4];
                                                unsafe {
                                                    Self::matmul_microkernel_4x1_avx2(
                                                        a_rows,
                                                        b_col,
                                                        &mut partial_dots,
                                                    );
                                                }

                                                result.data[i * result.cols + j] += partial_dots[0];
                                                result.data[(i + 1) * result.cols + j] +=
                                                    partial_dots[1];
                                                result.data[(i + 2) * result.cols + j] +=
                                                    partial_dots[2];
                                                result.data[(i + 3) * result.cols + j] +=
                                                    partial_dots[3];
                                            }

                                            i += 4;
                                        }

                                        // Handle remaining rows (< 4)
                                        for i in i..i_end {
                                            let row_start = i * self.cols + kk;
                                            let a_row =
                                                &self.data[row_start..row_start + block_size];

                                            for j in jj..j_end {
                                                let col_start = j * b_transposed.cols + kk;
                                                let b_col = &b_transposed.data
                                                    [col_start..col_start + block_size];

                                                let partial_dot =
                                                    unsafe { Avx2Backend::dot(a_row, b_col) };
                                                result.data[i * result.cols + j] += partial_dot;
                                            }
                                        }
                                    } else {
                                        // Non-AVX2 path
                                        #[allow(unused_variables)]
                                        for i in ii..i_end {
                                            let row_start = i * self.cols + kk;
                                            let a_row =
                                                &self.data[row_start..row_start + block_size];

                                            for j in jj..j_end {
                                                let col_start = j * b_transposed.cols + kk;
                                                let b_col = &b_transposed.data
                                                    [col_start..col_start + block_size];

                                                let partial_dot = unsafe {
                                                    match self.backend {
                                                        Backend::Scalar => {
                                                            ScalarBackend::dot(a_row, b_col)
                                                        }
                                                        #[cfg(target_arch = "x86_64")]
                                                        Backend::SSE2 | Backend::AVX => {
                                                            Sse2Backend::dot(a_row, b_col)
                                                        }
                                                        #[cfg(not(target_arch = "x86_64"))]
                                                        Backend::SSE2
                                                        | Backend::AVX
                                                        | Backend::AVX2
                                                        | Backend::AVX512 => {
                                                            ScalarBackend::dot(a_row, b_col)
                                                        }
                                                        #[cfg(any(
                                                            target_arch = "aarch64",
                                                            target_arch = "arm"
                                                        ))]
                                                        Backend::NEON => {
                                                            use crate::backends::neon::NeonBackend;
                                                            NeonBackend::dot(a_row, b_col)
                                                        }
                                                        #[cfg(not(any(
                                                            target_arch = "aarch64",
                                                            target_arch = "arm"
                                                        )))]
                                                        Backend::NEON => {
                                                            ScalarBackend::dot(a_row, b_col)
                                                        }
                                                        #[cfg(target_arch = "wasm32")]
                                                        Backend::WasmSIMD => {
                                                            use crate::backends::wasm::WasmBackend;
                                                            WasmBackend::dot(a_row, b_col)
                                                        }
                                                        #[cfg(not(target_arch = "wasm32"))]
                                                        Backend::WasmSIMD => {
                                                            ScalarBackend::dot(a_row, b_col)
                                                        }
                                                        Backend::GPU
                                                        | Backend::Auto
                                                        | Backend::AVX2
                                                        | Backend::AVX512 => {
                                                            ScalarBackend::dot(a_row, b_col)
                                                        }
                                                    }
                                                };

                                                result.data[i * result.cols + j] += partial_dot;
                                            }
                                        }
                                    }

                                    // Non-x86_64 platforms
                                    #[cfg(not(target_arch = "x86_64"))]
                                    for i in ii..i_end {
                                        let row_start = i * self.cols + kk;
                                        let a_row = &self.data[row_start..row_start + block_size];

                                        for j in jj..j_end {
                                            let col_start = j * b_transposed.cols + kk;
                                            let b_col = &b_transposed.data
                                                [col_start..col_start + block_size];

                                            let partial_dot = unsafe {
                                                match self.backend {
                                                    Backend::Scalar => {
                                                        ScalarBackend::dot(a_row, b_col)
                                                    }
                                                    #[cfg(any(
                                                        target_arch = "aarch64",
                                                        target_arch = "arm"
                                                    ))]
                                                    Backend::NEON => {
                                                        use crate::backends::neon::NeonBackend;
                                                        NeonBackend::dot(a_row, b_col)
                                                    }
                                                    #[cfg(not(any(
                                                        target_arch = "aarch64",
                                                        target_arch = "arm"
                                                    )))]
                                                    Backend::NEON => {
                                                        ScalarBackend::dot(a_row, b_col)
                                                    }
                                                    #[cfg(target_arch = "wasm32")]
                                                    Backend::WasmSIMD => {
                                                        use crate::backends::wasm::WasmBackend;
                                                        WasmBackend::dot(a_row, b_col)
                                                    }
                                                    #[cfg(not(target_arch = "wasm32"))]
                                                    Backend::WasmSIMD => {
                                                        ScalarBackend::dot(a_row, b_col)
                                                    }
                                                    _ => ScalarBackend::dot(a_row, b_col),
                                                }
                                            };

                                            result.data[i * result.cols + j] += partial_dot;
                                        }
                                    }
                                }
                            }
                        }
                    }
                }
            }
        } else {
            // ===== Phase 1/2: 2-Level Cache Blocking (L2 → micro-kernel) =====
            // For medium matrices (32-512), use original 2-level blocking
            //
            // This path preserves the fast performance for 256×256 and smaller matrices
            // by avoiding the overhead of 3-level loop nesting

            for ii in (0..self.rows).step_by(L2_BLOCK_SIZE) {
                let i_end = (ii + L2_BLOCK_SIZE).min(self.rows);

                for jj in (0..other.cols).step_by(L2_BLOCK_SIZE) {
                    let j_end = (jj + L2_BLOCK_SIZE).min(other.cols);

                    for kk in (0..self.cols).step_by(L2_BLOCK_SIZE) {
                        let k_end = (kk + L2_BLOCK_SIZE).min(self.cols);
                        let block_size = k_end - kk;

                        // Inner loops: Process L2 block with micro-kernel (Phase 2) or SIMD
                        #[cfg(target_arch = "x86_64")]
                        let use_microkernel =
                            matches!(self.backend, Backend::AVX2 | Backend::AVX512);

                        #[cfg(target_arch = "x86_64")]
                        if use_microkernel {
                            // Phase 2: Use 4×1 micro-kernel for AVX2/AVX512
                            let mut i = ii;

                            // Process 4 rows at a time with micro-kernel
                            while i + 4 <= i_end {
                                // Get 4 consecutive rows of A
                                let row0_start = i * self.cols + kk;
                                let row1_start = (i + 1) * self.cols + kk;
                                let row2_start = (i + 2) * self.cols + kk;
                                let row3_start = (i + 3) * self.cols + kk;

                                let a_rows = [
                                    &self.data[row0_start..row0_start + block_size],
                                    &self.data[row1_start..row1_start + block_size],
                                    &self.data[row2_start..row2_start + block_size],
                                    &self.data[row3_start..row3_start + block_size],
                                ];

                                // Process each column of B with the micro-kernel
                                for j in jj..j_end {
                                    let col_start = j * b_transposed.cols + kk;
                                    let b_col =
                                        &b_transposed.data[col_start..col_start + block_size];

                                    // Compute 4 dot products simultaneously
                                    let mut partial_dots = [0.0f32; 4];
                                    unsafe {
                                        Self::matmul_microkernel_4x1_avx2(
                                            a_rows,
                                            b_col,
                                            &mut partial_dots,
                                        );
                                    }

                                    // Accumulate results
                                    result.data[i * result.cols + j] += partial_dots[0];
                                    result.data[(i + 1) * result.cols + j] += partial_dots[1];
                                    result.data[(i + 2) * result.cols + j] += partial_dots[2];
                                    result.data[(i + 3) * result.cols + j] += partial_dots[3];
                                }

                                i += 4;
                            }

                            // Handle remaining rows (< 4) with standard path
                            for i in i..i_end {
                                let row_start = i * self.cols + kk;
                                let a_row = &self.data[row_start..row_start + block_size];

                                for j in jj..j_end {
                                    let col_start = j * b_transposed.cols + kk;
                                    let b_col =
                                        &b_transposed.data[col_start..col_start + block_size];

                                    let partial_dot = unsafe { Avx2Backend::dot(a_row, b_col) };
                                    result.data[i * result.cols + j] += partial_dot;
                                }
                            }
                        } else {
                            // Phase 1: Standard SIMD path (non-AVX2 backends)
                            #[allow(unused_variables)]
                            for i in ii..i_end {
                                let row_start = i * self.cols + kk;
                                let a_row = &self.data[row_start..row_start + block_size];

                                for j in jj..j_end {
                                    let col_start = j * b_transposed.cols + kk;
                                    let b_col =
                                        &b_transposed.data[col_start..col_start + block_size];

                                    let partial_dot = unsafe {
                                        match self.backend {
                                            Backend::Scalar => ScalarBackend::dot(a_row, b_col),
                                            #[cfg(target_arch = "x86_64")]
                                            Backend::SSE2 | Backend::AVX => {
                                                Sse2Backend::dot(a_row, b_col)
                                            }
                                            #[cfg(not(target_arch = "x86_64"))]
                                            Backend::SSE2
                                            | Backend::AVX
                                            | Backend::AVX2
                                            | Backend::AVX512 => ScalarBackend::dot(a_row, b_col),
                                            #[cfg(any(
                                                target_arch = "aarch64",
                                                target_arch = "arm"
                                            ))]
                                            Backend::NEON => {
                                                use crate::backends::neon::NeonBackend;
                                                NeonBackend::dot(a_row, b_col)
                                            }
                                            #[cfg(not(any(
                                                target_arch = "aarch64",
                                                target_arch = "arm"
                                            )))]
                                            Backend::NEON => ScalarBackend::dot(a_row, b_col),
                                            #[cfg(target_arch = "wasm32")]
                                            Backend::WasmSIMD => {
                                                use crate::backends::wasm::WasmBackend;
                                                WasmBackend::dot(a_row, b_col)
                                            }
                                            #[cfg(not(target_arch = "wasm32"))]
                                            Backend::WasmSIMD => ScalarBackend::dot(a_row, b_col),
                                            Backend::GPU
                                            | Backend::Auto
                                            | Backend::AVX2
                                            | Backend::AVX512 => ScalarBackend::dot(a_row, b_col),
                                        }
                                    };

                                    result.data[i * result.cols + j] += partial_dot;
                                }
                            }
                        }

                        // Non-x86_64 platforms: Use standard SIMD path
                        #[cfg(not(target_arch = "x86_64"))]
                        for i in ii..i_end {
                            let row_start = i * self.cols + kk;
                            let a_row = &self.data[row_start..row_start + block_size];

                            for j in jj..j_end {
                                let col_start = j * b_transposed.cols + kk;
                                let b_col = &b_transposed.data[col_start..col_start + block_size];

                                let partial_dot = unsafe {
                                    match self.backend {
                                        Backend::Scalar => ScalarBackend::dot(a_row, b_col),
                                        #[cfg(any(target_arch = "aarch64", target_arch = "arm"))]
                                        Backend::NEON => {
                                            use crate::backends::neon::NeonBackend;
                                            NeonBackend::dot(a_row, b_col)
                                        }
                                        #[cfg(not(any(
                                            target_arch = "aarch64",
                                            target_arch = "arm"
                                        )))]
                                        Backend::NEON => ScalarBackend::dot(a_row, b_col),
                                        #[cfg(target_arch = "wasm32")]
                                        Backend::WasmSIMD => {
                                            use crate::backends::wasm::WasmBackend;
                                            WasmBackend::dot(a_row, b_col)
                                        }
                                        #[cfg(not(target_arch = "wasm32"))]
                                        Backend::WasmSIMD => ScalarBackend::dot(a_row, b_col),
                                        _ => ScalarBackend::dot(a_row, b_col),
                                    }
                                };

                                result.data[i * result.cols + j] += partial_dot;
                            }
                        }
                    }
                }
            }
        }

        Ok(())
    }

    /// Simple SIMD matrix multiplication without blocking (for small matrices)
    ///
    /// This is the pre-blocking implementation that works well for small matrices
    /// where cache blocking overhead exceeds benefits.
    fn matmul_simd_simple(
        &self,
        other: &Matrix<f32>,
        result: &mut Matrix<f32>,
    ) -> Result<(), TruenoError> {
        #[cfg(target_arch = "x86_64")]
        use crate::backends::{avx2::Avx2Backend, sse2::Sse2Backend};
        use crate::backends::{scalar::ScalarBackend, VectorBackend};

        // Pre-transpose B for better cache locality
        let b_transposed = other.transpose();

        for i in 0..self.rows {
            let row_start = i * self.cols;
            let row_end = row_start + self.cols;
            let a_row = &self.data[row_start..row_end];

            for j in 0..other.cols {
                let col_start = j * b_transposed.cols;
                let col_end = col_start + b_transposed.cols;
                let b_col = &b_transposed.data[col_start..col_end];

                // Compute dot product using SIMD backend directly
                // SAFETY: Backend dot() maintains safety invariants
                let dot_result = unsafe {
                    match self.backend {
                        Backend::Scalar => ScalarBackend::dot(a_row, b_col),
                        #[cfg(target_arch = "x86_64")]
                        Backend::SSE2 | Backend::AVX => Sse2Backend::dot(a_row, b_col),
                        #[cfg(target_arch = "x86_64")]
                        Backend::AVX2 | Backend::AVX512 => Avx2Backend::dot(a_row, b_col),
                        #[cfg(not(target_arch = "x86_64"))]
                        Backend::SSE2 | Backend::AVX | Backend::AVX2 | Backend::AVX512 => {
                            ScalarBackend::dot(a_row, b_col)
                        }
                        #[cfg(any(target_arch = "aarch64", target_arch = "arm"))]
                        Backend::NEON => {
                            use crate::backends::neon::NeonBackend;
                            NeonBackend::dot(a_row, b_col)
                        }
                        #[cfg(not(any(target_arch = "aarch64", target_arch = "arm")))]
                        Backend::NEON => ScalarBackend::dot(a_row, b_col),
                        #[cfg(target_arch = "wasm32")]
                        Backend::WasmSIMD => {
                            use crate::backends::wasm::WasmBackend;
                            WasmBackend::dot(a_row, b_col)
                        }
                        #[cfg(not(target_arch = "wasm32"))]
                        Backend::WasmSIMD => ScalarBackend::dot(a_row, b_col),
                        Backend::GPU | Backend::Auto => ScalarBackend::dot(a_row, b_col),
                    }
                };

                result.data[i * result.cols + j] = dot_result;
            }
        }

        Ok(())
    }

    /// WASM-optimized tiled matrix multiplication with SIMD inner loop
    ///
    /// Key optimizations:
    /// 1. NO transpose - avoids O(n²) memory allocation and copy
    /// 2. Tiled blocking with SIMD-aligned tile widths
    /// 3. Inner j-loop uses SIMD (B rows are contiguous in memory)
    /// 4. Register accumulation to minimize memory traffic
    ///
    /// Performance: Targets <30ms for 384×74×384 (Whisper encoder attention)
    fn matmul_wasm_tiled(
        &self,
        other: &Matrix<f32>,
        result: &mut Matrix<f32>,
    ) -> Result<(), TruenoError> {
        let m = self.rows;
        let k = self.cols;
        let n = other.cols;

        // For each row of A
        for i in 0..m {
            let a_row_start = i * k;
            let result_row_start = i * n;

            // For each column of B, compute dot product A[i,:] · B[:,j]
            // BUT: B[:,j] is not contiguous. Instead, iterate over k and accumulate.
            //
            // C[i,j] = Σ_k A[i,k] * B[k,j]
            //
            // For efficiency, broadcast A[i,k] and multiply with B[k, j0:j0+width]
            // This uses SIMD on the contiguous B row segment.

            // Process output columns in SIMD-width chunks
            let simd_width = 8; // AVX2 processes 8 f32s
            let n_simd = (n / simd_width) * simd_width;

            // SIMD portion: columns 0..n_simd
            // Note: Explicit indexing is intentional for LLVM auto-vectorization.
            // Iterator patterns prevent the compiler from recognizing the SIMD pattern.
            #[allow(clippy::needless_range_loop)]
            for j0 in (0..n_simd).step_by(simd_width) {
                let mut acc = [0.0f32; 8];

                for kk in 0..k {
                    let a_val = self.data[a_row_start + kk];
                    let b_row_start = kk * n + j0;

                    // Multiply a_val with B[kk, j0:j0+8]
                    for jj in 0..simd_width {
                        acc[jj] += a_val * other.data[b_row_start + jj];
                    }
                }

                // Write accumulated results
                for jj in 0..simd_width {
                    result.data[result_row_start + j0 + jj] = acc[jj];
                }
            }

            // Remainder columns (non-SIMD)
            for j in n_simd..n {
                let mut sum = 0.0f32;
                for kk in 0..k {
                    sum += self.data[a_row_start + kk] * other.data[kk * n + j];
                }
                result.data[result_row_start + j] = sum;
            }
        }

        Ok(())
    }

    /// GPU-accelerated matrix multiplication (very large matrices only)
    #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    fn matmul_gpu(&self, other: &Matrix<f32>) -> Result<Matrix<f32>, TruenoError> {
        use crate::backends::gpu::GpuBackend;

        // Check if GPU is available
        if !GpuBackend::is_available() {
            return Err(TruenoError::InvalidInput("GPU not available".to_string()));
        }

        // Create GPU backend
        let mut gpu = GpuBackend::new();

        // Execute GPU matmul
        let result_data = gpu
            .matmul(&self.data, &other.data, self.rows, self.cols, other.cols)
            .map_err(|e| TruenoError::InvalidInput(format!("GPU matmul failed: {}", e)))?;

        // Create result matrix
        let mut result = Matrix::zeros(self.rows, other.cols);
        result.data = result_data;

        Ok(result)
    }

    /// Transpose the matrix (swap rows and columns)
    ///
    /// Returns a new matrix where element `(i, j)` of the original becomes
    /// element `(j, i)` in the result.
    ///
    /// # Returns
    ///
    /// A new matrix with dimensions swapped: if input is `m×n`, output is `n×m`
    ///
    /// # Example
    ///
    /// ```
    /// use trueno::Matrix;
    ///
    /// let m = Matrix::from_vec(2, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    /// let t = m.transpose();
    ///
    /// // [[1, 2, 3],     [[1, 4],
    /// //  [4, 5, 6]]  →   [2, 5],
    /// //                  [3, 6]]
    /// assert_eq!(t.rows(), 3);
    /// assert_eq!(t.cols(), 2);
    /// assert_eq!(t.get(0, 0), Some(&1.0));
    /// assert_eq!(t.get(0, 1), Some(&4.0));
    /// assert_eq!(t.get(1, 0), Some(&2.0));
    /// ```
    #[cfg_attr(feature = "tracing", instrument(skip(self), fields(dims = %format!("{}x{}", self.rows, self.cols))))]
    pub fn transpose(&self) -> Matrix<f32> {
        let mut result = Matrix::zeros_with_backend(self.cols, self.rows, self.backend);

        // Use block-wise transpose for better cache locality
        // Block size of 64 fits well in L1 cache (64*64*4 = 16KB for f32)
        const BLOCK_SIZE: usize = 64;

        // Process matrix in BLOCK_SIZE x BLOCK_SIZE blocks
        for i_block in (0..self.rows).step_by(BLOCK_SIZE) {
            for j_block in (0..self.cols).step_by(BLOCK_SIZE) {
                // Process elements within this block
                let i_end = (i_block + BLOCK_SIZE).min(self.rows);
                let j_end = (j_block + BLOCK_SIZE).min(self.cols);

                for i in i_block..i_end {
                    // Direct slice access within row for better performance
                    let src_row_start = i * self.cols;
                    for j in j_block..j_end {
                        // result[j, i] = self[i, j]
                        // Use direct indexing instead of get/get_mut for speed
                        result.data[j * result.cols + i] = self.data[src_row_start + j];
                    }
                }
            }
        }

        result
    }

    /// Matrix-vector multiplication (column vector): A × v
    ///
    /// Multiplies this matrix by a column vector, computing `A × v` where the result
    /// is a column vector with length equal to the number of rows in `A`.
    ///
    /// # Mathematical Definition
    ///
    /// For an m×n matrix A and an n-dimensional vector v:
    /// ```text
    /// result[i] = Σ(j=0 to n-1) A[i,j] × v[j]
    /// ```
    ///
    /// # Arguments
    ///
    /// * `v` - Column vector with length equal to `self.cols()`
    ///
    /// # Returns
    ///
    /// A new vector with length `self.rows()`
    ///
    /// # Errors
    ///
    /// Returns `InvalidInput` if `v.len() != self.cols()`
    ///
    /// # Example
    ///
    /// ```
    /// use trueno::{Matrix, Vector};
    ///
    /// let m = Matrix::from_vec(2, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    /// let v = Vector::from_slice(&[1.0, 2.0, 3.0]);
    /// let result = m.matvec(&v).unwrap();
    ///
    /// // [[1, 2, 3]   [1]   [1×1 + 2×2 + 3×3]   [14]
    /// //  [4, 5, 6]] × [2] = [4×1 + 5×2 + 6×3] = [32]
    /// //               [3]
    /// assert_eq!(result.as_slice(), &[14.0, 32.0]);
    /// ```
    pub fn matvec(&self, v: &Vector<f32>) -> Result<Vector<f32>, TruenoError> {
        if v.len() != self.cols {
            return Err(TruenoError::InvalidInput(format!(
                "Vector length {} does not match matrix columns {} for matrix-vector multiplication",
                v.len(),
                self.cols
            )));
        }

        #[cfg(target_arch = "x86_64")]
        use crate::backends::{avx2::Avx2Backend, sse2::Sse2Backend};
        use crate::backends::{scalar::ScalarBackend, VectorBackend};

        let v_slice = v.as_slice();

        let mut result_data = vec![0.0; self.rows];

        // Parallel execution for very large matrices (≥4096 rows)
        // Note: Thread overhead dominates for smaller matrices
        #[cfg(feature = "parallel")]
        {
            const PARALLEL_THRESHOLD: usize = 4096;

            if self.rows >= PARALLEL_THRESHOLD {
                use rayon::prelude::*;
                use std::sync::atomic::{AtomicPtr, Ordering};
                use std::sync::Arc;

                let result_ptr = Arc::new(AtomicPtr::new(result_data.as_mut_ptr()));

                // Process rows in parallel - each row computes an independent dot product
                (0..self.rows).into_par_iter().for_each(|i| {
                    let row_start = i * self.cols;
                    let row = &self.data[row_start..(row_start + self.cols)];

                    let dot_result = unsafe {
                        #[cfg(target_arch = "x86_64")]
                        {
                            match self.backend {
                                Backend::AVX2 | Backend::AVX512 => Avx2Backend::dot(row, v_slice),
                                Backend::SSE2 | Backend::AVX => Sse2Backend::dot(row, v_slice),
                                _ => ScalarBackend::dot(row, v_slice),
                            }
                        }
                        #[cfg(not(target_arch = "x86_64"))]
                        {
                            ScalarBackend::dot(row, v_slice)
                        }
                    };

                    // Write to non-overlapping memory location (thread-safe)
                    unsafe {
                        let ptr = result_ptr.load(Ordering::Relaxed);
                        *ptr.add(i) = dot_result;
                    }
                });

                return Ok(Vector::from_slice(&result_data));
            }
        }

        // SIMD-optimized execution: each row-vector product is a dot product
        for (i, result) in result_data.iter_mut().enumerate() {
            let row_start = i * self.cols;
            let row = &self.data[row_start..(row_start + self.cols)];

            // Use SIMD dot product for each row
            *result = unsafe {
                #[cfg(target_arch = "x86_64")]
                {
                    match self.backend {
                        Backend::AVX2 | Backend::AVX512 => Avx2Backend::dot(row, v_slice),
                        Backend::SSE2 | Backend::AVX => Sse2Backend::dot(row, v_slice),
                        _ => ScalarBackend::dot(row, v_slice),
                    }
                }
                #[cfg(not(target_arch = "x86_64"))]
                {
                    ScalarBackend::dot(row, v_slice)
                }
            };
        }

        Ok(Vector::from_slice(&result_data))
    }

    /// Vector-matrix multiplication (row vector): v^T × A
    ///
    /// Multiplies a row vector by this matrix, computing `v^T × A` where the result
    /// is a row vector with length equal to the number of columns in `A`.
    ///
    /// # Mathematical Definition
    ///
    /// For an m-dimensional vector v and an m×n matrix A:
    /// ```text
    /// result[j] = Σ(i=0 to m-1) v[i] × A[i,j]
    /// ```
    ///
    /// # Arguments
    ///
    /// * `v` - Row vector with length equal to `m.rows()`
    /// * `m` - Matrix to multiply
    ///
    /// # Returns
    ///
    /// A new vector with length `m.cols()`
    ///
    /// # Errors
    ///
    /// Returns `InvalidInput` if `v.len() != m.rows()`
    ///
    /// # Example
    ///
    /// ```
    /// use trueno::{Matrix, Vector};
    ///
    /// let m = Matrix::from_vec(2, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
    /// let v = Vector::from_slice(&[1.0, 2.0]);
    /// let result = Matrix::vecmat(&v, &m).unwrap();
    ///
    /// // [1, 2] × [[1, 2, 3]  = [1×1 + 2×4, 1×2 + 2×5, 1×3 + 2×6]
    /// //           [4, 5, 6]]
    /// //         = [9, 12, 15]
    /// assert_eq!(result.as_slice(), &[9.0, 12.0, 15.0]);
    /// ```
    pub fn vecmat(v: &Vector<f32>, m: &Matrix<f32>) -> Result<Vector<f32>, TruenoError> {
        if v.len() != m.rows {
            return Err(TruenoError::InvalidInput(format!(
                "Vector length {} does not match matrix rows {} for vector-matrix multiplication",
                v.len(),
                m.rows
            )));
        }

        // SIMD-optimized implementation using row-wise accumulation
        // Instead of column-wise access (cache-unfriendly), we compute:
        // result = Σ(i) v[i] * row_i (cache-friendly, vectorizable)
        //
        // This approach:
        // 1. Sequential row access (cache-friendly vs strided column access)
        // 2. Uses SIMD scale and add operations
        // 3. Leverages existing optimized Vector operations

        let mut result = Vector::from_slice(&vec![0.0; m.cols]);
        let v_slice = v.as_slice();

        // Accumulate each scaled row into result
        for (i, &scalar) in v_slice.iter().enumerate().take(m.rows) {
            let row_start = i * m.cols;
            let row = &m.data[row_start..(row_start + m.cols)];

            // Create vector for this row
            let row_vec = Vector::from_slice(row);

            // result += scalar * row (using SIMD scale and add)
            let scaled_row = row_vec.scale(scalar)?;
            result = result.add(&scaled_row)?;
        }

        Ok(result)
    }

    /// Perform 2D convolution with a kernel
    ///
    /// Applies a 2D convolution operation using "valid" padding (no padding),
    /// resulting in an output smaller than the input.
    ///
    /// # Arguments
    ///
    /// * `kernel` - Convolution kernel (filter) to apply
    ///
    /// # Returns
    ///
    /// Convolved matrix with dimensions:
    /// - rows: `input.rows - kernel.rows + 1`
    /// - cols: `input.cols - kernel.cols + 1`
    ///
    /// # Errors
    ///
    /// Returns `InvalidInput` if:
    /// - Kernel is larger than input in any dimension
    /// - Kernel has even dimensions (center pixel ambiguous)
    ///
    /// # Example
    ///
    /// ```
    /// use trueno::Matrix;
    ///
    /// // 5x5 input image
    /// let input = Matrix::from_vec(
    ///     5, 5,
    ///     vec![
    ///         0.0, 0.0, 0.0, 0.0, 0.0,
    ///         0.0, 0.0, 0.0, 0.0, 0.0,
    ///         0.0, 0.0, 9.0, 0.0, 0.0,
    ///         0.0, 0.0, 0.0, 0.0, 0.0,
    ///         0.0, 0.0, 0.0, 0.0, 0.0,
    ///     ]
    /// ).unwrap();
    ///
    /// // 3x3 averaging kernel
    /// let kernel_val = 1.0 / 9.0;
    /// let kernel = Matrix::from_vec(
    ///     3, 3,
    ///     vec![kernel_val; 9]
    /// ).unwrap();
    ///
    /// let result = input.convolve2d(&kernel).unwrap();
    /// assert_eq!(result.rows(), 3); // 5 - 3 + 1
    /// assert_eq!(result.cols(), 3);
    /// ```
    pub fn convolve2d(&self, kernel: &Matrix<f32>) -> Result<Matrix<f32>, TruenoError> {
        // Validate kernel size
        if kernel.rows > self.rows || kernel.cols > self.cols {
            return Err(TruenoError::InvalidInput(format!(
                "Kernel size ({}x{}) larger than input ({}x{})",
                kernel.rows, kernel.cols, self.rows, self.cols
            )));
        }

        // Calculate output dimensions (valid padding)
        let output_rows = self.rows - kernel.rows + 1;
        let output_cols = self.cols - kernel.cols + 1;

        // Initialize output matrix (reuse parent's backend)
        let mut result = Matrix::zeros_with_backend(output_rows, output_cols, self.backend);

        // Backend selection strategy:
        // OpComplexity::High - GPU beneficial at >10K elements
        // GPU for large images (output > 10K elements)
        // Scalar for smaller images

        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
        const GPU_THRESHOLD: usize = 10_000;

        // Try GPU first for large convolutions
        #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
        {
            if output_rows * output_cols >= GPU_THRESHOLD {
                use crate::backends::gpu::GpuBackend;

                if GpuBackend::is_available() {
                    if let Ok(gpu_result) =
                        self.convolve2d_gpu(kernel, &mut result, output_rows, output_cols)
                    {
                        return Ok(gpu_result);
                    }
                    // Fall through to scalar if GPU fails
                }
            }
        }

        // Scalar baseline implementation
        // Bounds: output_rows = self.rows - kernel.rows + 1, output_cols = self.cols - kernel.cols + 1
        for out_row in 0..output_rows {
            for out_col in 0..output_cols {
                let mut sum = 0.0;

                // Apply kernel
                for k_row in 0..kernel.rows {
                    for k_col in 0..kernel.cols {
                        let in_row = out_row + k_row;
                        let in_col = out_col + k_col;

                        // Bounds guaranteed: in_row < self.rows, in_col < self.cols
                        let input_val = self
                            .get(in_row, in_col)
                            .expect("convolve2d: input bounds validated by output dimensions");
                        let kernel_val = kernel
                            .get(k_row, k_col)
                            .expect("convolve2d: kernel bounds validated by loop");

                        sum += input_val * kernel_val;
                    }
                }

                *result
                    .get_mut(out_row, out_col)
                    .expect("convolve2d: output bounds validated by allocation") = sum;
            }
        }

        Ok(result)
    }

    /// GPU-accelerated 2D convolution helper
    #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    fn convolve2d_gpu(
        &self,
        kernel: &Matrix<f32>,
        result: &mut Matrix<f32>,
        _output_rows: usize,
        _output_cols: usize,
    ) -> Result<Matrix<f32>, TruenoError> {
        use crate::backends::gpu::GpuDevice;

        let gpu = GpuDevice::new().map_err(TruenoError::InvalidInput)?;

        gpu.convolve2d(
            self.as_slice(),
            kernel.as_slice(),
            result.data.as_mut_slice(),
            self.rows,
            self.cols,
            kernel.rows,
            kernel.cols,
        )
        .map_err(TruenoError::InvalidInput)?;

        Ok(result.clone())
    }

    /// Lookup embeddings by indices (Issue #61: ML primitives)
    ///
    /// Performs embedding lookup where self is the embedding table with shape
    /// `[vocab_size, embed_dim]` and indices specify which rows to select.
    ///
    /// # Arguments
    ///
    /// * `indices` - Slice of indices into the embedding table
    ///
    /// # Returns
    ///
    /// A matrix with shape `[indices.len(), embed_dim]` containing the selected rows
    ///
    /// # Errors
    ///
    /// Returns `InvalidInput` if any index is out of bounds
    ///
    /// # Example
    ///
    /// ```
    /// use trueno::Matrix;
    ///
    /// // Create embedding table: 4 words, 3-dimensional embeddings
    /// let embeddings = Matrix::from_vec(4, 3, vec![
    ///     1.0, 2.0, 3.0,   // word 0
    ///     4.0, 5.0, 6.0,   // word 1
    ///     7.0, 8.0, 9.0,   // word 2
    ///     10.0, 11.0, 12.0 // word 3
    /// ]).unwrap();
    ///
    /// // Lookup embeddings for indices [1, 3, 0]
    /// let result = embeddings.embedding_lookup(&[1, 3, 0]).unwrap();
    ///
    /// assert_eq!(result.rows(), 3);
    /// assert_eq!(result.cols(), 3);
    /// assert_eq!(result.get(0, 0), Some(&4.0)); // word 1
    /// assert_eq!(result.get(1, 0), Some(&10.0)); // word 3
    /// assert_eq!(result.get(2, 0), Some(&1.0)); // word 0
    /// ```
    pub fn embedding_lookup(&self, indices: &[usize]) -> Result<Matrix<f32>, TruenoError> {
        // Validate indices
        for (i, &idx) in indices.iter().enumerate() {
            if idx >= self.rows {
                return Err(TruenoError::InvalidInput(format!(
                    "Index {} at position {} is out of bounds for embedding table with {} rows",
                    idx, i, self.rows
                )));
            }
        }

        // Handle empty indices
        if indices.is_empty() {
            return Ok(Matrix::zeros_with_backend(0, self.cols, self.backend));
        }

        // Allocate output matrix: [seq_len, embed_dim]
        let seq_len = indices.len();
        let embed_dim = self.cols;
        let mut result = Matrix::zeros_with_backend(seq_len, embed_dim, self.backend);

        // Copy rows from embedding table to result
        for (out_row, &idx) in indices.iter().enumerate() {
            let src_start = idx * embed_dim;
            let dst_start = out_row * embed_dim;

            // Copy entire row
            result.data[dst_start..dst_start + embed_dim]
                .copy_from_slice(&self.data[src_start..src_start + embed_dim]);
        }

        Ok(result)
    }

    /// Lookup embeddings with gradient tracking support (for training)
    ///
    /// Returns both the embeddings and a sparse gradient accumulator.
    /// This is useful for sparse gradient updates in training.
    ///
    /// # Arguments
    ///
    /// * `indices` - Slice of indices into the embedding table
    ///
    /// # Returns
    ///
    /// Tuple of (embeddings, unique_indices) where unique_indices can be used
    /// for sparse gradient updates
    ///
    /// # Errors
    ///
    /// Returns `InvalidInput` if any index is out of bounds
    pub fn embedding_lookup_sparse(
        &self,
        indices: &[usize],
    ) -> Result<(Matrix<f32>, Vec<usize>), TruenoError> {
        let embeddings = self.embedding_lookup(indices)?;

        // Get unique indices for sparse gradient updates
        let mut unique: Vec<usize> = indices.to_vec();
        unique.sort_unstable();
        unique.dedup();

        Ok((embeddings, unique))
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_matrix_new() {
        let m = Matrix::new(3, 4);
        assert_eq!(m.rows(), 3);
        assert_eq!(m.cols(), 4);
        assert_eq!(m.shape(), (3, 4));
        assert_eq!(m.as_slice().len(), 12);
    }

    #[test]
    fn test_matrix_from_vec() {
        let data = vec![1.0, 2.0, 3.0, 4.0];
        let m = Matrix::from_vec(2, 2, data).unwrap();
        assert_eq!(m.rows(), 2);
        assert_eq!(m.cols(), 2);
        assert_eq!(m.get(0, 0), Some(&1.0));
        assert_eq!(m.get(0, 1), Some(&2.0));
        assert_eq!(m.get(1, 0), Some(&3.0));
        assert_eq!(m.get(1, 1), Some(&4.0));
    }

    #[test]
    fn test_matrix_from_vec_invalid_size() {
        let data = vec![1.0, 2.0, 3.0];
        let result = Matrix::from_vec(2, 2, data);
        assert!(matches!(result, Err(TruenoError::InvalidInput(_))));
    }

    #[test]
    fn test_matrix_from_slice() {
        // TRUENO-SPEC-014 coverage test
        let data = [1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
        let m = Matrix::from_slice(2, 3, &data).unwrap();
        assert_eq!(m.rows(), 2);
        assert_eq!(m.cols(), 3);
        assert_eq!(m.get(0, 0), Some(&1.0));
        assert_eq!(m.get(1, 2), Some(&6.0));
    }

    #[test]
    fn test_matrix_from_slice_invalid() {
        // TRUENO-SPEC-014 coverage test - error path
        let data = [1.0, 2.0, 3.0];
        let result = Matrix::from_slice(2, 2, &data);
        assert!(matches!(result, Err(TruenoError::InvalidInput(_))));
    }

    #[test]
    fn test_matrix_zeros() {
        let m = Matrix::zeros(2, 3);
        assert_eq!(m.rows(), 2);
        assert_eq!(m.cols(), 3);
        for &val in m.as_slice() {
            assert_eq!(val, 0.0);
        }
    }

    #[test]
    fn test_matrix_identity() {
        let m = Matrix::identity(3);
        assert_eq!(m.rows(), 3);
        assert_eq!(m.cols(), 3);

        // Check diagonal
        assert_eq!(m.get(0, 0), Some(&1.0));
        assert_eq!(m.get(1, 1), Some(&1.0));
        assert_eq!(m.get(2, 2), Some(&1.0));

        // Check off-diagonal
        assert_eq!(m.get(0, 1), Some(&0.0));
        assert_eq!(m.get(0, 2), Some(&0.0));
        assert_eq!(m.get(1, 0), Some(&0.0));
        assert_eq!(m.get(1, 2), Some(&0.0));
        assert_eq!(m.get(2, 0), Some(&0.0));
        assert_eq!(m.get(2, 1), Some(&0.0));
    }

    #[test]
    fn test_matrix_get_out_of_bounds() {
        let m = Matrix::new(2, 2);
        assert_eq!(m.get(2, 0), None);
        assert_eq!(m.get(0, 2), None);
        assert_eq!(m.get(2, 2), None);
    }

    // ===== Matrix Multiplication Tests =====

    #[test]
    fn test_matmul_basic() {
        // [[1, 2],   [[5, 6],   [[19, 22],
        //  [3, 4]] ×  [7, 8]] =  [43, 50]]
        let a = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]).unwrap();
        let b = Matrix::from_vec(2, 2, vec![5.0, 6.0, 7.0, 8.0]).unwrap();
        let c = a.matmul(&b).unwrap();

        assert_eq!(c.rows(), 2);
        assert_eq!(c.cols(), 2);
        assert_eq!(c.get(0, 0), Some(&19.0));
        assert_eq!(c.get(0, 1), Some(&22.0));
        assert_eq!(c.get(1, 0), Some(&43.0));
        assert_eq!(c.get(1, 1), Some(&50.0));
    }

    #[test]
    fn test_matmul_identity() {
        // A × I = A
        let a = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]).unwrap();
        let identity = Matrix::identity(2);
        let result = a.matmul(&identity).unwrap();

        assert_eq!(result.get(0, 0), Some(&1.0));
        assert_eq!(result.get(0, 1), Some(&2.0));
        assert_eq!(result.get(1, 0), Some(&3.0));
        assert_eq!(result.get(1, 1), Some(&4.0));
    }

    #[test]
    fn test_matmul_zeros() {
        // A × 0 = 0
        let a = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]).unwrap();
        let zeros = Matrix::zeros(2, 2);
        let result = a.matmul(&zeros).unwrap();

        for &val in result.as_slice() {
            assert_eq!(val, 0.0);
        }
    }

    #[test]
    fn test_matmul_dimension_mismatch() {
        // 2×3 matrix cannot multiply with 2×2 matrix (inner dimensions don't match)
        let a = Matrix::from_vec(2, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        let b = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]).unwrap();
        let result = a.matmul(&b);

        assert!(matches!(result, Err(TruenoError::InvalidInput(_))));
    }

    #[test]
    fn test_matmul_non_square() {
        // 2×3 × 3×2 = 2×2
        // [[1, 2, 3],   [[7,  8],    [[58,  64],
        //  [4, 5, 6]] ×  [9, 10],  =  [139, 154]]
        //                [11, 12]]
        let a = Matrix::from_vec(2, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        let b = Matrix::from_vec(3, 2, vec![7.0, 8.0, 9.0, 10.0, 11.0, 12.0]).unwrap();
        let c = a.matmul(&b).unwrap();

        assert_eq!(c.rows(), 2);
        assert_eq!(c.cols(), 2);
        assert_eq!(c.get(0, 0), Some(&58.0));
        assert_eq!(c.get(0, 1), Some(&64.0));
        assert_eq!(c.get(1, 0), Some(&139.0));
        assert_eq!(c.get(1, 1), Some(&154.0));
    }

    #[test]
    fn test_matmul_single_element() {
        // 1×1 × 1×1 = 1×1
        let a = Matrix::from_vec(1, 1, vec![3.0]).unwrap();
        let b = Matrix::from_vec(1, 1, vec![4.0]).unwrap();
        let c = a.matmul(&b).unwrap();

        assert_eq!(c.rows(), 1);
        assert_eq!(c.cols(), 1);
        assert_eq!(c.get(0, 0), Some(&12.0));
    }

    #[test]
    fn test_matmul_remainder_rows() {
        // TRUENO-SPEC-014: Test matmul with rows not divisible by 4
        // This exercises the remainder handling path in SIMD matmul
        // 5×8 × 8×6 = 5×6 (5 % 4 = 1 remainder row)
        let a = Matrix::from_vec(5, 8, (0..40).map(|i| (i + 1) as f32).collect()).unwrap();
        let b = Matrix::from_vec(8, 6, (0..48).map(|i| (i + 1) as f32).collect()).unwrap();
        let c = a.matmul(&b).unwrap();

        assert_eq!(c.rows(), 5);
        assert_eq!(c.cols(), 6);

        // Verify using naive calculation for first and last row
        // First row: [1,2,3,4,5,6,7,8] . columns of B
        let expected_00 = (1..=8)
            .zip((0..48).step_by(6).map(|i| (i + 1) as f32))
            .map(|(a, b)| a as f32 * b)
            .sum::<f32>();
        assert!((c.get(0, 0).unwrap() - expected_00).abs() < 1.0);
    }

    #[test]
    fn test_matmul_remainder_rows_7() {
        // TRUENO-SPEC-014: 7×8 × 8×5 = 7×5 (7 % 4 = 3 remainder rows)
        let a = Matrix::from_vec(7, 8, (0..56).map(|_| 1.0f32).collect()).unwrap();
        let b = Matrix::from_vec(8, 5, (0..40).map(|_| 1.0f32).collect()).unwrap();
        let c = a.matmul(&b).unwrap();

        assert_eq!(c.rows(), 7);
        assert_eq!(c.cols(), 5);
        // Each element should be 8.0 (dot product of 8 ones)
        for &val in c.as_slice() {
            assert!((val - 8.0).abs() < 1e-5);
        }
    }

    // ===== Backend Equivalence Tests =====

    #[test]
    fn test_matmul_simd_equivalence_small() {
        // Small matrix (below SIMD threshold) - verify both paths work
        let a = Matrix::from_vec(8, 8, (0..64).map(|i| i as f32).collect()).unwrap();
        let b = Matrix::from_vec(8, 8, (0..64).map(|i| (i * 2) as f32).collect()).unwrap();

        let mut result_naive = Matrix::zeros(8, 8);
        let mut result_simd = Matrix::zeros(8, 8);

        a.matmul_naive(&b, &mut result_naive).unwrap();
        a.matmul_simd(&b, &mut result_simd).unwrap();

        // Results should be identical
        for i in 0..8 {
            for j in 0..8 {
                let naive_val = result_naive.get(i, j).unwrap();
                let simd_val = result_simd.get(i, j).unwrap();
                assert!(
                    (naive_val - simd_val).abs() < 1e-5,
                    "Mismatch at ({}, {}): naive={}, simd={}",
                    i,
                    j,
                    naive_val,
                    simd_val
                );
            }
        }
    }

    #[test]
    fn test_matmul_simd_equivalence_large() {
        // Large matrix (above SIMD threshold) - verify SIMD correctness
        let size = 128;
        let a = Matrix::from_vec(
            size,
            size,
            (0..size * size).map(|i| (i % 100) as f32).collect(),
        )
        .unwrap();
        let b = Matrix::from_vec(
            size,
            size,
            (0..size * size).map(|i| ((i * 2) % 100) as f32).collect(),
        )
        .unwrap();

        let mut result_naive = Matrix::zeros(size, size);
        let mut result_simd = Matrix::zeros(size, size);

        a.matmul_naive(&b, &mut result_naive).unwrap();
        a.matmul_simd(&b, &mut result_simd).unwrap();

        // Results should be identical (within floating-point tolerance)
        for i in 0..size {
            for j in 0..size {
                let naive_val = result_naive.get(i, j).unwrap();
                let simd_val = result_simd.get(i, j).unwrap();
                assert!(
                    (naive_val - simd_val).abs() < 1e-3,
                    "Mismatch at ({}, {}): naive={}, simd={}",
                    i,
                    j,
                    naive_val,
                    simd_val
                );
            }
        }
    }

    #[test]
    fn test_matmul_simd_equivalence_rectangular() {
        // Rectangular matrices
        let a = Matrix::from_vec(64, 128, (0..64 * 128).map(|i| i as f32).collect()).unwrap();
        let b = Matrix::from_vec(128, 32, (0..128 * 32).map(|i| (i * 3) as f32).collect()).unwrap();

        let mut result_naive = Matrix::zeros(64, 32);
        let mut result_simd = Matrix::zeros(64, 32);

        a.matmul_naive(&b, &mut result_naive).unwrap();
        a.matmul_simd(&b, &mut result_simd).unwrap();

        // Results should be identical (use relative tolerance for large values)
        for i in 0..64 {
            for j in 0..32 {
                let naive_val = result_naive.get(i, j).unwrap();
                let simd_val = result_simd.get(i, j).unwrap();
                let diff = (naive_val - simd_val).abs();
                let tolerance = if naive_val.abs() > 1.0 {
                    naive_val.abs() * 1e-5 // Relative tolerance for large values
                } else {
                    1e-5 // Absolute tolerance for small values
                };
                assert!(
                    diff < tolerance,
                    "Mismatch at ({}, {}): naive={}, simd={}, diff={}",
                    i,
                    j,
                    naive_val,
                    simd_val,
                    diff
                );
            }
        }
    }

    // ===== Cache-Aware Blocking Tests (Issue #10) =====

    #[test]
    fn test_matmul_blocking_small_matrices() {
        // Small matrices (≤32) should use simple path (no blocking overhead)
        let sizes = vec![8, 16, 32];
        for size in sizes {
            let a =
                Matrix::from_vec(size, size, (0..size * size).map(|i| i as f32).collect()).unwrap();
            let b = Matrix::from_vec(
                size,
                size,
                (0..size * size).map(|i| (i * 2) as f32).collect(),
            )
            .unwrap();

            let mut result_naive = Matrix::zeros(size, size);
            let mut result_simd = Matrix::zeros(size, size);

            a.matmul_naive(&b, &mut result_naive).unwrap();
            a.matmul_simd(&b, &mut result_simd).unwrap();

            // Verify correctness
            for i in 0..size {
                for j in 0..size {
                    let naive_val = result_naive.get(i, j).unwrap();
                    let simd_val = result_simd.get(i, j).unwrap();
                    let diff = (naive_val - simd_val).abs();
                    let tolerance = if naive_val.abs() > 1.0 {
                        naive_val.abs() * 1e-4
                    } else {
                        1e-4
                    };
                    assert!(
                        diff < tolerance,
                        "Size {}: Mismatch at ({}, {}): naive={}, simd={}, diff={}",
                        size,
                        i,
                        j,
                        naive_val,
                        simd_val,
                        diff
                    );
                }
            }
        }
    }

    #[test]
    fn test_matmul_blocking_medium_matrices() {
        // Medium matrices (>32, <512) should benefit from L2 blocking
        let sizes = vec![64, 128, 256];
        for size in sizes {
            let a = Matrix::from_vec(
                size,
                size,
                (0..size * size).map(|i| (i % 100) as f32).collect(),
            )
            .unwrap();
            let b = Matrix::from_vec(
                size,
                size,
                (0..size * size).map(|i| ((i * 3) % 100) as f32).collect(),
            )
            .unwrap();

            let mut result_naive = Matrix::zeros(size, size);
            let mut result_simd = Matrix::zeros(size, size);

            a.matmul_naive(&b, &mut result_naive).unwrap();
            a.matmul_simd(&b, &mut result_simd).unwrap();

            // Verify correctness with relative tolerance for large accumulated values
            for i in 0..size {
                for j in 0..size {
                    let naive_val = result_naive.get(i, j).unwrap();
                    let simd_val = result_simd.get(i, j).unwrap();
                    let diff = (naive_val - simd_val).abs();
                    let tolerance = if naive_val.abs() > 1.0 {
                        naive_val.abs() * 1e-3 // More relaxed for large values
                    } else {
                        1e-3
                    };
                    assert!(
                        diff < tolerance,
                        "Size {}: Mismatch at ({}, {}): naive={}, simd={}, diff={}",
                        size,
                        i,
                        j,
                        naive_val,
                        simd_val,
                        diff
                    );
                }
            }
        }
    }

    #[test]
    fn test_matmul_blocking_non_aligned_sizes() {
        // Test matrices with sizes not aligned to block boundaries
        let test_cases = vec![
            (33, 33, 33),    // Just over small threshold
            (65, 65, 65),    // Just over L2 block size
            (100, 100, 100), // Middle of L2 block
            (127, 127, 127), // Just under 2× L2 block size
        ];

        for (m, k, n) in test_cases {
            let a = Matrix::from_vec(m, k, (0..m * k).map(|i| (i % 50) as f32).collect()).unwrap();
            let b = Matrix::from_vec(k, n, (0..k * n).map(|i| ((i * 2) % 50) as f32).collect())
                .unwrap();

            let mut result_naive = Matrix::zeros(m, n);
            let mut result_simd = Matrix::zeros(m, n);

            a.matmul_naive(&b, &mut result_naive).unwrap();
            a.matmul_simd(&b, &mut result_simd).unwrap();

            // Verify correctness
            for i in 0..m {
                for j in 0..n {
                    let naive_val = result_naive.get(i, j).unwrap();
                    let simd_val = result_simd.get(i, j).unwrap();
                    let diff = (naive_val - simd_val).abs();
                    let tolerance = if naive_val.abs() > 1.0 {
                        naive_val.abs() * 1e-3
                    } else {
                        1e-3
                    };
                    assert!(
                        diff < tolerance,
                        "Size {}×{}×{}: Mismatch at ({}, {}): naive={}, simd={}, diff={}",
                        m,
                        k,
                        n,
                        i,
                        j,
                        naive_val,
                        simd_val,
                        diff
                    );
                }
            }
        }
    }

    #[test]
    fn test_matmul_blocking_large_matrices() {
        // Large matrix to verify blocking algorithm correctness
        // Keep size manageable for test speed but large enough to trigger blocking
        let size = 256;
        let a = Matrix::from_vec(
            size,
            size,
            (0..size * size)
                .map(|i| ((i % 100) as f32) / 10.0)
                .collect(),
        )
        .unwrap();
        let b = Matrix::from_vec(
            size,
            size,
            (0..size * size)
                .map(|i| (((i * 7) % 100) as f32) / 10.0)
                .collect(),
        )
        .unwrap();

        let mut result_naive = Matrix::zeros(size, size);
        let mut result_simd = Matrix::zeros(size, size);

        a.matmul_naive(&b, &mut result_naive).unwrap();
        a.matmul_simd(&b, &mut result_simd).unwrap();

        // Verify correctness with appropriate tolerance for accumulated floating-point errors
        let mut max_diff = 0.0f32;
        let mut mismatches = 0;
        for i in 0..size {
            for j in 0..size {
                let naive_val = result_naive.get(i, j).unwrap();
                let simd_val = result_simd.get(i, j).unwrap();
                let diff = (naive_val - simd_val).abs();
                let tolerance = if naive_val.abs() > 1.0 {
                    naive_val.abs() * 1e-2 // Relaxed tolerance for large accumulated values
                } else {
                    1e-2
                };
                if diff >= tolerance {
                    mismatches += 1;
                    if mismatches <= 5 {
                        eprintln!(
                            "Mismatch at ({}, {}): naive={}, simd={}, diff={}, tolerance={}",
                            i, j, naive_val, simd_val, diff, tolerance
                        );
                    }
                }
                max_diff = max_diff.max(diff);
            }
        }
        assert_eq!(
            mismatches, 0,
            "Found {} mismatches in {}×{} matmul, max_diff={}",
            mismatches, size, size, max_diff
        );
    }

    #[test]
    fn test_matmul_3level_blocking() {
        // Phase 3: Test 3-level cache blocking for very large matrices (≥512×512)
        // This test ensures the L3 → L2 → micro-kernel hierarchy works correctly
        let size = 512; // Triggers 3-level blocking (L3_THRESHOLD = 512)
        let a = Matrix::from_vec(
            size,
            size,
            (0..size * size)
                .map(|i| ((i % 100) as f32) / 10.0)
                .collect(),
        )
        .unwrap();
        let b = Matrix::from_vec(
            size,
            size,
            (0..size * size)
                .map(|i| (((i * 7) % 100) as f32) / 10.0)
                .collect(),
        )
        .unwrap();

        let mut result_naive = Matrix::zeros(size, size);
        let mut result_simd = Matrix::zeros(size, size);

        a.matmul_naive(&b, &mut result_naive).unwrap();
        a.matmul_simd(&b, &mut result_simd).unwrap();

        // Verify correctness with appropriate tolerance
        let mut max_diff = 0.0f32;
        let mut mismatches = 0;
        for i in 0..size {
            for j in 0..size {
                let naive_val = result_naive.get(i, j).unwrap();
                let simd_val = result_simd.get(i, j).unwrap();
                let diff = (naive_val - simd_val).abs();
                let tolerance = if naive_val.abs() > 1.0 {
                    naive_val.abs() * 1e-2
                } else {
                    1e-2
                };
                if diff >= tolerance {
                    mismatches += 1;
                    if mismatches <= 5 {
                        eprintln!(
                            "Mismatch at ({}, {}): naive={}, simd={}, diff={}, tolerance={}",
                            i, j, naive_val, simd_val, diff, tolerance
                        );
                    }
                }
                max_diff = max_diff.max(diff);
            }
        }
        assert_eq!(
            mismatches, 0,
            "Found {} mismatches in {}×{} matmul (3-level blocking), max_diff={}",
            mismatches, size, size, max_diff
        );
    }

    #[test]
    #[cfg(feature = "parallel")]
    fn test_matmul_parallel_1024() {
        // Phase 4: Test parallel matmul for 1024×1024 matrices
        // This triggers the parallel path (PARALLEL_THRESHOLD = 1024)
        let size = 1024;
        let a = Matrix::from_vec(
            size,
            size,
            (0..size * size)
                .map(|i| ((i % 100) as f32) / 10.0)
                .collect(),
        )
        .unwrap();
        let b = Matrix::from_vec(
            size,
            size,
            (0..size * size)
                .map(|i| (((i * 7) % 100) as f32) / 10.0)
                .collect(),
        )
        .unwrap();

        let mut result_naive = Matrix::zeros(size, size);
        let mut result_parallel = Matrix::zeros(size, size);

        a.matmul_naive(&b, &mut result_naive).unwrap();
        a.matmul_simd(&b, &mut result_parallel).unwrap(); // Uses parallel path with 'parallel' feature

        // Verify correctness with appropriate tolerance
        let mut max_diff = 0.0f32;
        let mut mismatches = 0;
        for i in 0..size {
            for j in 0..size {
                let naive_val = result_naive.get(i, j).unwrap();
                let parallel_val = result_parallel.get(i, j).unwrap();
                let diff = (naive_val - parallel_val).abs();
                let tolerance = if naive_val.abs() > 1.0 {
                    naive_val.abs() * 1e-2
                } else {
                    1e-2
                };
                if diff >= tolerance {
                    mismatches += 1;
                    if mismatches <= 5 {
                        eprintln!(
                            "Mismatch at ({}, {}): naive={}, parallel={}, diff={}, tolerance={}",
                            i, j, naive_val, parallel_val, diff, tolerance
                        );
                    }
                }
                max_diff = max_diff.max(diff);
            }
        }
        assert_eq!(
            mismatches, 0,
            "Found {} mismatches in {}×{} parallel matmul, max_diff={}",
            mismatches, size, size, max_diff
        );
    }

    #[test]
    #[cfg(feature = "parallel")]
    fn test_matvec_parallel_4096() {
        // Test parallel matvec for very large matrices (≥4096 rows)
        // This triggers the parallel path (PARALLEL_THRESHOLD = 4096)
        let rows = 4096;
        let cols = 512;

        let matrix = Matrix::from_vec(
            rows,
            cols,
            (0..rows * cols)
                .map(|i| ((i % 100) as f32) / 10.0)
                .collect(),
        )
        .unwrap();

        let vector = Vector::from_slice(
            &(0..cols)
                .map(|i| ((i % 50) as f32) / 5.0)
                .collect::<Vec<f32>>(),
        );

        // Compute result (should use parallel path)
        let result = matrix.matvec(&vector).unwrap();

        // Verify result shape
        assert_eq!(result.len(), rows);

        // Verify correctness by comparing with manual dot product calculation
        // Check a few sample rows
        for sample_row in [0, 1024, 2048, 3072, 4095] {
            let row_start = sample_row * cols;
            let row = &matrix.data[row_start..(row_start + cols)];

            // Manual dot product
            let expected: f32 = row
                .iter()
                .zip(vector.as_slice().iter())
                .map(|(a, b)| a * b)
                .sum();

            let actual = result.as_slice()[sample_row];
            let diff = (expected - actual).abs();
            let tolerance = if expected.abs() > 1.0 {
                expected.abs() * 1e-3
            } else {
                1e-3
            };

            assert!(
                diff < tolerance,
                "Mismatch at row {}: expected={}, actual={}, diff={}",
                sample_row,
                expected,
                actual,
                diff
            );
        }
    }

    // ===== Phase 2 Micro-kernel Tests (Issue #10) =====

    #[test]
    #[cfg(target_arch = "x86_64")]
    fn test_horizontal_sum_avx2() {
        // Test the AVX2 horizontal sum helper function
        if !is_x86_feature_detected!("avx2") {
            println!("Skipping AVX2 horizontal sum test (CPU doesn't support AVX2)");
            return;
        }

        use std::arch::x86_64::*;

        unsafe {
            // Test case 1: All ones
            let v = _mm256_set1_ps(1.0);
            let sum = Matrix::<f32>::horizontal_sum_avx2(v);
            assert!((sum - 8.0).abs() < 1e-6, "Expected 8.0, got {}", sum);

            // Test case 2: Sequence 1..8
            let v = _mm256_setr_ps(1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0);
            let sum = Matrix::<f32>::horizontal_sum_avx2(v);
            assert!((sum - 36.0).abs() < 1e-6, "Expected 36.0, got {}", sum);

            // Test case 3: Alternating signs
            let v = _mm256_setr_ps(1.0, -1.0, 1.0, -1.0, 1.0, -1.0, 1.0, -1.0);
            let sum = Matrix::<f32>::horizontal_sum_avx2(v);
            assert!(sum.abs() < 1e-6, "Expected ~0.0, got {}", sum);

            // Test case 4: Large values
            let v = _mm256_setr_ps(100.0, 200.0, 300.0, 400.0, 500.0, 600.0, 700.0, 800.0);
            let sum = Matrix::<f32>::horizontal_sum_avx2(v);
            assert!((sum - 3600.0).abs() < 1e-3, "Expected 3600.0, got {}", sum);

            // Test case 5: Mixed positive/negative
            let v = _mm256_setr_ps(10.5, -5.25, 3.75, -8.0, 12.0, -6.5, 4.25, -2.75);
            let expected = 10.5 - 5.25 + 3.75 - 8.0 + 12.0 - 6.5 + 4.25 - 2.75;
            let sum = Matrix::<f32>::horizontal_sum_avx2(v);
            assert!(
                (sum - expected).abs() < 1e-5,
                "Expected {}, got {}",
                expected,
                sum
            );
        }
    }

    #[test]
    #[cfg(target_arch = "x86_64")]
    fn test_matmul_microkernel_4x1_avx2() {
        // Test the 4×1 AVX2 micro-kernel for matrix multiplication
        if !is_x86_feature_detected!("avx2") || !is_x86_feature_detected!("fma") {
            println!("Skipping AVX2 micro-kernel test (CPU doesn't support AVX2/FMA)");
            return;
        }

        // Test case 1: Simple dot products
        // A rows: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]
        // B col:  [1, 1, 1, 1, 1, 1, 1, 1, 1,  1,  1,  1,  1,  1,  1,  1]
        // Expected: Row sums
        {
            let row0: Vec<f32> = (1..=16).map(|x| x as f32).collect();
            let row1: Vec<f32> = (17..=32).map(|x| x as f32).collect();
            let row2: Vec<f32> = (33..=48).map(|x| x as f32).collect();
            let row3: Vec<f32> = (49..=64).map(|x| x as f32).collect();
            let b_col = vec![1.0f32; 16];

            let a_rows = [
                row0.as_slice(),
                row1.as_slice(),
                row2.as_slice(),
                row3.as_slice(),
            ];
            let mut results = [0.0f32; 4];

            unsafe {
                Matrix::<f32>::matmul_microkernel_4x1_avx2(a_rows, &b_col, &mut results);
            }

            // Expected: sum(1..16), sum(17..32), sum(33..48), sum(49..64)
            let expected = [
                (1..=16).sum::<i32>() as f32,
                (17..=32).sum::<i32>() as f32,
                (33..=48).sum::<i32>() as f32,
                (49..=64).sum::<i32>() as f32,
            ];

            for i in 0..4 {
                assert!(
                    (results[i] - expected[i]).abs() < 1e-3,
                    "Row {}: expected {}, got {}",
                    i,
                    expected[i],
                    results[i]
                );
            }
        }

        // Test case 2: Identity-like pattern
        // Each row is all zeros except one 1.0
        {
            let row0 = vec![
                1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
            ];
            let row1 = vec![
                0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
            ];
            let row2 = vec![
                0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
            ];
            let row3 = vec![
                0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0,
            ];
            let b_col: Vec<f32> = (1..=16).map(|x| x as f32).collect();

            let a_rows = [
                row0.as_slice(),
                row1.as_slice(),
                row2.as_slice(),
                row3.as_slice(),
            ];
            let mut results = [0.0f32; 4];

            unsafe {
                Matrix::<f32>::matmul_microkernel_4x1_avx2(a_rows, &b_col, &mut results);
            }

            // Expected: Each result picks one element from b_col
            let expected = [1.0, 2.0, 3.0, 4.0];
            for i in 0..4 {
                assert!(
                    (results[i] - expected[i]).abs() < 1e-6,
                    "Row {}: expected {}, got {}",
                    i,
                    expected[i],
                    results[i]
                );
            }
        }

        // Test case 3: Non-aligned size (not multiple of 8)
        // Size 10 (8 + 2 remainder)
        {
            let row0: Vec<f32> = (1..=10).map(|x| x as f32).collect();
            let row1: Vec<f32> = (11..=20).map(|x| x as f32).collect();
            let row2: Vec<f32> = (21..=30).map(|x| x as f32).collect();
            let row3: Vec<f32> = (31..=40).map(|x| x as f32).collect();
            let b_col = vec![2.0f32; 10];

            let a_rows = [
                row0.as_slice(),
                row1.as_slice(),
                row2.as_slice(),
                row3.as_slice(),
            ];
            let mut results = [0.0f32; 4];

            unsafe {
                Matrix::<f32>::matmul_microkernel_4x1_avx2(a_rows, &b_col, &mut results);
            }

            // Expected: 2× each row sum
            let expected = [
                2.0 * (1..=10).sum::<i32>() as f32,
                2.0 * (11..=20).sum::<i32>() as f32,
                2.0 * (21..=30).sum::<i32>() as f32,
                2.0 * (31..=40).sum::<i32>() as f32,
            ];

            for i in 0..4 {
                assert!(
                    (results[i] - expected[i]).abs() < 1e-3,
                    "Row {}: expected {}, got {}",
                    i,
                    expected[i],
                    results[i]
                );
            }
        }

        // Test case 4: Mixed positive/negative values
        {
            let row0 = vec![
                1.0, -2.0, 3.0, -4.0, 5.0, -6.0, 7.0, -8.0, 9.0, -10.0, 11.0, -12.0, 13.0, -14.0,
                15.0, -16.0,
            ];
            let row1 = vec![
                2.0, -4.0, 6.0, -8.0, 10.0, -12.0, 14.0, -16.0, 18.0, -20.0, 22.0, -24.0, 26.0,
                -28.0, 30.0, -32.0,
            ];
            let row2 = vec![
                0.5, -1.0, 1.5, -2.0, 2.5, -3.0, 3.5, -4.0, 4.5, -5.0, 5.5, -6.0, 6.5, -7.0, 7.5,
                -8.0,
            ];
            let row3 = vec![
                10.0, -10.0, 10.0, -10.0, 10.0, -10.0, 10.0, -10.0, 10.0, -10.0, 10.0, -10.0, 10.0,
                -10.0, 10.0, -10.0,
            ];
            let b_col = vec![
                1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0,
            ];

            let a_rows = [
                row0.as_slice(),
                row1.as_slice(),
                row2.as_slice(),
                row3.as_slice(),
            ];
            let mut results = [0.0f32; 4];

            unsafe {
                Matrix::<f32>::matmul_microkernel_4x1_avx2(a_rows, &b_col, &mut results);
            }

            // Compute expected manually
            let expected = [
                row0.iter().sum::<f32>(),
                row1.iter().sum::<f32>(),
                row2.iter().sum::<f32>(),
                row3.iter().sum::<f32>(),
            ];

            for i in 0..4 {
                assert!(
                    (results[i] - expected[i]).abs() < 1e-4,
                    "Row {}: expected {}, got {}",
                    i,
                    expected[i],
                    results[i]
                );
            }
        }

        // Test case 5: Zero accumulation
        {
            let row0 = vec![0.0f32; 16];
            let row1 = vec![0.0f32; 16];
            let row2 = vec![0.0f32; 16];
            let row3 = vec![0.0f32; 16];
            let b_col: Vec<f32> = (1..=16).map(|x| x as f32).collect();

            let a_rows = [
                row0.as_slice(),
                row1.as_slice(),
                row2.as_slice(),
                row3.as_slice(),
            ];
            let mut results = [0.0f32; 4];

            unsafe {
                Matrix::<f32>::matmul_microkernel_4x1_avx2(a_rows, &b_col, &mut results);
            }

            for (i, &result) in results.iter().enumerate() {
                assert!(
                    result.abs() < 1e-6,
                    "Row {}: expected 0.0, got {}",
                    i,
                    result
                );
            }
        }

        // Test case 6: Verify FMA correctness (a * b + c pattern)
        // Micro-kernel computes: sum(a[i] * b[i])
        {
            let row0 = vec![
                1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0,
                16.0,
            ];
            let row1 = vec![
                2.0, 4.0, 6.0, 8.0, 10.0, 12.0, 14.0, 16.0, 18.0, 20.0, 22.0, 24.0, 26.0, 28.0,
                30.0, 32.0,
            ];
            let row2 = vec![
                0.5, 1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5, 5.0, 5.5, 6.0, 6.5, 7.0, 7.5, 8.0,
            ];
            let row3 = vec![
                3.0, 6.0, 9.0, 12.0, 15.0, 18.0, 21.0, 24.0, 27.0, 30.0, 33.0, 36.0, 39.0, 42.0,
                45.0, 48.0,
            ];
            let b_col = vec![
                0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.5,
            ];

            let a_rows = [
                row0.as_slice(),
                row1.as_slice(),
                row2.as_slice(),
                row3.as_slice(),
            ];
            let mut results = [0.0f32; 4];

            unsafe {
                Matrix::<f32>::matmul_microkernel_4x1_avx2(a_rows, &b_col, &mut results);
            }

            // Expected: 0.5 × each row sum
            let expected = [
                0.5 * row0.iter().sum::<f32>(),
                0.5 * row1.iter().sum::<f32>(),
                0.5 * row2.iter().sum::<f32>(),
                0.5 * row3.iter().sum::<f32>(),
            ];

            for i in 0..4 {
                assert!(
                    (results[i] - expected[i]).abs() < 1e-3,
                    "Row {}: expected {}, got {}",
                    i,
                    expected[i],
                    results[i]
                );
            }
        }
    }

    // ===== GPU Tests =====

    #[test]
    #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    fn test_gpu_availability() {
        use crate::backends::gpu::GpuBackend;
        // Just test that we can check GPU availability without crashing
        let _available = GpuBackend::is_available();
        // Note: We don't assert availability since CI may not have GPU
    }

    #[test]
    #[cfg(all(feature = "gpu", not(target_arch = "wasm32")))]
    #[ignore] // Ignore by default since CI may not have GPU
    fn test_gpu_matmul_basic() {
        use crate::backends::gpu::GpuBackend;

        if !GpuBackend::is_available() {
            eprintln!("GPU not available, skipping test");
            return;
        }

        // Small test matrix (will use GPU if threshold is low enough)
        let a = Matrix::from_vec(
            4,
            4,
            vec![
                1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0, 15.0,
                16.0,
            ],
        )
        .unwrap();

        let b = Matrix::from_vec(
            4,
            4,
            vec![
                16.0, 15.0, 14.0, 13.0, 12.0, 11.0, 10.0, 9.0, 8.0, 7.0, 6.0, 5.0, 4.0, 3.0, 2.0,
                1.0,
            ],
        )
        .unwrap();

        // Try GPU matmul directly
        let result = a.matmul_gpu(&b);

        if let Ok(c) = result {
            // Verify some basic properties
            assert_eq!(c.rows(), 4);
            assert_eq!(c.cols(), 4);

            // Verify against known result (first element)
            // [1,2,3,4] · [16,12,8,4] = 16+24+24+16 = 80
            assert!((c.get(0, 0).unwrap() - 80.0).abs() < 1e-4);
        } else {
            eprintln!("GPU matmul failed: {:?}", result);
        }
    }

    // ===== Transpose Tests =====

    #[test]
    fn test_transpose_basic() {
        // [[1, 2, 3],     [[1, 4],
        //  [4, 5, 6]]  →   [2, 5],
        //                  [3, 6]]
        let m = Matrix::from_vec(2, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        let t = m.transpose();

        assert_eq!(t.rows(), 3);
        assert_eq!(t.cols(), 2);
        assert_eq!(t.get(0, 0), Some(&1.0));
        assert_eq!(t.get(0, 1), Some(&4.0));
        assert_eq!(t.get(1, 0), Some(&2.0));
        assert_eq!(t.get(1, 1), Some(&5.0));
        assert_eq!(t.get(2, 0), Some(&3.0));
        assert_eq!(t.get(2, 1), Some(&6.0));
    }

    #[test]
    fn test_transpose_square() {
        // [[1, 2],     [[1, 3],
        //  [3, 4]]  →   [2, 4]]
        let m = Matrix::from_vec(2, 2, vec![1.0, 2.0, 3.0, 4.0]).unwrap();
        let t = m.transpose();

        assert_eq!(t.rows(), 2);
        assert_eq!(t.cols(), 2);
        assert_eq!(t.get(0, 0), Some(&1.0));
        assert_eq!(t.get(0, 1), Some(&3.0));
        assert_eq!(t.get(1, 0), Some(&2.0));
        assert_eq!(t.get(1, 1), Some(&4.0));
    }

    #[test]
    fn test_transpose_single_row() {
        // [[1, 2, 3]] → [[1],
        //                 [2],
        //                 [3]]
        let m = Matrix::from_vec(1, 3, vec![1.0, 2.0, 3.0]).unwrap();
        let t = m.transpose();

        assert_eq!(t.rows(), 3);
        assert_eq!(t.cols(), 1);
        assert_eq!(t.get(0, 0), Some(&1.0));
        assert_eq!(t.get(1, 0), Some(&2.0));
        assert_eq!(t.get(2, 0), Some(&3.0));
    }

    #[test]
    fn test_transpose_single_col() {
        // [[1],        [[1, 2, 3]]
        //  [2],   →
        //  [3]]
        let m = Matrix::from_vec(3, 1, vec![1.0, 2.0, 3.0]).unwrap();
        let t = m.transpose();

        assert_eq!(t.rows(), 1);
        assert_eq!(t.cols(), 3);
        assert_eq!(t.get(0, 0), Some(&1.0));
        assert_eq!(t.get(0, 1), Some(&2.0));
        assert_eq!(t.get(0, 2), Some(&3.0));
    }

    #[test]
    fn test_transpose_single_element() {
        // [[5]] → [[5]]
        let m = Matrix::from_vec(1, 1, vec![5.0]).unwrap();
        let t = m.transpose();

        assert_eq!(t.rows(), 1);
        assert_eq!(t.cols(), 1);
        assert_eq!(t.get(0, 0), Some(&5.0));
    }

    #[test]
    fn test_transpose_identity() {
        // I^T = I
        let identity = Matrix::identity(3);
        let t = identity.transpose();

        assert_eq!(t.rows(), 3);
        assert_eq!(t.cols(), 3);

        // Check it's still identity
        for i in 0..3 {
            for j in 0..3 {
                let expected = if i == j { 1.0 } else { 0.0 };
                assert_eq!(t.get(i, j), Some(&expected));
            }
        }
    }
}

// Property-based tests for matmul
#[cfg(test)]
mod property_tests {
    use super::*;
    use proptest::prelude::*;

    /// Generate a matrix of given dimensions with random values
    fn matrix_strategy(rows: usize, cols: usize) -> impl Strategy<Value = Matrix<f32>> {
        proptest::collection::vec(-100.0f32..100.0, rows * cols)
            .prop_map(move |data| Matrix::from_vec(rows, cols, data).unwrap())
    }

    proptest! {
        #![proptest_config(ProptestConfig::with_cases(100))]

        /// Property: Matrix multiplication is associative
        /// (A × B) × C = A × (B × C)
        #[test]
        fn test_matmul_associative(
            a in matrix_strategy(3, 4),
            b in matrix_strategy(4, 5),
            c in matrix_strategy(5, 3)
        ) {
            let ab = a.matmul(&b).unwrap();
            let ab_c = ab.matmul(&c).unwrap();

            let bc = b.matmul(&c).unwrap();
            let a_bc = a.matmul(&bc).unwrap();

            // Check dimensions
            prop_assert_eq!(ab_c.rows(), a_bc.rows());
            prop_assert_eq!(ab_c.cols(), a_bc.cols());

            // Check values with tolerance for floating-point errors
            // Use relative tolerance for large values, absolute for small values
            for i in 0..ab_c.rows() {
                for j in 0..ab_c.cols() {
                    let val1 = ab_c.get(i, j).unwrap();
                    let val2 = a_bc.get(i, j).unwrap();
                    let diff = (val1 - val2).abs();
                    let max_val = val1.abs().max(val2.abs());

                    // Use hybrid tolerance: absolute for small values, relative for large
                    // Matrix multiplication accumulates rounding errors across multiple operations
                    // Different evaluation orders (A×B)×C vs A×(B×C) produce different rounding
                    // AVX512 FMA instructions accumulate errors differently than scalar operations
                    // Tolerance must account for:
                    //   - 3-way matrix multiplication (more accumulation than 2-way)
                    //   - SIMD reordering (AVX512, AVX2, SSE2 all have different patterns)
                    //   - FMA vs separate multiply+add
                    let tolerance = if max_val < 1.0 {
                        1e-3  // Absolute tolerance for small values
                    } else {
                        max_val * 5e-2  // Relative tolerance (5%) for large values
                        // Increased from 1e-2 (1%) to 5e-2 (5%) for AVX512 FMA
                        // AVX512 FMA instructions have different rounding behavior:
                        //   (A×B)×C: Different op count than A×(B×C)
                        //   3-way matmul accumulates 4.3x more error than expected
                        //   Empirical: proptest regression shows 4.28% error
                        //   Industry standard: 1-5% for accumulated FP operations
                    };

                    prop_assert!(
                        diff < tolerance,
                        "Associativity failed at ({}, {}): {} != {} (diff: {}, tolerance: {})",
                        i, j, val1, val2, diff, tolerance
                    );
                }
            }
        }

        /// Property: Multiplying by identity matrix preserves the matrix
        /// A × I = A
        #[test]
        fn test_matmul_identity_property(
            rows in 1usize..10,
            cols in 1usize..10,
            data in proptest::collection::vec(-100.0f32..100.0, 1..100)
        ) {
            // Ensure data length matches dimensions
            let size = rows * cols;
            if data.len() < size {
                return Ok(());
            }
            let matrix_data = data[0..size].to_vec();

            let a = Matrix::from_vec(rows, cols, matrix_data).unwrap();
            let identity = Matrix::identity(cols);
            let result = a.matmul(&identity).unwrap();

            // Check dimensions
            prop_assert_eq!(result.rows(), a.rows());
            prop_assert_eq!(result.cols(), a.cols());

            // Check values (should be identical)
            for i in 0..rows {
                for j in 0..cols {
                    let original = a.get(i, j).unwrap();
                    let multiplied = result.get(i, j).unwrap();
                    let diff = (original - multiplied).abs();
                    prop_assert!(
                        diff < 1e-5,
                        "Identity property failed at ({}, {}): {} != {} (diff: {})",
                        i, j, original, multiplied, diff
                    );
                }
            }
        }

        /// Property: Dimension property
        /// If A is m×n and B is n×p, then A×B is m×p
        #[test]
        fn test_matmul_dimension_property(
            m in 1usize..10,
            n in 1usize..10,
            p in 1usize..10
        ) {
            let a = Matrix::zeros(m, n);
            let b = Matrix::zeros(n, p);
            let c = a.matmul(&b).unwrap();

            prop_assert_eq!(c.rows(), m);
            prop_assert_eq!(c.cols(), p);
        }

        /// Property: Double transpose returns original
        /// (A^T)^T = A
        #[test]
        fn test_transpose_double_transpose(
            a in matrix_strategy(5, 7)
        ) {
            let t = a.transpose();
            let tt = t.transpose();

            prop_assert_eq!(tt.rows(), a.rows());
            prop_assert_eq!(tt.cols(), a.cols());

            for i in 0..a.rows() {
                for j in 0..a.cols() {
                    prop_assert_eq!(tt.get(i, j), a.get(i, j));
                }
            }
        }

        /// Property: Transpose swaps dimensions
        /// If A is m×n, then A^T is n×m
        #[test]
        fn test_transpose_dimension_swap(
            m in 1usize..20,
            n in 1usize..20
        ) {
            let a = Matrix::zeros(m, n);
            let t = a.transpose();

            prop_assert_eq!(t.rows(), n);
            prop_assert_eq!(t.cols(), m);
        }

        /// Property: Transpose of product
        /// (A×B)^T = B^T×A^T
        #[test]
        fn test_transpose_of_product(
            a in matrix_strategy(3, 4),
            b in matrix_strategy(4, 5)
        ) {
            let ab = a.matmul(&b).unwrap();
            let ab_t = ab.transpose();

            let b_t = b.transpose();
            let a_t = a.transpose();
            let bt_at = b_t.matmul(&a_t).unwrap();

            prop_assert_eq!(ab_t.rows(), bt_at.rows());
            prop_assert_eq!(ab_t.cols(), bt_at.cols());

            // Check values with tolerance for floating-point errors
            for i in 0..ab_t.rows() {
                for j in 0..ab_t.cols() {
                    let val1 = ab_t.get(i, j).unwrap();
                    let val2 = bt_at.get(i, j).unwrap();
                    let diff = (val1 - val2).abs();
                    let max_val = val1.abs().max(val2.abs());

                    let tolerance = if max_val < 1.0 {
                        1e-3
                    } else {
                        max_val * 1e-3
                    };

                    prop_assert!(
                        diff < tolerance,
                        "Transpose of product failed at ({}, {}): {} != {} (diff: {}, tolerance: {})",
                        i, j, val1, val2, diff, tolerance
                    );
                }
            }
        }

        /// Matrix-vector multiplication: (A×B)×v = A×(B×v)
        #[test]
        fn test_matvec_associativity(
            a in matrix_strategy(3, 4),
            b in matrix_strategy(4, 5),
            v_data in prop::collection::vec(-10.0f32..10.0, 5)
        ) {
            let v = Vector::from_slice(&v_data);

            let ab = a.matmul(&b).unwrap();
            let ab_v = ab.matvec(&v).unwrap();

            let b_v = b.matvec(&v).unwrap();
            let a_bv = a.matvec(&b_v).unwrap();

            prop_assert_eq!(ab_v.len(), a_bv.len());

            for i in 0..ab_v.len() {
                let diff = (ab_v.as_slice()[i] - a_bv.as_slice()[i]).abs();
                let max_val = ab_v.as_slice()[i].abs().max(a_bv.as_slice()[i].abs());
                // Relaxed tolerance for SIMD backends (AVX512 accumulates more rounding error)
                let tolerance = if max_val < 1.0 { 1e-2 } else { max_val * 2e-2 };

                prop_assert!(
                    diff < tolerance,
                    "Associativity failed at index {}: {} != {} (diff: {}, tolerance: {})",
                    i, ab_v.as_slice()[i], a_bv.as_slice()[i], diff, tolerance
                );
            }
        }

        /// Vector-matrix multiplication: v×(A×B) = (v×A)×B
        #[test]
        fn test_vecmat_associativity(
            a in matrix_strategy(3, 4),
            b in matrix_strategy(4, 5),
            v_data in prop::collection::vec(-10.0f32..10.0, 3)
        ) {
            let v = Vector::from_slice(&v_data);

            let ab = a.matmul(&b).unwrap();
            let v_ab = Matrix::vecmat(&v, &ab).unwrap();

            let v_a = Matrix::vecmat(&v, &a).unwrap();
            let va_b = Matrix::vecmat(&v_a, &b).unwrap();

            prop_assert_eq!(v_ab.len(), va_b.len());

            for i in 0..v_ab.len() {
                let diff = (v_ab.as_slice()[i] - va_b.as_slice()[i]).abs();
                let max_val = v_ab.as_slice()[i].abs().max(va_b.as_slice()[i].abs());
                let tolerance = if max_val < 1.0 { 1e-2 } else { max_val * 1e-2 };

                prop_assert!(
                    diff < tolerance,
                    "Associativity failed at index {}: {} != {} (diff: {}, tolerance: {})",
                    i, v_ab.as_slice()[i], va_b.as_slice()[i], diff, tolerance
                );
            }
        }
    }

    // Unit tests for matrix-vector operations
    #[test]
    fn test_matvec_basic() {
        let m = Matrix::from_vec(2, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        let v = Vector::from_slice(&[1.0, 2.0, 3.0]);
        let result = m.matvec(&v).unwrap();

        // [[1, 2, 3]   [1]   [14]
        //  [4, 5, 6]] × [2] = [32]
        //               [3]
        assert_eq!(result.len(), 2);
        assert!((result.as_slice()[0] - 14.0).abs() < 1e-6);
        assert!((result.as_slice()[1] - 32.0).abs() < 1e-6);
    }

    #[test]
    fn test_matvec_identity() {
        let m = Matrix::identity(3);
        let v = Vector::from_slice(&[1.0, 2.0, 3.0]);
        let result = m.matvec(&v).unwrap();

        // I×v = v
        assert_eq!(result.as_slice(), v.as_slice());
    }

    #[test]
    fn test_matvec_dimension_mismatch() {
        let m = Matrix::from_vec(2, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        let v = Vector::from_slice(&[1.0, 2.0]); // Wrong size

        assert!(m.matvec(&v).is_err());
    }

    #[test]
    fn test_vecmat_basic() {
        let m = Matrix::from_vec(2, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        let v = Vector::from_slice(&[1.0, 2.0]);
        let result = Matrix::vecmat(&v, &m).unwrap();

        // [1, 2] × [[1, 2, 3]  = [9, 12, 15]
        //           [4, 5, 6]]
        assert_eq!(result.len(), 3);
        assert!((result.as_slice()[0] - 9.0).abs() < 1e-6);
        assert!((result.as_slice()[1] - 12.0).abs() < 1e-6);
        assert!((result.as_slice()[2] - 15.0).abs() < 1e-6);
    }

    #[test]
    fn test_vecmat_identity() {
        let m = Matrix::identity(3);
        let v = Vector::from_slice(&[1.0, 2.0, 3.0]);
        let result = Matrix::vecmat(&v, &m).unwrap();

        // v×I = v
        assert_eq!(result.as_slice(), v.as_slice());
    }

    #[test]
    fn test_vecmat_dimension_mismatch() {
        let m = Matrix::from_vec(2, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        let v = Vector::from_slice(&[1.0, 2.0, 3.0]); // Wrong size

        assert!(Matrix::vecmat(&v, &m).is_err());
    }

    #[test]
    fn test_matvec_zero_vector() {
        let m = Matrix::from_vec(2, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        let v = Vector::from_slice(&[0.0, 0.0, 0.0]);
        let result = m.matvec(&v).unwrap();

        // A×0 = 0
        assert_eq!(result.as_slice(), &[0.0, 0.0]);
    }

    #[test]
    fn test_vecmat_zero_vector() {
        let m = Matrix::from_vec(2, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        let v = Vector::from_slice(&[0.0, 0.0]);
        let result = Matrix::vecmat(&v, &m).unwrap();

        // 0×A = 0
        assert_eq!(result.as_slice(), &[0.0, 0.0, 0.0]);
    }

    #[test]
    fn test_matvec_transpose_equivalence() {
        // v^T × A = (A^T × v)^T
        // If A is m×n and v is m-dimensional, then:
        // - v^T × A is n-dimensional
        // - A^T is n×m, so A^T × v needs v to be n-dimensional
        // Actually, this is wrong. Let me use correct equivalence:
        // If A is m×n, v is n-dimensional:
        // - A × v is m-dimensional (matrix-vector)
        // - A^T is n×m, u is m-dimensional:
        // - u^T × A is n-dimensional (vector-matrix)
        // These are equivalent when u = A × v

        let m = Matrix::from_vec(3, 2, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();
        let v = Vector::from_slice(&[1.0, 2.0]); // 2-dimensional

        // A × v (3×2 times 2D = 3D result)
        let av = m.matvec(&v).unwrap();

        // v^T × A^T (2D times 2×3 = 3D result)
        let m_t = m.transpose(); // Now 2×3
        let v_mt = Matrix::vecmat(&v, &m_t).unwrap();

        // (A × v)^T = v^T × A^T
        assert_eq!(av.as_slice(), v_mt.as_slice());
    }

    // ===== 2D Convolution Tests =====

    #[test]
    fn test_convolve2d_basic_3x3() {
        // Simple 3x3 convolution with identity kernel (should preserve input)
        let input =
            Matrix::from_vec(3, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0]).unwrap();

        // 1x1 identity kernel (should return center pixel)
        let kernel = Matrix::from_vec(1, 1, vec![1.0]).unwrap();

        let result = input.convolve2d(&kernel).unwrap();

        // Result should be 3x3 (same input size with valid padding)
        assert_eq!(result.rows(), 3);
        assert_eq!(result.cols(), 3);
        assert_eq!(result.as_slice(), input.as_slice());
    }

    #[test]
    fn test_convolve2d_edge_detection() {
        // Test edge detection with Sobel-like kernel
        let input = Matrix::from_vec(
            4,
            4,
            vec![
                1.0, 1.0, 1.0, 1.0, //
                1.0, 2.0, 2.0, 1.0, //
                1.0, 2.0, 2.0, 1.0, //
                1.0, 1.0, 1.0, 1.0, //
            ],
        )
        .unwrap();

        // Simple 3x3 horizontal edge detection kernel
        #[rustfmt::skip]
        let kernel = Matrix::from_vec(
            3,
            3,
            vec![
                -1.0, -1.0, -1.0,
                 0.0,  0.0,  0.0,
                 1.0,  1.0,  1.0,
            ],
        )
        .unwrap();

        let result = input.convolve2d(&kernel).unwrap();

        // Result should be 2x2 (4-3+1 = 2)
        assert_eq!(result.rows(), 2);
        assert_eq!(result.cols(), 2);
    }

    #[test]
    fn test_convolve2d_averaging_filter() {
        // Test averaging filter (blur)
        let input = Matrix::from_vec(
            5,
            5,
            vec![
                0.0, 0.0, 0.0, 0.0, 0.0, //
                0.0, 0.0, 0.0, 0.0, 0.0, //
                0.0, 0.0, 9.0, 0.0, 0.0, // Center pixel
                0.0, 0.0, 0.0, 0.0, 0.0, //
                0.0, 0.0, 0.0, 0.0, 0.0, //
            ],
        )
        .unwrap();

        // 3x3 averaging kernel (all 1/9)
        let kernel_val = 1.0 / 9.0;
        let kernel = Matrix::from_vec(
            3,
            3,
            vec![
                kernel_val, kernel_val, kernel_val, //
                kernel_val, kernel_val, kernel_val, //
                kernel_val, kernel_val, kernel_val, //
            ],
        )
        .unwrap();

        let result = input.convolve2d(&kernel).unwrap();

        // Result should be 3x3
        assert_eq!(result.rows(), 3);
        assert_eq!(result.cols(), 3);

        // Center should be 1.0 (9/9)
        assert!((result.get(1, 1).unwrap() - 1.0).abs() < 1e-5);
    }

    #[test]
    fn test_convolve2d_invalid_kernel() {
        let input = Matrix::from_vec(3, 3, vec![1.0; 9]).unwrap();

        // Kernel larger than input
        let kernel = Matrix::from_vec(4, 4, vec![1.0; 16]).unwrap();

        assert!(input.convolve2d(&kernel).is_err());
    }

    // ===== Embedding Lookup Tests (Issue #61) =====

    #[test]
    fn test_embedding_lookup_basic() {
        // Create embedding table: 4 words, 3-dimensional embeddings
        let embeddings = Matrix::from_vec(
            4,
            3,
            vec![
                1.0, 2.0, 3.0, // word 0
                4.0, 5.0, 6.0, // word 1
                7.0, 8.0, 9.0, // word 2
                10.0, 11.0, 12.0, // word 3
            ],
        )
        .unwrap();

        // Lookup embeddings for indices [1, 3, 0]
        let result = embeddings.embedding_lookup(&[1, 3, 0]).unwrap();

        assert_eq!(result.rows(), 3);
        assert_eq!(result.cols(), 3);

        // Check word 1 embedding
        assert_eq!(result.get(0, 0), Some(&4.0));
        assert_eq!(result.get(0, 1), Some(&5.0));
        assert_eq!(result.get(0, 2), Some(&6.0));

        // Check word 3 embedding
        assert_eq!(result.get(1, 0), Some(&10.0));
        assert_eq!(result.get(1, 1), Some(&11.0));
        assert_eq!(result.get(1, 2), Some(&12.0));

        // Check word 0 embedding
        assert_eq!(result.get(2, 0), Some(&1.0));
        assert_eq!(result.get(2, 1), Some(&2.0));
        assert_eq!(result.get(2, 2), Some(&3.0));
    }

    #[test]
    fn test_embedding_lookup_single_index() {
        let embeddings = Matrix::from_vec(3, 2, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();

        let result = embeddings.embedding_lookup(&[1]).unwrap();

        assert_eq!(result.rows(), 1);
        assert_eq!(result.cols(), 2);
        assert_eq!(result.get(0, 0), Some(&3.0));
        assert_eq!(result.get(0, 1), Some(&4.0));
    }

    #[test]
    fn test_embedding_lookup_repeated_indices() {
        let embeddings = Matrix::from_vec(2, 3, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();

        // Same index can appear multiple times
        let result = embeddings.embedding_lookup(&[0, 0, 1, 0]).unwrap();

        assert_eq!(result.rows(), 4);
        assert_eq!(result.cols(), 3);

        // All index-0 rows should be identical
        assert_eq!(result.get(0, 0), result.get(1, 0));
        assert_eq!(result.get(0, 0), result.get(3, 0));
    }

    #[test]
    fn test_embedding_lookup_empty_indices() {
        let embeddings = Matrix::from_vec(3, 2, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();

        let result = embeddings.embedding_lookup(&[]).unwrap();

        assert_eq!(result.rows(), 0);
        assert_eq!(result.cols(), 2);
    }

    #[test]
    fn test_embedding_lookup_out_of_bounds() {
        let embeddings = Matrix::from_vec(3, 2, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0]).unwrap();

        // Index 5 is out of bounds for 3-row table
        let result = embeddings.embedding_lookup(&[0, 5, 1]);

        assert!(result.is_err());
        let err = result.unwrap_err();
        assert!(err.to_string().contains("out of bounds"));
    }

    #[test]
    fn test_embedding_lookup_sparse() {
        let embeddings =
            Matrix::from_vec(4, 2, vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0]).unwrap();

        // Lookup with repeated indices
        let (result, unique) = embeddings
            .embedding_lookup_sparse(&[1, 3, 1, 0, 3])
            .unwrap();

        assert_eq!(result.rows(), 5);
        assert_eq!(result.cols(), 2);

        // Unique indices should be sorted and deduplicated
        assert_eq!(unique, vec![0, 1, 3]);
    }

    #[test]
    fn test_embedding_lookup_large_embeddings() {
        // Test with realistic NLP dimensions
        let vocab_size = 1000;
        let embed_dim = 256;
        let data: Vec<f32> = (0..vocab_size * embed_dim).map(|i| i as f32).collect();
        let embeddings = Matrix::from_vec(vocab_size, embed_dim, data).unwrap();

        // Lookup a sequence
        let indices: Vec<usize> = vec![0, 500, 999, 42, 100];
        let result = embeddings.embedding_lookup(&indices).unwrap();

        assert_eq!(result.rows(), 5);
        assert_eq!(result.cols(), embed_dim);

        // Verify first element of each row
        assert_eq!(result.get(0, 0), Some(&0.0)); // word 0
        assert_eq!(result.get(1, 0), Some(&(500.0 * 256.0))); // word 500
        assert_eq!(result.get(2, 0), Some(&(999.0 * 256.0))); // word 999
    }

    // ===== Batched Matrix Multiplication Tests =====

    #[test]
    fn test_batched_matmul_basic() {
        // [batch=2, m=2, k=3] @ [batch=2, k=3, n=2] -> [batch=2, m=2, n=2]
        let batch = 2;
        let m = 2;
        let k = 3;
        let n = 2;

        // Batch 0: [[1,2,3],[4,5,6]] @ [[1,2],[3,4],[5,6]] = [[22,28],[49,64]]
        // Batch 1: [[7,8,9],[10,11,12]] @ [[7,8],[9,10],[11,12]] = [[184,202],[265,292]]
        let a_data: Vec<f32> = vec![
            1.0, 2.0, 3.0, 4.0, 5.0, 6.0, // Batch 0
            7.0, 8.0, 9.0, 10.0, 11.0, 12.0, // Batch 1
        ];
        let b_data: Vec<f32> = vec![
            1.0, 2.0, 3.0, 4.0, 5.0, 6.0, // Batch 0
            7.0, 8.0, 9.0, 10.0, 11.0, 12.0, // Batch 1
        ];

        let result = Matrix::batched_matmul(&a_data, &b_data, batch, m, k, n).unwrap();

        assert_eq!(result.len(), batch * m * n);

        // Verify batch 0
        assert!((result[0] - 22.0).abs() < 1e-5);
        assert!((result[1] - 28.0).abs() < 1e-5);
        assert!((result[2] - 49.0).abs() < 1e-5);
        assert!((result[3] - 64.0).abs() < 1e-5);

        // Verify batch 1: [[7,8,9],[10,11,12]] @ [[7,8],[9,10],[11,12]]
        // C[0,0] = 7*7 + 8*9 + 9*11 = 49 + 72 + 99 = 220
        // C[0,1] = 7*8 + 8*10 + 9*12 = 56 + 80 + 108 = 244
        // C[1,0] = 10*7 + 11*9 + 12*11 = 70 + 99 + 132 = 301
        // C[1,1] = 10*8 + 11*10 + 12*12 = 80 + 110 + 144 = 334
        assert!((result[4] - 220.0).abs() < 1e-5);
        assert!((result[5] - 244.0).abs() < 1e-5);
        assert!((result[6] - 301.0).abs() < 1e-5);
        assert!((result[7] - 334.0).abs() < 1e-5);
    }

    #[test]
    fn test_batched_matmul_single_batch() {
        let batch = 1;
        let m = 2;
        let k = 2;
        let n = 2;

        let a_data = vec![1.0, 0.0, 0.0, 1.0]; // Identity
        let b_data = vec![5.0, 6.0, 7.0, 8.0];

        let result = Matrix::batched_matmul(&a_data, &b_data, batch, m, k, n).unwrap();

        // Identity @ B = B
        assert!((result[0] - 5.0).abs() < 1e-5);
        assert!((result[1] - 6.0).abs() < 1e-5);
        assert!((result[2] - 7.0).abs() < 1e-5);
        assert!((result[3] - 8.0).abs() < 1e-5);
    }

    #[test]
    fn test_batched_matmul_a_size_mismatch() {
        let batch = 2;
        let m = 2;
        let k = 3;
        let n = 2;

        let a_data = vec![1.0; 10]; // Wrong size (should be 12)
        let b_data = vec![1.0; batch * k * n];

        let result = Matrix::batched_matmul(&a_data, &b_data, batch, m, k, n);
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("A data size mismatch"));
    }

    #[test]
    fn test_batched_matmul_b_size_mismatch() {
        let batch = 2;
        let m = 2;
        let k = 3;
        let n = 2;

        let a_data = vec![1.0; batch * m * k];
        let b_data = vec![1.0; 10]; // Wrong size (should be 12)

        let result = Matrix::batched_matmul(&a_data, &b_data, batch, m, k, n);
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("B data size mismatch"));
    }

    #[test]
    fn test_batched_matmul_4d_basic() {
        // [batch=1, heads=2, m=2, k=2] @ [batch=1, heads=2, k=2, n=2]
        let batch = 1;
        let heads = 2;
        let m = 2;
        let k = 2;
        let n = 2;

        // Head 0: [[1,2],[3,4]] @ [[1,0],[0,1]] = [[1,2],[3,4]]
        // Head 1: [[5,6],[7,8]] @ [[1,0],[0,1]] = [[5,6],[7,8]]
        let a_data: Vec<f32> = vec![
            1.0, 2.0, 3.0, 4.0, // Head 0
            5.0, 6.0, 7.0, 8.0, // Head 1
        ];
        let b_data: Vec<f32> = vec![
            1.0, 0.0, 0.0, 1.0, // Head 0 (identity)
            1.0, 0.0, 0.0, 1.0, // Head 1 (identity)
        ];

        let result = Matrix::batched_matmul_4d(&a_data, &b_data, batch, heads, m, k, n).unwrap();

        assert_eq!(result.len(), batch * heads * m * n);

        // Head 0: A @ I = A
        assert!((result[0] - 1.0).abs() < 1e-5);
        assert!((result[1] - 2.0).abs() < 1e-5);
        assert!((result[2] - 3.0).abs() < 1e-5);
        assert!((result[3] - 4.0).abs() < 1e-5);

        // Head 1: A @ I = A
        assert!((result[4] - 5.0).abs() < 1e-5);
        assert!((result[5] - 6.0).abs() < 1e-5);
        assert!((result[6] - 7.0).abs() < 1e-5);
        assert!((result[7] - 8.0).abs() < 1e-5);
    }

    #[test]
    fn test_batched_matmul_4d_attention_pattern() {
        // Simulate Q @ K^T for attention: [batch=1, heads=2, seq=4, head_dim=8]
        let batch = 1;
        let heads = 2;
        let seq_len = 4;
        let head_dim = 8;

        let q_data: Vec<f32> = (0..batch * heads * seq_len * head_dim)
            .map(|i| (i as f32) * 0.01)
            .collect();
        let kt_data: Vec<f32> = (0..batch * heads * head_dim * seq_len)
            .map(|i| (i as f32) * 0.01)
            .collect();

        let result = Matrix::batched_matmul_4d(
            &q_data,
            &kt_data,
            batch,
            heads,
            seq_len,
            head_dim,
            seq_len,
        )
        .unwrap();

        // Output should be [batch, heads, seq, seq] = 1 * 2 * 4 * 4 = 32 elements
        assert_eq!(result.len(), batch * heads * seq_len * seq_len);
    }

    #[test]
    fn test_batched_matmul_4d_a_size_mismatch() {
        let batch = 1;
        let heads = 2;
        let m = 4;
        let k = 8;
        let n = 4;

        let a_data = vec![1.0; 50]; // Wrong size
        let b_data = vec![1.0; batch * heads * k * n];

        let result = Matrix::batched_matmul_4d(&a_data, &b_data, batch, heads, m, k, n);
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("A data size mismatch"));
    }

    #[test]
    fn test_batched_matmul_4d_b_size_mismatch() {
        let batch = 1;
        let heads = 2;
        let m = 4;
        let k = 8;
        let n = 4;

        let a_data = vec![1.0; batch * heads * m * k];
        let b_data = vec![1.0; 50]; // Wrong size

        let result = Matrix::batched_matmul_4d(&a_data, &b_data, batch, heads, m, k, n);
        assert!(result.is_err());
        assert!(result.unwrap_err().to_string().contains("B data size mismatch"));
    }

    // ===== Property-Based Tests for Convolution =====

    #[cfg(test)]
    mod conv_property_tests {
        use super::*;

        proptest! {
            #[test]
            fn test_convolve2d_output_size(
                input_rows in 3usize..20,
                input_cols in 3usize..20,
                kernel_rows in 1usize..5,
                kernel_cols in 1usize..5,
            ) {
                // Property: Output size is always (input - kernel + 1) for valid padding
                if kernel_rows <= input_rows && kernel_cols <= input_cols {
                    let input = Matrix::from_vec(input_rows, input_cols, vec![1.0; input_rows * input_cols]).unwrap();
                    let kernel = Matrix::from_vec(kernel_rows, kernel_cols, vec![1.0; kernel_rows * kernel_cols]).unwrap();

                    let result = input.convolve2d(&kernel).unwrap();

                    prop_assert_eq!(result.rows(), input_rows - kernel_rows + 1);
                    prop_assert_eq!(result.cols(), input_cols - kernel_cols + 1);
                }
            }

            #[test]
            fn test_convolve2d_identity_kernel(
                input_rows in 3usize..10,
                input_cols in 3usize..10,
                values in prop::collection::vec(-100.0f32..100.0, 9..100)
            ) {
                // Property: 1x1 identity kernel preserves input
                if values.len() >= input_rows * input_cols {
                    let data: Vec<f32> = values.iter().take(input_rows * input_cols).copied().collect();
                    let input = Matrix::from_vec(input_rows, input_cols, data.clone()).unwrap();
                    let kernel = Matrix::from_vec(1, 1, vec![1.0]).unwrap();

                    let result = input.convolve2d(&kernel).unwrap();

                    prop_assert_eq!(result.rows(), input_rows);
                    prop_assert_eq!(result.cols(), input_cols);
                    prop_assert_eq!(result.as_slice(), input.as_slice());
                }
            }

            #[test]
            fn test_convolve2d_zero_kernel(
                input_rows in 3usize..10,
                input_cols in 3usize..10,
                kernel_rows in 1usize..4,
                kernel_cols in 1usize..4,
            ) {
                // Property: Zero kernel produces zero output
                if kernel_rows <= input_rows && kernel_cols <= input_cols {
                    let input = Matrix::from_vec(input_rows, input_cols, vec![5.0; input_rows * input_cols]).unwrap();
                    let kernel = Matrix::from_vec(kernel_rows, kernel_cols, vec![0.0; kernel_rows * kernel_cols]).unwrap();

                    let result = input.convolve2d(&kernel).unwrap();

                    for &val in result.as_slice() {
                        prop_assert!((val - 0.0).abs() < 1e-5);
                    }
                }
            }

            #[test]
            fn test_convolve2d_scalar_multiplication(
                input_rows in 3usize..10,
                input_cols in 3usize..10,
                scalar in -10.0f32..10.0,
            ) {
                // Property: Convolving with scalar * kernel = scalar * (convolve with kernel)
                let input = Matrix::from_vec(input_rows, input_cols, vec![2.0; input_rows * input_cols]).unwrap();
                let kernel = Matrix::from_vec(3, 3, vec![1.0; 9]).unwrap();
                let kernel_scaled = Matrix::from_vec(3, 3, vec![scalar; 9]).unwrap();

                let result1 = input.convolve2d(&kernel).unwrap();
                let result2 = input.convolve2d(&kernel_scaled).unwrap();

                for (v1, v2) in result1.as_slice().iter().zip(result2.as_slice().iter()) {
                    prop_assert!((v1 * scalar - v2).abs() < 1e-3);
                }
            }
        }
    }
}