1use 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#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
29pub enum SpMVAlgo {
30 Scalar,
32 Vector,
35 Adaptive,
37}
38
39const SPMV_SCALAR_BLOCK: u32 = 256;
41
42const SPMV_VECTOR_BLOCK: u32 = 256;
44
45const VECTOR_THRESHOLD: f64 = 4.0;
47
48#[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#[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 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 spmv_scalar(handle, alpha, a, x_ptr, beta, y_ptr)
107 }
108 }
109}
110
111fn 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 ¶ms,
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
146fn 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 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 ¶ms,
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
183fn 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 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 let acc = load_float_imm::<T>(b, 0.0);
226
227 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 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 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 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 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 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 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 b.raw_ptx(&format!("add.u32 {k}, {k}, 1;"));
271 b.branch(&loop_label);
272 b.label(&done_label);
273
274 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
291fn 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 let tid_global = b.global_thread_id_x();
308 let num_rows = b.load_param_u32("num_rows");
309
310 let lane = b.alloc_reg(PtxType::U32);
312 b.raw_ptx(&format!("and.b32 {lane}, {tid_global}, 31;"));
313
314 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 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 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 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 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 b.raw_ptx(&format!("add.u32 {k}, {k}, 32;"));
382 b.branch(&loop_label);
383 b.label(&done_label);
384
385 let reduced = emit_warp_reduce_sum::<T>(b, acc);
387
388 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 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 Err(_) => None,
440 }
441 }
442
443 #[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 assert!(
456 ptx.contains(".reg .b64 %f"),
457 "f64 {name} kernel must declare 64-bit %f value registers:\n{ptx}"
458 );
459 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 assert!(
471 ptx.contains("0D0000000000000000"),
472 "f64 {name} kernel must materialize the f64 0.0 immediate (0D…):\n{ptx}"
473 );
474 assert!(
476 !ptx.contains("shfl.sync.down.b64"),
477 "f64 {name} kernel must not emit shfl.sync.down.b64:\n{ptx}"
478 );
479 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 assert!(
488 vector.contains("shfl.sync.down.b32"),
489 "f64 vector kernel must reduce via b32 shuffles:\n{vector}"
490 );
491
492 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 }
502 }
503 }
504 }
505
506 #[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 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 #[test]
565 fn test_spmv_selects_scalar_for_very_sparse() {
566 let avg = 150.0_f64 / 100.0;
568 assert!(avg < VECTOR_THRESHOLD);
569 assert_eq!(resolve_adaptive(avg), SpMVAlgo::Scalar);
570 }
571
572 #[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 #[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 #[test]
590 fn test_spmv_selection_boundary_conditions() {
591 let just_below = VECTOR_THRESHOLD - f64::EPSILON * VECTOR_THRESHOLD;
593 assert_eq!(resolve_adaptive(just_below), SpMVAlgo::Scalar);
594
595 assert_eq!(resolve_adaptive(VECTOR_THRESHOLD), SpMVAlgo::Vector);
597
598 let just_above = VECTOR_THRESHOLD + f64::EPSILON * VECTOR_THRESHOLD;
600 assert_eq!(resolve_adaptive(just_above), SpMVAlgo::Vector);
601 }
602
603 #[test]
605 fn test_spmv_selection_empty_matrix() {
606 assert_eq!(resolve_adaptive(0.0), SpMVAlgo::Scalar);
607 }
608
609 #[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 #[test]
628 fn test_spmv_scalar_for_diagonal_matrix() {
629 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 #[test]
643 fn test_spmv_scalar_for_tridiagonal_matrix() {
644 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 #[test]
658 fn test_spmv_vector_for_5pt_stencil() {
659 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 #[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 #[test]
686 fn test_spmv_vector_at_exact_threshold() {
687 let avg = VECTOR_THRESHOLD; assert_eq!(
689 resolve_adaptive(avg),
690 SpMVAlgo::Vector,
691 "avg == VECTOR_THRESHOLD should select Vector (inclusive boundary)"
692 );
693 let below = VECTOR_THRESHOLD - f64::MIN_POSITIVE;
695 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 #[test]
711 fn test_spmv_vector_for_high_density_rows() {
712 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 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 #[test]
733 fn test_spmv_adaptive_algo_is_not_concrete() {
734 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 #[test]
758 fn spmv_warp_reduction_sim_32_threads() {
759 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 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 #[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 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 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 #[test]
845 fn spmv_numerical_accuracy_identity_4x4() {
846 let n = 4_usize;
847 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 #[test]
877 fn spmv_very_sparse_0_1_percent_1000x1000() {
878 let n = 1000_usize;
879 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]; 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 #[test]
902 fn spmv_moderate_10_percent_100x100() {
903 let n = 100_usize;
904 let bandwidth = 5_usize; 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; let x = vec![1.0_f64; n];
923 let y_csr = csr_spmv_sim(n, &row_ptr, &col_idx, &values, &x);
924
925 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 #[test]
957 fn spmv_format_selection_three_brackets() {
958 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 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 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 #[test]
999 fn spmv_suitesparse_proxy_throughput_10k() {
1000 let grid = 100_usize;
1003 let n = grid * grid; 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 if r > 0 {
1015 col_idx.push(row - grid);
1016 values.push(-1.0);
1017 }
1018 if c > 0 {
1020 col_idx.push(row - 1);
1021 values.push(-1.0);
1022 }
1023 col_idx.push(row);
1025 values.push(4.0);
1026 if c + 1 < grid {
1028 col_idx.push(row + 1);
1029 values.push(-1.0);
1030 }
1031 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 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 let total_flops = 2.0 * nnz as f64 * ITERS as f64;
1055 let gflops = total_flops / elapsed_ns; println!(
1058 "SpMV SuiteSparse proxy (10k×10k 5-pt stencil, {} nnz, {} iters): {:.3} GFLOPS (CPU reference)",
1059 nnz, ITERS, gflops
1060 );
1061
1062 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#[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 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 #[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 fn matrix_4x4() -> (u32, u32, Vec<i32>, Vec<i32>, Vec<f64>) {
1157 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 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 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}