oxicuda-sparse 0.4.0

//! Batched sparse operations.
//!
//! This module provides routines for executing many small sparse operations
//! in a single logical batch, amortising kernel launch overhead and enabling
//! fusion when beneficial.
//!
//! ## Provided types
//!
//! - [`BatchedSpMV`] -- batched sparse matrix-vector multiply (host-side loop).
//! - [`BatchedSpMVPlan`] -- precomputed plan with concatenated CSR arrays.
//! - [`UniformBatchedSpMV`] -- optimised path when all matrices share structure.
//! - [`BatchedSpGEMM`] -- batched sparse-sparse matrix multiply.
//! - [`BatchedTriSolve`] -- batched sparse triangular solve.
//! - [`BatchScheduler`] -- heuristic-based execution strategy selector.

use std::ops::{Add, AddAssign, Mul};

use oxicuda_blas::GpuFloat;

use crate::error::{SparseError, SparseResult};
use crate::format::CsrMatrix;

/// Result type for a single SpGEMM: (row_ptr, col_idx, values, rows, cols).
type SpGEMMResultU32<T> = (Vec<i32>, Vec<i32>, Vec<T>, u32, u32);

/// Result type for a single SpGEMM from host arrays: (row_ptr, col_idx, values, rows, cols).
type SpGEMMResultUsize<T> = (Vec<i32>, Vec<i32>, Vec<T>, usize, usize);

// ---------------------------------------------------------------------------
// BatchScheduler -- strategy selection
// ---------------------------------------------------------------------------

/// Execution strategy chosen by [`BatchScheduler`].
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum Strategy {
    /// Process each operation sequentially on a single stream.
    Sequential,
    /// Issue operations on `num_streams` concurrent CUDA streams.
    Concurrent(usize),
    /// Fuse all operations into a single kernel launch.
    Fused,
}

/// Heuristic-based scheduler that decides how to execute a batch.
///
/// The strategy depends on batch size and average matrix complexity:
///
/// | Condition | Strategy |
/// |-----------|----------|
/// | batch <= 4 *and* avg\_nnz >= 10 000 | [`Sequential`](Strategy::Sequential) |
/// | batch >= 64 *and* avg\_nnz < 256 | [`Fused`](Strategy::Fused) |
/// | otherwise | [`Concurrent`](Strategy::Concurrent) with up to 8 streams |
#[derive(Debug, Clone)]
pub struct BatchScheduler {
    _private: (),
}

impl BatchScheduler {
    /// Create a new scheduler.
    #[inline]
    pub fn new() -> Self {
        Self { _private: () }
    }

    /// Select an execution strategy for the given batch characteristics.
    ///
    /// # Arguments
    ///
    /// * `batch_size` -- Number of independent operations in the batch.
    /// * `avg_nnz` -- Average number of non-zeros per matrix in the batch.
    pub fn select_strategy(&self, batch_size: usize, avg_nnz: usize) -> Strategy {
        Self::select_strategy_static(batch_size, avg_nnz)
    }

    /// Static version of [`select_strategy`](Self::select_strategy) for use
    /// without constructing a scheduler instance.
    pub fn select_strategy_static(batch_size: usize, avg_nnz: usize) -> Strategy {
        // Small batches of large matrices -> sequential is fine
        if batch_size <= 4 && avg_nnz >= 10_000 {
            return Strategy::Sequential;
        }
        // Many tiny matrices -> fuse into single kernel
        if batch_size >= 64 && avg_nnz < 256 {
            return Strategy::Fused;
        }
        // Default: concurrent streams (capped at 8)
        let streams = batch_size.clamp(1, 8);
        Strategy::Concurrent(streams)
    }
}

impl Default for BatchScheduler {
    fn default() -> Self {
        Self::new()
    }
}

// ---------------------------------------------------------------------------
// BatchedSpMVPlan -- precomputed concatenated arrays
// ---------------------------------------------------------------------------

/// Precomputed plan for batched SpMV execution.
///
/// Concatenates all CSR arrays (row\_ptr, col\_indices, values) from multiple
/// matrices into contiguous host-side arrays with offset tables, enabling
/// efficient data transfer and potential kernel fusion.
///
/// # Type parameters
///
/// * `T` -- Element type (must satisfy [`GpuFloat`]).
#[derive(Debug, Clone)]
pub struct BatchedSpMVPlan<T> {
    /// Concatenated row pointer array across all matrices.
    pub concat_row_ptr: Vec<i32>,
    /// Concatenated column index array across all matrices.
    pub concat_col_idx: Vec<i32>,
    /// Concatenated values array across all matrices.
    pub concat_values: Vec<T>,
    /// Start offset of each matrix's row\_ptr in `concat_row_ptr`.
    pub batch_offsets_row_ptr: Vec<usize>,
    /// Start offset of each matrix's col\_idx / values in the concatenated arrays.
    pub batch_offsets_nnz: Vec<usize>,
    /// Number of rows in each matrix.
    pub row_counts: Vec<usize>,
    /// Number of columns in each matrix.
    pub col_counts: Vec<usize>,
    /// Number of non-zeros in each matrix.
    pub nnz_counts: Vec<usize>,
    /// Total number of matrices in the batch.
    pub batch_size: usize,
    /// Recommended execution strategy.
    pub strategy: Strategy,
}

impl<T: GpuFloat> BatchedSpMVPlan<T> {
    /// Build a plan from a slice of CSR matrices.
    ///
    /// Downloads each matrix to host memory, concatenates the arrays, and
    /// records offset metadata.
    ///
    /// # Errors
    ///
    /// Returns [`SparseError::InvalidArgument`] if `matrices` is empty.
    /// Returns [`SparseError::Cuda`] on download failure.
    pub fn from_matrices(matrices: &[CsrMatrix<T>]) -> SparseResult<Self> {
        if matrices.is_empty() {
            return Err(SparseError::InvalidArgument(
                "batch must contain at least one matrix".to_string(),
            ));
        }

        let batch_size = matrices.len();
        let mut concat_row_ptr = Vec::new();
        let mut concat_col_idx = Vec::new();
        let mut concat_values: Vec<T> = Vec::new();
        let mut batch_offsets_row_ptr = Vec::with_capacity(batch_size);
        let mut batch_offsets_nnz = Vec::with_capacity(batch_size);
        let mut row_counts = Vec::with_capacity(batch_size);
        let mut col_counts = Vec::with_capacity(batch_size);
        let mut nnz_counts = Vec::with_capacity(batch_size);

        for mat in matrices {
            let (h_rp, h_ci, h_vals) = mat.to_host()?;

            batch_offsets_row_ptr.push(concat_row_ptr.len());
            batch_offsets_nnz.push(concat_col_idx.len());
            row_counts.push(mat.rows() as usize);
            col_counts.push(mat.cols() as usize);
            nnz_counts.push(mat.nnz() as usize);

            concat_row_ptr.extend_from_slice(&h_rp);
            concat_col_idx.extend_from_slice(&h_ci);
            concat_values.extend_from_slice(&h_vals);
        }

        let total_nnz = nnz_counts.iter().copied().sum::<usize>();
        let avg_nnz = total_nnz.checked_div(batch_size).unwrap_or(0);
        let strategy = BatchScheduler::select_strategy_static(batch_size, avg_nnz);

        Ok(Self {
            concat_row_ptr,
            concat_col_idx,
            concat_values,
            batch_offsets_row_ptr,
            batch_offsets_nnz,
            row_counts,
            col_counts,
            nnz_counts,
            batch_size,
            strategy,
        })
    }

    /// Build a plan from host-side arrays directly (no GPU download required).
    ///
    /// Each entry in the parallel arrays describes one matrix in the batch.
    ///
    /// # Errors
    ///
    /// Returns [`SparseError::InvalidArgument`] if any of the slices are empty
    /// or if their lengths do not match.
    pub fn from_host_arrays(
        row_ptrs: &[Vec<i32>],
        col_indices: &[Vec<i32>],
        values: &[Vec<T>],
        rows: &[usize],
        cols: &[usize],
    ) -> SparseResult<Self> {
        let batch_size = row_ptrs.len();
        if batch_size == 0 {
            return Err(SparseError::InvalidArgument(
                "batch must contain at least one matrix".to_string(),
            ));
        }
        if col_indices.len() != batch_size
            || values.len() != batch_size
            || rows.len() != batch_size
            || cols.len() != batch_size
        {
            return Err(SparseError::InvalidArgument(
                "all input slices must have the same length".to_string(),
            ));
        }

        let mut concat_row_ptr = Vec::new();
        let mut concat_col_idx = Vec::new();
        let mut concat_values: Vec<T> = Vec::new();
        let mut batch_offsets_row_ptr = Vec::with_capacity(batch_size);
        let mut batch_offsets_nnz = Vec::with_capacity(batch_size);
        let mut row_counts = Vec::with_capacity(batch_size);
        let mut col_counts = Vec::with_capacity(batch_size);
        let mut nnz_counts = Vec::with_capacity(batch_size);

        for i in 0..batch_size {
            batch_offsets_row_ptr.push(concat_row_ptr.len());
            batch_offsets_nnz.push(concat_col_idx.len());
            row_counts.push(rows[i]);
            col_counts.push(cols[i]);
            nnz_counts.push(values[i].len());

            concat_row_ptr.extend_from_slice(&row_ptrs[i]);
            concat_col_idx.extend_from_slice(&col_indices[i]);
            concat_values.extend_from_slice(&values[i]);
        }

        let total_nnz = nnz_counts.iter().copied().sum::<usize>();
        let avg_nnz = total_nnz.checked_div(batch_size).unwrap_or(0);
        let strategy = BatchScheduler::select_strategy_static(batch_size, avg_nnz);

        Ok(Self {
            concat_row_ptr,
            concat_col_idx,
            concat_values,
            batch_offsets_row_ptr,
            batch_offsets_nnz,
            row_counts,
            col_counts,
            nnz_counts,
            batch_size,
            strategy,
        })
    }

    /// Returns the total number of non-zeros across the entire batch.
    #[inline]
    pub fn total_nnz(&self) -> usize {
        self.nnz_counts.iter().copied().sum()
    }

    /// Returns the total number of rows across the entire batch.
    #[inline]
    pub fn total_rows(&self) -> usize {
        self.row_counts.iter().copied().sum()
    }

    /// Returns the average non-zeros per matrix.
    #[inline]
    pub fn avg_nnz(&self) -> usize {
        if self.batch_size == 0 {
            return 0;
        }
        self.total_nnz() / self.batch_size
    }
}

// ---------------------------------------------------------------------------
// BatchedSpMV -- batched sparse matrix-vector multiply
// ---------------------------------------------------------------------------

/// Batched sparse matrix-vector multiplication.
///
/// Computes `y_i = alpha * A_i * x_i + beta * y_i` for a batch of sparse
/// matrices `A_i`, dense vectors `x_i`, and output vectors `y_i`.
///
/// The baseline implementation iterates sequentially on the host. For GPU
/// execution, use [`generate_batched_spmv_ptx`] to create a single fused
/// kernel that processes all matrices via offset arrays.
#[derive(Debug)]
pub struct BatchedSpMV<T: GpuFloat> {
    /// The batch of sparse matrices (stored as host-side CSR data).
    matrices: Vec<HostCsr<T>>,
}

/// Host-side CSR representation for batched operations.
#[derive(Debug, Clone)]
struct HostCsr<T> {
    rows: usize,
    cols: usize,
    row_ptr: Vec<i32>,
    col_idx: Vec<i32>,
    values: Vec<T>,
}

impl<T: GpuFloat + Add<Output = T> + Mul<Output = T> + AddAssign> BatchedSpMV<T> {
    /// Create a `BatchedSpMV` by downloading matrix data from device memory.
    ///
    /// # Errors
    ///
    /// Returns [`SparseError::InvalidArgument`] if `matrices` is empty.
    /// Returns [`SparseError::Cuda`] on download failure.
    pub fn from_device(matrices: &[CsrMatrix<T>]) -> SparseResult<Self> {
        if matrices.is_empty() {
            return Err(SparseError::InvalidArgument(
                "batch must contain at least one matrix".to_string(),
            ));
        }

        let mut host_mats = Vec::with_capacity(matrices.len());
        for mat in matrices {
            let (rp, ci, vals) = mat.to_host()?;
            host_mats.push(HostCsr {
                rows: mat.rows() as usize,
                cols: mat.cols() as usize,
                row_ptr: rp,
                col_idx: ci,
                values: vals,
            });
        }

        Ok(Self {
            matrices: host_mats,
        })
    }

    /// Create a `BatchedSpMV` from host-side CSR arrays.
    ///
    /// # Errors
    ///
    /// Returns [`SparseError::InvalidArgument`] on dimension mismatch.
    pub fn from_host(
        row_ptrs: Vec<Vec<i32>>,
        col_indices: Vec<Vec<i32>>,
        values: Vec<Vec<T>>,
        rows: Vec<usize>,
        cols: Vec<usize>,
    ) -> SparseResult<Self> {
        let n = row_ptrs.len();
        if n == 0 {
            return Err(SparseError::InvalidArgument(
                "batch must contain at least one matrix".to_string(),
            ));
        }
        if col_indices.len() != n || values.len() != n || rows.len() != n || cols.len() != n {
            return Err(SparseError::InvalidArgument(
                "all input vectors must have the same length".to_string(),
            ));
        }

        let mut host_mats = Vec::with_capacity(n);
        for i in 0..n {
            host_mats.push(HostCsr {
                rows: rows[i],
                cols: cols[i],
                row_ptr: row_ptrs[i].clone(),
                col_idx: col_indices[i].clone(),
                values: values[i].clone(),
            });
        }

        Ok(Self {
            matrices: host_mats,
        })
    }

    /// Returns the number of matrices in the batch.
    #[inline]
    pub fn batch_size(&self) -> usize {
        self.matrices.len()
    }

    /// Execute the batched SpMV on the host (sequential baseline).
    ///
    /// For each matrix `A_i`:
    ///   `y_i[r] = alpha * sum(A_i[r,c] * x_i[c]) + beta * y_i[r]`
    ///
    /// # Arguments
    ///
    /// * `xs` -- Slice of input vectors, one per matrix.
    /// * `ys` -- Slice of output vectors, one per matrix (modified in place).
    /// * `alpha` -- Scalar multiplier for `A * x`.
    /// * `beta` -- Scalar multiplier for existing `y`.
    ///
    /// # Errors
    ///
    /// Returns [`SparseError::DimensionMismatch`] if slice lengths do not match
    /// the batch size, or if individual vector dimensions are wrong.
    pub fn execute(&self, xs: &[Vec<T>], ys: &mut [Vec<T>], alpha: T, beta: T) -> SparseResult<()> {
        let n = self.matrices.len();
        if xs.len() != n || ys.len() != n {
            return Err(SparseError::DimensionMismatch(format!(
                "expected {} vectors, got xs={}, ys={}",
                n,
                xs.len(),
                ys.len()
            )));
        }

        for (i, mat) in self.matrices.iter().enumerate() {
            if xs[i].len() < mat.cols {
                return Err(SparseError::DimensionMismatch(format!(
                    "matrix {} has {} cols but x has {} elements",
                    i,
                    mat.cols,
                    xs[i].len()
                )));
            }
            if ys[i].len() < mat.rows {
                return Err(SparseError::DimensionMismatch(format!(
                    "matrix {} has {} rows but y has {} elements",
                    i,
                    mat.rows,
                    ys[i].len()
                )));
            }

            host_csr_spmv(mat, &xs[i], &mut ys[i], alpha, beta);
        }

        Ok(())
    }
}

/// Perform a single CSR SpMV on the host: y = alpha*A*x + beta*y.
fn host_csr_spmv<T: GpuFloat + Add<Output = T> + Mul<Output = T> + AddAssign>(
    mat: &HostCsr<T>,
    x: &[T],
    y: &mut [T],
    alpha: T,
    beta: T,
) {
    for (row, y_row) in y.iter_mut().enumerate().take(mat.rows) {
        let start = mat.row_ptr[row] as usize;
        let end = mat.row_ptr[row + 1] as usize;

        let mut acc = T::gpu_zero();
        for j in start..end {
            let col = mat.col_idx[j] as usize;
            acc += mat.values[j] * x[col];
        }

        *y_row = alpha * acc + beta * *y_row;
    }
}

// ---------------------------------------------------------------------------
// UniformBatchedSpMV -- shared structure, only values differ
// ---------------------------------------------------------------------------

/// Optimised batched SpMV when all matrices share the same sparsity pattern.
///
/// Only the non-zero values differ between matrices; the row pointers and
/// column indices are shared. This enables more efficient memory layout and
/// potential vectorisation across the batch dimension.
#[derive(Debug, Clone)]
pub struct UniformBatchedSpMV<T> {
    /// Number of rows in every matrix.
    rows: usize,
    /// Number of columns in every matrix.
    cols: usize,
    /// Shared row pointer array (length `rows + 1`).
    row_ptr: Vec<i32>,
    /// Shared column index array (length `nnz`).
    col_idx: Vec<i32>,
    /// Per-matrix values. `batch_values[i]` has length `nnz`.
    batch_values: Vec<Vec<T>>,
}

impl<T: GpuFloat + Add<Output = T> + Mul<Output = T> + AddAssign> UniformBatchedSpMV<T> {
    /// Create a uniform batch from a pattern matrix and per-batch values.
    ///
    /// The pattern matrix supplies the shared structure (row\_ptr, col\_idx);
    /// `batch_values[i]` supplies the non-zero values for the i-th matrix.
    ///
    /// # Errors
    ///
    /// Returns [`SparseError::InvalidArgument`] if `batch_values` is empty or
    /// if any entry has the wrong length.
    /// Returns [`SparseError::Cuda`] on download failure.
    pub fn from_pattern(pattern: &CsrMatrix<T>, batch_values: Vec<Vec<T>>) -> SparseResult<Self> {
        if batch_values.is_empty() {
            return Err(SparseError::InvalidArgument(
                "batch_values must not be empty".to_string(),
            ));
        }

        let nnz = pattern.nnz() as usize;
        for (i, vals) in batch_values.iter().enumerate() {
            if vals.len() != nnz {
                return Err(SparseError::InvalidArgument(format!(
                    "batch_values[{}] has length {} but pattern nnz is {}",
                    i,
                    vals.len(),
                    nnz
                )));
            }
        }

        let (rp, ci, _) = pattern.to_host()?;

        Ok(Self {
            rows: pattern.rows() as usize,
            cols: pattern.cols() as usize,
            row_ptr: rp,
            col_idx: ci,
            batch_values,
        })
    }

    /// Create a uniform batch from host-side arrays (no GPU required).
    ///
    /// # Errors
    ///
    /// Returns [`SparseError::InvalidArgument`] on dimension / length mismatch.
    pub fn from_host_arrays(
        rows: usize,
        cols: usize,
        row_ptr: Vec<i32>,
        col_idx: Vec<i32>,
        batch_values: Vec<Vec<T>>,
    ) -> SparseResult<Self> {
        if batch_values.is_empty() {
            return Err(SparseError::InvalidArgument(
                "batch_values must not be empty".to_string(),
            ));
        }
        if row_ptr.len() != rows + 1 {
            return Err(SparseError::InvalidArgument(format!(
                "row_ptr length {} != rows + 1 ({})",
                row_ptr.len(),
                rows + 1
            )));
        }
        let nnz = col_idx.len();
        for (i, vals) in batch_values.iter().enumerate() {
            if vals.len() != nnz {
                return Err(SparseError::InvalidArgument(format!(
                    "batch_values[{}] length {} != nnz {}",
                    i,
                    vals.len(),
                    nnz
                )));
            }
        }

        Ok(Self {
            rows,
            cols,
            row_ptr,
            col_idx,
            batch_values,
        })
    }

    /// Number of matrices in the batch.
    #[inline]
    pub fn batch_size(&self) -> usize {
        self.batch_values.len()
    }

    /// Execute batched SpMV on the host.
    ///
    /// # Errors
    ///
    /// Returns [`SparseError::DimensionMismatch`] on vector size mismatch.
    pub fn execute(&self, xs: &[Vec<T>], ys: &mut [Vec<T>], alpha: T, beta: T) -> SparseResult<()> {
        let n = self.batch_values.len();
        if xs.len() != n || ys.len() != n {
            return Err(SparseError::DimensionMismatch(format!(
                "expected {} vectors, got xs={}, ys={}",
                n,
                xs.len(),
                ys.len()
            )));
        }

        for i in 0..n {
            if xs[i].len() < self.cols {
                return Err(SparseError::DimensionMismatch(format!(
                    "x[{}] length {} < cols {}",
                    i,
                    xs[i].len(),
                    self.cols
                )));
            }
            if ys[i].len() < self.rows {
                return Err(SparseError::DimensionMismatch(format!(
                    "y[{}] length {} < rows {}",
                    i,
                    ys[i].len(),
                    self.rows
                )));
            }

            let mat = HostCsr {
                rows: self.rows,
                cols: self.cols,
                row_ptr: self.row_ptr.clone(),
                col_idx: self.col_idx.clone(),
                values: self.batch_values[i].clone(),
            };
            host_csr_spmv(&mat, &xs[i], &mut ys[i], alpha, beta);
        }

        Ok(())
    }
}

// ---------------------------------------------------------------------------
// BatchedSpGEMM -- batched sparse-sparse matrix multiply
// ---------------------------------------------------------------------------

/// Batched sparse-sparse matrix multiplication.
///
/// Computes `C_i = A_i * B_i` for a batch of sparse matrix pairs.
/// The baseline runs each SpGEMM independently on the host using a
/// hash-table accumulation approach.
#[derive(Debug)]
pub struct BatchedSpGEMM {
    _private: (),
}

impl BatchedSpGEMM {
    /// Create a new batched SpGEMM executor.
    #[inline]
    pub fn new() -> Self {
        Self { _private: () }
    }

    /// Execute batched SpGEMM on the host.
    ///
    /// For each pair `(A_i, B_i)`, computes `C_i = A_i * B_i` using a
    /// row-by-row hash-table accumulation.
    ///
    /// Returns host-side CSR triples `(row_ptr, col_idx, values)` for each
    /// result matrix.
    ///
    /// # Errors
    ///
    /// Returns [`SparseError::InvalidArgument`] if batch sizes differ or are empty.
    /// Returns [`SparseError::DimensionMismatch`] if dimensions are incompatible.
    /// Returns [`SparseError::Cuda`] on download failure.
    pub fn execute<T: GpuFloat + Add<Output = T> + Mul<Output = T> + AddAssign>(
        a_batch: &[CsrMatrix<T>],
        b_batch: &[CsrMatrix<T>],
    ) -> SparseResult<Vec<SpGEMMResultU32<T>>> {
        if a_batch.is_empty() {
            return Err(SparseError::InvalidArgument(
                "batch must not be empty".to_string(),
            ));
        }
        if a_batch.len() != b_batch.len() {
            return Err(SparseError::InvalidArgument(format!(
                "a_batch length {} != b_batch length {}",
                a_batch.len(),
                b_batch.len()
            )));
        }

        let mut results = Vec::with_capacity(a_batch.len());

        for (i, (a, b)) in a_batch.iter().zip(b_batch.iter()).enumerate() {
            if a.cols() != b.rows() {
                return Err(SparseError::DimensionMismatch(format!(
                    "pair {}: A.cols ({}) != B.rows ({})",
                    i,
                    a.cols(),
                    b.rows()
                )));
            }

            let (a_rp, a_ci, a_vals) = a.to_host()?;
            let (b_rp, b_ci, b_vals) = b.to_host()?;

            let (c_rp, c_ci, c_vals) = host_spgemm(
                &a_rp,
                &a_ci,
                &a_vals,
                a.rows() as usize,
                &b_rp,
                &b_ci,
                &b_vals,
                b.cols() as usize,
            );

            results.push((c_rp, c_ci, c_vals, a.rows(), b.cols()));
        }

        Ok(results)
    }

    /// Execute batched SpGEMM from host-side CSR arrays (no GPU required).
    ///
    /// # Errors
    ///
    /// Returns [`SparseError::InvalidArgument`] on length mismatch.
    /// Returns [`SparseError::DimensionMismatch`] on dimension mismatch.
    #[allow(clippy::too_many_arguments)]
    pub fn execute_host<T: GpuFloat + Add<Output = T> + Mul<Output = T> + AddAssign>(
        a_row_ptrs: &[Vec<i32>],
        a_col_indices: &[Vec<i32>],
        a_values: &[Vec<T>],
        a_rows: &[usize],
        a_cols: &[usize],
        b_row_ptrs: &[Vec<i32>],
        b_col_indices: &[Vec<i32>],
        b_values: &[Vec<T>],
        b_cols: &[usize],
    ) -> SparseResult<Vec<SpGEMMResultUsize<T>>> {
        let n = a_row_ptrs.len();
        if n == 0 {
            return Err(SparseError::InvalidArgument(
                "batch must not be empty".to_string(),
            ));
        }
        if b_row_ptrs.len() != n
            || a_col_indices.len() != n
            || a_values.len() != n
            || a_rows.len() != n
            || a_cols.len() != n
            || b_col_indices.len() != n
            || b_values.len() != n
            || b_cols.len() != n
        {
            return Err(SparseError::InvalidArgument(
                "all input slices must have the same length".to_string(),
            ));
        }

        let mut results = Vec::with_capacity(n);
        for i in 0..n {
            let b_rows_i = a_cols[i]; // A.cols == B.rows
            let (c_rp, c_ci, c_vals) = host_spgemm(
                &a_row_ptrs[i],
                &a_col_indices[i],
                &a_values[i],
                a_rows[i],
                &b_row_ptrs[i],
                &b_col_indices[i],
                &b_values[i],
                b_cols[i],
            );
            let _ = b_rows_i; // validated by caller
            results.push((c_rp, c_ci, c_vals, a_rows[i], b_cols[i]));
        }

        Ok(results)
    }
}

impl Default for BatchedSpGEMM {
    fn default() -> Self {
        Self::new()
    }
}

/// Host-side SpGEMM for a single pair: C = A * B using hash accumulation.
#[allow(clippy::too_many_arguments)]
fn host_spgemm<T: GpuFloat + Add<Output = T> + Mul<Output = T> + AddAssign>(
    a_rp: &[i32],
    a_ci: &[i32],
    a_vals: &[T],
    a_rows: usize,
    b_rp: &[i32],
    b_ci: &[i32],
    b_vals: &[T],
    b_cols: usize,
) -> (Vec<i32>, Vec<i32>, Vec<T>) {
    use std::collections::BTreeMap;

    let mut c_row_ptr = vec![0i32; a_rows + 1];
    let mut c_col_idx = Vec::new();
    let mut c_values: Vec<T> = Vec::new();

    for row in 0..a_rows {
        let a_start = a_rp[row] as usize;
        let a_end = a_rp[row + 1] as usize;

        // BTreeMap gives sorted column indices automatically
        let mut acc: BTreeMap<usize, T> = BTreeMap::new();

        for ja in a_start..a_end {
            let a_col = a_ci[ja] as usize;
            let a_val = a_vals[ja];

            let b_start = b_rp[a_col] as usize;
            let b_end = b_rp[a_col + 1] as usize;

            for jb in b_start..b_end {
                let b_col = b_ci[jb] as usize;
                if b_col < b_cols {
                    let product = a_val * b_vals[jb];
                    acc.entry(b_col)
                        .and_modify(|v| *v += product)
                        .or_insert(product);
                }
            }
        }

        for (&col, &val) in &acc {
            c_col_idx.push(col as i32);
            c_values.push(val);
        }
        c_row_ptr[row + 1] = c_col_idx.len() as i32;
    }

    (c_row_ptr, c_col_idx, c_values)
}

// ---------------------------------------------------------------------------
// BatchedTriSolve -- batched sparse triangular solve
// ---------------------------------------------------------------------------

/// Batched sparse triangular solve.
///
/// Solves `L_i * x_i = b_i` for each lower-triangular CSR matrix `L_i`
/// and right-hand side `b_i` in the batch.
#[derive(Debug)]
pub struct BatchedTriSolve {
    _private: (),
}

impl BatchedTriSolve {
    /// Create a new batched triangular solve executor.
    #[inline]
    pub fn new() -> Self {
        Self { _private: () }
    }

    /// Execute batched forward-substitution on the host.
    ///
    /// Each `L_i` must be square and lower-triangular with non-zero diagonal.
    ///
    /// # Errors
    ///
    /// Returns [`SparseError::InvalidArgument`] on batch size mismatch.
    /// Returns [`SparseError::DimensionMismatch`] if matrices are not square.
    /// Returns [`SparseError::SingularMatrix`] if a zero diagonal is encountered.
    /// Returns [`SparseError::Cuda`] on download failure.
    pub fn execute<T: GpuFloat + Add<Output = T> + Mul<Output = T> + AddAssign>(
        l_batch: &[CsrMatrix<T>],
        b_batch: &[Vec<T>],
    ) -> SparseResult<Vec<Vec<T>>> {
        if l_batch.is_empty() {
            return Err(SparseError::InvalidArgument(
                "batch must not be empty".to_string(),
            ));
        }
        if l_batch.len() != b_batch.len() {
            return Err(SparseError::InvalidArgument(format!(
                "l_batch length {} != b_batch length {}",
                l_batch.len(),
                b_batch.len()
            )));
        }

        let mut results = Vec::with_capacity(l_batch.len());

        for (i, (l, b)) in l_batch.iter().zip(b_batch.iter()).enumerate() {
            if l.rows() != l.cols() {
                return Err(SparseError::DimensionMismatch(format!(
                    "matrix {} is not square: {}x{}",
                    i,
                    l.rows(),
                    l.cols()
                )));
            }
            if b.len() < l.rows() as usize {
                return Err(SparseError::DimensionMismatch(format!(
                    "matrix {} has {} rows but rhs has {} elements",
                    i,
                    l.rows(),
                    b.len()
                )));
            }

            let (rp, ci, vals) = l.to_host()?;
            let x = host_forward_solve(&rp, &ci, &vals, l.rows() as usize, b)?;
            results.push(x);
        }

        Ok(results)
    }

    /// Execute batched forward-substitution from host-side arrays.
    ///
    /// # Errors
    ///
    /// Returns [`SparseError::InvalidArgument`] on length mismatch.
    /// Returns [`SparseError::SingularMatrix`] on zero diagonal.
    pub fn execute_host<T: GpuFloat + Add<Output = T> + Mul<Output = T> + AddAssign>(
        row_ptrs: &[Vec<i32>],
        col_indices: &[Vec<i32>],
        values: &[Vec<T>],
        sizes: &[usize],
        rhs: &[Vec<T>],
    ) -> SparseResult<Vec<Vec<T>>> {
        let n = row_ptrs.len();
        if n == 0 {
            return Err(SparseError::InvalidArgument(
                "batch must not be empty".to_string(),
            ));
        }
        if col_indices.len() != n || values.len() != n || sizes.len() != n || rhs.len() != n {
            return Err(SparseError::InvalidArgument(
                "all input slices must have the same length".to_string(),
            ));
        }

        let mut results = Vec::with_capacity(n);
        for i in 0..n {
            let x =
                host_forward_solve(&row_ptrs[i], &col_indices[i], &values[i], sizes[i], &rhs[i])?;
            results.push(x);
        }

        Ok(results)
    }
}

impl Default for BatchedTriSolve {
    fn default() -> Self {
        Self::new()
    }
}

/// Host-side lower-triangular forward substitution: L * x = b.
fn host_forward_solve<T: GpuFloat + Add<Output = T> + Mul<Output = T> + AddAssign>(
    row_ptr: &[i32],
    col_idx: &[i32],
    values: &[T],
    n: usize,
    b: &[T],
) -> SparseResult<Vec<T>> {
    let mut x = vec![T::gpu_zero(); n];

    for row in 0..n {
        let start = row_ptr[row] as usize;
        let end = row_ptr[row + 1] as usize;

        // Accumulate off-diagonal contributions and find diagonal entry.
        // x[row] = (b[row] - sum(L[row,j]*x[j] for j<row)) / L[row,row]
        let mut off_diag_sum = T::gpu_zero();
        let mut diag = T::gpu_zero();

        for j in start..end {
            let col = col_idx[j] as usize;
            if col < row {
                off_diag_sum += values[j] * x[col];
            } else if col == row {
                diag = values[j];
            }
        }

        // Check for zero diagonal
        if diag == T::gpu_zero() {
            return Err(SparseError::SingularMatrix);
        }

        // x[row] = (b[row] - off_diag_sum) / diag
        // We cannot subtract directly without Sub trait. Use bit manipulation
        // to negate for f32/f64.
        let neg_sum = negate_float(off_diag_sum);
        let numerator = b[row] + neg_sum;
        x[row] = divide_float(numerator, diag);
    }

    Ok(x)
}

/// Negate a GpuFloat value by flipping the sign bit.
#[inline]
fn negate_float<T: GpuFloat>(val: T) -> T {
    let bits = val.to_bits_u64();
    if T::SIZE == 4 {
        // f32: flip bit 31
        T::from_bits_u64(bits ^ (1u64 << 31))
    } else {
        // f64: flip bit 63
        T::from_bits_u64(bits ^ (1u64 << 63))
    }
}

/// Divide two GpuFloat values by converting through f64.
#[inline]
fn divide_float<T: GpuFloat>(num: T, den: T) -> T {
    // Convert to f64, divide, convert back. Works for f32 and f64.
    let n_bits = num.to_bits_u64();
    let d_bits = den.to_bits_u64();
    if T::SIZE == 4 {
        let n = f32::from_bits(n_bits as u32) as f64;
        let d = f32::from_bits(d_bits as u32) as f64;
        let result = (n / d) as f32;
        T::from_bits_u64(u64::from(result.to_bits()))
    } else {
        let n = f64::from_bits(n_bits);
        let d = f64::from_bits(d_bits);
        let result = n / d;
        T::from_bits_u64(result.to_bits())
    }
}

// ---------------------------------------------------------------------------
// PTX generation stub for batched SpMV
// ---------------------------------------------------------------------------

/// Generate a batched SpMV PTX kernel that processes all matrices via offset
/// arrays in a single launch.
///
/// The generated kernel expects the following arguments:
/// - Concatenated row\_ptr, col\_idx, values arrays
/// - batch\_offsets\_row\_ptr, batch\_offsets\_nnz arrays
/// - row\_counts array
/// - x pointers array, y pointers array
/// - alpha, beta scalars
/// - batch\_size
///
/// Each thread block processes one matrix in the batch; within a block,
/// threads cooperate on rows using the scalar SpMV strategy.
///
/// # Note
///
/// This generates valid PTX source but actual GPU execution requires a
/// CUDA-capable device.
pub fn generate_batched_spmv_ptx<T: GpuFloat>() -> String {
    let type_name = T::NAME;
    let ptx_type = match T::SIZE {
        4 => ".f32",
        8 => ".f64",
        _ => ".f32",
    };
    let elem_size = match T::SIZE {
        8 => 8,
        _ => 4,
    };
    let zero_literal = match T::SIZE {
        8 => "0d0000000000000000",
        _ => "0f00000000",
    };

    format!(
        r#"//
// Batched SpMV kernel for {type_name}
// Generated by oxicuda-sparse batched module
//
.version 7.0
.target sm_70
.address_size 64

.visible .entry batched_spmv_{type_name}(
    .param .u64 concat_row_ptr,
    .param .u64 concat_col_idx,
    .param .u64 concat_values,
    .param .u64 batch_offsets_rp,
    .param .u64 batch_offsets_nnz,
    .param .u64 row_counts,
    .param .u64 x_ptrs,
    .param .u64 y_ptrs,
    .param {ptx_type} alpha,
    .param {ptx_type} beta,
    .param .u32 batch_size
)
{{
    .reg .u32 %r<16>;
    .reg .u64 %rd<32>;
    .reg {ptx_type} %f<8>;
    .reg .pred %p<4>;

    // blockIdx.x = matrix index in batch
    mov.u32 %r0, %ctaid.x;
    // Early exit if blockIdx >= batch_size
    ld.param.u32 %r1, [batch_size];
    setp.ge.u32 %p0, %r0, %r1;
    @%p0 ret;

    // threadIdx.x = local row within this matrix
    mov.u32 %r2, %tid.x;

    // Load row_count for this matrix
    ld.param.u64 %rd0, [row_counts];
    cvt.u64.u32 %rd1, %r0;
    mad.wide.u32 %rd2, %r0, 4, %rd0;
    ld.global.u32 %r3, [%rd2];

    // Early exit if tid >= row_count
    setp.ge.u32 %p1, %r2, %r3;
    @%p1 ret;

    // Load row_ptr offset for this matrix
    ld.param.u64 %rd3, [batch_offsets_rp];
    mad.wide.u32 %rd4, %r0, 4, %rd3;
    ld.global.u32 %r4, [%rd4];

    // row_start = concat_row_ptr[rp_offset + tid], row_end = next entry
    ld.param.u64 %rd5, [concat_row_ptr];
    add.u32 %r11, %r4, %r2;
    mul.wide.u32 %rd6, %r11, 4;
    add.u64 %rd5, %rd5, %rd6;
    ld.global.u32 %r5, [%rd5];
    add.u64 %rd6, %rd5, 4;
    ld.global.u32 %r6, [%rd6];

    // Load nnz offset for this matrix
    ld.param.u64 %rd7, [batch_offsets_nnz];
    mad.wide.u32 %rd8, %r0, 4, %rd7;
    ld.global.u32 %r7, [%rd8];

    // Load x and y pointers for this matrix
    ld.param.u64 %rd9, [x_ptrs];
    mad.wide.u32 %rd10, %r0, 8, %rd9;
    ld.global.u64 %rd10, [%rd10];
    ld.param.u64 %rd11, [y_ptrs];
    mad.wide.u32 %rd12, %r0, 8, %rd11;
    ld.global.u64 %rd11, [%rd12];

    // Load concatenated col_idx / values bases and alpha / beta scalars
    ld.param.u64 %rd12, [concat_col_idx];
    ld.param.u64 %rd13, [concat_values];
    ld.param{ptx_type} %f6, [alpha];
    ld.param{ptx_type} %f7, [beta];

    // acc = 0; iterate row_start .. row_end
    mov{ptx_type} %f0, {zero_literal};
    mov.u32 %r8, %r5;

$ROW_LOOP:
    setp.lt.u32 %p2, %r8, %r6;
    @!%p2 bra $ROW_DONE;

    // k = nnz_offset + absolute row nnz index
    add.u32 %r10, %r7, %r8;

    // col = concat_col_idx[k]
    mul.wide.u32 %rd14, %r10, 4;
    add.u64 %rd15, %rd12, %rd14;
    ld.global.u32 %r9, [%rd15];

    // val = concat_values[k]
    mul.wide.u32 %rd16, %r10, {elem_size};
    add.u64 %rd17, %rd13, %rd16;
    ld.global{ptx_type} %f1, [%rd17];

    // x_val = x[col]
    mul.wide.u32 %rd18, %r9, {elem_size};
    add.u64 %rd19, %rd10, %rd18;
    ld.global{ptx_type} %f2, [%rd19];

    // acc += val * x_val
    fma.rn{ptx_type} %f0, %f1, %f2, %f0;

    add.u32 %r8, %r8, 1;
    bra $ROW_LOOP;

$ROW_DONE:
    // y[tid] = alpha * acc + beta * y[tid]
    mul.wide.u32 %rd20, %r2, {elem_size};
    add.u64 %rd21, %rd11, %rd20;
    ld.global{ptx_type} %f3, [%rd21];
    mul.rn{ptx_type} %f4, %f6, %f0;
    mul.rn{ptx_type} %f5, %f7, %f3;
    add.rn{ptx_type} %f4, %f4, %f5;
    st.global{ptx_type} [%rd21], %f4;

    ret;
}}
"#
    )
}

// ---------------------------------------------------------------------------
// CPU reference: batched SpMV for multiple right-hand sides
// ---------------------------------------------------------------------------

/// Batched SpMV (CPU reference): compute `y_i = A * x_i` for all `b` RHS vectors.
///
/// The shared sparse matrix `A` is given in CSR format. `x_batch` and `y_batch`
/// are stored in **column-major** order: element `(row_or_col, b)` is at index
/// `row_or_col * batch_size + b`.
///
/// # Arguments
///
/// * `n_rows` -- Number of rows in `A`.
/// * `n_cols` -- Number of columns in `A` (used for bounds; not strictly needed here).
/// * `row_ptr` -- CSR row pointer of length `n_rows + 1`.
/// * `col_idx` -- CSR column indices of length `nnz`.
/// * `values` -- CSR values of length `nnz`.
/// * `x_batch` -- Column-major matrix of shape `[n_cols × batch_size]`.
/// * `batch_size` -- Number of simultaneous right-hand sides.
///
/// # Returns
///
/// Output matrix `y` of shape `[n_rows × batch_size]` in column-major order.
pub fn batched_spmv_cpu(
    n_rows: usize,
    _n_cols: usize,
    row_ptr: &[u32],
    col_idx: &[u32],
    values: &[f32],
    x_batch: &[f32],
    batch_size: usize,
) -> Vec<f32> {
    let mut y = vec![0.0f32; n_rows * batch_size];
    for row in 0..n_rows {
        let start = row_ptr[row] as usize;
        let end = row_ptr[row + 1] as usize;
        for idx in start..end {
            let col = col_idx[idx] as usize;
            let val = values[idx];
            for b in 0..batch_size {
                y[row * batch_size + b] += val * x_batch[col * batch_size + b];
            }
        }
    }
    y
}

// ---------------------------------------------------------------------------
// CPU reference: mixed-precision SpMV (FP16 storage, FP64 accumulation -> FP32)
// ---------------------------------------------------------------------------

/// Mixed-precision SpMV (CPU reference): matrix values stored in FP16 range,
/// accumulation performed in FP64, result cast to FP32.
///
/// On actual GPU hardware, values would be loaded as FP16 (2 bytes each) and
/// accumulated in FP32 registers. This CPU reference simulates that by treating
/// `values_fp16` as FP32 scalars already quantised to FP16 precision, and
/// accumulating in FP64 to avoid catastrophic cancellation in the reference.
///
/// # Arguments
///
/// * `n_rows` -- Number of rows.
/// * `row_ptr` -- CSR row pointer of length `n_rows + 1`.
/// * `col_idx` -- CSR column indices of length `nnz`.
/// * `values_fp16` -- Per-entry matrix values (FP16-range, stored as f32).
/// * `x` -- Dense input vector of length `n_cols`.
///
/// # Returns
///
/// Output vector `y` of length `n_rows` as FP32.
pub fn mixed_precision_spmv_cpu(
    n_rows: usize,
    row_ptr: &[u32],
    col_idx: &[u32],
    values_fp16: &[f32],
    x: &[f32],
) -> Vec<f32> {
    let mut y = vec![0.0f32; n_rows];
    for row in 0..n_rows {
        let start = row_ptr[row] as usize;
        let end = row_ptr[row + 1] as usize;
        // Accumulate in f64 for numerical stability of the reference
        let mut acc = 0.0f64;
        for idx in start..end {
            let col = col_idx[idx] as usize;
            let val = values_fp16[idx] as f64;
            acc += val * (x[col] as f64);
        }
        y[row] = acc as f32;
    }
    y
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests;