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    ptx_suffix, 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 elem_bytes = T::size_u32();
186
187    KernelBuilder::new("spmv_scalar")
188        .target(sm)
189        .param("row_ptr", PtxType::U64)
190        .param("col_idx", PtxType::U64)
191        .param("values", PtxType::U64)
192        .param("x_ptr", PtxType::U64)
193        .param("y_ptr", PtxType::U64)
194        .param("alpha_bits", PtxType::U64)
195        .param("beta_bits", PtxType::U64)
196        .param("num_rows", PtxType::U32)
197        .body(move |b| {
198            let gid = b.global_thread_id_x();
199            let num_rows = b.load_param_u32("num_rows");
200
201            let gid_inner = gid.clone();
202            b.if_lt_u32(gid, num_rows, move |b| {
203                let row = gid_inner;
204                let row_ptr_base = b.load_param_u64("row_ptr");
205                let col_idx_base = b.load_param_u64("col_idx");
206                let values_base = b.load_param_u64("values");
207                let x_ptr = b.load_param_u64("x_ptr");
208                let y_ptr = b.load_param_u64("y_ptr");
209                let alpha_bits = b.load_param_u64("alpha_bits");
210                let beta_bits = b.load_param_u64("beta_bits");
211
212                let alpha = reinterpret_bits_to_float::<T>(b, alpha_bits);
213                let beta = reinterpret_bits_to_float::<T>(b, beta_bits);
214
215                // Load row_ptr[row] and row_ptr[row+1] (i32 = 4 bytes)
216                let rp_addr = b.byte_offset_addr(row_ptr_base.clone(), row.clone(), 4);
217                let row_start = b.load_global_i32(rp_addr);
218
219                let row_plus_1 = b.alloc_reg(PtxType::U32);
220                b.raw_ptx(&format!("add.u32 {row_plus_1}, {row}, 1;"));
221                let rp_addr_next = b.byte_offset_addr(row_ptr_base, row_plus_1, 4);
222                let row_end = b.load_global_i32(rp_addr_next);
223
224                // Initialize accumulator
225                let acc = load_float_imm::<T>(b, 0.0);
226
227                // Loop over non-zeros in this row
228                let loop_label = b.fresh_label("spmv_loop");
229                let done_label = b.fresh_label("spmv_done");
230
231                let k = b.alloc_reg(PtxType::U32);
232                // Convert row_start (i32) to u32
233                let rs_u32 = b.alloc_reg(PtxType::U32);
234                b.raw_ptx(&format!("mov.b32 {rs_u32}, {row_start};"));
235                b.raw_ptx(&format!("mov.u32 {k}, {rs_u32};"));
236
237                let re_u32 = b.alloc_reg(PtxType::U32);
238                b.raw_ptx(&format!("mov.b32 {re_u32}, {row_end};"));
239
240                b.label(&loop_label);
241                // Exit the loop when k >= row_end. Use the structured `branch_if`
242                // so the branch target is emitted with the same `$`-prefix
243                // convention as `b.label`/`b.branch`; a raw `bra L__...` would not
244                // match the `$L__...:` label definition and `ptxas` would reject it
245                // ("Unknown symbol").
246                let pred = b.alloc_reg(PtxType::Pred);
247                b.raw_ptx(&format!("setp.hs.u32 {pred}, {k}, {re_u32};"));
248                b.branch_if(pred, &done_label);
249
250                // Load col_idx[k] (i32 = 4 bytes)
251                let ci_addr = b.byte_offset_addr(col_idx_base.clone(), k.clone(), 4);
252                let col = b.load_global_i32(ci_addr);
253                let col_u32 = b.alloc_reg(PtxType::U32);
254                b.raw_ptx(&format!("mov.b32 {col_u32}, {col};"));
255
256                // Load values[k]
257                let v_addr = b.byte_offset_addr(values_base.clone(), k.clone(), elem_bytes);
258                let val = load_global_float::<T>(b, v_addr);
259
260                // Load x[col]
261                let x_addr = b.byte_offset_addr(x_ptr.clone(), col_u32, elem_bytes);
262                let x_val = load_global_float::<T>(b, x_addr);
263
264                // acc += val * x_val
265                let new_acc = fma_float::<T>(b, val, x_val, acc.clone());
266                let mov_suffix = ptx_suffix::<T>();
267                b.raw_ptx(&format!("mov.{mov_suffix} {acc}, {new_acc};"));
268
269                // k++
270                b.raw_ptx(&format!("add.u32 {k}, {k}, 1;"));
271                b.branch(&loop_label);
272                b.label(&done_label);
273
274                // Compute y = alpha * acc + beta * y_old
275                let y_addr = b.byte_offset_addr(y_ptr, row, elem_bytes);
276                let y_old = load_global_float::<T>(b, y_addr.clone());
277
278                let alpha_acc = mul_float::<T>(b, alpha, acc);
279                let beta_y = mul_float::<T>(b, beta, y_old);
280                let result = add_float::<T>(b, alpha_acc, beta_y);
281
282                store_global_float::<T>(b, y_addr, result);
283            });
284
285            b.ret();
286        })
287        .build()
288        .map_err(|e| SparseError::PtxGeneration(e.to_string()))
289}
290
291/// Generates PTX for vector SpMV (one warp per row).
292fn emit_spmv_vector<T: GpuFloat>(sm: SmVersion) -> SparseResult<String> {
293    let elem_bytes = T::size_u32();
294
295    KernelBuilder::new("spmv_vector")
296        .target(sm)
297        .param("row_ptr", PtxType::U64)
298        .param("col_idx", PtxType::U64)
299        .param("values", PtxType::U64)
300        .param("x_ptr", PtxType::U64)
301        .param("y_ptr", PtxType::U64)
302        .param("alpha_bits", PtxType::U64)
303        .param("beta_bits", PtxType::U64)
304        .param("num_rows", PtxType::U32)
305        .body(move |b| {
306            // Each warp handles one row. Warp ID = global_thread_id / 32
307            let tid_global = b.global_thread_id_x();
308            let num_rows = b.load_param_u32("num_rows");
309
310            // Lane within warp (0..31)
311            let lane = b.alloc_reg(PtxType::U32);
312            b.raw_ptx(&format!("and.b32 {lane}, {tid_global}, 31;"));
313
314            // Warp ID = tid_global >> 5
315            let warp_id = b.alloc_reg(PtxType::U32);
316            b.raw_ptx(&format!("shr.u32 {warp_id}, {tid_global}, 5;"));
317
318            let warp_id_inner = warp_id.clone();
319            let lane_inner = lane.clone();
320            b.if_lt_u32(warp_id, num_rows, move |b| {
321                let row = warp_id_inner;
322                let lane = lane_inner;
323
324                let row_ptr_base = b.load_param_u64("row_ptr");
325                let col_idx_base = b.load_param_u64("col_idx");
326                let values_base = b.load_param_u64("values");
327                let x_ptr = b.load_param_u64("x_ptr");
328                let y_ptr = b.load_param_u64("y_ptr");
329                let alpha_bits = b.load_param_u64("alpha_bits");
330                let beta_bits = b.load_param_u64("beta_bits");
331
332                let alpha = reinterpret_bits_to_float::<T>(b, alpha_bits);
333                let beta = reinterpret_bits_to_float::<T>(b, beta_bits);
334
335                // Load row bounds
336                let rp_addr = b.byte_offset_addr(row_ptr_base.clone(), row.clone(), 4);
337                let row_start_i32 = b.load_global_i32(rp_addr);
338                let row_start = b.alloc_reg(PtxType::U32);
339                b.raw_ptx(&format!("mov.b32 {row_start}, {row_start_i32};"));
340
341                let row_plus_1 = b.alloc_reg(PtxType::U32);
342                b.raw_ptx(&format!("add.u32 {row_plus_1}, {row}, 1;"));
343                let rp_addr_next = b.byte_offset_addr(row_ptr_base, row_plus_1, 4);
344                let row_end_i32 = b.load_global_i32(rp_addr_next);
345                let row_end = b.alloc_reg(PtxType::U32);
346                b.raw_ptx(&format!("mov.b32 {row_end}, {row_end_i32};"));
347
348                // Each lane starts at row_start + lane, stride 32
349                let acc = load_float_imm::<T>(b, 0.0);
350
351                let k = b.alloc_reg(PtxType::U32);
352                b.raw_ptx(&format!("add.u32 {k}, {row_start}, {lane};"));
353
354                let loop_label = b.fresh_label("spmv_vloop");
355                let done_label = b.fresh_label("spmv_vdone");
356
357                b.label(&loop_label);
358                // Exit the loop when k >= row_end via the structured `branch_if`
359                // (see the scalar kernel for why a raw `bra` would be rejected).
360                let pred = b.alloc_reg(PtxType::Pred);
361                b.raw_ptx(&format!("setp.hs.u32 {pred}, {k}, {row_end};"));
362                b.branch_if(pred, &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 = ptx_suffix::<T>();
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                // Only lane 0 writes the row result; every other lane skips ahead.
389                // Branch when lane != 0 using the structured `branch_if` so the
390                // target matches the `$`-prefixed label definition.
391                let not_lane_0 = b.alloc_reg(PtxType::Pred);
392                b.raw_ptx(&format!("setp.ne.u32 {not_lane_0}, {lane}, 0;"));
393                let skip_label = b.fresh_label("spmv_skip");
394                b.branch_if(not_lane_0, &skip_label);
395
396                let y_addr = b.byte_offset_addr(y_ptr, row, elem_bytes);
397                let y_old = load_global_float::<T>(b, y_addr.clone());
398
399                let alpha_acc = mul_float::<T>(b, alpha, reduced);
400                let beta_y = mul_float::<T>(b, beta, y_old);
401                let result = add_float::<T>(b, alpha_acc, beta_y);
402                store_global_float::<T>(b, y_addr, result);
403
404                b.label(&skip_label);
405            });
406
407            b.ret();
408        })
409        .build()
410        .map_err(|e| SparseError::PtxGeneration(e.to_string()))
411}
412
413#[cfg(test)]
414mod tests {
415    use super::*;
416
417    /// Runs `ptxas -arch=sm_86` on the given PTX, writing it through
418    /// [`std::env::temp_dir`]. Returns `Some(())` on success, `Some` error text
419    /// on assembler rejection, or `None` if `ptxas` is not on `PATH` (so callers
420    /// can skip gracefully on machines without the CUDA toolkit).
421    fn try_ptxas(name: &str, ptx: &str) -> Option<Result<(), String>> {
422        use std::io::Write;
423        let path = std::env::temp_dir().join(format!("oxicuda_spmv_{name}.ptx"));
424        {
425            let mut f = std::fs::File::create(&path).expect("test: create temp PTX file");
426            f.write_all(ptx.as_bytes())
427                .expect("test: write temp PTX file");
428        }
429        match std::process::Command::new("ptxas")
430            .arg("-arch=sm_86")
431            .arg(&path)
432            .arg("-o")
433            .arg("/dev/null")
434            .output()
435        {
436            Ok(out) if out.status.success() => Some(Ok(())),
437            Ok(out) => Some(Err(String::from_utf8_lossy(&out.stderr).into_owned())),
438            // ptxas missing (no CUDA toolkit): skip gracefully.
439            Err(_) => None,
440        }
441    }
442
443    /// The f64 scalar and vector SpMV kernels must be well-formed double-precision
444    /// PTX: 64-bit value registers, `0D`-prefixed f64 immediates (never an f32
445    /// `0F00000000` zero), no illegal `shfl.sync.down.b64`, and — when `ptxas` is
446    /// available — they must assemble for `sm_86` (RTX A4000). Regression guard for
447    /// the "CUDA: invalid PTX" failure of `cuda_spmv_csr`.
448    #[test]
449    fn spmv_f64_ptx_well_formed_and_assembles() {
450        let scalar = emit_spmv_scalar::<f64>(SmVersion::Sm86).expect("f64 scalar PTX");
451        let vector = emit_spmv_vector::<f64>(SmVersion::Sm86).expect("f64 vector PTX");
452
453        for (name, ptx) in [("scalar", &scalar), ("vector", &vector)] {
454            // f64 value registers are declared 64-bit (`.b64` is the reg-class of f64).
455            assert!(
456                ptx.contains(".reg .b64 %f"),
457                "f64 {name} kernel must declare 64-bit %f value registers:\n{ptx}"
458            );
459            // No f32 zero immediate and no f32-typed float math may leak into the
460            // f64 value path (the GEMM-class antipattern this guards against).
461            assert!(
462                !ptx.contains("0F00000000"),
463                "f64 {name} kernel must not materialize an f32 0.0 immediate:\n{ptx}"
464            );
465            assert!(
466                !ptx.contains(".f32"),
467                "f64 {name} kernel must not contain any .f32-typed instruction:\n{ptx}"
468            );
469            // The f64 zero immediate uses the 16-hex-digit 0D form.
470            assert!(
471                ptx.contains("0D0000000000000000"),
472                "f64 {name} kernel must materialize the f64 0.0 immediate (0D…):\n{ptx}"
473            );
474            // shfl only supports b32; a b64 shuffle is rejected by ptxas.
475            assert!(
476                !ptx.contains("shfl.sync.down.b64"),
477                "f64 {name} kernel must not emit shfl.sync.down.b64:\n{ptx}"
478            );
479            // Genuine double-precision arithmetic must be present.
480            assert!(
481                ptx.contains("fma.rn.f64") && ptx.contains("ld.global.f64"),
482                "f64 {name} kernel must use f64 fma/load instructions:\n{ptx}"
483            );
484        }
485
486        // The vector kernel performs the warp reduction via paired b32 shuffles.
487        assert!(
488            vector.contains("shfl.sync.down.b32"),
489            "f64 vector kernel must reduce via b32 shuffles:\n{vector}"
490        );
491
492        // Assemble with ptxas when present; both kernels must be accepted by sm_86.
493        for (name, ptx) in [("scalar", &scalar), ("vector", &vector)] {
494            match try_ptxas(&format!("f64_{name}"), ptx) {
495                Some(Ok(())) => {}
496                Some(Err(stderr)) => {
497                    panic!("ptxas rejected the f64 {name} SpMV kernel:\n{stderr}\nPTX:\n{ptx}")
498                }
499                None => {
500                    // ptxas unavailable: textual well-formedness checks above stand in.
501                }
502            }
503        }
504    }
505
506    /// The f32 SpMV kernels must keep assembling too (the label fix is shared).
507    #[test]
508    fn spmv_f32_ptx_assembles() {
509        let scalar = emit_spmv_scalar::<f32>(SmVersion::Sm86).expect("f32 scalar PTX");
510        let vector = emit_spmv_vector::<f32>(SmVersion::Sm86).expect("f32 vector PTX");
511        for (name, ptx) in [("scalar", &scalar), ("vector", &vector)] {
512            assert!(
513                ptx.contains(".reg .b32 %f"),
514                "f32 {name} kernel must declare 32-bit %f value registers:\n{ptx}"
515            );
516            if let Some(Err(stderr)) = try_ptxas(&format!("f32_{name}"), ptx) {
517                panic!("ptxas rejected the f32 {name} SpMV kernel:\n{stderr}\nPTX:\n{ptx}");
518            }
519        }
520    }
521
522    #[test]
523    fn spmv_algo_auto_select() {
524        // avg_nnz < threshold => Scalar
525        // Verify VECTOR_THRESHOLD is set to a reasonable value for algorithm selection.
526        let threshold = VECTOR_THRESHOLD;
527        assert!(threshold > 3.0);
528    }
529
530    #[test]
531    fn spmv_scalar_ptx_generates() {
532        let ptx = emit_spmv_scalar::<f32>(SmVersion::Sm80);
533        assert!(ptx.is_ok());
534        let ptx = ptx.expect("test: PTX gen should succeed");
535        assert!(ptx.contains(".entry spmv_scalar"));
536        assert!(ptx.contains(".target sm_80"));
537    }
538
539    #[test]
540    fn spmv_vector_ptx_generates() {
541        let ptx = emit_spmv_vector::<f32>(SmVersion::Sm80);
542        assert!(ptx.is_ok());
543        let ptx = ptx.expect("test: PTX gen should succeed");
544        assert!(ptx.contains(".entry spmv_vector"));
545    }
546
547    #[test]
548    fn spmv_scalar_ptx_f64() {
549        let ptx = emit_spmv_scalar::<f64>(SmVersion::Sm80);
550        assert!(ptx.is_ok());
551    }
552
553    #[test]
554    fn spmv_vector_ptx_f64() {
555        let ptx = emit_spmv_vector::<f64>(SmVersion::Sm80);
556        assert!(ptx.is_ok());
557    }
558
559    // -----------------------------------------------------------------------
560    // Task 5a: Auto-selection heuristic tests (CPU-only, no GPU required)
561    // -----------------------------------------------------------------------
562
563    /// Very sparse rows (avg_nnz ≈ 1.5, well below threshold 4.0) → Scalar.
564    #[test]
565    fn test_spmv_selects_scalar_for_very_sparse() {
566        // 100 rows, 150 nnz → avg = 1.5
567        let avg = 150.0_f64 / 100.0;
568        assert!(avg < VECTOR_THRESHOLD);
569        assert_eq!(resolve_adaptive(avg), SpMVAlgo::Scalar);
570    }
571
572    /// Moderate density (avg_nnz = 32, above threshold 4.0) → Vector.
573    #[test]
574    fn test_spmv_selects_vector_for_moderate_density() {
575        let avg = 32.0_f64;
576        assert!(avg >= VECTOR_THRESHOLD);
577        assert_eq!(resolve_adaptive(avg), SpMVAlgo::Vector);
578    }
579
580    /// Dense rows (avg_nnz = 128, well above threshold) → Vector.
581    #[test]
582    fn test_spmv_selects_vector_for_dense() {
583        let avg = 128.0_f64;
584        assert!(avg >= VECTOR_THRESHOLD);
585        assert_eq!(resolve_adaptive(avg), SpMVAlgo::Vector);
586    }
587
588    /// Boundary: just below threshold → Scalar; at threshold → Vector.
589    #[test]
590    fn test_spmv_selection_boundary_conditions() {
591        // Just below threshold (3.9999…)
592        let just_below = VECTOR_THRESHOLD - f64::EPSILON * VECTOR_THRESHOLD;
593        assert_eq!(resolve_adaptive(just_below), SpMVAlgo::Scalar);
594
595        // Exactly at threshold
596        assert_eq!(resolve_adaptive(VECTOR_THRESHOLD), SpMVAlgo::Vector);
597
598        // Slightly above threshold
599        let just_above = VECTOR_THRESHOLD + f64::EPSILON * VECTOR_THRESHOLD;
600        assert_eq!(resolve_adaptive(just_above), SpMVAlgo::Vector);
601    }
602
603    /// Empty matrix (0.0 avg_nnz) is handled gracefully → Scalar (no Vector wasted).
604    #[test]
605    fn test_spmv_selection_empty_matrix() {
606        assert_eq!(resolve_adaptive(0.0), SpMVAlgo::Scalar);
607    }
608
609    /// VECTOR_THRESHOLD sanity: must equal 4.0 (the spec-defined boundary).
610    #[test]
611    fn test_vector_threshold_sanity() {
612        assert_eq!(
613            VECTOR_THRESHOLD, 4.0,
614            "VECTOR_THRESHOLD must be 4.0 per spec"
615        );
616        assert!(VECTOR_THRESHOLD.is_finite());
617    }
618
619    // -----------------------------------------------------------------------
620    // Deepening: explicit avg_nnz_per_row bracket tests matching sparse
621    // matrix categories from estimation.md and architecture notes.
622    // -----------------------------------------------------------------------
623
624    /// avg_nnz_per_row ≤ 2 (diagonal / identity matrices) → Scalar kernel.
625    ///
626    /// Models a diagonal matrix (1 nnz/row) — the most sparse real-world case.
627    #[test]
628    fn test_spmv_scalar_for_diagonal_matrix() {
629        // 1000-row diagonal → avg = 1.0
630        let avg = 1000.0_f64 / 1000.0;
631        assert!(avg <= 2.0, "avg={avg} should be ≤ 2");
632        assert_eq!(
633            resolve_adaptive(avg),
634            SpMVAlgo::Scalar,
635            "diagonal matrices (avg ≤ 2) should use Scalar SpMV"
636        );
637    }
638
639    /// avg_nnz_per_row ≤ 2, fractional (near-diagonal) → Scalar kernel.
640    ///
641    /// Models a tridiagonal-like matrix with ~2 nnz/row.
642    #[test]
643    fn test_spmv_scalar_for_tridiagonal_matrix() {
644        // 1000 rows, 2000 nnz → avg = 2.0 (tridiagonal boundary)
645        let avg = 2000.0_f64 / 1000.0;
646        assert!(avg <= 2.0, "avg={avg} should be ≤ 2");
647        assert_eq!(
648            resolve_adaptive(avg),
649            SpMVAlgo::Scalar,
650            "near-diagonal matrices (avg ≤ 2) should use Scalar SpMV"
651        );
652    }
653
654    /// avg_nnz_per_row in (2, 32] (moderate stencil / FEM) → Vector kernel.
655    ///
656    /// Models a 5-point 2D finite-difference stencil (avg ≈ 5 nnz/row).
657    #[test]
658    fn test_spmv_vector_for_5pt_stencil() {
659        // 1000×1000 grid → 5_000_000 rows with ~5 nnz each
660        let avg = 5.0_f64;
661        assert!(avg > 2.0 && avg <= 32.0, "avg={avg} should be in (2, 32]");
662        assert_eq!(
663            resolve_adaptive(avg),
664            SpMVAlgo::Vector,
665            "5-point stencil (avg ≈ 5) should use Vector SpMV"
666        );
667    }
668
669    /// avg_nnz_per_row ≈ 16 (7-point 3D stencil) → Vector kernel.
670    #[test]
671    fn test_spmv_vector_for_7pt_3d_stencil() {
672        let avg = 7.0_f64;
673        assert!(avg <= 32.0, "avg={avg} should be ≤ 32");
674        assert_eq!(
675            resolve_adaptive(avg),
676            SpMVAlgo::Vector,
677            "7-point 3D stencil (avg ≈ 7) should use Vector SpMV"
678        );
679    }
680
681    /// avg_nnz_per_row exactly at VECTOR_THRESHOLD boundary (4.0) → Vector.
682    ///
683    /// Tests that the boundary is inclusive: avg = VECTOR_THRESHOLD selects
684    /// Vector, not Scalar (i.e., `>=` rather than `>`).
685    #[test]
686    fn test_spmv_vector_at_exact_threshold() {
687        let avg = VECTOR_THRESHOLD; // 4.0
688        assert_eq!(
689            resolve_adaptive(avg),
690            SpMVAlgo::Vector,
691            "avg == VECTOR_THRESHOLD should select Vector (inclusive boundary)"
692        );
693        // One ULP below threshold → Scalar
694        let below = VECTOR_THRESHOLD - f64::MIN_POSITIVE;
695        // May still be 4.0 due to float precision, so only check if strictly below
696        if below < VECTOR_THRESHOLD {
697            assert_eq!(
698                resolve_adaptive(below),
699                SpMVAlgo::Scalar,
700                "avg strictly below VECTOR_THRESHOLD should select Scalar"
701            );
702        }
703    }
704
705    /// avg_nnz_per_row > 32 (dense row, graph networks) → Vector kernel.
706    ///
707    /// In the current two-class model, any avg ≥ VECTOR_THRESHOLD selects
708    /// Vector regardless of whether avg is 5 or 500. This confirms that the
709    /// "Adaptive" algorithm resolves correctly for highly dense rows.
710    #[test]
711    fn test_spmv_vector_for_high_density_rows() {
712        // avg = 64: above the ≤ 32 bracket, still selects Vector
713        let avg_64 = 64.0_f64;
714        assert_eq!(
715            resolve_adaptive(avg_64),
716            SpMVAlgo::Vector,
717            "high-density rows (avg = 64) should use Vector SpMV via Adaptive"
718        );
719
720        // avg = 256: very dense (near-dense matrix)
721        let avg_256 = 256.0_f64;
722        assert_eq!(
723            resolve_adaptive(avg_256),
724            SpMVAlgo::Vector,
725            "near-dense rows (avg = 256) should use Vector SpMV via Adaptive"
726        );
727    }
728
729    /// Adaptive algo resolves to the same result as calling resolve_adaptive
730    /// directly for various avg_nnz values. Confirms SpMVAlgo::Adaptive is
731    /// not accidentally treated as a concrete kernel variant.
732    #[test]
733    fn test_spmv_adaptive_algo_is_not_concrete() {
734        // SpMVAlgo::Adaptive is a selection hint, not a concrete kernel.
735        // resolve_adaptive must return Scalar or Vector, never Adaptive.
736        let test_avgs = [0.0, 0.5, 1.0, 2.0, 3.99, 4.0, 4.01, 32.0, 64.0, 128.0];
737        for avg in test_avgs {
738            let resolved = resolve_adaptive(avg);
739            assert!(
740                matches!(resolved, SpMVAlgo::Scalar | SpMVAlgo::Vector),
741                "resolve_adaptive({avg}) returned {resolved:?}, expected Scalar or Vector"
742            );
743        }
744    }
745
746    // -----------------------------------------------------------------------
747    // Quality gate: CSR-Vector warp shuffle reduction simulation (CPU)
748    // -----------------------------------------------------------------------
749
750    /// Simulate a single-warp (32 threads) tree reduction of partial dot-products.
751    ///
752    /// In the Vector SpMV kernel each warp computes partial sums for the row
753    /// elements it handles, then performs a binary tree (warp-shuffle) reduction
754    /// to sum all 32 partial sums into a single row result.
755    ///
756    /// This test verifies the correctness of that reduction algorithm on the CPU.
757    #[test]
758    fn spmv_warp_reduction_sim_32_threads() {
759        // 32 partial sums (one per thread in a warp)
760        let partial: Vec<f64> = (0..32_u32).map(|i| f64::from(i * i + 1)).collect();
761        let naive_sum: f64 = partial.iter().sum();
762
763        // Simulate binary tree reduction (warp shuffle pattern):
764        // stride 16, 8, 4, 2, 1
765        let mut sums = partial.clone();
766        let mut active = 32_usize;
767        while active > 1 {
768            let half = active / 2;
769            for lane in 0..half {
770                sums[lane] += sums[lane + half];
771            }
772            active = half;
773        }
774        let tree_sum = sums[0];
775
776        assert!(
777            (tree_sum - naive_sum).abs() < 1e-9,
778            "Warp tree reduction ({tree_sum}) must match naive sum ({naive_sum})"
779        );
780    }
781
782    /// Simulate a half-warp (16 threads) tree reduction.
783    ///
784    /// Verifies reduction correctness for the half-warp code path used when
785    /// the row is shorter than a full warp.
786    #[test]
787    fn spmv_half_warp_reduction_sim_16_threads() {
788        let partial: Vec<f64> = (0..16_u32).map(|i| f64::from(2 * i + 3)).collect();
789        let naive_sum: f64 = partial.iter().sum();
790
791        let mut sums = partial.clone();
792        let mut active = 16_usize;
793        while active > 1 {
794            let half = active / 2;
795            for lane in 0..half {
796                sums[lane] += sums[lane + half];
797            }
798            active = half;
799        }
800        let tree_sum = sums[0];
801
802        assert!(
803            (tree_sum - naive_sum).abs() < 1e-9,
804            "Half-warp tree reduction ({tree_sum}) must match naive sum ({naive_sum})"
805        );
806    }
807
808    // -----------------------------------------------------------------------
809    // Quality gate: SpMV numerical accuracy vs dense reference (CPU simulation)
810    // -----------------------------------------------------------------------
811
812    /// Dense-reference SpMV: computes y = A * x for a general dense matrix.
813    fn dense_spmv(a_rows: usize, a_cols: usize, a: &[f64], x: &[f64]) -> Vec<f64> {
814        let mut y = vec![0.0_f64; a_rows];
815        for i in 0..a_rows {
816            for j in 0..a_cols {
817                y[i] += a[i * a_cols + j] * x[j];
818            }
819        }
820        y
821    }
822
823    /// CSR SpMV simulation: computes y = A_csr * x on the CPU.
824    fn csr_spmv_sim(
825        nrows: usize,
826        row_ptr: &[usize],
827        col_idx: &[usize],
828        values: &[f64],
829        x: &[f64],
830    ) -> Vec<f64> {
831        let mut y = vec![0.0_f64; nrows];
832        for i in 0..nrows {
833            for idx in row_ptr[i]..row_ptr[i + 1] {
834                y[i] += values[idx] * x[col_idx[idx]];
835            }
836        }
837        y
838    }
839
840    /// SpMV for 4×4 identity matrix: y = I * x must equal x.
841    ///
842    /// This is the simplest correctness test: the identity provides a known
843    /// reference where every output equals the corresponding input.
844    #[test]
845    fn spmv_numerical_accuracy_identity_4x4() {
846        let n = 4_usize;
847        // Identity matrix in CSR format
848        let row_ptr = vec![0, 1, 2, 3, 4];
849        let col_idx = vec![0, 1, 2, 3];
850        let values = vec![1.0_f64; n];
851        let x = vec![1.0_f64, 2.0, 3.0, 4.0];
852
853        let y_csr = csr_spmv_sim(n, &row_ptr, &col_idx, &values, &x);
854        let y_dense = dense_spmv(
855            n,
856            n,
857            &[
858                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,
859            ],
860            &x,
861        );
862        for i in 0..n {
863            assert!(
864                (y_csr[i] - y_dense[i]).abs() < 1e-13,
865                "SpMV I×x: y_csr[{i}]={} != y_dense[{i}]={}",
866                y_csr[i],
867                y_dense[i],
868            );
869        }
870    }
871
872    /// SpMV for a 0.1% sparse 1000×1000 matrix with a known diagonal pattern.
873    ///
874    /// Only diagonal entries are set (1000 out of 1_000_000 possible entries = 0.1%).
875    /// Result must equal x (diagonal matrix with ones = identity).
876    #[test]
877    fn spmv_very_sparse_0_1_percent_1000x1000() {
878        let n = 1000_usize;
879        // Diagonal matrix (0.1% density)
880        let row_ptr: Vec<usize> = (0..=n).collect();
881        let col_idx: Vec<usize> = (0..n).collect();
882        let values: Vec<f64> = vec![2.0; n]; // diagonal value = 2
883        let x: Vec<f64> = (0..n).map(|i| i as f64 * 0.001 + 1.0).collect();
884
885        let y = csr_spmv_sim(n, &row_ptr, &col_idx, &values, &x);
886
887        for i in 0..n {
888            let expected = 2.0 * x[i];
889            assert!(
890                (y[i] - expected).abs() < 1e-10,
891                "0.1% sparse SpMV row {i}: got {}, expected {expected}",
892                y[i],
893            );
894        }
895    }
896
897    /// SpMV for a 10% sparse 100×100 matrix: banded structure.
898    ///
899    /// Uses a banded matrix with bandwidth 5 (5 non-zeros per row on average),
900    /// giving approximately 10% density for a 100×100 system.
901    #[test]
902    fn spmv_moderate_10_percent_100x100() {
903        let n = 100_usize;
904        let bandwidth = 5_usize; // ±2 off-diagonal + diagonal
905
906        let mut row_ptr = vec![0_usize; n + 1];
907        let mut col_idx = Vec::new();
908        let mut values = Vec::new();
909
910        for i in 0..n {
911            let start = i.saturating_sub(2);
912            let end = (i + 3).min(n);
913            for j in start..end {
914                col_idx.push(j);
915                values.push(if i == j { 4.0_f64 } else { -1.0 });
916            }
917            row_ptr[i + 1] = col_idx.len();
918        }
919        let _ = bandwidth; // document variable used in comments
920
921        // x = [1, 1, 1, ..., 1]
922        let x = vec![1.0_f64; n];
923        let y_csr = csr_spmv_sim(n, &row_ptr, &col_idx, &values, &x);
924
925        // Build dense matrix and compute reference
926        let mut a_dense = vec![0.0_f64; n * n];
927        for i in 0..n {
928            let start = i.saturating_sub(2);
929            let end = (i + 3).min(n);
930            for j in start..end {
931                a_dense[i * n + j] = if i == j { 4.0 } else { -1.0 };
932            }
933        }
934        let y_dense = dense_spmv(n, n, &a_dense, &x);
935
936        for i in 0..n {
937            assert!(
938                (y_csr[i] - y_dense[i]).abs() < 1e-10,
939                "10% sparse SpMV row {i}: got {}, expected {}",
940                y_csr[i],
941                y_dense[i],
942            );
943        }
944    }
945
946    // -----------------------------------------------------------------------
947    // Quality gate: auto-format selection thresholds
948    // -----------------------------------------------------------------------
949
950    /// Verify format selection thresholds cover the three density regimes.
951    ///
952    /// - avg_nnz < VECTOR_THRESHOLD (4.0) → Scalar
953    /// - avg_nnz >= VECTOR_THRESHOLD       → Vector
954    ///
955    /// The test explicitly checks the three named brackets from the spec.
956    #[test]
957    fn spmv_format_selection_three_brackets() {
958        // Very sparse (≤ 2): diagonal-like
959        assert_eq!(
960            resolve_adaptive(1.0),
961            SpMVAlgo::Scalar,
962            "avg_nnz=1.0 (≤ 2 bracket) must select Scalar"
963        );
964        assert_eq!(
965            resolve_adaptive(2.0),
966            SpMVAlgo::Scalar,
967            "avg_nnz=2.0 (≤ 2 bracket) must select Scalar"
968        );
969        // Moderate (≤ 64): stencil-like (above VECTOR_THRESHOLD)
970        assert_eq!(
971            resolve_adaptive(5.0),
972            SpMVAlgo::Vector,
973            "avg_nnz=5.0 (≤ 64 bracket) must select Vector"
974        );
975        assert_eq!(
976            resolve_adaptive(32.0),
977            SpMVAlgo::Vector,
978            "avg_nnz=32.0 (≤ 64 bracket) must select Vector"
979        );
980        // Dense (> 64): near-dense graph
981        assert_eq!(
982            resolve_adaptive(65.0),
983            SpMVAlgo::Vector,
984            "avg_nnz=65.0 (> 64 bracket) must select Vector (binary model)"
985        );
986        assert_eq!(
987            resolve_adaptive(256.0),
988            SpMVAlgo::Vector,
989            "avg_nnz=256.0 (> 64 bracket) must select Vector"
990        );
991    }
992
993    /// CPU-proxy throughput benchmark: SpMV on a synthetic 10k×10k 5-point stencil matrix.
994    ///
995    /// Simulates the type of sparse matrix found in the SuiteSparse collection.
996    /// Measures CPU reference throughput and reports GFLOPS as a structural
997    /// lower-bound for the GPU target (cuSPARSE comparison requires real hardware).
998    #[test]
999    fn spmv_suitesparse_proxy_throughput_10k() {
1000        // 2D 5-point Laplacian stencil on a 100×100 grid → 10k×10k sparse matrix.
1001        // Each interior row has 5 non-zeros; boundary rows have 3–4.
1002        let grid = 100_usize;
1003        let n = grid * grid; // 10_000 rows
1004
1005        let mut row_ptr: Vec<usize> = Vec::with_capacity(n + 1);
1006        let mut col_idx: Vec<usize> = Vec::new();
1007        let mut values: Vec<f64> = Vec::new();
1008
1009        row_ptr.push(0);
1010        for row in 0..n {
1011            let r = row / grid;
1012            let c = row % grid;
1013            // North neighbour
1014            if r > 0 {
1015                col_idx.push(row - grid);
1016                values.push(-1.0);
1017            }
1018            // West neighbour
1019            if c > 0 {
1020                col_idx.push(row - 1);
1021                values.push(-1.0);
1022            }
1023            // Self (diagonal = 4)
1024            col_idx.push(row);
1025            values.push(4.0);
1026            // East neighbour
1027            if c + 1 < grid {
1028                col_idx.push(row + 1);
1029                values.push(-1.0);
1030            }
1031            // South neighbour
1032            if r + 1 < grid {
1033                col_idx.push(row + grid);
1034                values.push(-1.0);
1035            }
1036            row_ptr.push(col_idx.len());
1037        }
1038
1039        let nnz = col_idx.len();
1040        let x: Vec<f64> = (0..n).map(|i| (i as f64) * 0.0001 + 1.0).collect();
1041
1042        // Warm-up pass
1043        let _ = csr_spmv_sim(n, &row_ptr, &col_idx, &values, &x);
1044
1045        const ITERS: usize = 10;
1046        let start = std::time::Instant::now();
1047        let mut y = vec![0.0_f64; n];
1048        for _ in 0..ITERS {
1049            y = csr_spmv_sim(n, &row_ptr, &col_idx, &values, &x);
1050        }
1051        let elapsed_ns = start.elapsed().as_nanos() as f64;
1052
1053        // 2 flops per non-zero (one multiply + one add)
1054        let total_flops = 2.0 * nnz as f64 * ITERS as f64;
1055        let gflops = total_flops / elapsed_ns; // (flops) / (ns) = GFlops/s
1056
1057        println!(
1058            "SpMV SuiteSparse proxy (10k×10k 5-pt stencil, {} nnz, {} iters): {:.3} GFLOPS (CPU reference)",
1059            nnz, ITERS, gflops
1060        );
1061
1062        // Sanity: result must be non-zero and throughput must be measurable
1063        assert!(y[n / 2] != 0.0, "SpMV result must be non-zero");
1064        assert!(
1065            gflops > 0.001,
1066            "SpMV CPU reference throughput unrealistically low: {:.6} GFLOPS",
1067            gflops
1068        );
1069    }
1070}
1071
1072// ---------------------------------------------------------------------------
1073// On-device numeric validation (feature = "gpu-tests")
1074// ---------------------------------------------------------------------------
1075
1076#[cfg(all(test, feature = "gpu-tests"))]
1077mod gpu_device_tests {
1078    use super::*;
1079    use crate::gpu_test_support::{assert_close, gpu_handle};
1080    use crate::host_csr::{f64_to_gpu, gpu_to_f64};
1081    use oxicuda_memory::DeviceBuffer;
1082
1083    /// CPU oracle for `y = alpha * A * x + beta * y0` over a CSR matrix
1084    /// (row count is derived from `row_ptr`).
1085    fn cpu_csr_spmv(
1086        row_ptr: &[i32],
1087        col_idx: &[i32],
1088        values: &[f64],
1089        x: &[f64],
1090        y0: &[f64],
1091        alpha: f64,
1092        beta: f64,
1093    ) -> Vec<f64> {
1094        let rows = row_ptr.len() - 1;
1095        let mut y = vec![0.0_f64; rows];
1096        for (i, slot) in y.iter_mut().enumerate() {
1097            let start = row_ptr[i] as usize;
1098            let end = row_ptr[i + 1] as usize;
1099            let mut acc = 0.0_f64;
1100            for k in start..end {
1101                acc += values[k] * x[col_idx[k] as usize];
1102            }
1103            *slot = alpha * acc + beta * y0[i];
1104        }
1105        y
1106    }
1107
1108    /// Drive the production `spmv` op for one element type and compare to the
1109    /// CPU oracle.
1110    #[allow(clippy::too_many_arguments)]
1111    fn run_spmv<T: GpuFloat>(
1112        algo: SpMVAlgo,
1113        rows: u32,
1114        cols: u32,
1115        row_ptr: &[i32],
1116        col_idx: &[i32],
1117        values: &[f64],
1118        x: &[f64],
1119        y0: &[f64],
1120        alpha: f64,
1121        beta: f64,
1122        tol: f64,
1123        tag: &str,
1124    ) {
1125        let Some(handle) = gpu_handle() else {
1126            return;
1127        };
1128        let dev_values: Vec<T> = values.iter().map(|&v| f64_to_gpu::<T>(v)).collect();
1129        let a = CsrMatrix::<T>::from_host(rows, cols, row_ptr, col_idx, &dev_values)
1130            .expect("test: build CSR");
1131        let dev_x: Vec<T> = x.iter().map(|&v| f64_to_gpu::<T>(v)).collect();
1132        let dev_y: Vec<T> = y0.iter().map(|&v| f64_to_gpu::<T>(v)).collect();
1133        let x_buf = DeviceBuffer::from_host(&dev_x).expect("test: upload x");
1134        let y_buf = DeviceBuffer::from_host(&dev_y).expect("test: upload y");
1135
1136        spmv::<T>(
1137            &handle,
1138            algo,
1139            f64_to_gpu::<T>(alpha),
1140            &a,
1141            x_buf.as_device_ptr(),
1142            f64_to_gpu::<T>(beta),
1143            y_buf.as_device_ptr(),
1144        )
1145        .expect("test: spmv launch");
1146        handle.stream().synchronize().expect("test: sync");
1147
1148        let mut out = vec![T::gpu_zero(); rows as usize];
1149        y_buf.copy_to_host(&mut out).expect("test: download y");
1150        let got: Vec<f64> = out.iter().map(|&v| gpu_to_f64(v)).collect();
1151        let want = cpu_csr_spmv(row_ptr, col_idx, values, x, y0, alpha, beta);
1152        assert_close(&got, &want, tol, tag);
1153    }
1154
1155    /// 4x4 symmetric tridiagonal-ish matrix with a dense-ish last row.
1156    fn matrix_4x4() -> (u32, u32, Vec<i32>, Vec<i32>, Vec<f64>) {
1157        // [ 2 -1  0  0]
1158        // [-1  2 -1  0]
1159        // [ 0 -1  2 -1]
1160        // [ 3  0 -1  4]
1161        let row_ptr = vec![0, 2, 5, 8, 11];
1162        let col_idx = vec![0, 1, 0, 1, 2, 1, 2, 3, 0, 2, 3];
1163        let values = vec![2.0, -1.0, -1.0, 2.0, -1.0, -1.0, 2.0, -1.0, 3.0, -1.0, 4.0];
1164        (4, 4, row_ptr, col_idx, values)
1165    }
1166
1167    /// Wider matrix (5x5, ~8 nnz/row average via banded structure) to exercise
1168    /// the warp/vector kernel path.
1169    fn matrix_6x6_banded() -> (u32, u32, Vec<i32>, Vec<i32>, Vec<f64>) {
1170        let n = 6usize;
1171        let mut row_ptr = vec![0i32];
1172        let mut col_idx = Vec::new();
1173        let mut values = Vec::new();
1174        for i in 0..n {
1175            let lo = i.saturating_sub(2);
1176            let hi = (i + 3).min(n);
1177            for j in lo..hi {
1178                col_idx.push(j as i32);
1179                values.push(if i == j { 5.0 } else { -1.0 + 0.1 * (i as f64) });
1180            }
1181            row_ptr.push(col_idx.len() as i32);
1182        }
1183        (n as u32, n as u32, row_ptr, col_idx, values)
1184    }
1185
1186    #[test]
1187    fn spmv_scalar_f64_alpha_beta() {
1188        let (r, c, rp, ci, v) = matrix_4x4();
1189        let x = vec![1.0, 2.0, 3.0, 4.0];
1190        let y0 = vec![10.0, 20.0, 30.0, 40.0];
1191        run_spmv::<f64>(
1192            SpMVAlgo::Scalar,
1193            r,
1194            c,
1195            &rp,
1196            &ci,
1197            &v,
1198            &x,
1199            &y0,
1200            2.5,
1201            -0.75,
1202            1e-10,
1203            "spmv_scalar_f64",
1204        );
1205    }
1206
1207    #[test]
1208    fn spmv_vector_f64_alpha_beta() {
1209        let (r, c, rp, ci, v) = matrix_6x6_banded();
1210        let x: Vec<f64> = (0..r as usize).map(|i| 0.5 + i as f64).collect();
1211        let y0: Vec<f64> = (0..r as usize).map(|i| 100.0 - i as f64).collect();
1212        run_spmv::<f64>(
1213            SpMVAlgo::Vector,
1214            r,
1215            c,
1216            &rp,
1217            &ci,
1218            &v,
1219            &x,
1220            &y0,
1221            1.5,
1222            0.25,
1223            1e-10,
1224            "spmv_vector_f64",
1225        );
1226    }
1227
1228    #[test]
1229    fn spmv_scalar_f32_alpha_beta() {
1230        let (r, c, rp, ci, v) = matrix_4x4();
1231        let x = vec![1.0, 2.0, 3.0, 4.0];
1232        let y0 = vec![10.0, 20.0, 30.0, 40.0];
1233        run_spmv::<f32>(
1234            SpMVAlgo::Scalar,
1235            r,
1236            c,
1237            &rp,
1238            &ci,
1239            &v,
1240            &x,
1241            &y0,
1242            2.0,
1243            0.5,
1244            1e-4,
1245            "spmv_scalar_f32",
1246        );
1247    }
1248
1249    #[test]
1250    fn spmv_vector_f32_alpha_beta() {
1251        let (r, c, rp, ci, v) = matrix_6x6_banded();
1252        let x: Vec<f64> = (0..r as usize).map(|i| 0.5 + i as f64).collect();
1253        let y0: Vec<f64> = (0..r as usize).map(|i| 7.0 + i as f64).collect();
1254        run_spmv::<f32>(
1255            SpMVAlgo::Vector,
1256            r,
1257            c,
1258            &rp,
1259            &ci,
1260            &v,
1261            &x,
1262            &y0,
1263            1.25,
1264            -0.5,
1265            1e-4,
1266            "spmv_vector_f32",
1267        );
1268    }
1269
1270    #[test]
1271    fn spmv_beta_zero_overwrites_garbage() {
1272        // beta = 0 must fully overwrite the prior y (incl. any stale content).
1273        let (r, c, rp, ci, v) = matrix_4x4();
1274        let x = vec![1.0, 1.0, 1.0, 1.0];
1275        let y0 = vec![1e9, -1e9, 5e8, -5e8];
1276        run_spmv::<f64>(
1277            SpMVAlgo::Scalar,
1278            r,
1279            c,
1280            &rp,
1281            &ci,
1282            &v,
1283            &x,
1284            &y0,
1285            1.0,
1286            0.0,
1287            1e-10,
1288            "spmv_beta_zero",
1289        );
1290    }
1291}