#![cfg(feature = "cuda")]
use crate::error::{GraphError, Result as GraphResult};
use oxicuda_sparse::ops::{spmv, SpMVAlgo};
use oxicuda_sparse::{CsrMatrix, SparseHandle};
pub fn cuda_is_available() -> bool {
oxicuda_driver::init().is_ok()
&& oxicuda_driver::device::Device::count()
.map(|c| c > 0)
.unwrap_or(false)
}
fn sparse_err(e: oxicuda_sparse::error::SparseError) -> GraphError {
GraphError::ComputationError(format!("oxicuda-sparse: {e}"))
}
fn cuda_err(e: oxicuda_driver::CudaError) -> GraphError {
GraphError::ComputationError(format!("oxicuda CUDA driver: {e}"))
}
fn build_context() -> GraphResult<std::sync::Arc<oxicuda_driver::Context>> {
oxicuda_driver::init()
.map_err(|e| GraphError::ComputationError(format!("CUDA unavailable: {e}")))?;
let count = oxicuda_driver::device::Device::count()
.map_err(|e| GraphError::ComputationError(format!("device count: {e}")))?;
if count <= 0 {
return Err(GraphError::ComputationError(
"no NVIDIA CUDA device available".into(),
));
}
let dev = oxicuda_driver::device::Device::get(0).map_err(cuda_err)?;
Ok(std::sync::Arc::new(
oxicuda_driver::Context::new(&dev).map_err(cuda_err)?,
))
}
pub fn cuda_spmv_csr(
row_offsets: &[i32],
col_indices: &[i32],
values: &[f64],
n_rows: usize,
n_cols: usize,
x: &[f64],
) -> GraphResult<Vec<f64>> {
if n_rows == 0 {
return Ok(Vec::new());
}
if row_offsets.len() != n_rows + 1 {
return Err(GraphError::InvalidGraph(format!(
"cuda_spmv_csr: row_offsets length {} does not match n_rows + 1 = {}",
row_offsets.len(),
n_rows + 1
)));
}
if col_indices.len() != values.len() {
return Err(GraphError::InvalidGraph(format!(
"cuda_spmv_csr: col_indices length {} != values length {}",
col_indices.len(),
values.len()
)));
}
if x.len() != n_cols {
return Err(GraphError::InvalidGraph(format!(
"cuda_spmv_csr: x length {} != n_cols {}",
x.len(),
n_cols
)));
}
let nnz = values.len();
let last = *row_offsets.last().unwrap_or(&0);
if i64::from(last) != nnz as i64 {
return Err(GraphError::InvalidGraph(format!(
"cuda_spmv_csr: row_offsets last element {last} != nnz {nnz}"
)));
}
if nnz == 0 {
return Ok(vec![0.0f64; n_rows]);
}
let ctx = build_context()?;
let handle = SparseHandle::new(&ctx).map_err(sparse_err)?;
let csr = CsrMatrix::<f64>::from_host(
n_rows as u32,
n_cols as u32,
row_offsets,
col_indices,
values,
)
.map_err(sparse_err)?;
let d_x = oxicuda_memory::DeviceBuffer::from_host(x).map_err(cuda_err)?;
let d_y = oxicuda_memory::DeviceBuffer::from_host(&vec![0.0f64; n_rows]).map_err(cuda_err)?;
spmv::<f64>(
&handle,
SpMVAlgo::Adaptive,
1.0,
&csr,
d_x.as_device_ptr(),
0.0,
d_y.as_device_ptr(),
)
.map_err(sparse_err)?;
let mut y = vec![0.0f64; n_rows];
d_y.copy_to_host(&mut y).map_err(cuda_err)?;
Ok(y)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn shape_mismatch_is_detected_without_gpu() {
let res = cuda_spmv_csr(&[0, 1], &[0], &[1.0], 3, 3, &[1.0, 2.0, 3.0]);
assert!(res.is_err());
}
#[test]
fn x_length_mismatch_is_detected_without_gpu() {
let res = cuda_spmv_csr(&[0, 1, 1, 1], &[0], &[1.0], 3, 3, &[1.0, 2.0]);
assert!(res.is_err());
}
#[test]
fn empty_matrix_is_ok_without_gpu() {
let y = cuda_spmv_csr(&[], &[], &[], 0, 0, &[]).expect("empty spmv should be Ok");
assert!(y.is_empty());
}
#[test]
fn no_edges_yields_zero_vector_without_gpu() {
let y = cuda_spmv_csr(&[0, 0, 0, 0], &[], &[], 3, 3, &[1.0, 2.0, 3.0])
.expect("no-edge spmv should be Ok");
assert_eq!(y, vec![0.0, 0.0, 0.0]);
}
#[test]
fn cuda_spmv_path_graph_or_skip() {
if !cuda_is_available() {
eprintln!("skipping: no NVIDIA CUDA device");
assert!(!cuda_is_available());
return;
}
use crate::compressed::CsrGraph;
let row_ptr_us = vec![0usize, 1, 3, 4];
let col_us = vec![1usize, 0, 2, 1];
let vals = vec![1.0f64, 1.0, 1.0, 1.0];
let g = CsrGraph::from_raw(3, row_ptr_us.clone(), col_us.clone(), vals.clone(), false)
.expect("from_raw");
let x = vec![1.0f64, 2.0, 3.0];
let cpu = g.spmv(&x).expect("cpu spmv");
assert_eq!(cpu, vec![2.0, 4.0, 2.0]);
let row_ptr_i: Vec<i32> = row_ptr_us.iter().map(|&v| v as i32).collect();
let col_i: Vec<i32> = col_us.iter().map(|&v| v as i32).collect();
let gpu = cuda_spmv_csr(&row_ptr_i, &col_i, &vals, 3, 3, &x).expect("gpu spmv");
let max_diff = gpu
.iter()
.zip(cpu.iter())
.map(|(gv, cv)| (gv - cv).abs())
.fold(0.0f64, f64::max);
assert!(max_diff < 1e-9, "max abs diff {max_diff} exceeds 1e-9");
}
#[test]
fn cuda_spmv_vector_kernel_directed_or_skip() {
if !cuda_is_available() {
eprintln!("skipping: no NVIDIA CUDA device");
assert!(!cuda_is_available());
return;
}
use crate::compressed::CsrGraph;
let n: usize = 8;
#[rustfmt::skip]
let row_ptr_us: Vec<usize> = vec![0, 5, 9, 14, 18, 23, 27, 32, 37];
#[rustfmt::skip]
let col_us: Vec<usize> = vec![
1, 2, 3, 4, 5,
0, 3, 5, 6,
0, 1, 4, 5, 7,
0, 2, 6, 7,
1, 2, 3, 5, 6,
0, 3, 4, 7,
1, 2, 3, 4, 5,
0, 2, 4, 5, 6,
];
#[rustfmt::skip]
let vals: Vec<f64> = vec![
0.5, 1.5, 2.0, 0.3, 1.1,
0.7, 1.2, 0.9, 2.3,
0.4, 1.8, 0.6, 2.1, 1.4,
0.8, 1.3, 0.2, 1.7,
0.5, 0.9, 1.6, 2.2, 0.7,
1.1, 0.4, 1.9, 0.6,
0.3, 1.5, 0.8, 2.4, 1.0,
0.6, 1.2, 0.8, 1.4, 0.9,
];
let nnz = vals.len();
let avg_nnz_per_row = nnz as f64 / n as f64;
assert!(
avg_nnz_per_row >= 4.0,
"avg nnz/row {avg_nnz_per_row:.3} must be >= 4.0 to select the VECTOR kernel"
);
let g = CsrGraph::from_raw(n, row_ptr_us.clone(), col_us.clone(), vals.clone(), true)
.expect("from_raw directed graph");
let x: Vec<f64> = (1..=n).map(|v| v as f64).collect(); let cpu = g.spmv(&x).expect("cpu spmv");
let row_ptr_i: Vec<i32> = row_ptr_us.iter().map(|&v| v as i32).collect();
let col_i: Vec<i32> = col_us.iter().map(|&v| v as i32).collect();
let gpu = cuda_spmv_csr(&row_ptr_i, &col_i, &vals, n, n, &x)
.expect("cuda_spmv_csr (vector kernel)");
assert_eq!(gpu.len(), n, "output length must match n_rows");
let max_diff = gpu
.iter()
.zip(cpu.iter())
.map(|(gv, cv)| (gv - cv).abs())
.fold(0.0f64, f64::max);
assert!(
max_diff < 1e-9,
"VECTOR SpMV max abs diff {max_diff:.3e} exceeds 1e-9 \
(avg_nnz_per_row={avg_nnz_per_row:.3}, nnz={nnz})"
);
}
}