singe_cusparse/
matrix.rs

1use std::{marker::PhantomData, mem::MaybeUninit, mem::size_of, ptr};
2
3use singe_cuda::{data_type::DataTypeLike, memory::DeviceMemory, types::DevicePtr};
4
5use crate::{
6    context::Context,
7    error::{Error, Result},
8    sys, try_ffi,
9    types::{
10        CsrToCscAlgorithm, DenseToSparseAlgorithm, DiagonalType, FillMode, Format, IndexBase,
11        IndexTypeLike, MatrixType, Order, SparseMatrixAttribute, SparseToDenseAlgorithm,
12    },
13    utility::{to_i32, to_i64, to_usize},
14};
15
16#[derive(Debug, Clone, Copy, PartialEq, Eq)]
17pub struct SparseMatrixShape {
18    pub rows: usize,
19    pub cols: usize,
20    pub nonzero_count: usize,
21}
22
23#[derive(Debug, Clone, Copy, PartialEq, Eq)]
24pub struct DenseMatrixInfo {
25    pub rows: usize,
26    pub cols: usize,
27    pub leading_dimension: usize,
28    pub values: DevicePtr,
29    pub order: Order,
30}
31
32#[derive(Debug)]
33pub struct MatrixDescriptor {
34    handle: sys::cusparseMatDescr_t,
35}
36
37#[derive(Debug)]
38pub struct SparseMatrixDescriptor<'a> {
39    handle: sys::cusparseSpMatDescr_t,
40    _buffers: PhantomData<&'a mut ()>,
41}
42
43#[derive(Debug)]
44pub struct DenseMatrixDescriptor<'a> {
45    handle: sys::cusparseDnMatDescr_t,
46    _buffers: PhantomData<&'a mut ()>,
47}
48
49// Legacy matrix descriptors contain descriptor attributes only, so immutable
50// sharing follows the cuSPARSE descriptor contract. Data-bearing descriptors
51// can rebind borrowed pointers and therefore require exclusive access.
52unsafe impl Send for MatrixDescriptor {}
53unsafe impl Sync for MatrixDescriptor {}
54unsafe impl Send for SparseMatrixDescriptor<'_> {}
55unsafe impl Send for DenseMatrixDescriptor<'_> {}
56
57impl MatrixDescriptor {
58    /// Initializes a matrix descriptor.
59    /// It sets [`MatrixType`] and [`IndexBase`] to [`MatrixType::General`] and
60    /// [`IndexBase::Zero`], respectively, while leaving other fields uninitialized.
61    ///
62    /// # Errors
63    ///
64    /// Returns an error if cuSPARSE cannot create the descriptor or if it returns
65    /// a null handle.
66    pub fn create() -> Result<Self> {
67        let mut handle = ptr::null_mut();
68        unsafe {
69            try_ffi!(sys::cusparseCreateMatDescr(&raw mut handle))?;
70        }
71
72        if handle.is_null() {
73            return Err(Error::NullHandle);
74        }
75
76        Ok(Self { handle })
77    }
78
79    /// Returns the [`MatrixType`] field of this matrix descriptor.
80    pub fn matrix_type(&self) -> MatrixType {
81        unsafe { sys::cusparseGetMatType(self.as_raw()).into() }
82    }
83
84    /// Sets the [`MatrixType`] field of this matrix descriptor.
85    ///
86    /// # Errors
87    ///
88    /// Returns an error if cuSPARSE rejects `matrix_type`.
89    pub fn set_matrix_type(&mut self, matrix_type: MatrixType) -> Result<()> {
90        unsafe {
91            try_ffi!(sys::cusparseSetMatType(self.as_raw(), matrix_type.into()))?;
92        }
93        Ok(())
94    }
95
96    /// Returns the [`FillMode`] field of this matrix descriptor.
97    pub fn fill_mode(&self) -> FillMode {
98        unsafe { sys::cusparseGetMatFillMode(self.as_raw()).into() }
99    }
100
101    /// Sets the [`FillMode`] field of this matrix descriptor.
102    ///
103    /// # Errors
104    ///
105    /// Returns an error if cuSPARSE rejects `fill_mode`.
106    pub fn set_fill_mode(&mut self, fill_mode: FillMode) -> Result<()> {
107        unsafe {
108            try_ffi!(sys::cusparseSetMatFillMode(self.as_raw(), fill_mode.into()))?;
109        }
110        Ok(())
111    }
112
113    /// Returns the [`DiagonalType`] field of this matrix descriptor.
114    pub fn diagonal_type(&self) -> DiagonalType {
115        unsafe { sys::cusparseGetMatDiagType(self.as_raw()).into() }
116    }
117
118    /// Sets the [`DiagonalType`] field of this matrix descriptor.
119    ///
120    /// # Errors
121    ///
122    /// Returns an error if cuSPARSE rejects `diagonal_type`.
123    pub fn set_diagonal_type(&mut self, diagonal_type: DiagonalType) -> Result<()> {
124        unsafe {
125            try_ffi!(sys::cusparseSetMatDiagType(
126                self.as_raw(),
127                diagonal_type.into(),
128            ))?;
129        }
130        Ok(())
131    }
132
133    /// Returns the [`IndexBase`] field of this matrix descriptor.
134    pub fn index_base(&self) -> IndexBase {
135        unsafe { sys::cusparseGetMatIndexBase(self.as_raw()).into() }
136    }
137
138    /// Sets the [`IndexBase`] field of this matrix descriptor.
139    ///
140    /// # Errors
141    ///
142    /// Returns an error if cuSPARSE rejects `index_base`.
143    pub fn set_index_base(&mut self, index_base: IndexBase) -> Result<()> {
144        unsafe {
145            try_ffi!(sys::cusparseSetMatIndexBase(
146                self.as_raw(),
147                index_base.into(),
148            ))?;
149        }
150        Ok(())
151    }
152
153    pub fn as_raw(&self) -> sys::cusparseMatDescr_t {
154        self.handle
155    }
156}
157
158impl Drop for MatrixDescriptor {
159    fn drop(&mut self) {
160        unsafe {
161            if let Err(err) = try_ffi!(sys::cusparseDestroyMatDescr(self.handle)) {
162                #[cfg(debug_assertions)]
163                eprintln!("failed to destroy cusparse matrix descriptor: {err}");
164            }
165        }
166    }
167}
168
169impl<'a> SparseMatrixDescriptor<'a> {
170    /// Initializes a sparse matrix descriptor in the CSR format.
171    ///
172    /// The descriptor borrows the row-offset, column-index, and value buffers by
173    /// pointer; those buffers must remain valid while the descriptor uses them.
174    pub fn create_csr<Offsets: IndexTypeLike, Columns: IndexTypeLike, T: DataTypeLike>(
175        rows: usize,
176        cols: usize,
177        nonzero_count: usize,
178        row_offsets: &'a mut DeviceMemory<Offsets>,
179        column_indices: &'a mut DeviceMemory<Columns>,
180        values: &'a mut DeviceMemory<T>,
181        index_base: IndexBase,
182    ) -> Result<Self> {
183        let mut raw = ptr::null_mut();
184        unsafe {
185            try_ffi!(sys::cusparseCreateCsr(
186                &raw mut raw,
187                to_i64(rows, "rows")?,
188                to_i64(cols, "cols")?,
189                to_i64(nonzero_count, "nonzero_count")?,
190                row_offsets.as_mut_ptr().cast(),
191                column_indices.as_mut_ptr().cast(),
192                values.as_mut_ptr().cast(),
193                Offsets::index_type().into(),
194                Columns::index_type().into(),
195                index_base.into(),
196                T::data_type().into(),
197            ))?;
198        }
199        ensure_spmat(raw)
200    }
201
202    /// Initializes a sparse matrix descriptor in the CSC format.
203    ///
204    /// The descriptor borrows the column-offset, row-index, and value buffers by
205    /// pointer; those buffers must remain valid while the descriptor uses them.
206    pub fn create_csc<Offsets: IndexTypeLike, Rows: IndexTypeLike, T: DataTypeLike>(
207        rows: usize,
208        cols: usize,
209        nonzero_count: usize,
210        column_offsets: &'a mut DeviceMemory<Offsets>,
211        row_indices: &'a mut DeviceMemory<Rows>,
212        values: &'a mut DeviceMemory<T>,
213        index_base: IndexBase,
214    ) -> Result<Self> {
215        let mut raw = ptr::null_mut();
216        unsafe {
217            try_ffi!(sys::cusparseCreateCsc(
218                &raw mut raw,
219                to_i64(rows, "rows")?,
220                to_i64(cols, "cols")?,
221                to_i64(nonzero_count, "nonzero_count")?,
222                column_offsets.as_mut_ptr().cast(),
223                row_indices.as_mut_ptr().cast(),
224                values.as_mut_ptr().cast(),
225                Offsets::index_type().into(),
226                Rows::index_type().into(),
227                index_base.into(),
228                T::data_type().into(),
229            ))?;
230        }
231        ensure_spmat(raw)
232    }
233
234    /// Initializes a sparse matrix descriptor in the COO format (Structure of Arrays layout).
235    ///
236    /// The descriptor borrows the row-index, column-index, and value buffers by
237    /// pointer; those buffers must remain valid while the descriptor uses them.
238    pub fn create_coo<I: IndexTypeLike, T: DataTypeLike>(
239        rows: usize,
240        cols: usize,
241        nonzero_count: usize,
242        row_indices: &'a mut DeviceMemory<I>,
243        column_indices: &'a mut DeviceMemory<I>,
244        values: &'a mut DeviceMemory<T>,
245        index_base: IndexBase,
246    ) -> Result<Self> {
247        let mut raw = ptr::null_mut();
248        unsafe {
249            try_ffi!(sys::cusparseCreateCoo(
250                &raw mut raw,
251                to_i64(rows, "rows")?,
252                to_i64(cols, "cols")?,
253                to_i64(nonzero_count, "nonzero_count")?,
254                row_indices.as_mut_ptr().cast(),
255                column_indices.as_mut_ptr().cast(),
256                values.as_mut_ptr().cast(),
257                I::index_type().into(),
258                index_base.into(),
259                T::data_type().into(),
260            ))?;
261        }
262        ensure_spmat(raw)
263    }
264
265    /// Initializes a sparse matrix descriptor for the Block Compressed Row (BSR) format.
266    ///
267    /// The descriptor borrows the block row-offset, column-index, and value
268    /// buffers by pointer; those buffers must remain valid while the descriptor uses them.
269    pub fn create_bsr<Offsets: IndexTypeLike, Columns: IndexTypeLike, T: DataTypeLike>(
270        block_rows: usize,
271        block_cols: usize,
272        block_nnz: usize,
273        row_block_size: usize,
274        col_block_size: usize,
275        row_offsets: &'a mut DeviceMemory<Offsets>,
276        column_indices: &'a mut DeviceMemory<Columns>,
277        values: &'a mut DeviceMemory<T>,
278        index_base: IndexBase,
279        order: Order,
280    ) -> Result<Self> {
281        let mut raw = ptr::null_mut();
282        unsafe {
283            try_ffi!(sys::cusparseCreateBsr(
284                &raw mut raw,
285                to_i64(block_rows, "block_rows")?,
286                to_i64(block_cols, "block_cols")?,
287                to_i64(block_nnz, "block_nnz")?,
288                to_i64(row_block_size, "row_block_size")?,
289                to_i64(col_block_size, "col_block_size")?,
290                row_offsets.as_mut_ptr().cast(),
291                column_indices.as_mut_ptr().cast(),
292                values.as_mut_ptr().cast(),
293                Offsets::index_type().into(),
294                Columns::index_type().into(),
295                index_base.into(),
296                T::data_type().into(),
297                order.into(),
298            ))?;
299        }
300        ensure_spmat(raw)
301    }
302
303    /// Initializes a sparse matrix descriptor for the Blocked-Ellpack (ELL) format.
304    ///
305    /// Blocked-ELL column indices are in the range `0..cols / ell_block_size`.
306    /// The array can contain `-1` values for indicating empty blocks.
307    /// The descriptor borrows the column-index and value buffers by pointer;
308    /// those buffers must remain valid while the descriptor uses them.
309    pub fn create_blocked_ell<I: IndexTypeLike, T: DataTypeLike>(
310        rows: usize,
311        cols: usize,
312        block_size: usize,
313        ell_cols: usize,
314        column_indices: &'a mut DeviceMemory<I>,
315        values: &'a mut DeviceMemory<T>,
316        index_base: IndexBase,
317    ) -> Result<Self> {
318        let mut raw = ptr::null_mut();
319        unsafe {
320            try_ffi!(sys::cusparseCreateBlockedEll(
321                &raw mut raw,
322                to_i64(rows, "rows")?,
323                to_i64(cols, "cols")?,
324                to_i64(block_size, "block_size")?,
325                to_i64(ell_cols, "ell_cols")?,
326                column_indices.as_mut_ptr().cast(),
327                values.as_mut_ptr().cast(),
328                I::index_type().into(),
329                index_base.into(),
330                T::data_type().into(),
331            ))?;
332        }
333        ensure_spmat(raw)
334    }
335
336    /// Initializes a sparse matrix descriptor for the Sliced Ellpack (SELL) format.
337    ///
338    /// The descriptor borrows the slice-offset, column-index, and value buffers
339    /// by pointer; those buffers must remain valid while the descriptor uses them.
340    pub fn create_sliced_ell<
341        SliceOffsets: IndexTypeLike,
342        Columns: IndexTypeLike,
343        T: DataTypeLike,
344    >(
345        rows: usize,
346        cols: usize,
347        nonzero_count: usize,
348        values_size: usize,
349        slice_size: usize,
350        slice_offsets: &'a mut DeviceMemory<SliceOffsets>,
351        column_indices: &'a mut DeviceMemory<Columns>,
352        values: &'a mut DeviceMemory<T>,
353        index_base: IndexBase,
354    ) -> Result<Self> {
355        let mut raw = ptr::null_mut();
356        unsafe {
357            try_ffi!(sys::cusparseCreateSlicedEll(
358                &raw mut raw,
359                to_i64(rows, "rows")?,
360                to_i64(cols, "cols")?,
361                to_i64(nonzero_count, "nonzero_count")?,
362                to_i64(values_size, "values_size")?,
363                to_i64(slice_size, "slice_size")?,
364                slice_offsets.as_mut_ptr().cast(),
365                column_indices.as_mut_ptr().cast(),
366                values.as_mut_ptr().cast(),
367                SliceOffsets::index_type().into(),
368                Columns::index_type().into(),
369                index_base.into(),
370                T::data_type().into(),
371            ))?;
372        }
373        ensure_spmat(raw)
374    }
375
376    /// Returns the format of this sparse matrix descriptor.
377    ///
378    /// # Errors
379    ///
380    /// Returns an error if cuSPARSE cannot report the sparse matrix format.
381    pub fn format(&self) -> Result<Format> {
382        let mut format = sys::cusparseFormat_t::CUSPARSE_FORMAT_CSR;
383        unsafe {
384            try_ffi!(sys::cusparseSpMatGetFormat(
385                self.as_raw_const(),
386                &raw mut format,
387            ))?;
388        }
389        Ok(format.into())
390    }
391
392    /// Returns the index base of this sparse matrix descriptor.
393    ///
394    /// # Errors
395    ///
396    /// Returns an error if cuSPARSE cannot report the sparse matrix index base.
397    pub fn index_base(&self) -> Result<IndexBase> {
398        let mut index_base = sys::cusparseIndexBase_t::CUSPARSE_INDEX_BASE_ZERO;
399        unsafe {
400            try_ffi!(sys::cusparseSpMatGetIndexBase(
401                self.as_raw_const(),
402                &raw mut index_base,
403            ))?;
404        }
405        Ok(index_base.into())
406    }
407
408    /// Returns the values pointer of this sparse matrix descriptor.
409    ///
410    /// # Errors
411    ///
412    /// Returns an error if cuSPARSE cannot report the values pointer.
413    pub fn values(&self) -> Result<DevicePtr> {
414        let mut values = ptr::null_mut();
415        unsafe {
416            try_ffi!(sys::cusparseSpMatGetValues(self.as_raw(), &raw mut values))?;
417        }
418        Ok(values.into())
419    }
420
421    /// Sets the values pointer of this sparse matrix descriptor.
422    ///
423    /// The pointed-to device buffer must remain valid while the descriptor uses it.
424    ///
425    /// # Errors
426    ///
427    /// Returns an error if cuSPARSE rejects the values pointer.
428    pub fn set_values<T: DataTypeLike>(&mut self, values: &'a mut DeviceMemory<T>) -> Result<()> {
429        unsafe {
430            try_ffi!(sys::cusparseSpMatSetValues(
431                self.as_raw(),
432                values.as_mut_ptr().cast(),
433            ))?;
434        }
435        Ok(())
436    }
437
438    /// Returns the shape of this sparse matrix descriptor.
439    ///
440    /// # Errors
441    ///
442    /// Returns an error if cuSPARSE cannot report the sparse matrix shape or if
443    /// the reported values cannot be represented as `usize`.
444    pub fn shape(&self) -> Result<SparseMatrixShape> {
445        let mut rows = 0_i64;
446        let mut cols = 0_i64;
447        let mut nnz = 0_i64;
448        unsafe {
449            try_ffi!(sys::cusparseSpMatGetSize(
450                self.as_raw_const(),
451                &raw mut rows,
452                &raw mut cols,
453                &raw mut nnz,
454            ))?;
455        }
456        Ok(SparseMatrixShape {
457            rows: to_usize(rows, "rows")?,
458            cols: to_usize(cols, "cols")?,
459            nonzero_count: to_usize(nnz, "nonzero_count")?,
460        })
461    }
462
463    /// Returns the batch count of this sparse matrix descriptor.
464    ///
465    /// # Errors
466    ///
467    /// Returns an error if cuSPARSE cannot report the strided batch settings.
468    pub fn strided_batch_count(&self) -> Result<i32> {
469        let mut batch_count = 0;
470        unsafe {
471            try_ffi!(sys::cusparseSpMatGetStridedBatch(
472                self.as_raw_const(),
473                &raw mut batch_count,
474            ))?;
475        }
476        Ok(batch_count)
477    }
478
479    /// Sets the batch count and batch stride fields of this COO sparse matrix descriptor.
480    ///
481    /// # Errors
482    ///
483    /// Returns an error if either argument cannot be represented by cuSPARSE or
484    /// if cuSPARSE rejects the strided batch settings.
485    pub fn set_coo_strided_batch(&mut self, batch_count: usize, batch_stride: usize) -> Result<()> {
486        unsafe {
487            try_ffi!(sys::cusparseCooSetStridedBatch(
488                self.as_raw(),
489                to_i32(batch_count, "batch_count")?,
490                to_i64(batch_stride, "batch_stride")?,
491            ))?;
492        }
493        Ok(())
494    }
495
496    /// Sets the batch count and batch stride fields of this CSR sparse matrix descriptor.
497    pub fn set_csr_strided_batch(
498        &mut self,
499        batch_count: usize,
500        offsets_batch_stride: usize,
501        columns_values_batch_stride: usize,
502    ) -> Result<()> {
503        unsafe {
504            try_ffi!(sys::cusparseCsrSetStridedBatch(
505                self.as_raw(),
506                to_i32(batch_count, "batch_count")?,
507                to_i64(offsets_batch_stride, "offsets_batch_stride")?,
508                to_i64(columns_values_batch_stride, "columns_values_batch_stride")?,
509            ))?;
510        }
511        Ok(())
512    }
513
514    /// Sets the batch count and batch stride fields of this BSR sparse matrix descriptor.
515    pub fn set_bsr_strided_batch(
516        &mut self,
517        batch_count: usize,
518        offsets_batch_stride: usize,
519        columns_batch_stride: usize,
520        values_batch_stride: usize,
521    ) -> Result<()> {
522        unsafe {
523            try_ffi!(sys::cusparseBsrSetStridedBatch(
524                self.as_raw(),
525                to_i32(batch_count, "batch_count")?,
526                to_i64(offsets_batch_stride, "offsets_batch_stride")?,
527                to_i64(columns_batch_stride, "columns_batch_stride")?,
528                to_i64(values_batch_stride, "values_batch_stride")?,
529            ))?;
530        }
531        Ok(())
532    }
533
534    pub fn fill_mode(&self) -> Result<FillMode> {
535        let raw: sys::cusparseFillMode_t = self.raw_attribute(SparseMatrixAttribute::FillMode)?;
536        Ok(raw.into())
537    }
538
539    pub fn set_fill_mode(&mut self, fill_mode: FillMode) -> Result<()> {
540        self.set_attribute_raw::<sys::cusparseFillMode_t>(
541            SparseMatrixAttribute::FillMode,
542            fill_mode.into(),
543        )
544    }
545
546    pub fn diagonal_type(&self) -> Result<DiagonalType> {
547        let raw: sys::cusparseDiagType_t =
548            self.raw_attribute(SparseMatrixAttribute::DiagonalType)?;
549        Ok(raw.into())
550    }
551
552    pub fn set_diagonal_type(&mut self, diagonal_type: DiagonalType) -> Result<()> {
553        self.set_attribute_raw::<sys::cusparseDiagType_t>(
554            SparseMatrixAttribute::DiagonalType,
555            diagonal_type.into(),
556        )
557    }
558
559    /// Sets the pointers of this CSR sparse matrix descriptor.
560    ///
561    /// The pointed-to device buffers must remain valid while the descriptor uses them.
562    pub fn set_csr_pointers(
563        &mut self,
564        row_offsets: &'a mut DeviceMemory<impl IndexTypeLike>,
565        column_indices: &'a mut DeviceMemory<impl IndexTypeLike>,
566        values: &'a mut DeviceMemory<impl DataTypeLike>,
567    ) -> Result<()> {
568        unsafe {
569            try_ffi!(sys::cusparseCsrSetPointers(
570                self.as_raw(),
571                row_offsets.as_mut_ptr().cast(),
572                column_indices.as_mut_ptr().cast(),
573                values.as_mut_ptr().cast(),
574            ))?;
575        }
576        Ok(())
577    }
578
579    /// Sets the pointers of this CSC sparse matrix descriptor.
580    ///
581    /// The pointed-to device buffers must remain valid while the descriptor uses them.
582    pub fn set_csc_pointers(
583        &mut self,
584        column_offsets: &'a mut DeviceMemory<impl IndexTypeLike>,
585        row_indices: &'a mut DeviceMemory<impl IndexTypeLike>,
586        values: &'a mut DeviceMemory<impl DataTypeLike>,
587    ) -> Result<()> {
588        unsafe {
589            try_ffi!(sys::cusparseCscSetPointers(
590                self.as_raw(),
591                column_offsets.as_mut_ptr().cast(),
592                row_indices.as_mut_ptr().cast(),
593                values.as_mut_ptr().cast(),
594            ))?;
595        }
596        Ok(())
597    }
598
599    /// Sets the pointers of this COO sparse matrix descriptor.
600    ///
601    /// The pointed-to device buffers must remain valid while the descriptor uses them.
602    pub fn set_coo_pointers(
603        &mut self,
604        row_indices: &'a mut DeviceMemory<impl IndexTypeLike>,
605        column_indices: &'a mut DeviceMemory<impl IndexTypeLike>,
606        values: &'a mut DeviceMemory<impl DataTypeLike>,
607    ) -> Result<()> {
608        unsafe {
609            try_ffi!(sys::cusparseCooSetPointers(
610                self.as_raw(),
611                row_indices.as_mut_ptr().cast(),
612                column_indices.as_mut_ptr().cast(),
613                values.as_mut_ptr().cast(),
614            ))?;
615        }
616        Ok(())
617    }
618
619    pub fn as_raw(&self) -> sys::cusparseSpMatDescr_t {
620        self.handle
621    }
622
623    pub fn as_raw_const(&self) -> sys::cusparseConstSpMatDescr_t {
624        self.handle.cast_const()
625    }
626
627    fn raw_attribute<T: Copy>(&self, attribute: SparseMatrixAttribute) -> Result<T> {
628        let mut value = MaybeUninit::<T>::uninit();
629        unsafe {
630            try_ffi!(sys::cusparseSpMatGetAttribute(
631                self.as_raw_const(),
632                attribute.into(),
633                value.as_mut_ptr().cast(),
634                u64::try_from(size_of::<T>()).map_err(|_| Error::OutOfRange {
635                    name: "attribute size".into(),
636                })?,
637            ))?;
638            Ok(value.assume_init())
639        }
640    }
641
642    fn set_attribute_raw<T>(&mut self, attribute: SparseMatrixAttribute, value: T) -> Result<()> {
643        let mut value = value;
644        unsafe {
645            try_ffi!(sys::cusparseSpMatSetAttribute(
646                self.as_raw(),
647                attribute.into(),
648                &raw mut value as *mut _,
649                u64::try_from(size_of::<T>()).map_err(|_| Error::OutOfRange {
650                    name: "attribute size".into(),
651                })?,
652            ))?;
653        }
654        Ok(())
655    }
656}
657
658impl Drop for SparseMatrixDescriptor<'_> {
659    fn drop(&mut self) {
660        unsafe {
661            if let Err(err) = try_ffi!(sys::cusparseDestroySpMat(self.as_raw_const())) {
662                #[cfg(debug_assertions)]
663                eprintln!("failed to destroy cusparse sparse matrix descriptor: {err}");
664            }
665        }
666    }
667}
668
669impl<'a> DenseMatrixDescriptor<'a> {
670    /// Initializes a dense matrix descriptor.
671    ///
672    /// The descriptor borrows the device value buffer by pointer; the buffer must
673    /// remain valid while the descriptor uses it.
674    pub fn create<T: DataTypeLike>(
675        rows: usize,
676        cols: usize,
677        leading_dimension: usize,
678        values: &'a mut DeviceMemory<T>,
679        order: Order,
680    ) -> Result<Self> {
681        let mut handle = ptr::null_mut();
682        unsafe {
683            try_ffi!(sys::cusparseCreateDnMat(
684                &raw mut handle,
685                to_i64(rows, "rows")?,
686                to_i64(cols, "cols")?,
687                to_i64(leading_dimension, "leading_dimension")?,
688                values.as_mut_ptr().cast(),
689                T::data_type().into(),
690                order.into(),
691            ))?;
692        }
693
694        if handle.is_null() {
695            return Err(Error::NullHandle);
696        }
697
698        Ok(Self {
699            handle,
700            _buffers: PhantomData,
701        })
702    }
703
704    /// Returns the values pointer of this dense matrix descriptor.
705    ///
706    /// # Errors
707    ///
708    /// Returns an error if cuSPARSE cannot report the values pointer.
709    pub fn values(&self) -> Result<DevicePtr> {
710        let mut values = ptr::null_mut();
711        unsafe {
712            try_ffi!(sys::cusparseDnMatGetValues(self.as_raw(), &raw mut values))?;
713        }
714        Ok(values.into())
715    }
716
717    /// Sets the values pointer of this dense matrix descriptor.
718    ///
719    /// The pointed-to device buffer must remain valid while the descriptor uses it.
720    ///
721    /// # Errors
722    ///
723    /// Returns an error if cuSPARSE rejects the values pointer.
724    pub fn set_values<T: DataTypeLike>(&mut self, values: &'a mut DeviceMemory<T>) -> Result<()> {
725        unsafe {
726            try_ffi!(sys::cusparseDnMatSetValues(
727                self.as_raw(),
728                values.as_mut_ptr().cast(),
729            ))?;
730        }
731        Ok(())
732    }
733
734    /// Returns the fields of this dense matrix descriptor.
735    ///
736    /// # Errors
737    ///
738    /// Returns an error if cuSPARSE cannot report the descriptor fields or if
739    /// reported dimensions cannot be represented as `usize`.
740    pub fn info(&self) -> Result<DenseMatrixInfo> {
741        let mut rows = 0_i64;
742        let mut cols = 0_i64;
743        let mut leading_dimension = 0_i64;
744        let mut values = ptr::null_mut();
745        let mut value_type = MaybeUninit::uninit();
746        let mut order = sys::cusparseOrder_t::CUSPARSE_ORDER_ROW;
747        unsafe {
748            try_ffi!(sys::cusparseDnMatGet(
749                self.as_raw(),
750                &raw mut rows,
751                &raw mut cols,
752                &raw mut leading_dimension,
753                &raw mut values,
754                value_type.as_mut_ptr(),
755                &raw mut order,
756            ))?;
757        }
758        Ok(DenseMatrixInfo {
759            rows: to_usize(rows, "rows")?,
760            cols: to_usize(cols, "cols")?,
761            leading_dimension: to_usize(leading_dimension, "leading_dimension")?,
762            values: values.into(),
763            order: order.into(),
764        })
765    }
766
767    /// Sets the number of batches and the batch stride of this dense matrix descriptor.
768    ///
769    /// # Errors
770    ///
771    /// Returns an error if either argument cannot be represented by cuSPARSE or
772    /// if cuSPARSE rejects the strided batch settings.
773    pub fn set_strided_batch(&mut self, batch_count: usize, batch_stride: usize) -> Result<()> {
774        unsafe {
775            try_ffi!(sys::cusparseDnMatSetStridedBatch(
776                self.as_raw(),
777                to_i32(batch_count, "batch_count")?,
778                to_i64(batch_stride, "batch_stride")?,
779            ))?;
780        }
781        Ok(())
782    }
783
784    /// Returns the number of batches and the batch stride of this dense matrix descriptor.
785    ///
786    /// # Errors
787    ///
788    /// Returns an error if cuSPARSE cannot report the strided batch settings or
789    /// if the reported stride cannot be represented as `usize`.
790    pub fn strided_batch(&self) -> Result<(i32, usize)> {
791        let mut batch_count = 0;
792        let mut batch_stride = 0_i64;
793        unsafe {
794            try_ffi!(sys::cusparseDnMatGetStridedBatch(
795                self.as_raw_const(),
796                &raw mut batch_count,
797                &raw mut batch_stride,
798            ))?;
799        }
800        Ok((batch_count, to_usize(batch_stride, "batch_stride")?))
801    }
802
803    pub fn as_raw(&self) -> sys::cusparseDnMatDescr_t {
804        self.handle
805    }
806
807    pub fn as_raw_const(&self) -> sys::cusparseConstDnMatDescr_t {
808        self.handle.cast_const()
809    }
810}
811
812impl Drop for DenseMatrixDescriptor<'_> {
813    fn drop(&mut self) {
814        unsafe {
815            if let Err(err) = try_ffi!(sys::cusparseDestroyDnMat(self.as_raw_const())) {
816                #[cfg(debug_assertions)]
817                eprintln!("failed to destroy cusparse dense matrix descriptor: {err}");
818            }
819        }
820    }
821}
822
823fn ensure_spmat<'a>(handle: sys::cusparseSpMatDescr_t) -> Result<SparseMatrixDescriptor<'a>> {
824    if handle.is_null() {
825        Err(Error::NullHandle)
826    } else {
827        Ok(SparseMatrixDescriptor {
828            handle,
829            _buffers: PhantomData,
830        })
831    }
832}
833
834/// Converts a sparse matrix in CSR format into a sparse matrix in CSC format.
835/// The resulting matrix can also be seen as the transpose of the original sparse matrix.
836/// This can also convert a matrix in CSC format into a matrix in CSR format.
837///
838/// This operation requires extra storage proportional to `nonzero_count`.
839/// It always produces the same matrix values in the destination format.
840///
841/// The operation executes asynchronously with respect to the host and may return before the result is ready.
842///
843/// [`csr_to_csc_buffer_size`] returns the size of the workspace needed by [`csr_to_csc`].
844/// Pass a buffer of this size to [`csr_to_csc`].
845///
846/// If `nonzero_count` is `0`, then the CSR column indices, CSR values, CSC
847/// values, and CSC row indices may be zero-length buffers.
848/// In this case, all CSC column pointers are initialized to the configured index base.
849///
850/// If `m` or `n` is `0`, the pointers are not checked and the operation succeeds with `Ok`.
851///
852/// [`csr_to_csc`] supports the following data types:
853///
854/// | `X`/`Y` | Notes |
855/// | --- | --- |
856/// | [`DataType::I8`](singe_cuda::data_type::DataType::I8) |  |
857/// | [`DataType::F16`](singe_cuda::data_type::DataType::F16) |  |
858/// | [`DataType::Bf16`](singe_cuda::data_type::DataType::Bf16) |  |
859/// | [`DataType::F32`](singe_cuda::data_type::DataType::F32) |  |
860/// | [`DataType::F64`](singe_cuda::data_type::DataType::F64) |  |
861/// | [`DataType::ComplexF16`](singe_cuda::data_type::DataType::ComplexF16) | Deprecated. |
862/// | [`DataType::ComplexBf16`](singe_cuda::data_type::DataType::ComplexBf16) | Deprecated. |
863/// | [`DataType::ComplexF32`](singe_cuda::data_type::DataType::ComplexF32) |  |
864/// | [`DataType::ComplexF64`](singe_cuda::data_type::DataType::ComplexF64) |  |
865///
866/// [`csr_to_csc`] supports the following algorithms ([`CsrToCscAlgorithm`]):
867///
868/// | Algorithm | Notes |
869/// | --- | --- |
870/// | [`CsrToCscAlgorithm::Default`] | Default algorithm. |
871///
872/// | Action | Notes |
873/// | --- | --- |
874/// | [`Action::Symbolic`](crate::types::Action::Symbolic) | Compute the structure of the CSC output matrix (offset, row indices). |
875/// | [`Action::Numeric`](crate::types::Action::Numeric) | Compute the structure of the CSC output matrix and copy the values. |
876///
877/// [`csr_to_csc`] has the following properties:
878///
879/// * Requires no extra storage.
880/// * Supports asynchronous execution.
881///
882/// [`csr_to_csc`] supports the following optimizations:
883///
884/// * CUDA graph capture.
885/// * Hardware Memory Compression.
886pub fn csr_to_csc<T: DataTypeLike>(
887    ctx: &Context,
888    row_count: usize,
889    col_count: usize,
890    nonzero_count: usize,
891    csr_values: &DeviceMemory<T>,
892    csr_row_offsets: &DeviceMemory<i32>,
893    csr_column_indices: &DeviceMemory<i32>,
894    csc_values: &mut DeviceMemory<T>,
895    csc_column_offsets: &mut DeviceMemory<i32>,
896    csc_row_indices: &mut DeviceMemory<i32>,
897    action: crate::types::Action,
898    index_base: IndexBase,
899    algorithm: CsrToCscAlgorithm,
900    external_buffer: &mut DeviceMemory<u8>,
901) -> Result<()> {
902    ctx.bind()?;
903    unsafe {
904        try_ffi!(sys::cusparseCsr2cscEx2(
905            ctx.as_raw(),
906            to_i32(row_count, "row_count")?,
907            to_i32(col_count, "col_count")?,
908            to_i32(nonzero_count, "nonzero_count")?,
909            csr_values.as_ptr().cast(),
910            csr_row_offsets.as_ptr() as _,
911            csr_column_indices.as_ptr() as _,
912            csc_values.as_mut_ptr().cast(),
913            csc_column_offsets.as_mut_ptr(),
914            csc_row_indices.as_mut_ptr(),
915            T::data_type().into(),
916            action.into(),
917            index_base.into(),
918            algorithm.into(),
919            external_buffer.as_mut_ptr().cast(),
920        ))?;
921    }
922    Ok(())
923}
924
925/// Converts `matrix_a` in CSR, CSC, or COO format into its dense representation `matrix_b`.
926/// Blocked-ELL is not currently supported.
927///
928/// [`sparse_to_dense_buffer_size`] returns the size of the workspace needed by [`sparse_to_dense`].
929///
930/// [`sparse_to_dense`] supports the following index type for representing `matrix_a`:
931///
932/// * 32-bit indices ([`IndexType::I32`](crate::types::IndexType::I32))
933/// * 64-bit indices ([`IndexType::I64`](crate::types::IndexType::I64))
934///
935/// [`sparse_to_dense`] supports the following data types:
936///
937/// | `A`/`B` | Notes |
938/// | --- | --- |
939/// | [`DataType::I8`](singe_cuda::data_type::DataType::I8) |  |
940/// | [`DataType::F16`](singe_cuda::data_type::DataType::F16) |  |
941/// | [`DataType::Bf16`](singe_cuda::data_type::DataType::Bf16) |  |
942/// | [`DataType::F32`](singe_cuda::data_type::DataType::F32) |  |
943/// | [`DataType::F64`](singe_cuda::data_type::DataType::F64) |  |
944/// | [`DataType::ComplexF16`](singe_cuda::data_type::DataType::ComplexF16) | Deprecated. |
945/// | [`DataType::ComplexBf16`](singe_cuda::data_type::DataType::ComplexBf16) | Deprecated. |
946/// | [`DataType::ComplexF32`](singe_cuda::data_type::DataType::ComplexF32) |  |
947/// | [`DataType::ComplexF64`](singe_cuda::data_type::DataType::ComplexF64) |  |
948///
949/// [`sparse_to_dense`] supports the following algorithm:
950///
951/// | Algorithm | Notes |
952/// | --- | --- |
953/// | [`SparseToDenseAlgorithm::Default`] | Default algorithm. |
954///
955/// [`sparse_to_dense`] has the following properties:
956///
957/// * Requires no extra storage.
958/// * Supports asynchronous execution.
959/// * Provides deterministic (bitwise) results for each run.
960/// * Allows the indices of `matrix_a` to be unsorted.
961///
962/// [`sparse_to_dense`] supports the following optimizations:
963///
964/// * CUDA graph capture.
965/// * Hardware Memory Compression.
966pub fn sparse_to_dense(
967    ctx: &Context,
968    matrix_a: &SparseMatrixDescriptor,
969    matrix_b: &mut DenseMatrixDescriptor,
970    algorithm: SparseToDenseAlgorithm,
971    external_buffer: Option<DevicePtr>,
972) -> Result<()> {
973    ctx.bind()?;
974    unsafe {
975        try_ffi!(sys::cusparseSparseToDense(
976            ctx.as_raw(),
977            matrix_a.as_raw_const(),
978            matrix_b.as_raw(),
979            algorithm.into(),
980            external_buffer.unwrap_or(DevicePtr::null()).as_ptr() as _,
981        ))?;
982    }
983    Ok(())
984}
985
986pub fn dense_to_sparse_buffer_size(
987    ctx: &Context,
988    matrix_a: &DenseMatrixDescriptor,
989    matrix_b: &mut SparseMatrixDescriptor,
990    algorithm: DenseToSparseAlgorithm,
991) -> Result<usize> {
992    ctx.bind()?;
993    let mut size = 0;
994    unsafe {
995        try_ffi!(sys::cusparseDenseToSparse_bufferSize(
996            ctx.as_raw(),
997            matrix_a.as_raw_const(),
998            matrix_b.as_raw(),
999            algorithm.into(),
1000            &raw mut size,
1001        ))?;
1002    }
1003    to_usize(size, "dense-to-sparse buffer size")
1004}
1005
1006pub fn dense_to_sparse_analysis(
1007    ctx: &Context,
1008    matrix_a: &DenseMatrixDescriptor,
1009    matrix_b: &mut SparseMatrixDescriptor,
1010    algorithm: DenseToSparseAlgorithm,
1011    external_buffer: Option<DevicePtr>,
1012) -> Result<()> {
1013    ctx.bind()?;
1014    unsafe {
1015        try_ffi!(sys::cusparseDenseToSparse_analysis(
1016            ctx.as_raw(),
1017            matrix_a.as_raw_const(),
1018            matrix_b.as_raw(),
1019            algorithm.into(),
1020            external_buffer.unwrap_or(DevicePtr::null()).as_ptr() as _,
1021        ))?;
1022    }
1023    Ok(())
1024}
1025
1026pub fn dense_to_sparse_convert(
1027    ctx: &Context,
1028    matrix_a: &DenseMatrixDescriptor,
1029    matrix_b: &mut SparseMatrixDescriptor,
1030    algorithm: DenseToSparseAlgorithm,
1031    external_buffer: Option<DevicePtr>,
1032) -> Result<()> {
1033    ctx.bind()?;
1034    unsafe {
1035        try_ffi!(sys::cusparseDenseToSparse_convert(
1036            ctx.as_raw(),
1037            matrix_a.as_raw_const(),
1038            matrix_b.as_raw(),
1039            algorithm.into(),
1040            external_buffer.unwrap_or(DevicePtr::null()).as_ptr() as _,
1041        ))?;
1042    }
1043    Ok(())
1044}
1045
1046pub fn csr_to_csc_buffer_size<T: DataTypeLike>(
1047    ctx: &Context,
1048    row_count: usize,
1049    col_count: usize,
1050    nonzero_count: usize,
1051    csr_values: &DeviceMemory<T>,
1052    csr_row_offsets: &DeviceMemory<i32>,
1053    csr_column_indices: &DeviceMemory<i32>,
1054    csc_values: &mut DeviceMemory<T>,
1055    csc_column_offsets: &mut DeviceMemory<i32>,
1056    csc_row_indices: &mut DeviceMemory<i32>,
1057    action: crate::types::Action,
1058    index_base: IndexBase,
1059    algorithm: CsrToCscAlgorithm,
1060) -> Result<usize> {
1061    ctx.bind()?;
1062    let mut size = 0;
1063    unsafe {
1064        try_ffi!(sys::cusparseCsr2cscEx2_bufferSize(
1065            ctx.as_raw(),
1066            to_i32(row_count, "row_count")?,
1067            to_i32(col_count, "col_count")?,
1068            to_i32(nonzero_count, "nonzero_count")?,
1069            csr_values.as_ptr().cast(),
1070            csr_row_offsets.as_ptr() as _,
1071            csr_column_indices.as_ptr() as _,
1072            csc_values.as_mut_ptr().cast(),
1073            csc_column_offsets.as_mut_ptr(),
1074            csc_row_indices.as_mut_ptr(),
1075            T::data_type().into(),
1076            action.into(),
1077            index_base.into(),
1078            algorithm.into(),
1079            &raw mut size,
1080        ))?;
1081    }
1082    to_usize(size, "csr-to-csc buffer size")
1083}
1084
1085/// Converts compressed row pointers in CSR format into uncompressed row indices in COO format.
1086///
1087/// It can also be used to convert the array containing the compressed column indices (corresponding to CSC format) into an array of uncompressed column indices (corresponding to COO format).
1088///
1089/// * Requires no extra storage.
1090/// * Supports asynchronous execution.
1091/// * Supports CUDA graph capture.
1092pub fn csr_to_coo(
1093    ctx: &Context,
1094    csr_row_offsets: &DeviceMemory<i32>,
1095    nonzero_count: usize,
1096    row_count: usize,
1097    coo_row_indices: &mut DeviceMemory<i32>,
1098    index_base: IndexBase,
1099) -> Result<()> {
1100    ctx.bind()?;
1101    if csr_row_offsets.len() < row_count.saturating_add(1) || coo_row_indices.len() < nonzero_count
1102    {
1103        return Err(Error::OutOfRange {
1104            name: "csr/coo buffer length".into(),
1105        });
1106    }
1107    unsafe {
1108        try_ffi!(sys::cusparseXcsr2coo(
1109            ctx.as_raw(),
1110            csr_row_offsets.as_ptr() as _,
1111            to_i32(nonzero_count, "nonzero_count")?,
1112            to_i32(row_count, "row_count")?,
1113            coo_row_indices.as_mut_ptr(),
1114            index_base.into(),
1115        ))?;
1116    }
1117    Ok(())
1118}
1119
1120pub fn sparse_to_dense_buffer_size(
1121    ctx: &Context,
1122    matrix_a: &SparseMatrixDescriptor,
1123    matrix_b: &mut DenseMatrixDescriptor,
1124    algorithm: SparseToDenseAlgorithm,
1125) -> Result<usize> {
1126    ctx.bind()?;
1127    let mut size = 0;
1128    unsafe {
1129        try_ffi!(sys::cusparseSparseToDense_bufferSize(
1130            ctx.as_raw(),
1131            matrix_a.as_raw_const(),
1132            matrix_b.as_raw(),
1133            algorithm.into(),
1134            &raw mut size,
1135        ))?;
1136    }
1137    to_usize(size, "sparse-to-dense buffer size")
1138}
1139
1140/// Converts uncompressed row indices in COO format into compressed row pointers in CSR format.
1141///
1142/// It can also be used to convert the array containing the uncompressed column indices (corresponding to COO format) into an array of column pointers (corresponding to CSC format).
1143///
1144/// * Requires no extra storage.
1145/// * Supports asynchronous execution.
1146/// * Supports CUDA graph capture.
1147pub fn coo_to_csr(
1148    ctx: &Context,
1149    coo_row_indices: &DeviceMemory<i32>,
1150    row_count: usize,
1151    csr_row_offsets: &mut DeviceMemory<i32>,
1152    index_base: IndexBase,
1153) -> Result<()> {
1154    ctx.bind()?;
1155    let nonzero_count = coo_row_indices.len();
1156    if csr_row_offsets.len() < row_count.saturating_add(1) {
1157        return Err(Error::OutOfRange {
1158            name: "csr_row_offsets length".into(),
1159        });
1160    }
1161    unsafe {
1162        try_ffi!(sys::cusparseXcoo2csr(
1163            ctx.as_raw(),
1164            coo_row_indices.as_ptr() as _,
1165            to_i32(nonzero_count, "nonzero_count")?,
1166            to_i32(row_count, "row_count")?,
1167            csr_row_offsets.as_mut_ptr(),
1168            index_base.into(),
1169        ))?;
1170    }
1171    Ok(())
1172}
singe_cusparse/matrix.rs

singe_cusparse/
matrix.rs
