Skip to main content

oxicuda_sparse/ops/
spmv.rs

1//! Sparse matrix-vector multiplication (SpMV).
2//!
3//! Computes `y = alpha * A * x + beta * y` where `A` is a sparse CSR matrix
4//! and `x`, `y` are dense vectors stored as raw device pointers.
5//!
6//! Three kernel strategies are available:
7//! - **Scalar**: one thread per row (best for very sparse rows, < 4 nnz/row)
8//! - **Vector**: one warp per row with shuffle reduction (best for moderate sparsity)
9//! - **Adaptive**: auto-selects based on average nnz per row
10
11use std::sync::Arc;
12
13use oxicuda_blas::GpuFloat;
14use oxicuda_driver::Module;
15use oxicuda_driver::ffi::CUdeviceptr;
16use oxicuda_launch::{Kernel, LaunchParams, grid_size_for};
17use oxicuda_ptx::prelude::*;
18
19use crate::error::{SparseError, SparseResult};
20use crate::format::CsrMatrix;
21use crate::handle::SparseHandle;
22use crate::ptx_helpers::{
23    add_float, emit_warp_reduce_sum, fma_float, load_float_imm, load_global_float, mul_float,
24    reinterpret_bits_to_float, store_global_float,
25};
26
27/// Algorithm selection for SpMV.
28#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
29pub enum SpMVAlgo {
30    /// One thread per row. Best for very sparse matrices (< 4 nnz/row).
31    Scalar,
32    /// One warp (32 threads) per row with warp shuffle reduction.
33    /// Best for moderate sparsity (4-64 nnz/row).
34    Vector,
35    /// Automatically selects Scalar or Vector based on the matrix structure.
36    Adaptive,
37}
38
39/// Default block size for scalar SpMV.
40const SPMV_SCALAR_BLOCK: u32 = 256;
41
42/// Default block size for vector SpMV (must be a multiple of warp size 32).
43const SPMV_VECTOR_BLOCK: u32 = 256;
44
45/// Threshold for auto-selecting Vector over Scalar (nnz per row).
46const VECTOR_THRESHOLD: f64 = 4.0;
47
48/// Resolve the [`SpMVAlgo::Adaptive`] selection to a concrete kernel given the
49/// matrix's average non-zeros per row.
50///
51/// This pure function contains the kernel-selection heuristic so it can be
52/// tested independently of GPU device memory.
53///
54/// * `avg_nnz_per_row < VECTOR_THRESHOLD` → [`SpMVAlgo::Scalar`]
55/// * `avg_nnz_per_row >= VECTOR_THRESHOLD` → [`SpMVAlgo::Vector`]
56#[inline]
57pub(crate) fn resolve_adaptive(avg_nnz_per_row: f64) -> SpMVAlgo {
58    if avg_nnz_per_row >= VECTOR_THRESHOLD {
59        SpMVAlgo::Vector
60    } else {
61        SpMVAlgo::Scalar
62    }
63}
64
65/// Sparse matrix-vector multiplication: `y = alpha * A * x + beta * y`.
66///
67/// # Arguments
68///
69/// * `handle` -- Sparse handle providing stream and device context.
70/// * `algo` -- Algorithm selection strategy.
71/// * `alpha` -- Scalar multiplier for `A * x`.
72/// * `a` -- Sparse CSR matrix `A`.
73/// * `x_ptr` -- Device pointer to dense vector `x` of length `A.cols()`.
74/// * `beta` -- Scalar multiplier for existing `y`.
75/// * `y_ptr` -- Device pointer to dense vector `y` of length `A.rows()`.
76///
77/// # Errors
78///
79/// Returns [`SparseError::PtxGeneration`] if kernel generation fails.
80/// Returns [`SparseError::Cuda`] on kernel launch failure.
81#[allow(clippy::too_many_arguments)]
82pub fn spmv<T: GpuFloat>(
83    handle: &SparseHandle,
84    algo: SpMVAlgo,
85    alpha: T,
86    a: &CsrMatrix<T>,
87    x_ptr: CUdeviceptr,
88    beta: T,
89    y_ptr: CUdeviceptr,
90) -> SparseResult<()> {
91    if a.rows() == 0 || a.cols() == 0 {
92        return Ok(());
93    }
94
95    // Resolve Adaptive algorithm
96    let effective_algo = match algo {
97        SpMVAlgo::Adaptive => resolve_adaptive(a.avg_nnz_per_row()),
98        other => other,
99    };
100
101    match effective_algo {
102        SpMVAlgo::Scalar => spmv_scalar(handle, alpha, a, x_ptr, beta, y_ptr),
103        SpMVAlgo::Vector => spmv_vector(handle, alpha, a, x_ptr, beta, y_ptr),
104        SpMVAlgo::Adaptive => {
105            // Already resolved above; unreachable
106            spmv_scalar(handle, alpha, a, x_ptr, beta, y_ptr)
107        }
108    }
109}
110
111/// Scalar SpMV: one thread per row.
112fn spmv_scalar<T: GpuFloat>(
113    handle: &SparseHandle,
114    alpha: T,
115    a: &CsrMatrix<T>,
116    x_ptr: CUdeviceptr,
117    beta: T,
118    y_ptr: CUdeviceptr,
119) -> SparseResult<()> {
120    let ptx = emit_spmv_scalar::<T>(handle.sm_version())?;
121    let module = Arc::new(Module::from_ptx(&ptx)?);
122    let kernel = Kernel::from_module(module, "spmv_scalar")?;
123
124    let block_size = SPMV_SCALAR_BLOCK;
125    let grid_size = grid_size_for(a.rows(), block_size);
126    let params = LaunchParams::new(grid_size, block_size);
127
128    kernel.launch(
129        &params,
130        handle.stream(),
131        &(
132            a.row_ptr().as_device_ptr(),
133            a.col_idx().as_device_ptr(),
134            a.values().as_device_ptr(),
135            x_ptr,
136            y_ptr,
137            alpha.to_bits_u64(),
138            beta.to_bits_u64(),
139            a.rows(),
140        ),
141    )?;
142
143    Ok(())
144}
145
146/// Vector SpMV: one warp (32 threads) per row.
147fn spmv_vector<T: GpuFloat>(
148    handle: &SparseHandle,
149    alpha: T,
150    a: &CsrMatrix<T>,
151    x_ptr: CUdeviceptr,
152    beta: T,
153    y_ptr: CUdeviceptr,
154) -> SparseResult<()> {
155    let ptx = emit_spmv_vector::<T>(handle.sm_version())?;
156    let module = Arc::new(Module::from_ptx(&ptx)?);
157    let kernel = Kernel::from_module(module, "spmv_vector")?;
158
159    let block_size = SPMV_VECTOR_BLOCK;
160    // Each warp handles one row; warps_per_block = block_size / 32
161    let warps_per_block = block_size / 32;
162    let grid_size = grid_size_for(a.rows(), warps_per_block);
163    let params = LaunchParams::new(grid_size, block_size);
164
165    kernel.launch(
166        &params,
167        handle.stream(),
168        &(
169            a.row_ptr().as_device_ptr(),
170            a.col_idx().as_device_ptr(),
171            a.values().as_device_ptr(),
172            x_ptr,
173            y_ptr,
174            alpha.to_bits_u64(),
175            beta.to_bits_u64(),
176            a.rows(),
177        ),
178    )?;
179
180    Ok(())
181}
182
183/// Generates PTX for scalar SpMV (one thread per row).
184fn emit_spmv_scalar<T: GpuFloat>(sm: SmVersion) -> SparseResult<String> {
185    let _ptx_ty = T::PTX_TYPE;
186    let elem_bytes = T::size_u32();
187    let is_f64 = T::SIZE == 8;
188
189    KernelBuilder::new("spmv_scalar")
190        .target(sm)
191        .param("row_ptr", PtxType::U64)
192        .param("col_idx", PtxType::U64)
193        .param("values", PtxType::U64)
194        .param("x_ptr", PtxType::U64)
195        .param("y_ptr", PtxType::U64)
196        .param("alpha_bits", PtxType::U64)
197        .param("beta_bits", PtxType::U64)
198        .param("num_rows", PtxType::U32)
199        .body(move |b| {
200            let gid = b.global_thread_id_x();
201            let num_rows = b.load_param_u32("num_rows");
202
203            let gid_inner = gid.clone();
204            b.if_lt_u32(gid, num_rows, move |b| {
205                let row = gid_inner;
206                let row_ptr_base = b.load_param_u64("row_ptr");
207                let col_idx_base = b.load_param_u64("col_idx");
208                let values_base = b.load_param_u64("values");
209                let x_ptr = b.load_param_u64("x_ptr");
210                let y_ptr = b.load_param_u64("y_ptr");
211                let alpha_bits = b.load_param_u64("alpha_bits");
212                let beta_bits = b.load_param_u64("beta_bits");
213
214                let alpha = reinterpret_bits_to_float::<T>(b, alpha_bits);
215                let beta = reinterpret_bits_to_float::<T>(b, beta_bits);
216
217                // Load row_ptr[row] and row_ptr[row+1] (i32 = 4 bytes)
218                let rp_addr = b.byte_offset_addr(row_ptr_base.clone(), row.clone(), 4);
219                let row_start = b.load_global_i32(rp_addr);
220
221                let row_plus_1 = b.alloc_reg(PtxType::U32);
222                b.raw_ptx(&format!("add.u32 {row_plus_1}, {row}, 1;"));
223                let rp_addr_next = b.byte_offset_addr(row_ptr_base, row_plus_1, 4);
224                let row_end = b.load_global_i32(rp_addr_next);
225
226                // Initialize accumulator
227                let acc = load_float_imm::<T>(b, 0.0);
228
229                // Loop over non-zeros in this row
230                let loop_label = b.fresh_label("spmv_loop");
231                let done_label = b.fresh_label("spmv_done");
232
233                let k = b.alloc_reg(PtxType::U32);
234                // Convert row_start (i32) to u32
235                let rs_u32 = b.alloc_reg(PtxType::U32);
236                b.raw_ptx(&format!("mov.b32 {rs_u32}, {row_start};"));
237                b.raw_ptx(&format!("mov.u32 {k}, {rs_u32};"));
238
239                let re_u32 = b.alloc_reg(PtxType::U32);
240                b.raw_ptx(&format!("mov.b32 {re_u32}, {row_end};"));
241
242                b.label(&loop_label);
243                let pred = b.alloc_reg(PtxType::Pred);
244                b.raw_ptx(&format!("setp.lo.u32 {pred}, {k}, {re_u32};"));
245                b.raw_ptx(&format!("@!{pred} bra {done_label};"));
246
247                // Load col_idx[k] (i32 = 4 bytes)
248                let ci_addr = b.byte_offset_addr(col_idx_base.clone(), k.clone(), 4);
249                let col = b.load_global_i32(ci_addr);
250                let col_u32 = b.alloc_reg(PtxType::U32);
251                b.raw_ptx(&format!("mov.b32 {col_u32}, {col};"));
252
253                // Load values[k]
254                let v_addr = b.byte_offset_addr(values_base.clone(), k.clone(), elem_bytes);
255                let val = load_global_float::<T>(b, v_addr);
256
257                // Load x[col]
258                let x_addr = b.byte_offset_addr(x_ptr.clone(), col_u32, elem_bytes);
259                let x_val = load_global_float::<T>(b, x_addr);
260
261                // acc += val * x_val
262                let new_acc = fma_float::<T>(b, val, x_val, acc.clone());
263                let mov_suffix = if is_f64 { "f64" } else { "f32" };
264                b.raw_ptx(&format!("mov.{mov_suffix} {acc}, {new_acc};"));
265
266                // k++
267                b.raw_ptx(&format!("add.u32 {k}, {k}, 1;"));
268                b.branch(&loop_label);
269                b.label(&done_label);
270
271                // Compute y = alpha * acc + beta * y_old
272                let y_addr = b.byte_offset_addr(y_ptr, row, elem_bytes);
273                let y_old = load_global_float::<T>(b, y_addr.clone());
274
275                let alpha_acc = mul_float::<T>(b, alpha, acc);
276                let beta_y = mul_float::<T>(b, beta, y_old);
277                let result = add_float::<T>(b, alpha_acc, beta_y);
278
279                store_global_float::<T>(b, y_addr, result);
280            });
281
282            b.ret();
283        })
284        .build()
285        .map_err(|e| SparseError::PtxGeneration(e.to_string()))
286}
287
288/// Generates PTX for vector SpMV (one warp per row).
289fn emit_spmv_vector<T: GpuFloat>(sm: SmVersion) -> SparseResult<String> {
290    let ptx_ty = T::PTX_TYPE;
291    let elem_bytes = T::size_u32();
292    let is_f64 = T::SIZE == 8;
293
294    // Suppress unused variable warnings for ptx_ty
295    let _ = ptx_ty;
296
297    KernelBuilder::new("spmv_vector")
298        .target(sm)
299        .param("row_ptr", PtxType::U64)
300        .param("col_idx", PtxType::U64)
301        .param("values", PtxType::U64)
302        .param("x_ptr", PtxType::U64)
303        .param("y_ptr", PtxType::U64)
304        .param("alpha_bits", PtxType::U64)
305        .param("beta_bits", PtxType::U64)
306        .param("num_rows", PtxType::U32)
307        .body(move |b| {
308            // Each warp handles one row. Warp ID = global_thread_id / 32
309            let tid_global = b.global_thread_id_x();
310            let num_rows = b.load_param_u32("num_rows");
311
312            // Lane within warp (0..31)
313            let lane = b.alloc_reg(PtxType::U32);
314            b.raw_ptx(&format!("and.b32 {lane}, {tid_global}, 31;"));
315
316            // Warp ID = tid_global >> 5
317            let warp_id = b.alloc_reg(PtxType::U32);
318            b.raw_ptx(&format!("shr.u32 {warp_id}, {tid_global}, 5;"));
319
320            let warp_id_inner = warp_id.clone();
321            let lane_inner = lane.clone();
322            b.if_lt_u32(warp_id, num_rows, move |b| {
323                let row = warp_id_inner;
324                let lane = lane_inner;
325
326                let row_ptr_base = b.load_param_u64("row_ptr");
327                let col_idx_base = b.load_param_u64("col_idx");
328                let values_base = b.load_param_u64("values");
329                let x_ptr = b.load_param_u64("x_ptr");
330                let y_ptr = b.load_param_u64("y_ptr");
331                let alpha_bits = b.load_param_u64("alpha_bits");
332                let beta_bits = b.load_param_u64("beta_bits");
333
334                let alpha = reinterpret_bits_to_float::<T>(b, alpha_bits);
335                let beta = reinterpret_bits_to_float::<T>(b, beta_bits);
336
337                // Load row bounds
338                let rp_addr = b.byte_offset_addr(row_ptr_base.clone(), row.clone(), 4);
339                let row_start_i32 = b.load_global_i32(rp_addr);
340                let row_start = b.alloc_reg(PtxType::U32);
341                b.raw_ptx(&format!("mov.b32 {row_start}, {row_start_i32};"));
342
343                let row_plus_1 = b.alloc_reg(PtxType::U32);
344                b.raw_ptx(&format!("add.u32 {row_plus_1}, {row}, 1;"));
345                let rp_addr_next = b.byte_offset_addr(row_ptr_base, row_plus_1, 4);
346                let row_end_i32 = b.load_global_i32(rp_addr_next);
347                let row_end = b.alloc_reg(PtxType::U32);
348                b.raw_ptx(&format!("mov.b32 {row_end}, {row_end_i32};"));
349
350                // Each lane starts at row_start + lane, stride 32
351                let acc = load_float_imm::<T>(b, 0.0);
352
353                let k = b.alloc_reg(PtxType::U32);
354                b.raw_ptx(&format!("add.u32 {k}, {row_start}, {lane};"));
355
356                let loop_label = b.fresh_label("spmv_vloop");
357                let done_label = b.fresh_label("spmv_vdone");
358
359                b.label(&loop_label);
360                let pred = b.alloc_reg(PtxType::Pred);
361                b.raw_ptx(&format!("setp.lo.u32 {pred}, {k}, {row_end};"));
362                b.raw_ptx(&format!("@!{pred} bra {done_label};"));
363
364                // Load col and value
365                let ci_addr = b.byte_offset_addr(col_idx_base.clone(), k.clone(), 4);
366                let col_i32 = b.load_global_i32(ci_addr);
367                let col_u32 = b.alloc_reg(PtxType::U32);
368                b.raw_ptx(&format!("mov.b32 {col_u32}, {col_i32};"));
369
370                let v_addr = b.byte_offset_addr(values_base.clone(), k.clone(), elem_bytes);
371                let val = load_global_float::<T>(b, v_addr);
372
373                let x_addr = b.byte_offset_addr(x_ptr.clone(), col_u32, elem_bytes);
374                let x_val = load_global_float::<T>(b, x_addr);
375
376                let new_acc = fma_float::<T>(b, val, x_val, acc.clone());
377                let mov_suffix = if is_f64 { "f64" } else { "f32" };
378                b.raw_ptx(&format!("mov.{mov_suffix} {acc}, {new_acc};"));
379
380                // k += 32 (warp width)
381                b.raw_ptx(&format!("add.u32 {k}, {k}, 32;"));
382                b.branch(&loop_label);
383                b.label(&done_label);
384
385                // Warp shuffle reduction
386                let reduced = emit_warp_reduce_sum::<T>(b, acc);
387
388                // Lane 0 writes the result
389                let is_lane_0 = b.alloc_reg(PtxType::Pred);
390                b.raw_ptx(&format!("setp.eq.u32 {is_lane_0}, {lane}, 0;"));
391
392                let write_label = b.fresh_label("spmv_write");
393                let skip_label = b.fresh_label("spmv_skip");
394                b.raw_ptx(&format!("@!{is_lane_0} bra {skip_label};"));
395
396                b.label(&write_label);
397                let y_addr = b.byte_offset_addr(y_ptr, row, elem_bytes);
398                let y_old = load_global_float::<T>(b, y_addr.clone());
399
400                let alpha_acc = mul_float::<T>(b, alpha, reduced);
401                let beta_y = mul_float::<T>(b, beta, y_old);
402                let result = add_float::<T>(b, alpha_acc, beta_y);
403                store_global_float::<T>(b, y_addr, result);
404
405                b.label(&skip_label);
406            });
407
408            b.ret();
409        })
410        .build()
411        .map_err(|e| SparseError::PtxGeneration(e.to_string()))
412}
413
414#[cfg(test)]
415mod tests {
416    use super::*;
417
418    #[test]
419    fn spmv_algo_auto_select() {
420        // avg_nnz < threshold => Scalar
421        // Verify VECTOR_THRESHOLD is set to a reasonable value for algorithm selection.
422        let threshold = VECTOR_THRESHOLD;
423        assert!(threshold > 3.0);
424    }
425
426    #[test]
427    fn spmv_scalar_ptx_generates() {
428        let ptx = emit_spmv_scalar::<f32>(SmVersion::Sm80);
429        assert!(ptx.is_ok());
430        let ptx = ptx.expect("test: PTX gen should succeed");
431        assert!(ptx.contains(".entry spmv_scalar"));
432        assert!(ptx.contains(".target sm_80"));
433    }
434
435    #[test]
436    fn spmv_vector_ptx_generates() {
437        let ptx = emit_spmv_vector::<f32>(SmVersion::Sm80);
438        assert!(ptx.is_ok());
439        let ptx = ptx.expect("test: PTX gen should succeed");
440        assert!(ptx.contains(".entry spmv_vector"));
441    }
442
443    #[test]
444    fn spmv_scalar_ptx_f64() {
445        let ptx = emit_spmv_scalar::<f64>(SmVersion::Sm80);
446        assert!(ptx.is_ok());
447    }
448
449    #[test]
450    fn spmv_vector_ptx_f64() {
451        let ptx = emit_spmv_vector::<f64>(SmVersion::Sm80);
452        assert!(ptx.is_ok());
453    }
454
455    // -----------------------------------------------------------------------
456    // Task 5a: Auto-selection heuristic tests (CPU-only, no GPU required)
457    // -----------------------------------------------------------------------
458
459    /// Very sparse rows (avg_nnz ≈ 1.5, well below threshold 4.0) → Scalar.
460    #[test]
461    fn test_spmv_selects_scalar_for_very_sparse() {
462        // 100 rows, 150 nnz → avg = 1.5
463        let avg = 150.0_f64 / 100.0;
464        assert!(avg < VECTOR_THRESHOLD);
465        assert_eq!(resolve_adaptive(avg), SpMVAlgo::Scalar);
466    }
467
468    /// Moderate density (avg_nnz = 32, above threshold 4.0) → Vector.
469    #[test]
470    fn test_spmv_selects_vector_for_moderate_density() {
471        let avg = 32.0_f64;
472        assert!(avg >= VECTOR_THRESHOLD);
473        assert_eq!(resolve_adaptive(avg), SpMVAlgo::Vector);
474    }
475
476    /// Dense rows (avg_nnz = 128, well above threshold) → Vector.
477    #[test]
478    fn test_spmv_selects_vector_for_dense() {
479        let avg = 128.0_f64;
480        assert!(avg >= VECTOR_THRESHOLD);
481        assert_eq!(resolve_adaptive(avg), SpMVAlgo::Vector);
482    }
483
484    /// Boundary: just below threshold → Scalar; at threshold → Vector.
485    #[test]
486    fn test_spmv_selection_boundary_conditions() {
487        // Just below threshold (3.9999…)
488        let just_below = VECTOR_THRESHOLD - f64::EPSILON * VECTOR_THRESHOLD;
489        assert_eq!(resolve_adaptive(just_below), SpMVAlgo::Scalar);
490
491        // Exactly at threshold
492        assert_eq!(resolve_adaptive(VECTOR_THRESHOLD), SpMVAlgo::Vector);
493
494        // Slightly above threshold
495        let just_above = VECTOR_THRESHOLD + f64::EPSILON * VECTOR_THRESHOLD;
496        assert_eq!(resolve_adaptive(just_above), SpMVAlgo::Vector);
497    }
498
499    /// Empty matrix (0.0 avg_nnz) is handled gracefully → Scalar (no Vector wasted).
500    #[test]
501    fn test_spmv_selection_empty_matrix() {
502        assert_eq!(resolve_adaptive(0.0), SpMVAlgo::Scalar);
503    }
504
505    /// VECTOR_THRESHOLD sanity: must equal 4.0 (the spec-defined boundary).
506    #[test]
507    fn test_vector_threshold_sanity() {
508        assert_eq!(
509            VECTOR_THRESHOLD, 4.0,
510            "VECTOR_THRESHOLD must be 4.0 per spec"
511        );
512        assert!(VECTOR_THRESHOLD.is_finite());
513    }
514
515    // -----------------------------------------------------------------------
516    // Deepening: explicit avg_nnz_per_row bracket tests matching sparse
517    // matrix categories from estimation.md and architecture notes.
518    // -----------------------------------------------------------------------
519
520    /// avg_nnz_per_row ≤ 2 (diagonal / identity matrices) → Scalar kernel.
521    ///
522    /// Models a diagonal matrix (1 nnz/row) — the most sparse real-world case.
523    #[test]
524    fn test_spmv_scalar_for_diagonal_matrix() {
525        // 1000-row diagonal → avg = 1.0
526        let avg = 1000.0_f64 / 1000.0;
527        assert!(avg <= 2.0, "avg={avg} should be ≤ 2");
528        assert_eq!(
529            resolve_adaptive(avg),
530            SpMVAlgo::Scalar,
531            "diagonal matrices (avg ≤ 2) should use Scalar SpMV"
532        );
533    }
534
535    /// avg_nnz_per_row ≤ 2, fractional (near-diagonal) → Scalar kernel.
536    ///
537    /// Models a tridiagonal-like matrix with ~2 nnz/row.
538    #[test]
539    fn test_spmv_scalar_for_tridiagonal_matrix() {
540        // 1000 rows, 2000 nnz → avg = 2.0 (tridiagonal boundary)
541        let avg = 2000.0_f64 / 1000.0;
542        assert!(avg <= 2.0, "avg={avg} should be ≤ 2");
543        assert_eq!(
544            resolve_adaptive(avg),
545            SpMVAlgo::Scalar,
546            "near-diagonal matrices (avg ≤ 2) should use Scalar SpMV"
547        );
548    }
549
550    /// avg_nnz_per_row in (2, 32] (moderate stencil / FEM) → Vector kernel.
551    ///
552    /// Models a 5-point 2D finite-difference stencil (avg ≈ 5 nnz/row).
553    #[test]
554    fn test_spmv_vector_for_5pt_stencil() {
555        // 1000×1000 grid → 5_000_000 rows with ~5 nnz each
556        let avg = 5.0_f64;
557        assert!(avg > 2.0 && avg <= 32.0, "avg={avg} should be in (2, 32]");
558        assert_eq!(
559            resolve_adaptive(avg),
560            SpMVAlgo::Vector,
561            "5-point stencil (avg ≈ 5) should use Vector SpMV"
562        );
563    }
564
565    /// avg_nnz_per_row ≈ 16 (7-point 3D stencil) → Vector kernel.
566    #[test]
567    fn test_spmv_vector_for_7pt_3d_stencil() {
568        let avg = 7.0_f64;
569        assert!(avg <= 32.0, "avg={avg} should be ≤ 32");
570        assert_eq!(
571            resolve_adaptive(avg),
572            SpMVAlgo::Vector,
573            "7-point 3D stencil (avg ≈ 7) should use Vector SpMV"
574        );
575    }
576
577    /// avg_nnz_per_row exactly at VECTOR_THRESHOLD boundary (4.0) → Vector.
578    ///
579    /// Tests that the boundary is inclusive: avg = VECTOR_THRESHOLD selects
580    /// Vector, not Scalar (i.e., `>=` rather than `>`).
581    #[test]
582    fn test_spmv_vector_at_exact_threshold() {
583        let avg = VECTOR_THRESHOLD; // 4.0
584        assert_eq!(
585            resolve_adaptive(avg),
586            SpMVAlgo::Vector,
587            "avg == VECTOR_THRESHOLD should select Vector (inclusive boundary)"
588        );
589        // One ULP below threshold → Scalar
590        let below = VECTOR_THRESHOLD - f64::MIN_POSITIVE;
591        // May still be 4.0 due to float precision, so only check if strictly below
592        if below < VECTOR_THRESHOLD {
593            assert_eq!(
594                resolve_adaptive(below),
595                SpMVAlgo::Scalar,
596                "avg strictly below VECTOR_THRESHOLD should select Scalar"
597            );
598        }
599    }
600
601    /// avg_nnz_per_row > 32 (dense row, graph networks) → Vector kernel.
602    ///
603    /// In the current two-class model, any avg ≥ VECTOR_THRESHOLD selects
604    /// Vector regardless of whether avg is 5 or 500. This confirms that the
605    /// "Adaptive" algorithm resolves correctly for highly dense rows.
606    #[test]
607    fn test_spmv_vector_for_high_density_rows() {
608        // avg = 64: above the ≤ 32 bracket, still selects Vector
609        let avg_64 = 64.0_f64;
610        assert_eq!(
611            resolve_adaptive(avg_64),
612            SpMVAlgo::Vector,
613            "high-density rows (avg = 64) should use Vector SpMV via Adaptive"
614        );
615
616        // avg = 256: very dense (near-dense matrix)
617        let avg_256 = 256.0_f64;
618        assert_eq!(
619            resolve_adaptive(avg_256),
620            SpMVAlgo::Vector,
621            "near-dense rows (avg = 256) should use Vector SpMV via Adaptive"
622        );
623    }
624
625    /// Adaptive algo resolves to the same result as calling resolve_adaptive
626    /// directly for various avg_nnz values. Confirms SpMVAlgo::Adaptive is
627    /// not accidentally treated as a concrete kernel variant.
628    #[test]
629    fn test_spmv_adaptive_algo_is_not_concrete() {
630        // SpMVAlgo::Adaptive is a selection hint, not a concrete kernel.
631        // resolve_adaptive must return Scalar or Vector, never Adaptive.
632        let test_avgs = [0.0, 0.5, 1.0, 2.0, 3.99, 4.0, 4.01, 32.0, 64.0, 128.0];
633        for avg in test_avgs {
634            let resolved = resolve_adaptive(avg);
635            assert!(
636                matches!(resolved, SpMVAlgo::Scalar | SpMVAlgo::Vector),
637                "resolve_adaptive({avg}) returned {resolved:?}, expected Scalar or Vector"
638            );
639        }
640    }
641
642    // -----------------------------------------------------------------------
643    // Quality gate: CSR-Vector warp shuffle reduction simulation (CPU)
644    // -----------------------------------------------------------------------
645
646    /// Simulate a single-warp (32 threads) tree reduction of partial dot-products.
647    ///
648    /// In the Vector SpMV kernel each warp computes partial sums for the row
649    /// elements it handles, then performs a binary tree (warp-shuffle) reduction
650    /// to sum all 32 partial sums into a single row result.
651    ///
652    /// This test verifies the correctness of that reduction algorithm on the CPU.
653    #[test]
654    fn spmv_warp_reduction_sim_32_threads() {
655        // 32 partial sums (one per thread in a warp)
656        let partial: Vec<f64> = (0..32_u32).map(|i| f64::from(i * i + 1)).collect();
657        let naive_sum: f64 = partial.iter().sum();
658
659        // Simulate binary tree reduction (warp shuffle pattern):
660        // stride 16, 8, 4, 2, 1
661        let mut sums = partial.clone();
662        let mut active = 32_usize;
663        while active > 1 {
664            let half = active / 2;
665            for lane in 0..half {
666                sums[lane] += sums[lane + half];
667            }
668            active = half;
669        }
670        let tree_sum = sums[0];
671
672        assert!(
673            (tree_sum - naive_sum).abs() < 1e-9,
674            "Warp tree reduction ({tree_sum}) must match naive sum ({naive_sum})"
675        );
676    }
677
678    /// Simulate a half-warp (16 threads) tree reduction.
679    ///
680    /// Verifies reduction correctness for the half-warp code path used when
681    /// the row is shorter than a full warp.
682    #[test]
683    fn spmv_half_warp_reduction_sim_16_threads() {
684        let partial: Vec<f64> = (0..16_u32).map(|i| f64::from(2 * i + 3)).collect();
685        let naive_sum: f64 = partial.iter().sum();
686
687        let mut sums = partial.clone();
688        let mut active = 16_usize;
689        while active > 1 {
690            let half = active / 2;
691            for lane in 0..half {
692                sums[lane] += sums[lane + half];
693            }
694            active = half;
695        }
696        let tree_sum = sums[0];
697
698        assert!(
699            (tree_sum - naive_sum).abs() < 1e-9,
700            "Half-warp tree reduction ({tree_sum}) must match naive sum ({naive_sum})"
701        );
702    }
703
704    // -----------------------------------------------------------------------
705    // Quality gate: SpMV numerical accuracy vs dense reference (CPU simulation)
706    // -----------------------------------------------------------------------
707
708    /// Dense-reference SpMV: computes y = A * x for a general dense matrix.
709    fn dense_spmv(a_rows: usize, a_cols: usize, a: &[f64], x: &[f64]) -> Vec<f64> {
710        let mut y = vec![0.0_f64; a_rows];
711        for i in 0..a_rows {
712            for j in 0..a_cols {
713                y[i] += a[i * a_cols + j] * x[j];
714            }
715        }
716        y
717    }
718
719    /// CSR SpMV simulation: computes y = A_csr * x on the CPU.
720    fn csr_spmv_sim(
721        nrows: usize,
722        row_ptr: &[usize],
723        col_idx: &[usize],
724        values: &[f64],
725        x: &[f64],
726    ) -> Vec<f64> {
727        let mut y = vec![0.0_f64; nrows];
728        for i in 0..nrows {
729            for idx in row_ptr[i]..row_ptr[i + 1] {
730                y[i] += values[idx] * x[col_idx[idx]];
731            }
732        }
733        y
734    }
735
736    /// SpMV for 4×4 identity matrix: y = I * x must equal x.
737    ///
738    /// This is the simplest correctness test: the identity provides a known
739    /// reference where every output equals the corresponding input.
740    #[test]
741    fn spmv_numerical_accuracy_identity_4x4() {
742        let n = 4_usize;
743        // Identity matrix in CSR format
744        let row_ptr = vec![0, 1, 2, 3, 4];
745        let col_idx = vec![0, 1, 2, 3];
746        let values = vec![1.0_f64; n];
747        let x = vec![1.0_f64, 2.0, 3.0, 4.0];
748
749        let y_csr = csr_spmv_sim(n, &row_ptr, &col_idx, &values, &x);
750        let y_dense = dense_spmv(
751            n,
752            n,
753            &[
754                1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 1.0,
755            ],
756            &x,
757        );
758        for i in 0..n {
759            assert!(
760                (y_csr[i] - y_dense[i]).abs() < 1e-13,
761                "SpMV I×x: y_csr[{i}]={} != y_dense[{i}]={}",
762                y_csr[i],
763                y_dense[i],
764            );
765        }
766    }
767
768    /// SpMV for a 0.1% sparse 1000×1000 matrix with a known diagonal pattern.
769    ///
770    /// Only diagonal entries are set (1000 out of 1_000_000 possible entries = 0.1%).
771    /// Result must equal x (diagonal matrix with ones = identity).
772    #[test]
773    fn spmv_very_sparse_0_1_percent_1000x1000() {
774        let n = 1000_usize;
775        // Diagonal matrix (0.1% density)
776        let row_ptr: Vec<usize> = (0..=n).collect();
777        let col_idx: Vec<usize> = (0..n).collect();
778        let values: Vec<f64> = vec![2.0; n]; // diagonal value = 2
779        let x: Vec<f64> = (0..n).map(|i| i as f64 * 0.001 + 1.0).collect();
780
781        let y = csr_spmv_sim(n, &row_ptr, &col_idx, &values, &x);
782
783        for i in 0..n {
784            let expected = 2.0 * x[i];
785            assert!(
786                (y[i] - expected).abs() < 1e-10,
787                "0.1% sparse SpMV row {i}: got {}, expected {expected}",
788                y[i],
789            );
790        }
791    }
792
793    /// SpMV for a 10% sparse 100×100 matrix: banded structure.
794    ///
795    /// Uses a banded matrix with bandwidth 5 (5 non-zeros per row on average),
796    /// giving approximately 10% density for a 100×100 system.
797    #[test]
798    fn spmv_moderate_10_percent_100x100() {
799        let n = 100_usize;
800        let bandwidth = 5_usize; // ±2 off-diagonal + diagonal
801
802        let mut row_ptr = vec![0_usize; n + 1];
803        let mut col_idx = Vec::new();
804        let mut values = Vec::new();
805
806        for i in 0..n {
807            let start = i.saturating_sub(2);
808            let end = (i + 3).min(n);
809            for j in start..end {
810                col_idx.push(j);
811                values.push(if i == j { 4.0_f64 } else { -1.0 });
812            }
813            row_ptr[i + 1] = col_idx.len();
814        }
815        let _ = bandwidth; // document variable used in comments
816
817        // x = [1, 1, 1, ..., 1]
818        let x = vec![1.0_f64; n];
819        let y_csr = csr_spmv_sim(n, &row_ptr, &col_idx, &values, &x);
820
821        // Build dense matrix and compute reference
822        let mut a_dense = vec![0.0_f64; n * n];
823        for i in 0..n {
824            let start = i.saturating_sub(2);
825            let end = (i + 3).min(n);
826            for j in start..end {
827                a_dense[i * n + j] = if i == j { 4.0 } else { -1.0 };
828            }
829        }
830        let y_dense = dense_spmv(n, n, &a_dense, &x);
831
832        for i in 0..n {
833            assert!(
834                (y_csr[i] - y_dense[i]).abs() < 1e-10,
835                "10% sparse SpMV row {i}: got {}, expected {}",
836                y_csr[i],
837                y_dense[i],
838            );
839        }
840    }
841
842    // -----------------------------------------------------------------------
843    // Quality gate: auto-format selection thresholds
844    // -----------------------------------------------------------------------
845
846    /// Verify format selection thresholds cover the three density regimes.
847    ///
848    /// - avg_nnz < VECTOR_THRESHOLD (4.0) → Scalar
849    /// - avg_nnz >= VECTOR_THRESHOLD       → Vector
850    ///
851    /// The test explicitly checks the three named brackets from the spec.
852    #[test]
853    fn spmv_format_selection_three_brackets() {
854        // Very sparse (≤ 2): diagonal-like
855        assert_eq!(
856            resolve_adaptive(1.0),
857            SpMVAlgo::Scalar,
858            "avg_nnz=1.0 (≤ 2 bracket) must select Scalar"
859        );
860        assert_eq!(
861            resolve_adaptive(2.0),
862            SpMVAlgo::Scalar,
863            "avg_nnz=2.0 (≤ 2 bracket) must select Scalar"
864        );
865        // Moderate (≤ 64): stencil-like (above VECTOR_THRESHOLD)
866        assert_eq!(
867            resolve_adaptive(5.0),
868            SpMVAlgo::Vector,
869            "avg_nnz=5.0 (≤ 64 bracket) must select Vector"
870        );
871        assert_eq!(
872            resolve_adaptive(32.0),
873            SpMVAlgo::Vector,
874            "avg_nnz=32.0 (≤ 64 bracket) must select Vector"
875        );
876        // Dense (> 64): near-dense graph
877        assert_eq!(
878            resolve_adaptive(65.0),
879            SpMVAlgo::Vector,
880            "avg_nnz=65.0 (> 64 bracket) must select Vector (binary model)"
881        );
882        assert_eq!(
883            resolve_adaptive(256.0),
884            SpMVAlgo::Vector,
885            "avg_nnz=256.0 (> 64 bracket) must select Vector"
886        );
887    }
888
889    /// CPU-proxy throughput benchmark: SpMV on a synthetic 10k×10k 5-point stencil matrix.
890    ///
891    /// Simulates the type of sparse matrix found in the SuiteSparse collection.
892    /// Measures CPU reference throughput and reports GFLOPS as a structural
893    /// lower-bound for the GPU target (cuSPARSE comparison requires real hardware).
894    #[test]
895    fn spmv_suitesparse_proxy_throughput_10k() {
896        // 2D 5-point Laplacian stencil on a 100×100 grid → 10k×10k sparse matrix.
897        // Each interior row has 5 non-zeros; boundary rows have 3–4.
898        let grid = 100_usize;
899        let n = grid * grid; // 10_000 rows
900
901        let mut row_ptr: Vec<usize> = Vec::with_capacity(n + 1);
902        let mut col_idx: Vec<usize> = Vec::new();
903        let mut values: Vec<f64> = Vec::new();
904
905        row_ptr.push(0);
906        for row in 0..n {
907            let r = row / grid;
908            let c = row % grid;
909            // North neighbour
910            if r > 0 {
911                col_idx.push(row - grid);
912                values.push(-1.0);
913            }
914            // West neighbour
915            if c > 0 {
916                col_idx.push(row - 1);
917                values.push(-1.0);
918            }
919            // Self (diagonal = 4)
920            col_idx.push(row);
921            values.push(4.0);
922            // East neighbour
923            if c + 1 < grid {
924                col_idx.push(row + 1);
925                values.push(-1.0);
926            }
927            // South neighbour
928            if r + 1 < grid {
929                col_idx.push(row + grid);
930                values.push(-1.0);
931            }
932            row_ptr.push(col_idx.len());
933        }
934
935        let nnz = col_idx.len();
936        let x: Vec<f64> = (0..n).map(|i| (i as f64) * 0.0001 + 1.0).collect();
937
938        // Warm-up pass
939        let _ = csr_spmv_sim(n, &row_ptr, &col_idx, &values, &x);
940
941        const ITERS: usize = 10;
942        let start = std::time::Instant::now();
943        let mut y = vec![0.0_f64; n];
944        for _ in 0..ITERS {
945            y = csr_spmv_sim(n, &row_ptr, &col_idx, &values, &x);
946        }
947        let elapsed_ns = start.elapsed().as_nanos() as f64;
948
949        // 2 flops per non-zero (one multiply + one add)
950        let total_flops = 2.0 * nnz as f64 * ITERS as f64;
951        let gflops = total_flops / elapsed_ns; // (flops) / (ns) = GFlops/s
952
953        println!(
954            "SpMV SuiteSparse proxy (10k×10k 5-pt stencil, {} nnz, {} iters): {:.3} GFLOPS (CPU reference)",
955            nnz, ITERS, gflops
956        );
957
958        // Sanity: result must be non-zero and throughput must be measurable
959        assert!(y[n / 2] != 0.0, "SpMV result must be non-zero");
960        assert!(
961            gflops > 0.001,
962            "SpMV CPU reference throughput unrealistically low: {:.6} GFLOPS",
963            gflops
964        );
965    }
966}