oxicuda_sparse/ops/

spmm.rs

//! 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};"));

                // We process 1 column (simplified approach for correctness)
                // In production, we'd unroll tile_cols times
                let col = col_start;
                let col_bound = b.alloc_reg(PtxType::Pred);
                b.raw_ptx(&format!("setp.lo.u32 {col_bound}, {col}, {n_param};"));

                let do_col = b.fresh_label("spmm_do_col");
                let skip_col = b.fresh_label("spmm_skip_col");
                b.raw_ptx(&format!("@!{col_bound} bra {skip_col};"));
                b.label(&do_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);
                let pred = b.alloc_reg(PtxType::Pred);
                b.raw_ptx(&format!("setp.lo.u32 {pred}, {k_reg}, {re};"));
                b.raw_ptx(&format!("@!{pred} bra {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());
                let mov_suffix = if is_f64 { "f64" } else { "f32" };
                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, c_idx, elem_bytes);
                let c_old = load_global_float::<T>(b, c_addr.clone());

                let alpha_acc = mul_float::<T>(b, alpha, acc);
                let beta_c = mul_float::<T>(b, beta, 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::*;

    // ---------------------------------------------------------------------------
    // 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"
            );
        }
    }
}