oxicuda_sparse/format/
bsr.rs

1//! Block Sparse Row (BSR) format.
2//!
3//! BSR stores sparse matrices as a collection of dense blocks of size
4//! `block_dim x block_dim`. This is efficient for matrices with structured
5//! sparsity patterns, such as those arising from finite-element
6//! discretizations.
7//!
8//! The storage is organized as:
9//! - `row_ptr[block_rows+1]`: indices into `col_idx`/`values` for each block row
10//! - `col_idx[nnz_blocks]`: block column index for each non-zero block
11//! - `values[nnz_blocks * block_dim * block_dim]`: dense block data (row-major)
12
13use oxicuda_blas::GpuFloat;
14use oxicuda_memory::DeviceBuffer;
15
16use crate::error::{SparseError, SparseResult};
17
18/// A sparse matrix in Block Sparse Row (BSR) format, stored on GPU.
19///
20/// The matrix has shape `(rows, cols)` where `rows` and `cols` are multiples
21/// of `block_dim`. It contains `nnz_blocks` non-zero dense blocks.
22pub struct BsrMatrix<T: GpuFloat> {
23    /// Number of rows (must be a multiple of `block_dim`).
24    rows: u32,
25    /// Number of columns (must be a multiple of `block_dim`).
26    cols: u32,
27    /// Number of non-zero blocks.
28    nnz_blocks: u32,
29    /// Size of each dense block (block_dim x block_dim).
30    block_dim: u32,
31    /// Block row pointer array of length `block_rows + 1`.
32    row_ptr: DeviceBuffer<i32>,
33    /// Block column indices of length `nnz_blocks`.
34    col_idx: DeviceBuffer<i32>,
35    /// Dense block values of length `nnz_blocks * block_dim * block_dim`.
36    values: DeviceBuffer<T>,
37}
38
39impl<T: GpuFloat> BsrMatrix<T> {
40    /// Creates a BSR matrix from host-side arrays, uploading to GPU.
41    ///
42    /// # Arguments
43    ///
44    /// * `rows` -- Number of rows (must be a multiple of `block_dim`).
45    /// * `cols` -- Number of columns (must be a multiple of `block_dim`).
46    /// * `block_dim` -- Size of each dense block.
47    /// * `row_ptr` -- Block row pointer array of length `block_rows + 1`.
48    /// * `col_idx` -- Block column indices of length `nnz_blocks`.
49    /// * `values` -- Dense block values (row-major within each block).
50    ///
51    /// # Errors
52    ///
53    /// Returns [`SparseError::InvalidFormat`] if dimensions or array lengths
54    /// are inconsistent.
55    #[allow(clippy::too_many_arguments)]
56    pub fn from_host(
57        rows: u32,
58        cols: u32,
59        block_dim: u32,
60        row_ptr: &[i32],
61        col_idx: &[i32],
62        values: &[T],
63    ) -> SparseResult<Self> {
64        if rows == 0 || cols == 0 || block_dim == 0 {
65            return Err(SparseError::InvalidFormat(
66                "rows, cols, and block_dim must be non-zero".to_string(),
67            ));
68        }
69        if rows % block_dim != 0 {
70            return Err(SparseError::InvalidFormat(format!(
71                "rows ({rows}) must be a multiple of block_dim ({block_dim})"
72            )));
73        }
74        if cols % block_dim != 0 {
75            return Err(SparseError::InvalidFormat(format!(
76                "cols ({cols}) must be a multiple of block_dim ({block_dim})"
77            )));
78        }
79
80        let block_rows = rows / block_dim;
81        let expected_row_ptr_len = block_rows as usize + 1;
82        if row_ptr.len() != expected_row_ptr_len {
83            return Err(SparseError::InvalidFormat(format!(
84                "row_ptr length ({}) must be block_rows + 1 ({})",
85                row_ptr.len(),
86                expected_row_ptr_len
87            )));
88        }
89
90        let nnz_blocks = col_idx.len() as u32;
91        if nnz_blocks == 0 {
92            return Err(SparseError::ZeroNnz);
93        }
94
95        let block_elems = block_dim as usize * block_dim as usize;
96        let expected_values_len = nnz_blocks as usize * block_elems;
97        if values.len() != expected_values_len {
98            return Err(SparseError::InvalidFormat(format!(
99                "values length ({}) must be nnz_blocks * block_dim^2 ({})",
100                values.len(),
101                expected_values_len
102            )));
103        }
104
105        // Validate row_ptr
106        if row_ptr[0] != 0 {
107            return Err(SparseError::InvalidFormat(
108                "row_ptr[0] must be 0".to_string(),
109            ));
110        }
111        if row_ptr[block_rows as usize] != nnz_blocks as i32 {
112            return Err(SparseError::InvalidFormat(format!(
113                "row_ptr[block_rows] ({}) must equal nnz_blocks ({})",
114                row_ptr[block_rows as usize], nnz_blocks
115            )));
116        }
117
118        // Validate block column indices are within [0, block_cols); an
119        // out-of-range col_idx would otherwise cause SpMV/SpMM kernels to
120        // read device memory out of bounds.
121        let block_cols = cols / block_dim;
122        for (k, &c) in col_idx.iter().enumerate() {
123            if c < 0 || c as u32 >= block_cols {
124                return Err(SparseError::InvalidFormat(format!(
125                    "col_idx[{k}] = {c} out of range [0, {block_cols})"
126                )));
127            }
128        }
129
130        let d_row_ptr = DeviceBuffer::from_host(row_ptr)?;
131        let d_col_idx = DeviceBuffer::from_host(col_idx)?;
132        let d_values = DeviceBuffer::from_host(values)?;
133
134        Ok(Self {
135            rows,
136            cols,
137            nnz_blocks,
138            block_dim,
139            row_ptr: d_row_ptr,
140            col_idx: d_col_idx,
141            values: d_values,
142        })
143    }
144
145    /// Downloads the BSR arrays from GPU to host memory.
146    ///
147    /// # Errors
148    ///
149    /// Returns [`SparseError::Cuda`] on transfer failure.
150    pub fn to_host(&self) -> SparseResult<(Vec<i32>, Vec<i32>, Vec<T>)> {
151        let mut h_row_ptr = vec![0i32; self.row_ptr.len()];
152        let mut h_col_idx = vec![0i32; self.col_idx.len()];
153        let mut h_values = vec![T::gpu_zero(); self.values.len()];
154
155        self.row_ptr.copy_to_host(&mut h_row_ptr)?;
156        self.col_idx.copy_to_host(&mut h_col_idx)?;
157        self.values.copy_to_host(&mut h_values)?;
158
159        Ok((h_row_ptr, h_col_idx, h_values))
160    }
161
162    /// Returns the number of rows.
163    #[inline]
164    pub fn rows(&self) -> u32 {
165        self.rows
166    }
167
168    /// Returns the number of columns.
169    #[inline]
170    pub fn cols(&self) -> u32 {
171        self.cols
172    }
173
174    /// Returns the number of non-zero blocks.
175    #[inline]
176    pub fn nnz_blocks(&self) -> u32 {
177        self.nnz_blocks
178    }
179
180    /// Returns the block dimension.
181    #[inline]
182    pub fn block_dim(&self) -> u32 {
183        self.block_dim
184    }
185
186    /// Returns the number of block rows.
187    #[inline]
188    pub fn block_rows(&self) -> u32 {
189        self.rows / self.block_dim
190    }
191
192    /// Returns the number of block columns.
193    #[inline]
194    pub fn block_cols(&self) -> u32 {
195        self.cols / self.block_dim
196    }
197
198    /// Returns the total scalar non-zero count.
199    #[inline]
200    pub fn scalar_nnz(&self) -> u32 {
201        self.nnz_blocks * self.block_dim * self.block_dim
202    }
203
204    /// Returns a reference to the block row pointer device buffer.
205    #[inline]
206    pub fn row_ptr(&self) -> &DeviceBuffer<i32> {
207        &self.row_ptr
208    }
209
210    /// Returns a reference to the block column index device buffer.
211    #[inline]
212    pub fn col_idx(&self) -> &DeviceBuffer<i32> {
213        &self.col_idx
214    }
215
216    /// Returns a reference to the values device buffer.
217    #[inline]
218    pub fn values(&self) -> &DeviceBuffer<T> {
219        &self.values
220    }
221}
222
223#[cfg(test)]
224mod tests {
225    use super::*;
226
227    #[test]
228    fn bsr_validation_block_alignment() {
229        // rows not a multiple of block_dim
230        let result = BsrMatrix::<f32>::from_host(
231            5,
232            4,
233            2,
234            &[0, 1, 2, 3], // would be wrong length too
235            &[0, 1, 0],
236            &[1.0; 12],
237        );
238        assert!(result.is_err());
239    }
240
241    #[test]
242    fn bsr_validation_col_idx_out_of_range() {
243        // 4x4 matrix, block_dim = 2 => block_cols = 2, so block col index 2
244        // is out of range.
245        let result = BsrMatrix::<f32>::from_host(4, 4, 2, &[0, 1, 1], &[2], &[1.0; 4]);
246        assert!(matches!(result, Err(SparseError::InvalidFormat(_))));
247    }
248
249    #[test]
250    fn bsr_validation_negative_col_idx() {
251        let result = BsrMatrix::<f32>::from_host(4, 4, 2, &[0, 1, 1], &[-1], &[1.0; 4]);
252        assert!(matches!(result, Err(SparseError::InvalidFormat(_))));
253    }
254
255    #[test]
256    fn bsr_block_counts() {
257        // 4x4 matrix with 2x2 blocks => 2 block rows, 2 block cols
258        let br = 4 / 2;
259        let bc = 4 / 2;
260        assert_eq!(br, 2);
261        assert_eq!(bc, 2);
262    }
263}
oxicuda_sparse/format/bsr.rs

oxicuda_sparse/format/
bsr.rs
