Skip to main content

oxicuda_sparse/format/
csr.rs

1//! Compressed Sparse Row (CSR) format.
2//!
3//! CSR is the most widely used sparse matrix format. It stores:
4//! - `row_ptr[rows+1]`: indices into `col_idx`/`values` for each row
5//! - `col_idx[nnz]`: column index for each non-zero
6//! - `values[nnz]`: the non-zero values
7//!
8//! This is the primary format for SpMV, SpMM, and most sparse operations.
9
10use oxicuda_blas::GpuFloat;
11use oxicuda_memory::DeviceBuffer;
12
13use crate::error::{SparseError, SparseResult};
14
15/// A sparse matrix in Compressed Sparse Row (CSR) format, stored on GPU.
16///
17/// The matrix has shape `(rows, cols)` with `nnz` non-zero elements.
18/// All index and value arrays reside in device (GPU) memory.
19pub struct CsrMatrix<T: GpuFloat> {
20    /// Number of rows.
21    rows: u32,
22    /// Number of columns.
23    cols: u32,
24    /// Number of non-zero elements.
25    nnz: u32,
26    /// Row pointer array of length `rows + 1`. `row_ptr[i]` is the index
27    /// into `col_idx`/`values` where row `i` begins.
28    row_ptr: DeviceBuffer<i32>,
29    /// Column indices of length `nnz`.
30    col_idx: DeviceBuffer<i32>,
31    /// Non-zero values of length `nnz`.
32    values: DeviceBuffer<T>,
33}
34
35impl<T: GpuFloat> CsrMatrix<T> {
36    /// Creates a CSR matrix from host-side arrays, uploading to GPU.
37    ///
38    /// # Arguments
39    ///
40    /// * `rows` -- Number of rows.
41    /// * `cols` -- Number of columns.
42    /// * `row_ptr` -- Row pointer array of length `rows + 1`.
43    /// * `col_idx` -- Column indices of length `nnz`.
44    /// * `values` -- Non-zero values of length `nnz`.
45    ///
46    /// # Errors
47    ///
48    /// Returns [`SparseError::InvalidFormat`] if array lengths are inconsistent.
49    /// Returns [`SparseError::ZeroNnz`] if `values` is empty.
50    /// Returns [`SparseError::Cuda`] on GPU memory allocation failure.
51    pub fn from_host(
52        rows: u32,
53        cols: u32,
54        row_ptr: &[i32],
55        col_idx: &[i32],
56        values: &[T],
57    ) -> SparseResult<Self> {
58        // Validate dimensions
59        if rows == 0 || cols == 0 {
60            return Err(SparseError::InvalidFormat(
61                "rows and cols must be non-zero".to_string(),
62            ));
63        }
64
65        // Validate row_ptr length
66        let expected_row_ptr_len = rows as usize + 1;
67        if row_ptr.len() != expected_row_ptr_len {
68            return Err(SparseError::InvalidFormat(format!(
69                "row_ptr length ({}) must be rows + 1 ({})",
70                row_ptr.len(),
71                expected_row_ptr_len
72            )));
73        }
74
75        // Validate nnz consistency
76        let nnz = values.len();
77        if nnz == 0 {
78            return Err(SparseError::ZeroNnz);
79        }
80        if col_idx.len() != nnz {
81            return Err(SparseError::InvalidFormat(format!(
82                "col_idx length ({}) must equal values length ({})",
83                col_idx.len(),
84                nnz
85            )));
86        }
87
88        // Validate row_ptr is non-decreasing and within bounds
89        if row_ptr[0] != 0 {
90            return Err(SparseError::InvalidFormat(
91                "row_ptr[0] must be 0".to_string(),
92            ));
93        }
94        if row_ptr[rows as usize] != nnz as i32 {
95            return Err(SparseError::InvalidFormat(format!(
96                "row_ptr[rows] ({}) must equal nnz ({})",
97                row_ptr[rows as usize], nnz
98            )));
99        }
100        for i in 0..rows as usize {
101            if row_ptr[i] > row_ptr[i + 1] {
102                return Err(SparseError::InvalidFormat(format!(
103                    "row_ptr must be non-decreasing: row_ptr[{}]={} > row_ptr[{}]={}",
104                    i,
105                    row_ptr[i],
106                    i + 1,
107                    row_ptr[i + 1]
108                )));
109            }
110        }
111
112        // Validate column indices are within [0, cols); an out-of-range
113        // col_idx would otherwise cause SpMV/SpMM kernels to read device
114        // memory out of bounds.
115        for (k, &c) in col_idx.iter().enumerate() {
116            if c < 0 || c as u32 >= cols {
117                return Err(SparseError::InvalidFormat(format!(
118                    "col_idx[{k}] = {c} out of range [0, {cols})"
119                )));
120            }
121        }
122
123        // Upload to GPU
124        let d_row_ptr = DeviceBuffer::from_host(row_ptr)?;
125        let d_col_idx = DeviceBuffer::from_host(col_idx)?;
126        let d_values = DeviceBuffer::from_host(values)?;
127
128        Ok(Self {
129            rows,
130            cols,
131            nnz: nnz as u32,
132            row_ptr: d_row_ptr,
133            col_idx: d_col_idx,
134            values: d_values,
135        })
136    }
137
138    /// Creates a CSR matrix from pre-allocated device buffers.
139    ///
140    /// This is the unchecked escape hatch: no validation of contents is
141    /// performed; only lengths are checked. In particular, `col_idx`
142    /// values are **not** range-checked against `cols`, and `row_ptr` is
143    /// not checked for monotonicity. Out-of-range column indices will
144    /// cause SpMV/SpMM kernels to read device memory out of bounds.
145    /// Callers are responsible for ensuring the contents are valid;
146    /// prefer [`from_host`](Self::from_host) when the data originates on
147    /// the host, as it validates both structure and index ranges.
148    ///
149    /// # Errors
150    ///
151    /// Returns [`SparseError::InvalidFormat`] if buffer lengths are inconsistent.
152    pub fn from_device(
153        rows: u32,
154        cols: u32,
155        nnz: u32,
156        row_ptr: DeviceBuffer<i32>,
157        col_idx: DeviceBuffer<i32>,
158        values: DeviceBuffer<T>,
159    ) -> SparseResult<Self> {
160        if row_ptr.len() != (rows as usize + 1) {
161            return Err(SparseError::InvalidFormat(format!(
162                "row_ptr length ({}) must be rows + 1 ({})",
163                row_ptr.len(),
164                rows as usize + 1
165            )));
166        }
167        if col_idx.len() != nnz as usize {
168            return Err(SparseError::InvalidFormat(format!(
169                "col_idx length ({}) must equal nnz ({})",
170                col_idx.len(),
171                nnz
172            )));
173        }
174        if values.len() != nnz as usize {
175            return Err(SparseError::InvalidFormat(format!(
176                "values length ({}) must equal nnz ({})",
177                values.len(),
178                nnz
179            )));
180        }
181        Ok(Self {
182            rows,
183            cols,
184            nnz,
185            row_ptr,
186            col_idx,
187            values,
188        })
189    }
190
191    /// Downloads the CSR arrays from GPU to host memory.
192    ///
193    /// # Errors
194    ///
195    /// Returns [`SparseError::Cuda`] on transfer failure.
196    pub fn to_host(&self) -> SparseResult<(Vec<i32>, Vec<i32>, Vec<T>)> {
197        let mut h_row_ptr = vec![0i32; self.row_ptr.len()];
198        let mut h_col_idx = vec![0i32; self.col_idx.len()];
199        let mut h_values = vec![T::gpu_zero(); self.values.len()];
200
201        self.row_ptr.copy_to_host(&mut h_row_ptr)?;
202        self.col_idx.copy_to_host(&mut h_col_idx)?;
203        self.values.copy_to_host(&mut h_values)?;
204
205        Ok((h_row_ptr, h_col_idx, h_values))
206    }
207
208    /// Returns the number of rows.
209    #[inline]
210    pub fn rows(&self) -> u32 {
211        self.rows
212    }
213
214    /// Returns the number of columns.
215    #[inline]
216    pub fn cols(&self) -> u32 {
217        self.cols
218    }
219
220    /// Returns the number of non-zero elements.
221    #[inline]
222    pub fn nnz(&self) -> u32 {
223        self.nnz
224    }
225
226    /// Returns the density of the matrix (nnz / (rows * cols)).
227    #[inline]
228    pub fn density(&self) -> f64 {
229        let total = self.rows as f64 * self.cols as f64;
230        if total == 0.0 {
231            return 0.0;
232        }
233        self.nnz as f64 / total
234    }
235
236    /// Checks whether the matrix metadata is structurally valid.
237    ///
238    /// This checks buffer lengths but does NOT download data to verify
239    /// content (e.g. sorted column indices). For full validation, use
240    /// [`to_host`](Self::to_host) and check on the CPU.
241    #[inline]
242    pub fn is_valid(&self) -> bool {
243        self.row_ptr.len() == (self.rows as usize + 1)
244            && self.col_idx.len() == self.nnz as usize
245            && self.values.len() == self.nnz as usize
246    }
247
248    /// Returns a reference to the row pointer device buffer.
249    #[inline]
250    pub fn row_ptr(&self) -> &DeviceBuffer<i32> {
251        &self.row_ptr
252    }
253
254    /// Returns a reference to the column index device buffer.
255    #[inline]
256    pub fn col_idx(&self) -> &DeviceBuffer<i32> {
257        &self.col_idx
258    }
259
260    /// Returns a reference to the values device buffer.
261    #[inline]
262    pub fn values(&self) -> &DeviceBuffer<T> {
263        &self.values
264    }
265
266    /// Returns a mutable reference to the values device buffer.
267    #[inline]
268    pub fn values_mut(&mut self) -> &mut DeviceBuffer<T> {
269        &mut self.values
270    }
271
272    /// Average number of non-zeros per row.
273    #[inline]
274    pub fn avg_nnz_per_row(&self) -> f64 {
275        if self.rows == 0 {
276            return 0.0;
277        }
278        self.nnz as f64 / self.rows as f64
279    }
280
281    /// Converts this CSR matrix to COO format on the host.
282    ///
283    /// Downloads data to host, expands row pointers to row indices,
284    /// then uploads the COO matrix back to GPU.
285    ///
286    /// # Errors
287    ///
288    /// Returns [`SparseError::Cuda`] on transfer failure.
289    pub fn to_coo(&self) -> SparseResult<super::CooMatrix<T>> {
290        let (h_row_ptr, h_col_idx, h_values) = self.to_host()?;
291
292        // Expand row_ptr to row_idx
293        let mut h_row_idx = Vec::with_capacity(self.nnz as usize);
294        for row in 0..self.rows {
295            let start = h_row_ptr[row as usize];
296            let end = h_row_ptr[row as usize + 1];
297            for _ in start..end {
298                h_row_idx.push(row as i32);
299            }
300        }
301
302        super::CooMatrix::from_host(self.rows, self.cols, &h_row_idx, &h_col_idx, &h_values)
303    }
304
305    /// Converts this CSR matrix to CSC format on the host.
306    ///
307    /// Downloads data, transposes the structure, then uploads.
308    ///
309    /// # Errors
310    ///
311    /// Returns [`SparseError::Cuda`] on transfer failure.
312    pub fn to_csc(&self) -> SparseResult<super::CscMatrix<T>> {
313        let (h_row_ptr, h_col_idx, h_values) = self.to_host()?;
314
315        // Build CSC by transposing CSR
316        let mut col_counts = vec![0i32; self.cols as usize];
317        for &c in &h_col_idx {
318            col_counts[c as usize] += 1;
319        }
320
321        let mut h_col_ptr = vec![0i32; self.cols as usize + 1];
322        for i in 0..self.cols as usize {
323            h_col_ptr[i + 1] = h_col_ptr[i] + col_counts[i];
324        }
325
326        let mut h_csc_row_idx = vec![0i32; self.nnz as usize];
327        let mut h_csc_values = vec![T::gpu_zero(); self.nnz as usize];
328        let mut write_pos = h_col_ptr.clone();
329
330        for row in 0..self.rows as usize {
331            let start = h_row_ptr[row] as usize;
332            let end = h_row_ptr[row + 1] as usize;
333            for j in start..end {
334                let col = h_col_idx[j] as usize;
335                let dest = write_pos[col] as usize;
336                h_csc_row_idx[dest] = row as i32;
337                h_csc_values[dest] = h_values[j];
338                write_pos[col] += 1;
339            }
340        }
341
342        super::CscMatrix::from_host(
343            self.rows,
344            self.cols,
345            &h_col_ptr,
346            &h_csc_row_idx,
347            &h_csc_values,
348        )
349    }
350}
351
352#[cfg(test)]
353mod tests {
354    use super::*;
355
356    #[test]
357    fn csr_validation_row_ptr_length() {
358        // row_ptr too short
359        let result = CsrMatrix::<f32>::from_host(3, 3, &[0, 2, 4], &[0, 1, 0, 2], &[1.0; 4]);
360        assert!(result.is_err());
361    }
362
363    #[test]
364    fn csr_validation_zero_nnz() {
365        let result = CsrMatrix::<f32>::from_host(2, 2, &[0, 0, 0], &[], &[]);
366        assert!(matches!(result, Err(SparseError::ZeroNnz)));
367    }
368
369    #[test]
370    fn csr_validation_mismatched_col_idx() {
371        let result = CsrMatrix::<f32>::from_host(2, 2, &[0, 1, 2], &[0], &[1.0, 2.0]);
372        assert!(result.is_err());
373    }
374
375    #[test]
376    fn csr_validation_col_idx_out_of_range() {
377        // cols = 2, so col index 2 is out of range
378        let result = CsrMatrix::<f32>::from_host(2, 2, &[0, 1, 2], &[0, 2], &[1.0, 2.0]);
379        assert!(matches!(result, Err(SparseError::InvalidFormat(_))));
380    }
381
382    #[test]
383    fn csr_validation_negative_col_idx() {
384        let result = CsrMatrix::<f32>::from_host(2, 2, &[0, 1, 2], &[0, -1], &[1.0, 2.0]);
385        assert!(matches!(result, Err(SparseError::InvalidFormat(_))));
386    }
387
388    #[test]
389    fn csr_density() {
390        // A 4x4 matrix with 4 nnz => density 0.25
391        // We cannot actually call from_host without a GPU, so test the formula
392        assert!((4.0_f64 / 16.0 - 0.25).abs() < 1e-10);
393    }
394}