1use oxicuda_blas::GpuFloat;
11use oxicuda_memory::DeviceBuffer;
12
13use crate::error::{SparseError, SparseResult};
14
15pub struct CsrMatrix<T: GpuFloat> {
20 rows: u32,
22 cols: u32,
24 nnz: u32,
26 row_ptr: DeviceBuffer<i32>,
29 col_idx: DeviceBuffer<i32>,
31 values: DeviceBuffer<T>,
33}
34
35impl<T: GpuFloat> CsrMatrix<T> {
36 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 if rows == 0 || cols == 0 {
60 return Err(SparseError::InvalidFormat(
61 "rows and cols must be non-zero".to_string(),
62 ));
63 }
64
65 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 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 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 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 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 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 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 #[inline]
210 pub fn rows(&self) -> u32 {
211 self.rows
212 }
213
214 #[inline]
216 pub fn cols(&self) -> u32 {
217 self.cols
218 }
219
220 #[inline]
222 pub fn nnz(&self) -> u32 {
223 self.nnz
224 }
225
226 #[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 #[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 #[inline]
250 pub fn row_ptr(&self) -> &DeviceBuffer<i32> {
251 &self.row_ptr
252 }
253
254 #[inline]
256 pub fn col_idx(&self) -> &DeviceBuffer<i32> {
257 &self.col_idx
258 }
259
260 #[inline]
262 pub fn values(&self) -> &DeviceBuffer<T> {
263 &self.values
264 }
265
266 #[inline]
268 pub fn values_mut(&mut self) -> &mut DeviceBuffer<T> {
269 &mut self.values
270 }
271
272 #[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 pub fn to_coo(&self) -> SparseResult<super::CooMatrix<T>> {
290 let (h_row_ptr, h_col_idx, h_values) = self.to_host()?;
291
292 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 pub fn to_csc(&self) -> SparseResult<super::CscMatrix<T>> {
313 let (h_row_ptr, h_col_idx, h_values) = self.to_host()?;
314
315 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 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 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 assert!((4.0_f64 / 16.0 - 0.25).abs() < 1e-10);
393 }
394}