oxicuda_sparse/format/
coo.rs

1//! Coordinate (COO) format.
2//!
3//! COO stores each non-zero as a `(row, col, value)` triplet. It is the
4//! simplest sparse format and is commonly used as an intermediate format
5//! for assembly before converting to CSR or CSC.
6//!
7//! The triplets can be in any order, but sorted COO enables more efficient
8//! conversion to CSR.
9
10use oxicuda_blas::GpuFloat;
11use oxicuda_memory::DeviceBuffer;
12
13use crate::error::{SparseError, SparseResult};
14
15/// A sparse matrix in Coordinate (COO) format, stored on GPU.
16///
17/// The matrix has shape `(rows, cols)` with `nnz` non-zero elements.
18/// Each non-zero is represented by its row index, column index, and value.
19pub struct CooMatrix<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 indices of length `nnz`.
27    row_idx: DeviceBuffer<i32>,
28    /// Column indices of length `nnz`.
29    col_idx: DeviceBuffer<i32>,
30    /// Non-zero values of length `nnz`.
31    values: DeviceBuffer<T>,
32    /// Whether the triplets are sorted by (row, col).
33    sorted: bool,
34}
35
36impl<T: GpuFloat> CooMatrix<T> {
37    /// Creates a COO matrix from host-side arrays, uploading to GPU.
38    ///
39    /// The triplets are assumed to be unsorted unless the caller knows
40    /// otherwise (use [`with_sorted`](Self::with_sorted) to override).
41    ///
42    /// # Arguments
43    ///
44    /// * `rows` -- Number of rows.
45    /// * `cols` -- Number of columns.
46    /// * `row_idx` -- Row indices of length `nnz`.
47    /// * `col_idx` -- Column indices of length `nnz`.
48    /// * `values` -- Non-zero values of length `nnz`.
49    ///
50    /// # Errors
51    ///
52    /// Returns [`SparseError::InvalidFormat`] if array lengths are inconsistent.
53    pub fn from_host(
54        rows: u32,
55        cols: u32,
56        row_idx: &[i32],
57        col_idx: &[i32],
58        values: &[T],
59    ) -> SparseResult<Self> {
60        if rows == 0 || cols == 0 {
61            return Err(SparseError::InvalidFormat(
62                "rows and cols must be non-zero".to_string(),
63            ));
64        }
65
66        let nnz = values.len();
67        if nnz == 0 {
68            return Err(SparseError::ZeroNnz);
69        }
70        if row_idx.len() != nnz || col_idx.len() != nnz {
71            return Err(SparseError::InvalidFormat(format!(
72                "row_idx ({}), col_idx ({}), and values ({}) must have equal length",
73                row_idx.len(),
74                col_idx.len(),
75                nnz
76            )));
77        }
78
79        for (k, &r) in row_idx.iter().enumerate() {
80            if r < 0 || r as u32 >= rows {
81                return Err(SparseError::InvalidFormat(format!(
82                    "row_idx[{k}] = {r} out of range [0, {rows})"
83                )));
84            }
85        }
86        for (k, &c) in col_idx.iter().enumerate() {
87            if c < 0 || c as u32 >= cols {
88                return Err(SparseError::InvalidFormat(format!(
89                    "col_idx[{k}] = {c} out of range [0, {cols})"
90                )));
91            }
92        }
93
94        let d_row_idx = DeviceBuffer::from_host(row_idx)?;
95        let d_col_idx = DeviceBuffer::from_host(col_idx)?;
96        let d_values = DeviceBuffer::from_host(values)?;
97
98        Ok(Self {
99            rows,
100            cols,
101            nnz: nnz as u32,
102            row_idx: d_row_idx,
103            col_idx: d_col_idx,
104            values: d_values,
105            sorted: false,
106        })
107    }
108
109    /// Creates a COO matrix from pre-allocated device buffers.
110    ///
111    /// This is the unchecked escape hatch: only buffer *lengths* are
112    /// validated. The contents of `row_idx`/`col_idx` are **not** range
113    /// checked against `rows`/`cols`. Passing indices outside
114    /// `[0, rows)`/`[0, cols)` will cause out-of-bounds device memory
115    /// access in downstream operations (e.g. [`to_csr`](Self::to_csr),
116    /// [`to_csc`](Self::to_csc), or SpMV/SpMM kernels). Callers are
117    /// responsible for ensuring indices are in range before calling this
118    /// constructor; prefer [`from_host`](Self::from_host) when the data
119    /// originates on the host, as it validates the full range.
120    ///
121    /// # Errors
122    ///
123    /// Returns [`SparseError::InvalidFormat`] if buffer lengths are inconsistent.
124    pub fn from_device(
125        rows: u32,
126        cols: u32,
127        nnz: u32,
128        row_idx: DeviceBuffer<i32>,
129        col_idx: DeviceBuffer<i32>,
130        values: DeviceBuffer<T>,
131    ) -> SparseResult<Self> {
132        if row_idx.len() != nnz as usize
133            || col_idx.len() != nnz as usize
134            || values.len() != nnz as usize
135        {
136            return Err(SparseError::InvalidFormat(
137                "all arrays must have length equal to nnz".to_string(),
138            ));
139        }
140        Ok(Self {
141            rows,
142            cols,
143            nnz,
144            row_idx,
145            col_idx,
146            values,
147            sorted: false,
148        })
149    }
150
151    /// Mark the COO matrix as sorted by (row, col).
152    ///
153    /// This is a hint for conversion routines; the caller is responsible
154    /// for ensuring correctness.
155    #[must_use]
156    pub fn with_sorted(mut self, sorted: bool) -> Self {
157        self.sorted = sorted;
158        self
159    }
160
161    /// Downloads the COO arrays from GPU to host memory.
162    ///
163    /// # Errors
164    ///
165    /// Returns [`SparseError::Cuda`] on transfer failure.
166    pub fn to_host(&self) -> SparseResult<(Vec<i32>, Vec<i32>, Vec<T>)> {
167        let mut h_row_idx = vec![0i32; self.row_idx.len()];
168        let mut h_col_idx = vec![0i32; self.col_idx.len()];
169        let mut h_values = vec![T::gpu_zero(); self.values.len()];
170
171        self.row_idx.copy_to_host(&mut h_row_idx)?;
172        self.col_idx.copy_to_host(&mut h_col_idx)?;
173        self.values.copy_to_host(&mut h_values)?;
174
175        Ok((h_row_idx, h_col_idx, h_values))
176    }
177
178    /// Converts this COO matrix to CSR format.
179    ///
180    /// Downloads to host, builds row pointers via histogram, then uploads.
181    ///
182    /// # Errors
183    ///
184    /// Returns [`SparseError::Cuda`] on transfer failure. Returns
185    /// [`SparseError::InvalidFormat`] if any downloaded row or column
186    /// index is out of range (this can happen for matrices constructed
187    /// via the unchecked [`from_device`](Self::from_device) constructor).
188    pub fn to_csr(&self) -> SparseResult<super::CsrMatrix<T>> {
189        let (h_row_idx, h_col_idx, h_values) = self.to_host()?;
190
191        // Build row pointers from row indices
192        let mut row_counts = vec![0i32; self.rows as usize];
193        for &r in &h_row_idx {
194            if r < 0 || r as u32 >= self.rows {
195                return Err(SparseError::InvalidFormat(format!(
196                    "row index {r} out of range for {} rows",
197                    self.rows
198                )));
199            }
200            row_counts[r as usize] += 1;
201        }
202        for &c in &h_col_idx {
203            if c < 0 || c as u32 >= self.cols {
204                return Err(SparseError::InvalidFormat(format!(
205                    "col index {c} out of range for {} cols",
206                    self.cols
207                )));
208            }
209        }
210
211        let mut h_row_ptr = vec![0i32; self.rows as usize + 1];
212        for i in 0..self.rows as usize {
213            h_row_ptr[i + 1] = h_row_ptr[i] + row_counts[i];
214        }
215
216        // Place entries in row order
217        let mut h_csr_col_idx = vec![0i32; self.nnz as usize];
218        let mut h_csr_values = vec![T::gpu_zero(); self.nnz as usize];
219        let mut write_pos = h_row_ptr.clone();
220
221        for i in 0..self.nnz as usize {
222            let row = h_row_idx[i] as usize;
223            let dest = write_pos[row] as usize;
224            h_csr_col_idx[dest] = h_col_idx[i];
225            h_csr_values[dest] = h_values[i];
226            write_pos[row] += 1;
227        }
228
229        super::CsrMatrix::from_host(
230            self.rows,
231            self.cols,
232            &h_row_ptr,
233            &h_csr_col_idx,
234            &h_csr_values,
235        )
236    }
237
238    /// Converts this COO matrix to CSC format.
239    ///
240    /// Downloads to host, builds column pointers, then uploads.
241    ///
242    /// # Errors
243    ///
244    /// Returns [`SparseError::Cuda`] on transfer failure. Returns
245    /// [`SparseError::InvalidFormat`] if any downloaded row or column
246    /// index is out of range (this can happen for matrices constructed
247    /// via the unchecked [`from_device`](Self::from_device) constructor).
248    pub fn to_csc(&self) -> SparseResult<super::CscMatrix<T>> {
249        let (h_row_idx, h_col_idx, h_values) = self.to_host()?;
250
251        for &r in &h_row_idx {
252            if r < 0 || r as u32 >= self.rows {
253                return Err(SparseError::InvalidFormat(format!(
254                    "row index {r} out of range for {} rows",
255                    self.rows
256                )));
257            }
258        }
259        let mut col_counts = vec![0i32; self.cols as usize];
260        for &c in &h_col_idx {
261            if c < 0 || c as u32 >= self.cols {
262                return Err(SparseError::InvalidFormat(format!(
263                    "col index {c} out of range for {} cols",
264                    self.cols
265                )));
266            }
267            col_counts[c as usize] += 1;
268        }
269
270        let mut h_col_ptr = vec![0i32; self.cols as usize + 1];
271        for i in 0..self.cols as usize {
272            h_col_ptr[i + 1] = h_col_ptr[i] + col_counts[i];
273        }
274
275        let mut h_csc_row_idx = vec![0i32; self.nnz as usize];
276        let mut h_csc_values = vec![T::gpu_zero(); self.nnz as usize];
277        let mut write_pos = h_col_ptr.clone();
278
279        for i in 0..self.nnz as usize {
280            let col = h_col_idx[i] as usize;
281            let dest = write_pos[col] as usize;
282            h_csc_row_idx[dest] = h_row_idx[i];
283            h_csc_values[dest] = h_values[i];
284            write_pos[col] += 1;
285        }
286
287        super::CscMatrix::from_host(
288            self.rows,
289            self.cols,
290            &h_col_ptr,
291            &h_csc_row_idx,
292            &h_csc_values,
293        )
294    }
295
296    /// Returns whether the triplets are sorted by (row, col).
297    #[inline]
298    pub fn is_sorted(&self) -> bool {
299        self.sorted
300    }
301
302    /// Returns the number of rows.
303    #[inline]
304    pub fn rows(&self) -> u32 {
305        self.rows
306    }
307
308    /// Returns the number of columns.
309    #[inline]
310    pub fn cols(&self) -> u32 {
311        self.cols
312    }
313
314    /// Returns the number of non-zero elements.
315    #[inline]
316    pub fn nnz(&self) -> u32 {
317        self.nnz
318    }
319
320    /// Returns a reference to the row index device buffer.
321    #[inline]
322    pub fn row_idx(&self) -> &DeviceBuffer<i32> {
323        &self.row_idx
324    }
325
326    /// Returns a reference to the column index device buffer.
327    #[inline]
328    pub fn col_idx(&self) -> &DeviceBuffer<i32> {
329        &self.col_idx
330    }
331
332    /// Returns a reference to the values device buffer.
333    #[inline]
334    pub fn values(&self) -> &DeviceBuffer<T> {
335        &self.values
336    }
337}
338
339#[cfg(test)]
340mod tests {
341    use super::*;
342
343    #[test]
344    fn coo_validation_mismatched_lengths() {
345        let result = CooMatrix::<f32>::from_host(3, 3, &[0, 1], &[0, 1, 2], &[1.0; 3]);
346        assert!(result.is_err());
347    }
348
349    #[test]
350    fn coo_validation_zero_nnz() {
351        let result = CooMatrix::<f32>::from_host(2, 2, &[], &[], &[]);
352        assert!(matches!(result, Err(SparseError::ZeroNnz)));
353    }
354
355    #[test]
356    fn coo_sorted_flag() {
357        // Just verify the flag API works (no GPU needed)
358    }
359
360    #[test]
361    fn coo_validation_row_idx_out_of_range() {
362        // rows = 2, so row index 2 is out of range
363        let result = CooMatrix::<f32>::from_host(2, 2, &[0, 2], &[0, 1], &[1.0, 1.0]);
364        assert!(matches!(result, Err(SparseError::InvalidFormat(_))));
365    }
366
367    #[test]
368    fn coo_validation_col_idx_out_of_range() {
369        // cols = 2, so col index -1 is out of range
370        let result = CooMatrix::<f32>::from_host(2, 2, &[0, 1], &[0, -1], &[1.0, 1.0]);
371        assert!(matches!(result, Err(SparseError::InvalidFormat(_))));
372    }
373}
374
375// ---------------------------------------------------------------------------
376// On-device validation of the unchecked `from_device` escape hatch
377// (feature = "gpu-tests")
378// ---------------------------------------------------------------------------
379
380#[cfg(all(test, feature = "gpu-tests"))]
381mod gpu_device_tests {
382    use super::*;
383    use crate::gpu_test_support::gpu_handle;
384
385    /// `to_csr`/`to_csc` must reject out-of-range indices even when the
386    /// matrix was built via the unchecked `from_device` constructor,
387    /// instead of indexing a host `Vec` out of bounds or silently
388    /// producing a corrupt CSR/CSC matrix.
389    #[test]
390    fn to_csr_rejects_out_of_range_row_idx_from_device() {
391        // Keep the handle (and the CUDA context it holds current) alive for
392        // the whole test; dropping it immediately would tear down the
393        // context before the `DeviceBuffer::from_host` calls below run.
394        let Some(_handle) = gpu_handle() else {
395            return;
396        };
397        let d_row_idx = DeviceBuffer::from_host(&[0i32, 5]).expect("test: upload row_idx");
398        let d_col_idx = DeviceBuffer::from_host(&[0i32, 1]).expect("test: upload col_idx");
399        let d_values = DeviceBuffer::from_host(&[1.0f32, 1.0]).expect("test: upload values");
400        let coo = CooMatrix::<f32>::from_device(2, 2, 2, d_row_idx, d_col_idx, d_values)
401            .expect("test: build COO from device buffers (lengths are consistent)");
402
403        let result = coo.to_csr();
404        assert!(matches!(result, Err(SparseError::InvalidFormat(_))));
405    }
406
407    #[test]
408    fn to_csc_rejects_out_of_range_col_idx_from_device() {
409        let Some(_handle) = gpu_handle() else {
410            return;
411        };
412        let d_row_idx = DeviceBuffer::from_host(&[0i32, 1]).expect("test: upload row_idx");
413        let d_col_idx = DeviceBuffer::from_host(&[0i32, 5]).expect("test: upload col_idx");
414        let d_values = DeviceBuffer::from_host(&[1.0f32, 1.0]).expect("test: upload values");
415        let coo = CooMatrix::<f32>::from_device(2, 2, 2, d_row_idx, d_col_idx, d_values)
416            .expect("test: build COO from device buffers (lengths are consistent)");
417
418        let result = coo.to_csc();
419        assert!(matches!(result, Err(SparseError::InvalidFormat(_))));
420    }
421}
oxicuda_sparse/format/coo.rs

oxicuda_sparse/format/
coo.rs
