oxicuda-sparse 0.4.1

//! Sparse matrix-dense matrix multiplication (SpMM).
//!
//! Computes `C = alpha * A * B + beta * C` where `A` is a sparse CSR matrix
//! and `B`, `C` are dense matrices.
//!
//! ## Strategy
//!
//! Each warp processes one row of `A` and tiles across columns of `B`.
//! Within each row, the warp collaboratively loads non-zero entries and
//! accumulates partial sums into the output columns.

use std::sync::Arc;

use oxicuda_blas::GpuFloat;
use oxicuda_blas::types::{MatrixDesc, MatrixDescMut};
use oxicuda_driver::Module;
use oxicuda_launch::{Kernel, LaunchParams, grid_size_for};
use oxicuda_ptx::prelude::*;

use crate::error::{SparseError, SparseResult};
use crate::format::CsrMatrix;
use crate::handle::SparseHandle;
use crate::ptx_helpers::{
    add_float, fma_float, load_float_imm, load_global_float, mul_float, reinterpret_bits_to_float,
    store_global_float,
};

/// Default block size for SpMM kernel.
const SPMM_BLOCK_SIZE: u32 = 256;

/// Number of columns of B processed per thread in one tile.
const SPMM_TILE_COLS: u32 = 4;

/// Sparse matrix-dense matrix multiplication: `C = alpha * A * B + beta * C`.
///
/// `A` is a sparse CSR matrix of shape `(m, k)`. `B` is a dense matrix of
/// shape `(k, n)`. `C` is a dense output matrix of shape `(m, n)`.
///
/// # Arguments
///
/// * `handle` -- Sparse handle.
/// * `alpha` -- Scalar multiplier for `A * B`.
/// * `a` -- Sparse CSR matrix `A`.
/// * `b` -- Dense matrix descriptor for `B`.
/// * `beta` -- Scalar multiplier for existing `C`.
/// * `c` -- Dense mutable matrix descriptor for `C`.
///
/// # Errors
///
/// Returns [`SparseError::DimensionMismatch`] if dimensions are incompatible.
/// Returns [`SparseError::PtxGeneration`] if kernel generation fails.
/// Returns [`SparseError::Cuda`] on kernel launch failure.
pub fn spmm<T: GpuFloat>(
    handle: &SparseHandle,
    alpha: T,
    a: &CsrMatrix<T>,
    b: &MatrixDesc<T>,
    beta: T,
    c: &mut MatrixDescMut<T>,
) -> SparseResult<()> {
    // Validate dimensions: A(m, k) * B(k, n) = C(m, n)
    if a.cols() != b.rows {
        return Err(SparseError::DimensionMismatch(format!(
            "A.cols ({}) != B.rows ({})",
            a.cols(),
            b.rows
        )));
    }
    if a.rows() != c.rows {
        return Err(SparseError::DimensionMismatch(format!(
            "A.rows ({}) != C.rows ({})",
            a.rows(),
            c.rows
        )));
    }
    if b.cols != c.cols {
        return Err(SparseError::DimensionMismatch(format!(
            "B.cols ({}) != C.cols ({})",
            b.cols, c.cols
        )));
    }

    if a.rows() == 0 || a.cols() == 0 || b.cols == 0 {
        return Ok(());
    }

    let ptx = emit_spmm_kernel::<T>(handle.sm_version())?;
    let module = Arc::new(Module::from_ptx(&ptx)?);
    let kernel = Kernel::from_module(module, "spmm")?;

    // Grid: one thread per (row, col_tile) pair
    let block_size = SPMM_BLOCK_SIZE;
    let total_work = a.rows() * b.cols.div_ceil(SPMM_TILE_COLS);
    let grid_size = grid_size_for(total_work, block_size);
    let params = LaunchParams::new(grid_size, block_size);

    kernel.launch(
        &params,
        handle.stream(),
        &(
            a.row_ptr().as_device_ptr(),
            a.col_idx().as_device_ptr(),
            a.values().as_device_ptr(),
            b.ptr,
            c.ptr,
            alpha.to_bits_u64(),
            beta.to_bits_u64(),
            a.rows(),
            b.cols,
            b.ld,
            c.ld,
        ),
    )?;

    Ok(())
}

/// Generates PTX for the SpMM kernel.
///
/// Each thread handles one row of A and a tile of SPMM_TILE_COLS columns of B.
/// For simplicity, we implement a scalar approach: each thread iterates over
/// the non-zeros of its row and accumulates into multiple output columns.
fn emit_spmm_kernel<T: GpuFloat>(sm: SmVersion) -> SparseResult<String> {
    let elem_bytes = T::size_u32();
    let is_f64 = T::SIZE == 8;
    let tile_cols = SPMM_TILE_COLS;

    KernelBuilder::new("spmm")
        .target(sm)
        .param("row_ptr", PtxType::U64)
        .param("col_idx", PtxType::U64)
        .param("values", PtxType::U64)
        .param("b_ptr", PtxType::U64)
        .param("c_ptr", PtxType::U64)
        .param("alpha_bits", PtxType::U64)
        .param("beta_bits", PtxType::U64)
        .param("m", PtxType::U32)
        .param("n", PtxType::U32)
        .param("ldb", PtxType::U32)
        .param("ldc", PtxType::U32)
        .body(move |b| {
            let gid = b.global_thread_id_x();

            // Compute which row and col tile we handle
            let n_param = b.load_param_u32("n");
            let m_param = b.load_param_u32("m");

            // tiles_per_row = (n + tile_cols - 1) / tile_cols
            let tiles_per_row = b.alloc_reg(PtxType::U32);
            let n_plus = b.alloc_reg(PtxType::U32);
            b.raw_ptx(&format!("add.u32 {n_plus}, {n_param}, {};", tile_cols - 1));
            b.raw_ptx(&format!(
                "div.u32 {tiles_per_row}, {n_plus}, {};",
                tile_cols
            ));

            // row = gid / tiles_per_row, tile_id = gid % tiles_per_row
            let row = b.alloc_reg(PtxType::U32);
            let tile_id = b.alloc_reg(PtxType::U32);
            b.raw_ptx(&format!("div.u32 {row}, {gid}, {tiles_per_row};"));
            b.raw_ptx(&format!("rem.u32 {tile_id}, {gid}, {tiles_per_row};"));

            let row_inner = row.clone();
            let tile_id_inner = tile_id.clone();
            b.if_lt_u32(row, m_param, move |b| {
                let row = row_inner;
                let tile_id = tile_id_inner;

                let row_ptr_base = b.load_param_u64("row_ptr");
                let col_idx_base = b.load_param_u64("col_idx");
                let values_base = b.load_param_u64("values");
                let b_ptr = b.load_param_u64("b_ptr");
                let c_ptr = b.load_param_u64("c_ptr");
                let alpha_bits = b.load_param_u64("alpha_bits");
                let beta_bits = b.load_param_u64("beta_bits");
                let n_param = b.load_param_u32("n");
                let ldb = b.load_param_u32("ldb");
                let ldc = b.load_param_u32("ldc");

                let alpha = reinterpret_bits_to_float::<T>(b, alpha_bits);
                let beta = reinterpret_bits_to_float::<T>(b, beta_bits);

                // col_start = tile_id * tile_cols
                let col_start = b.alloc_reg(PtxType::U32);
                b.raw_ptx(&format!(
                    "mul.lo.u32 {col_start}, {tile_id}, {};",
                    tile_cols
                ));

                // Load row bounds
                let rp_addr = b.byte_offset_addr(row_ptr_base.clone(), row.clone(), 4);
                let rs_i32 = b.load_global_i32(rp_addr);
                let rs = b.alloc_reg(PtxType::U32);
                b.raw_ptx(&format!("mov.b32 {rs}, {rs_i32};"));

                let row_p1 = b.alloc_reg(PtxType::U32);
                b.raw_ptx(&format!("add.u32 {row_p1}, {row}, 1;"));
                let rp_addr_next = b.byte_offset_addr(row_ptr_base, row_p1, 4);
                let re_i32 = b.load_global_i32(rp_addr_next);
                let re = b.alloc_reg(PtxType::U32);
                b.raw_ptx(&format!("mov.b32 {re}, {re_i32};"));

                // This thread owns the whole tile of `tile_cols` consecutive
                // output columns starting at `col_start`. The grid spawns
                // exactly one thread per (row, tile) pair, so *every* column in
                // the tile must be computed here -- handling only `col_start`
                // would leave any column whose index is not a multiple of
                // `tile_cols` permanently unwritten.
                let mov_suffix = if is_f64 { "f64" } else { "f32" };
                for tc in 0..tile_cols {
                    let col = b.alloc_reg(PtxType::U32);
                    b.raw_ptx(&format!("add.u32 {col}, {col_start}, {tc};"));

                    // Skip out-of-range columns (col >= n). Inverted skip-branch
                    // (`setp.lo` -> `setp.hs`) routed through the structured
                    // `branch_if` so the target matches the `$`-prefixed label.
                    let col_oob = b.alloc_reg(PtxType::Pred);
                    b.raw_ptx(&format!("setp.hs.u32 {col_oob}, {col}, {n_param};"));
                    let skip_col = b.fresh_label("spmm_skip_col");
                    b.branch_if(col_oob, &skip_col);

                    // Accumulate: acc = sum(A[row,k] * B[k, col])
                    let acc = load_float_imm::<T>(b, 0.0);
                    let k_reg = b.alloc_reg(PtxType::U32);
                    b.raw_ptx(&format!("mov.u32 {k_reg}, {rs};"));

                    let loop_label = b.fresh_label("spmm_loop");
                    let done_label = b.fresh_label("spmm_done");

                    b.label(&loop_label);
                    // Exit when k >= row_end (inverted skip-branch via branch_if).
                    let pred = b.alloc_reg(PtxType::Pred);
                    b.raw_ptx(&format!("setp.hs.u32 {pred}, {k_reg}, {re};"));
                    b.branch_if(pred, &done_label);

                    // Load A value and column
                    let ci_addr = b.byte_offset_addr(col_idx_base.clone(), k_reg.clone(), 4);
                    let a_col_i32 = b.load_global_i32(ci_addr);
                    let a_col = b.alloc_reg(PtxType::U32);
                    b.raw_ptx(&format!("mov.b32 {a_col}, {a_col_i32};"));

                    let v_addr = b.byte_offset_addr(values_base.clone(), k_reg.clone(), elem_bytes);
                    let a_val = load_global_float::<T>(b, v_addr);

                    // Load B[a_col, col] = b_ptr + (a_col * ldb + col) * elem_bytes
                    // Row-major: B[a_col][col] = b_ptr + a_col * ldb + col
                    let b_row_off = b.alloc_reg(PtxType::U32);
                    b.raw_ptx(&format!("mul.lo.u32 {b_row_off}, {a_col}, {ldb};"));
                    let b_idx = b.alloc_reg(PtxType::U32);
                    b.raw_ptx(&format!("add.u32 {b_idx}, {b_row_off}, {col};"));
                    let b_addr = b.byte_offset_addr(b_ptr.clone(), b_idx, elem_bytes);
                    let b_val = load_global_float::<T>(b, b_addr);

                    // acc += a_val * b_val
                    let new_acc = fma_float::<T>(b, a_val, b_val, acc.clone());
                    b.raw_ptx(&format!("mov.{mov_suffix} {acc}, {new_acc};"));

                    b.raw_ptx(&format!("add.u32 {k_reg}, {k_reg}, 1;"));
                    b.branch(&loop_label);
                    b.label(&done_label);

                    // Write C[row, col] = alpha * acc + beta * C_old
                    let c_row_off = b.alloc_reg(PtxType::U32);
                    b.raw_ptx(&format!("mul.lo.u32 {c_row_off}, {row}, {ldc};"));
                    let c_idx = b.alloc_reg(PtxType::U32);
                    b.raw_ptx(&format!("add.u32 {c_idx}, {c_row_off}, {col};"));
                    let c_addr = b.byte_offset_addr(c_ptr.clone(), c_idx, elem_bytes);
                    let c_old = load_global_float::<T>(b, c_addr.clone());

                    let alpha_acc = mul_float::<T>(b, alpha.clone(), acc);
                    let beta_c = mul_float::<T>(b, beta.clone(), c_old);
                    let result = add_float::<T>(b, alpha_acc, beta_c);
                    store_global_float::<T>(b, c_addr, result);

                    b.label(&skip_col);
                }
            });

            b.ret();
        })
        .build()
        .map_err(|e| SparseError::PtxGeneration(e.to_string()))
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ptx_helpers::test_support::assert_assembles_and_clean;

    /// The SpMM kernel must assemble for sm_86 in both precisions with
    /// `$`-prefixed branch targets and no `.b64` shuffle.
    #[test]
    fn spmm_f32_f64_assemble_sm86() {
        let f32_ptx = emit_spmm_kernel::<f32>(SmVersion::Sm86).expect("f32 SpMM PTX");
        assert_assembles_and_clean("spmm_f32", &f32_ptx);

        let f64_ptx = emit_spmm_kernel::<f64>(SmVersion::Sm86).expect("f64 SpMM PTX");
        assert_assembles_and_clean("spmm_f64", &f64_ptx);
        assert!(
            !f64_ptx.contains("0F00000000"),
            "f64 SpMM kernel must not materialize an f32 0.0 immediate:\n{f64_ptx}"
        );
    }

    // ---------------------------------------------------------------------------
    // CPU reference SpMM for numerical accuracy verification
    // ---------------------------------------------------------------------------

    /// CPU reference CSR SpMM: computes C = A * B (no alpha/beta scaling).
    ///
    /// * `row_ptr`, `col_idx`, `values` — CSR representation of A (m×k sparse).
    /// * `b` — row-major dense matrix B of shape (k, n) with leading dimension `ldb`.
    /// * `n` — number of columns in B (and C).
    ///
    /// Returns C as a row-major Vec<f32> of shape m×n (leading dimension n).
    fn cpu_csr_spmm(
        row_ptr: &[usize],
        col_idx: &[usize],
        values: &[f32],
        b: &[f32],
        n: usize,
        ldb: usize,
    ) -> Vec<f32> {
        let m = row_ptr.len() - 1;
        let mut c = vec![0.0_f32; m * n];
        for row in 0..m {
            for nnz_idx in row_ptr[row]..row_ptr[row + 1] {
                let a_col = col_idx[nnz_idx];
                let a_val = values[nnz_idx];
                // A[row, a_col] * B[a_col, col] for all cols
                for col in 0..n {
                    c[row * n + col] += a_val * b[a_col * ldb + col];
                }
            }
        }
        c
    }

    // ---------------------------------------------------------------------------
    // PTX generation tests
    // ---------------------------------------------------------------------------

    #[test]
    fn spmm_ptx_generates_f32() {
        let ptx = emit_spmm_kernel::<f32>(SmVersion::Sm80);
        assert!(ptx.is_ok());
        let ptx = ptx.expect("test: PTX gen should succeed");
        assert!(ptx.contains(".entry spmm"));
    }

    #[test]
    fn spmm_ptx_generates_f64() {
        let ptx = emit_spmm_kernel::<f64>(SmVersion::Sm80);
        assert!(ptx.is_ok());
    }

    #[test]
    fn spmm_ptx_contains_arithmetic_instructions() {
        let ptx = emit_spmm_kernel::<f32>(SmVersion::Sm80);
        assert!(ptx.is_ok());
        let ptx = ptx.expect("test: PTX gen should succeed");
        // Should contain FMA for the accumulation step
        assert!(
            ptx.contains("fma") || ptx.contains("mul"),
            "SpMM PTX should contain arithmetic instructions"
        );
    }

    // ---------------------------------------------------------------------------
    // CPU reference numerical accuracy tests
    // ---------------------------------------------------------------------------

    /// 4×4 identity sparse × 4×3 dense = 4×3 dense (same as dense matrix).
    ///
    /// A = I_4, B:
    ///   [1 2 3]
    ///   [4 5 6]
    ///   [7 8 9]
    ///   [10 11 12]
    ///
    /// C = A * B = B.
    #[test]
    fn spmm_identity_times_dense_equals_dense() {
        let row_ptr = vec![0usize, 1, 2, 3, 4];
        let col_idx = vec![0usize, 1, 2, 3];
        let values = vec![1.0_f32; 4];

        // B: 4×3 row-major
        let b = vec![
            1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0,
        ];
        let n = 3usize;
        let ldb = 3usize;

        let c = cpu_csr_spmm(&row_ptr, &col_idx, &values, &b, n, ldb);

        // C should equal B
        assert_eq!(c.len(), 4 * 3);
        for (i, (&ci, &bi)) in c.iter().zip(b.iter()).enumerate() {
            assert!((ci - bi).abs() < 1e-6, "C[{}] = {ci} expected {bi}", i);
        }
    }

    /// 2×3 sparse A × 3×2 dense B = 2×2 dense C with known values.
    ///
    /// A (CSR):
    ///   Row 0: A[0,0]=1, A[0,2]=3
    ///   Row 1: A[1,1]=2, A[1,2]=4
    ///
    /// B (row-major, 3×2):
    ///   [1 2]
    ///   [3 4]
    ///   [5 6]
    ///
    /// C = A*B:
    ///   C[0,0] = 1*1 + 3*5 = 16,  C[0,1] = 1*2 + 3*6 = 20
    ///   C[1,0] = 2*3 + 4*5 = 26,  C[1,1] = 2*4 + 4*6 = 32
    #[test]
    fn spmm_small_sparse_times_dense_known_values() {
        let row_ptr = vec![0usize, 2, 4];
        let col_idx = vec![0usize, 2, 1, 2];
        let values = vec![1.0_f32, 3.0, 2.0, 4.0];

        let b = vec![1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]; // 3×2 row-major
        let n = 2usize;
        let ldb = 2usize;

        let c = cpu_csr_spmm(&row_ptr, &col_idx, &values, &b, n, ldb);

        assert_eq!(c.len(), 4);
        assert!((c[0] - 16.0).abs() < 1e-5, "C[0,0] = {} expected 16", c[0]);
        assert!((c[1] - 20.0).abs() < 1e-5, "C[0,1] = {} expected 20", c[1]);
        assert!((c[2] - 26.0).abs() < 1e-5, "C[1,0] = {} expected 26", c[2]);
        assert!((c[3] - 32.0).abs() < 1e-5, "C[1,1] = {} expected 32", c[3]);
    }

    /// 4×4 diagonal sparse A × 4×3 dense B = 4×3 dense C.
    ///
    /// A = diag(2, 3, 4, 5), B rows are [1,0,0], [0,1,0], [0,0,1], [1,1,1].
    ///
    /// C[i] = A[i,i] * B[i] for each row i.
    #[test]
    fn spmm_diagonal_times_dense_row_scaling() {
        let row_ptr = vec![0usize, 1, 2, 3, 4];
        let col_idx = vec![0usize, 1, 2, 3];
        let values = vec![2.0_f32, 3.0, 4.0, 5.0];

        // B: 4×3, each row is a unit vector or all-ones
        let b = vec![
            1.0_f32, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0, 1.0, 1.0, 1.0,
        ];
        let n = 3usize;
        let ldb = 3usize;

        let c = cpu_csr_spmm(&row_ptr, &col_idx, &values, &b, n, ldb);

        // Row 0: 2 * [1,0,0] = [2,0,0]
        assert!((c[0] - 2.0).abs() < 1e-6, "C[0,0] = {}", c[0]);
        assert!(c[1].abs() < 1e-6, "C[0,1] = {}", c[1]);
        assert!(c[2].abs() < 1e-6, "C[0,2] = {}", c[2]);

        // Row 1: 3 * [0,1,0] = [0,3,0]
        assert!(c[3].abs() < 1e-6, "C[1,0] = {}", c[3]);
        assert!((c[4] - 3.0).abs() < 1e-6, "C[1,1] = {}", c[4]);
        assert!(c[5].abs() < 1e-6, "C[1,2] = {}", c[5]);

        // Row 2: 4 * [0,0,1] = [0,0,4]
        assert!(c[6].abs() < 1e-6, "C[2,0] = {}", c[6]);
        assert!(c[7].abs() < 1e-6, "C[2,1] = {}", c[7]);
        assert!((c[8] - 4.0).abs() < 1e-6, "C[2,2] = {}", c[8]);

        // Row 3: 5 * [1,1,1] = [5,5,5]
        assert!((c[9] - 5.0).abs() < 1e-6, "C[3,0] = {}", c[9]);
        assert!((c[10] - 5.0).abs() < 1e-6, "C[3,1] = {}", c[10]);
        assert!((c[11] - 5.0).abs() < 1e-6, "C[3,2] = {}", c[11]);
    }

    /// Verify SpMM with a zero sparse matrix produces an all-zero output.
    #[test]
    fn spmm_zero_sparse_matrix_produces_zero_output() {
        let row_ptr = vec![0usize, 0, 0, 0];
        let col_idx: Vec<usize> = vec![];
        let values: Vec<f32> = vec![];

        let b = vec![1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0]; // 3×2
        let n = 2usize;
        let ldb = 2usize;

        let c = cpu_csr_spmm(&row_ptr, &col_idx, &values, &b, n, ldb);

        assert_eq!(c.len(), 6);
        for (i, &ci) in c.iter().enumerate() {
            assert!(
                ci.abs() < 1e-6,
                "C[{i}] = {ci}, expected 0.0 for zero sparse matrix"
            );
        }
    }
}

// ---------------------------------------------------------------------------
// On-device numeric validation (feature = "gpu-tests")
// ---------------------------------------------------------------------------

#[cfg(all(test, feature = "gpu-tests"))]
mod gpu_device_tests {
    use super::*;
    use crate::gpu_test_support::{assert_close, gpu_handle};
    use crate::host_csr::{f64_to_gpu, gpu_to_f64};
    use oxicuda_blas::types::Layout;
    use oxicuda_memory::DeviceBuffer;

    /// CPU oracle for `C = alpha * A * B + beta * C0`, all row-major dense.
    #[allow(clippy::too_many_arguments)]
    fn cpu_spmm(
        m: usize,
        n: usize,
        row_ptr: &[i32],
        col_idx: &[i32],
        values: &[f64],
        b: &[f64],
        c0: &[f64],
        alpha: f64,
        beta: f64,
    ) -> Vec<f64> {
        let mut c = vec![0.0_f64; m * n];
        for row in 0..m {
            let mut acc = vec![0.0_f64; n];
            for k in row_ptr[row] as usize..row_ptr[row + 1] as usize {
                let a_col = col_idx[k] as usize;
                let a_val = values[k];
                for (col, slot) in acc.iter_mut().enumerate() {
                    *slot += a_val * b[a_col * n + col];
                }
            }
            for col in 0..n {
                c[row * n + col] = alpha * acc[col] + beta * c0[row * n + col];
            }
        }
        c
    }

    /// Drive the production `spmm` op and compare to the CPU oracle.
    #[allow(clippy::too_many_arguments)]
    fn run_spmm<T: GpuFloat>(
        m: u32,
        k: u32,
        n: u32,
        row_ptr: &[i32],
        col_idx: &[i32],
        values: &[f64],
        b_dense: &[f64],
        c0: &[f64],
        alpha: f64,
        beta: f64,
        tol: f64,
        tag: &str,
    ) {
        let Some(handle) = gpu_handle() else {
            return;
        };
        let dev_values: Vec<T> = values.iter().map(|&v| f64_to_gpu::<T>(v)).collect();
        let a = CsrMatrix::<T>::from_host(m, k, row_ptr, col_idx, &dev_values)
            .expect("test: build CSR");

        let dev_b: Vec<T> = b_dense.iter().map(|&v| f64_to_gpu::<T>(v)).collect();
        let dev_c: Vec<T> = c0.iter().map(|&v| f64_to_gpu::<T>(v)).collect();
        let b_buf = DeviceBuffer::from_host(&dev_b).expect("test: upload B");
        let c_buf = DeviceBuffer::from_host(&dev_c).expect("test: upload C");

        let b_desc = MatrixDesc::<T>::from_raw(b_buf.as_device_ptr(), k, n, n, Layout::RowMajor);
        let mut c_desc =
            MatrixDescMut::<T>::from_raw(c_buf.as_device_ptr(), m, n, n, Layout::RowMajor);

        spmm::<T>(
            &handle,
            f64_to_gpu::<T>(alpha),
            &a,
            &b_desc,
            f64_to_gpu::<T>(beta),
            &mut c_desc,
        )
        .expect("test: spmm launch");
        handle.stream().synchronize().expect("test: sync");

        let mut out = vec![T::gpu_zero(); (m * n) as usize];
        c_buf.copy_to_host(&mut out).expect("test: download C");
        let got: Vec<f64> = out.iter().map(|&v| gpu_to_f64(v)).collect();
        let want = cpu_spmm(
            m as usize, n as usize, row_ptr, col_idx, values, b_dense, c0, alpha, beta,
        );
        assert_close(&got, &want, tol, tag);
    }

    /// A = 3x4 CSR:
    /// [1 0 2 0]
    /// [0 3 0 4]
    /// [5 0 0 6]
    fn matrix_3x4() -> (u32, u32, Vec<i32>, Vec<i32>, Vec<f64>) {
        let row_ptr = vec![0, 2, 4, 6];
        let col_idx = vec![0, 2, 1, 3, 0, 3];
        let values = vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0];
        (3, 4, row_ptr, col_idx, values)
    }

    /// B is k x n row-major. Build a deterministic dense matrix.
    fn dense(k: usize, n: usize, base: f64) -> Vec<f64> {
        (0..k * n)
            .map(|idx| base + 0.5 * (idx as f64) - 0.1 * ((idx % 3) as f64))
            .collect()
    }

    #[test]
    fn spmm_f64_n3_alpha_beta() {
        // n = 3 is NOT a multiple of the internal tile width (4): this is the
        // exact configuration that exposes a one-column-per-tile kernel bug.
        let (m, k, rp, ci, v) = matrix_3x4();
        let n = 3usize;
        let b = dense(k as usize, n, 1.0);
        let c0 = dense(m as usize, n, 7.0);
        run_spmm::<f64>(
            m,
            k,
            n as u32,
            &rp,
            &ci,
            &v,
            &b,
            &c0,
            1.75,
            -0.5,
            1e-10,
            "spmm_f64_n3",
        );
    }

    #[test]
    fn spmm_f64_n5_alpha_beta() {
        // n = 5 spans more than one tile and is not a multiple of 4.
        let (m, k, rp, ci, v) = matrix_3x4();
        let n = 5usize;
        let b = dense(k as usize, n, -2.0);
        let c0 = dense(m as usize, n, 0.25);
        run_spmm::<f64>(
            m,
            k,
            n as u32,
            &rp,
            &ci,
            &v,
            &b,
            &c0,
            2.0,
            0.5,
            1e-10,
            "spmm_f64_n5",
        );
    }

    #[test]
    fn spmm_f32_n3_alpha_beta() {
        let (m, k, rp, ci, v) = matrix_3x4();
        let n = 3usize;
        let b = dense(k as usize, n, 0.5);
        let c0 = dense(m as usize, n, 3.0);
        run_spmm::<f32>(
            m,
            k,
            n as u32,
            &rp,
            &ci,
            &v,
            &b,
            &c0,
            1.25,
            -0.75,
            1e-4,
            "spmm_f32_n3",
        );
    }

    #[test]
    fn spmm_f64_single_column_beta_zero() {
        // n = 1 (the only width the old kernel handled) with beta = 0.
        let (m, k, rp, ci, v) = matrix_3x4();
        let n = 1usize;
        let b = dense(k as usize, n, 1.0);
        let c0 = vec![1e8; (m as usize) * n];
        run_spmm::<f64>(
            m,
            k,
            n as u32,
            &rp,
            &ci,
            &v,
            &b,
            &c0,
            1.0,
            0.0,
            1e-10,
            "spmm_f64_n1",
        );
    }
}