baracuda_cusparse/
lib.rs

1//! Safe Rust wrappers for NVIDIA cuSPARSE.
2//!
3//! Covers the modern generic-API surface: `Handle`, `SpMat` (CSR/CSC/COO/BSR),
4//! `DnMat`, `DnVec`, and the family of op algorithms — SpMV, SpMM, SpGEMM,
5//! SpSV, SpSM, SDDMM — plus CSR↔CSC and sparse↔dense conversions, and the
6//! sparse BLAS-1 helpers (`axpby`, `gather`, `scatter`, `rot`).
7//!
8//! All matrix/vector descriptors borrow the underlying [`DeviceBuffer`]s and
9//! tie their lifetime to them so the buffers can't be freed while cuSPARSE
10//! is still holding references.
11//!
12//! # Examples
13//!
14//! **SpMV (sparse-times-dense vector):** `y = A · x` for a 4×4 CSR matrix
15//! with 6 non-zeros.
16//!
17//! ```no_run
18//! use baracuda_cusparse::{spmv, spmv_buffer_size, DnVec, Handle, Op, SpMVAlg, SpMat};
19//! use baracuda_driver::{Context, Device, DeviceBuffer};
20//!
21//! # fn demo() -> Result<(), Box<dyn std::error::Error>> {
22//! baracuda_driver::init()?;
23//! let device = Device::get(0)?;
24//! let ctx = Context::new(&device)?;
25//! let handle = Handle::new()?;
26//!
27//! // A = [[1, 0, 0, 2], [0, 3, 0, 0], [0, 0, 0, 4], [5, 0, 6, 0]], nnz = 6
28//! let mut row_offsets = DeviceBuffer::from_slice(&ctx, &[0i32, 2, 3, 4, 6])?;
29//! let mut col_indices = DeviceBuffer::from_slice(&ctx, &[0i32, 3, 1, 3, 0, 2])?;
30//! let mut values     = DeviceBuffer::from_slice(&ctx, &[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0])?;
31//! let mut x          = DeviceBuffer::from_slice(&ctx, &[1.0f32, 2.0, 3.0, 4.0])?;
32//! let mut y: DeviceBuffer<f32> = DeviceBuffer::zeros(&ctx, 4)?;
33//!
34//! let a = SpMat::csr(4, 4, 6, &mut row_offsets, &mut col_indices, &mut values)?;
35//! let x_d = DnVec::new(&mut x)?;
36//! let mut y_d = DnVec::new(&mut y)?;
37//!
38//! let bytes = spmv_buffer_size::<f32>(
39//!     &handle, Op::N, &1.0, &a, &x_d, &0.0, &y_d, SpMVAlg::Default)?;
40//! let mut workspace: DeviceBuffer<u8> = DeviceBuffer::new(&ctx, bytes.max(1))?;
41//! spmv::<f32>(&handle, Op::N, &1.0, &a, &x_d, &0.0, &mut y_d, SpMVAlg::Default, &mut workspace)?;
42//! # Ok(()) }
43//! ```
44//!
45//! **COO descriptor build** — feed an unsorted `(row, col, value)` triple
46//! stream into cuSPARSE without pre-converting to CSR:
47//!
48//! ```no_run
49//! use baracuda_cusparse::SpMat;
50//! use baracuda_driver::{Context, Device, DeviceBuffer};
51//!
52//! # fn demo() -> Result<(), Box<dyn std::error::Error>> {
53//! baracuda_driver::init()?;
54//! let device = Device::get(0)?;
55//! let ctx = Context::new(&device)?;
56//!
57//! // Same 4×4 matrix as above, expressed as (row, col, value) triples.
58//! let mut row_idx = DeviceBuffer::from_slice(&ctx, &[0i32, 0, 1, 2, 3, 3])?;
59//! let mut col_idx = DeviceBuffer::from_slice(&ctx, &[0i32, 3, 1, 3, 0, 2])?;
60//! let mut values  = DeviceBuffer::from_slice(&ctx, &[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0])?;
61//!
62//! let _coo = SpMat::coo(4, 4, 6, &mut row_idx, &mut col_idx, &mut values)?;
63//! // Pass `_coo` to `spmv`, `spmm`, or `sparse_to_dense` just like a CSR descriptor.
64//! # Ok(()) }
65//! ```
66//!
67//! **SpMM (sparse-times-dense matrix):** `C = A · B` where `A` is sparse,
68//! `B` and `C` are dense column-major.
69//!
70//! ```no_run
71//! use baracuda_cusparse::{spmm, spmm_buffer_size, DnMat, Handle, Op, Order, SpMMAlg, SpMat};
72//! use baracuda_driver::{Context, Device, DeviceBuffer};
73//!
74//! # fn demo() -> Result<(), Box<dyn std::error::Error>> {
75//! baracuda_driver::init()?;
76//! let device = Device::get(0)?;
77//! let ctx = Context::new(&device)?;
78//! let handle = Handle::new()?;
79//!
80//! // 4×4 CSR A (same as the SpMV example) times a 4×2 dense B → 4×2 dense C.
81//! let mut row_offsets = DeviceBuffer::from_slice(&ctx, &[0i32, 2, 3, 4, 6])?;
82//! let mut col_indices = DeviceBuffer::from_slice(&ctx, &[0i32, 3, 1, 3, 0, 2])?;
83//! let mut values     = DeviceBuffer::from_slice(&ctx, &[1.0f32, 2.0, 3.0, 4.0, 5.0, 6.0])?;
84//! let mut b_data: DeviceBuffer<f32> = DeviceBuffer::zeros(&ctx, 4 * 2)?;
85//! let mut c_data: DeviceBuffer<f32> = DeviceBuffer::zeros(&ctx, 4 * 2)?;
86//!
87//! let a = SpMat::csr(4, 4, 6, &mut row_offsets, &mut col_indices, &mut values)?;
88//! let b = DnMat::new(4, 2, /* ld */ 4, Order::Col, &mut b_data)?;
89//! let mut c = DnMat::new(4, 2, /* ld */ 4, Order::Col, &mut c_data)?;
90//!
91//! let bytes = spmm_buffer_size::<f32>(
92//!     &handle, Op::N, Op::N, &1.0, &a, &b, &0.0, &c, SpMMAlg::Default)?;
93//! let mut workspace: DeviceBuffer<u8> = DeviceBuffer::new(&ctx, bytes.max(1))?;
94//! spmm::<f32>(&handle, Op::N, Op::N, &1.0, &a, &b, &0.0, &mut c, SpMMAlg::Default,
95//!     &mut workspace)?;
96//! # Ok(()) }
97//! ```
98
99#![warn(missing_debug_implementations)]
100
101use core::ffi::c_void;
102use std::marker::PhantomData;
103
104use baracuda_cusparse_sys::{
105    cudaDataType, cusparse, cusparseDiagType_t, cusparseDnMatDescr_t, cusparseDnVecDescr_t,
106    cusparseFillMode_t, cusparseHandle_t, cusparseIndexBase_t, cusparseIndexType_t,
107    cusparseOperation_t, cusparseOrder_t, cusparseSpGEMMDescr_t, cusparseSpMatAttribute_t,
108    cusparseSpMatDescr_t, cusparseSpSMDescr_t, cusparseSpSVDescr_t, cusparseStatus_t,
109};
110use baracuda_driver::{DeviceBuffer, Stream};
111use baracuda_types::{Complex32, Complex64};
112
113pub use baracuda_cusparse_sys::{
114    cusparseCsr2CscAlg_t as Csr2CscAlg, cusparseIndexBase_t as IndexBase,
115    cusparseSDDMMAlg_t as SDDMMAlg, cusparseSpGEMMAlg_t as SpGEMMAlg, cusparseSpMMAlg_t as SpMMAlg,
116    cusparseSpMVAlg_t as SpMVAlg, cusparseSpSMAlg_t as SpSMAlg, cusparseSpSVAlg_t as SpSVAlg,
117};
118
119/// Error type for cuSPARSE operations.
120pub type Error = baracuda_core::Error<cusparseStatus_t>;
121/// Result alias.
122pub type Result<T, E = Error> = core::result::Result<T, E>;
123
124#[inline]
125fn check(status: cusparseStatus_t) -> Result<()> {
126    Error::check(status)
127}
128
129// ---- scalar <-> cudaDataType --------------------------------------------
130
131/// Types that cuSPARSE's generic API accepts as element / compute type.
132pub trait SparseScalar: sealed::Sealed + Copy + 'static {
133    /// cuSPARSE / cuBLAS element-type tag.
134    fn data_type() -> cudaDataType;
135}
136
137impl SparseScalar for f32 {
138    fn data_type() -> cudaDataType {
139        cudaDataType::R_32F
140    }
141}
142impl SparseScalar for f64 {
143    fn data_type() -> cudaDataType {
144        cudaDataType::R_64F
145    }
146}
147impl SparseScalar for Complex32 {
148    fn data_type() -> cudaDataType {
149        cudaDataType::C_32F
150    }
151}
152impl SparseScalar for Complex64 {
153    fn data_type() -> cudaDataType {
154        cudaDataType::C_64F
155    }
156}
157
158mod sealed {
159    use baracuda_types::{Complex32, Complex64};
160    pub trait Sealed {}
161    impl Sealed for f32 {}
162    impl Sealed for f64 {}
163    impl Sealed for Complex32 {}
164    impl Sealed for Complex64 {}
165}
166
167// ---- Op / Order / Fill / Diag wrappers ----------------------------------
168
169/// Transpose / conjugate-transpose selector.
170#[derive(Copy, Clone, Debug, Eq, PartialEq, Default)]
171pub enum Op {
172    /// No transpose.
173    #[default]
174    N,
175    /// Transpose.
176    T,
177    /// Conjugate transpose (real types: same as `T`).
178    C,
179}
180
181impl Op {
182    fn raw(self) -> cusparseOperation_t {
183        match self {
184            Op::N => cusparseOperation_t::N,
185            Op::T => cusparseOperation_t::T,
186            Op::C => cusparseOperation_t::C,
187        }
188    }
189}
190
191/// Dense-matrix storage order.
192#[derive(Copy, Clone, Debug, Eq, PartialEq)]
193pub enum Order {
194    /// Row-major (C-style) storage.
195    Row,
196    /// Column-major (Fortran-style) storage.
197    Col,
198}
199
200impl Order {
201    fn raw(self) -> cusparseOrder_t {
202        match self {
203            Order::Row => cusparseOrder_t::Row,
204            Order::Col => cusparseOrder_t::Col,
205        }
206    }
207}
208
209/// Triangular storage mode for sparse matrices.
210#[derive(Copy, Clone, Debug, Eq, PartialEq)]
211pub enum Fill {
212    /// Lower triangle is the populated half.
213    Lower,
214    /// Upper triangle is the populated half.
215    Upper,
216}
217
218impl Fill {
219    fn raw(self) -> cusparseFillMode_t {
220        match self {
221            Fill::Lower => cusparseFillMode_t::Lower,
222            Fill::Upper => cusparseFillMode_t::Upper,
223        }
224    }
225}
226
227/// Unit-diagonal flag for triangular sparse matrices.
228#[derive(Copy, Clone, Debug, Eq, PartialEq)]
229pub enum Diag {
230    /// Diagonal entries are read from the matrix.
231    NonUnit,
232    /// Diagonal entries are implicitly 1 (not read).
233    Unit,
234}
235
236impl Diag {
237    fn raw(self) -> cusparseDiagType_t {
238        match self {
239            Diag::NonUnit => cusparseDiagType_t::NonUnit,
240            Diag::Unit => cusparseDiagType_t::Unit,
241        }
242    }
243}
244
245// ---- Handle -------------------------------------------------------------
246
247/// Owned cuSPARSE handle.
248pub struct Handle {
249    handle: cusparseHandle_t,
250}
251
252unsafe impl Send for Handle {}
253
254impl core::fmt::Debug for Handle {
255    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
256        f.debug_struct("cusparse::Handle")
257            .field("handle", &self.handle)
258            .finish()
259    }
260}
261
262impl Handle {
263    /// Create a new cuSPARSE handle (`cusparseCreate`).
264    pub fn new() -> Result<Self> {
265        let c = cusparse()?;
266        let cu = c.cusparse_create()?;
267        let mut h: cusparseHandle_t = core::ptr::null_mut();
268        check(unsafe { cu(&mut h) })?;
269        Ok(Self { handle: h })
270    }
271
272    /// Bind this handle to `stream`; subsequent cuSPARSE calls dispatch on
273    /// it (`cusparseSetStream`).
274    pub fn set_stream(&self, stream: &Stream) -> Result<()> {
275        let c = cusparse()?;
276        let cu = c.cusparse_set_stream()?;
277        check(unsafe { cu(self.handle, stream.as_raw() as _) })
278    }
279
280    /// cuSPARSE library version (`cusparseGetVersion`),
281    /// e.g. `12604` for cuSPARSE 12.6.4.
282    pub fn version(&self) -> Result<i32> {
283        let c = cusparse()?;
284        let cu = c.cusparse_get_version()?;
285        let mut v: core::ffi::c_int = 0;
286        check(unsafe { cu(self.handle, &mut v) })?;
287        Ok(v)
288    }
289
290    /// Raw `cusparseHandle_t`. Use with care.
291    #[inline]
292    pub fn as_raw(&self) -> cusparseHandle_t {
293        self.handle
294    }
295}
296
297impl Drop for Handle {
298    fn drop(&mut self) {
299        if let Ok(c) = cusparse()
300            && let Ok(cu) = c.cusparse_destroy()
301        {
302            let _ = unsafe { cu(self.handle) };
303        }
304    }
305}
306
307// ---- Sparse matrix descriptor -------------------------------------------
308
309/// A sparse-matrix descriptor (CSR / CSC / COO / BSR). The descriptor keeps
310/// pointers to externally-owned device buffers; the lifetime parameter ties
311/// those buffers to the descriptor.
312pub struct SpMat<'buf, T> {
313    descr: cusparseSpMatDescr_t,
314    _markers: PhantomData<&'buf mut T>,
315}
316
317unsafe impl<T> Send for SpMat<'_, T> {}
318
319impl<T> core::fmt::Debug for SpMat<'_, T> {
320    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
321        f.debug_struct("SpMat")
322            .field("descr", &self.descr)
323            .finish_non_exhaustive()
324    }
325}
326
327impl<'buf, T: SparseScalar + baracuda_types::DeviceRepr> SpMat<'buf, T> {
328    /// Build a CSR (compressed sparse row) descriptor.
329    ///
330    /// `row_offsets.len()` must equal `rows + 1`; `col_indices.len()` and
331    /// `values.len()` must equal `nnz`.
332    pub fn csr(
333        rows: i64,
334        cols: i64,
335        nnz: i64,
336        row_offsets: &'buf mut DeviceBuffer<i32>,
337        col_indices: &'buf mut DeviceBuffer<i32>,
338        values: &'buf mut DeviceBuffer<T>,
339    ) -> Result<Self> {
340        let c = cusparse()?;
341        let cu = c.cusparse_create_csr()?;
342        let mut descr: cusparseSpMatDescr_t = core::ptr::null_mut();
343        check(unsafe {
344            cu(
345                &mut descr,
346                rows,
347                cols,
348                nnz,
349                row_offsets.as_raw().0 as *mut c_void,
350                col_indices.as_raw().0 as *mut c_void,
351                values.as_raw().0 as *mut c_void,
352                cusparseIndexType_t::I32I,
353                cusparseIndexType_t::I32I,
354                cusparseIndexBase_t::Zero,
355                T::data_type(),
356            )
357        })?;
358        Ok(Self {
359            descr,
360            _markers: PhantomData,
361        })
362    }
363
364    /// Build a CSC (compressed sparse column) descriptor.
365    pub fn csc(
366        rows: i64,
367        cols: i64,
368        nnz: i64,
369        col_offsets: &'buf mut DeviceBuffer<i32>,
370        row_indices: &'buf mut DeviceBuffer<i32>,
371        values: &'buf mut DeviceBuffer<T>,
372    ) -> Result<Self> {
373        let c = cusparse()?;
374        let cu = c.cusparse_create_csc()?;
375        let mut descr: cusparseSpMatDescr_t = core::ptr::null_mut();
376        check(unsafe {
377            cu(
378                &mut descr,
379                rows,
380                cols,
381                nnz,
382                col_offsets.as_raw().0 as *mut c_void,
383                row_indices.as_raw().0 as *mut c_void,
384                values.as_raw().0 as *mut c_void,
385                cusparseIndexType_t::I32I,
386                cusparseIndexType_t::I32I,
387                cusparseIndexBase_t::Zero,
388                T::data_type(),
389            )
390        })?;
391        Ok(Self {
392            descr,
393            _markers: PhantomData,
394        })
395    }
396
397    /// Build a BSR (block-sparse-row) descriptor.
398    #[allow(clippy::too_many_arguments)]
399    pub fn bsr(
400        brows: i64,
401        bcols: i64,
402        bnnz: i64,
403        row_block_dim: i64,
404        col_block_dim: i64,
405        order: Order,
406        row_offsets: &'buf mut DeviceBuffer<i32>,
407        col_indices: &'buf mut DeviceBuffer<i32>,
408        values: &'buf mut DeviceBuffer<T>,
409    ) -> Result<Self> {
410        let c = cusparse()?;
411        let cu = c.cusparse_create_bsr()?;
412        let mut descr: cusparseSpMatDescr_t = core::ptr::null_mut();
413        check(unsafe {
414            cu(
415                &mut descr,
416                brows,
417                bcols,
418                bnnz,
419                row_block_dim,
420                col_block_dim,
421                row_offsets.as_raw().0 as *mut c_void,
422                col_indices.as_raw().0 as *mut c_void,
423                values.as_raw().0 as *mut c_void,
424                cusparseIndexType_t::I32I,
425                cusparseIndexType_t::I32I,
426                cusparseIndexBase_t::Zero,
427                T::data_type(),
428                order.raw(),
429            )
430        })?;
431        Ok(Self {
432            descr,
433            _markers: PhantomData,
434        })
435    }
436
437    /// Build a COO (coordinate) descriptor.
438    pub fn coo(
439        rows: i64,
440        cols: i64,
441        nnz: i64,
442        row_indices: &'buf mut DeviceBuffer<i32>,
443        col_indices: &'buf mut DeviceBuffer<i32>,
444        values: &'buf mut DeviceBuffer<T>,
445    ) -> Result<Self> {
446        let c = cusparse()?;
447        let cu = c.cusparse_create_coo()?;
448        let mut descr: cusparseSpMatDescr_t = core::ptr::null_mut();
449        check(unsafe {
450            cu(
451                &mut descr,
452                rows,
453                cols,
454                nnz,
455                row_indices.as_raw().0 as *mut c_void,
456                col_indices.as_raw().0 as *mut c_void,
457                values.as_raw().0 as *mut c_void,
458                cusparseIndexType_t::I32I,
459                cusparseIndexBase_t::Zero,
460                T::data_type(),
461            )
462        })?;
463        Ok(Self {
464            descr,
465            _markers: PhantomData,
466        })
467    }
468}
469
470impl<T> SpMat<'_, T> {
471    /// Sparse matrix dimensions: `(rows, cols, nnz)`.
472    pub fn shape(&self) -> Result<(i64, i64, i64)> {
473        let c = cusparse()?;
474        let cu = c.cusparse_sp_mat_get_size()?;
475        let (mut r, mut col, mut nz) = (0i64, 0i64, 0i64);
476        check(unsafe { cu(self.descr, &mut r, &mut col, &mut nz) })?;
477        Ok((r, col, nz))
478    }
479
480    /// Rebind a CSR descriptor's underlying device pointers without
481    /// rebuilding it. Saves descriptor recreation when the same shape
482    /// is reused with new data.
483    ///
484    /// # Safety
485    ///
486    /// All three pointers must be live device allocations matching the
487    /// original `(rows + 1, nnz, nnz)` element counts and the original
488    /// element types. They must stay valid until the next operation
489    /// on this descriptor completes.
490    pub unsafe fn set_csr_pointers(
491        &self,
492        row_offsets: *mut c_void,
493        col_indices: *mut c_void,
494        values: *mut c_void,
495    ) -> Result<()> {
496        unsafe {
497            let c = cusparse()?;
498            let cu = c.cusparse_csr_set_pointers()?;
499            check(cu(self.descr, row_offsets, col_indices, values))
500        }
501    }
502
503    /// Rebind a CSC descriptor's underlying device pointers.
504    ///
505    /// # Safety
506    ///
507    /// See [`Self::set_csr_pointers`].
508    pub unsafe fn set_csc_pointers(
509        &self,
510        col_offsets: *mut c_void,
511        row_indices: *mut c_void,
512        values: *mut c_void,
513    ) -> Result<()> {
514        unsafe {
515            let c = cusparse()?;
516            let cu = c.cusparse_csc_set_pointers()?;
517            check(cu(self.descr, col_offsets, row_indices, values))
518        }
519    }
520
521    /// Rebind a COO descriptor's underlying device pointers.
522    ///
523    /// # Safety
524    ///
525    /// See [`Self::set_csr_pointers`].
526    pub unsafe fn set_coo_pointers(
527        &self,
528        row_indices: *mut c_void,
529        col_indices: *mut c_void,
530        values: *mut c_void,
531    ) -> Result<()> {
532        unsafe {
533            let c = cusparse()?;
534            let cu = c.cusparse_coo_set_pointers()?;
535            check(cu(self.descr, row_indices, col_indices, values))
536        }
537    }
538
539    /// Set the fill-triangle attribute (for triangular solves).
540    pub fn set_fill(&self, fill: Fill) -> Result<()> {
541        let c = cusparse()?;
542        let cu = c.cusparse_sp_mat_set_attribute()?;
543        let raw = fill.raw();
544        check(unsafe {
545            cu(
546                self.descr,
547                cusparseSpMatAttribute_t::FillMode,
548                &raw as *const _ as *const c_void,
549                core::mem::size_of::<cusparseFillMode_t>(),
550            )
551        })
552    }
553
554    /// Set the diagonal-type attribute (unit vs non-unit, for triangular solves).
555    pub fn set_diag(&self, diag: Diag) -> Result<()> {
556        let c = cusparse()?;
557        let cu = c.cusparse_sp_mat_set_attribute()?;
558        let raw = diag.raw();
559        check(unsafe {
560            cu(
561                self.descr,
562                cusparseSpMatAttribute_t::DiagType,
563                &raw as *const _ as *const c_void,
564                core::mem::size_of::<cusparseDiagType_t>(),
565            )
566        })
567    }
568
569    /// Raw `cusparseSpMatDescr_t`. Use with care.
570    #[inline]
571    pub fn as_raw(&self) -> cusparseSpMatDescr_t {
572        self.descr
573    }
574}
575
576impl<T> Drop for SpMat<'_, T> {
577    fn drop(&mut self) {
578        if let Ok(c) = cusparse()
579            && let Ok(cu) = c.cusparse_destroy_sp_mat()
580        {
581            let _ = unsafe { cu(self.descr) };
582        }
583    }
584}
585
586// ---- Dense vector / matrix -----------------------------------------------
587
588/// Dense vector descriptor for cuSPARSE's generic API.
589///
590/// Borrows the underlying [`DeviceBuffer`]; the buffer cannot be freed while
591/// the descriptor is alive.
592pub struct DnVec<'buf, T> {
593    descr: cusparseDnVecDescr_t,
594    _marker: PhantomData<&'buf mut T>,
595}
596
597unsafe impl<T> Send for DnVec<'_, T> {}
598
599impl<T> core::fmt::Debug for DnVec<'_, T> {
600    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
601        f.debug_struct("DnVec")
602            .field("descr", &self.descr)
603            .finish_non_exhaustive()
604    }
605}
606
607impl<'buf, T: SparseScalar + baracuda_types::DeviceRepr> DnVec<'buf, T> {
608    /// Build a dense-vector descriptor (`cusparseCreateDnVec`)
609    /// over the buffer `values`.
610    pub fn new(values: &'buf mut DeviceBuffer<T>) -> Result<Self> {
611        let c = cusparse()?;
612        let cu = c.cusparse_create_dn_vec()?;
613        let mut descr: cusparseDnVecDescr_t = core::ptr::null_mut();
614        check(unsafe {
615            cu(
616                &mut descr,
617                values.len() as i64,
618                values.as_raw().0 as *mut c_void,
619                T::data_type(),
620            )
621        })?;
622        Ok(Self {
623            descr,
624            _marker: PhantomData,
625        })
626    }
627}
628
629impl<T> DnVec<'_, T> {
630    /// Raw `cusparseDnVecDescr_t`. Use with care.
631    #[inline]
632    pub fn as_raw(&self) -> cusparseDnVecDescr_t {
633        self.descr
634    }
635}
636
637impl<T> Drop for DnVec<'_, T> {
638    fn drop(&mut self) {
639        if let Ok(c) = cusparse()
640            && let Ok(cu) = c.cusparse_destroy_dn_vec()
641        {
642            let _ = unsafe { cu(self.descr) };
643        }
644    }
645}
646
647/// Dense matrix descriptor for cuSPARSE's generic API.
648///
649/// Borrows the underlying [`DeviceBuffer`]; the buffer cannot be freed while
650/// the descriptor is alive.
651pub struct DnMat<'buf, T> {
652    descr: cusparseDnMatDescr_t,
653    _marker: PhantomData<&'buf mut T>,
654}
655
656unsafe impl<T> Send for DnMat<'_, T> {}
657
658impl<T> core::fmt::Debug for DnMat<'_, T> {
659    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
660        f.debug_struct("DnMat")
661            .field("descr", &self.descr)
662            .finish_non_exhaustive()
663    }
664}
665
666impl<'buf, T: SparseScalar + baracuda_types::DeviceRepr> DnMat<'buf, T> {
667    /// Build a dense-matrix descriptor (`cusparseCreateDnMat`) for the
668    /// given shape, leading dimension, and storage order.
669    pub fn new(
670        rows: i64,
671        cols: i64,
672        ld: i64,
673        order: Order,
674        values: &'buf mut DeviceBuffer<T>,
675    ) -> Result<Self> {
676        let c = cusparse()?;
677        let cu = c.cusparse_create_dn_mat()?;
678        let mut descr: cusparseDnMatDescr_t = core::ptr::null_mut();
679        check(unsafe {
680            cu(
681                &mut descr,
682                rows,
683                cols,
684                ld,
685                values.as_raw().0 as *mut c_void,
686                T::data_type(),
687                order.raw(),
688            )
689        })?;
690        Ok(Self {
691            descr,
692            _marker: PhantomData,
693        })
694    }
695}
696
697impl<T> DnMat<'_, T> {
698    /// Raw `cusparseDnMatDescr_t`. Use with care.
699    #[inline]
700    pub fn as_raw(&self) -> cusparseDnMatDescr_t {
701        self.descr
702    }
703}
704
705impl<T> Drop for DnMat<'_, T> {
706    fn drop(&mut self) {
707        if let Ok(c) = cusparse()
708            && let Ok(cu) = c.cusparse_destroy_dn_mat()
709        {
710            let _ = unsafe { cu(self.descr) };
711        }
712    }
713}
714
715// ---- SpMV ---------------------------------------------------------------
716
717/// Query buffer-size for `y = alpha * op(A) * x + beta * y`.
718#[allow(clippy::too_many_arguments)]
719pub fn spmv_buffer_size<T: SparseScalar>(
720    handle: &Handle,
721    op: Op,
722    alpha: &T,
723    a: &SpMat<'_, T>,
724    x: &DnVec<'_, T>,
725    beta: &T,
726    y: &DnVec<'_, T>,
727    alg: SpMVAlg,
728) -> Result<usize> {
729    let c = cusparse()?;
730    let cu = c.cusparse_spmv_buffer_size()?;
731    let mut size: usize = 0;
732    check(unsafe {
733        cu(
734            handle.as_raw(),
735            op.raw(),
736            alpha as *const T as *const c_void,
737            a.descr,
738            x.descr,
739            beta as *const T as *const c_void,
740            y.descr,
741            T::data_type(),
742            alg,
743            &mut size,
744        )
745    })?;
746    Ok(size)
747}
748
749/// `y = alpha * op(A) * x + beta * y`.
750#[allow(clippy::too_many_arguments)]
751pub fn spmv<T: SparseScalar>(
752    handle: &Handle,
753    op: Op,
754    alpha: &T,
755    a: &SpMat<'_, T>,
756    x: &DnVec<'_, T>,
757    beta: &T,
758    y: &mut DnVec<'_, T>,
759    alg: SpMVAlg,
760    workspace: &mut DeviceBuffer<u8>,
761) -> Result<()> {
762    let c = cusparse()?;
763    let cu = c.cusparse_spmv()?;
764    check(unsafe {
765        cu(
766            handle.as_raw(),
767            op.raw(),
768            alpha as *const T as *const c_void,
769            a.descr,
770            x.descr,
771            beta as *const T as *const c_void,
772            y.descr,
773            T::data_type(),
774            alg,
775            workspace.as_raw().0 as *mut c_void,
776        )
777    })
778}
779
780// ---- SpMM ---------------------------------------------------------------
781
782/// Query buffer-size for `C = alpha * op(A) * op(B) + beta * C`, `A` sparse.
783#[allow(clippy::too_many_arguments)]
784pub fn spmm_buffer_size<T: SparseScalar>(
785    handle: &Handle,
786    op_a: Op,
787    op_b: Op,
788    alpha: &T,
789    a: &SpMat<'_, T>,
790    b: &DnMat<'_, T>,
791    beta: &T,
792    c: &DnMat<'_, T>,
793    alg: SpMMAlg,
794) -> Result<usize> {
795    let c_api = cusparse()?;
796    let cu = c_api.cusparse_spmm_buffer_size()?;
797    let mut size = 0usize;
798    check(unsafe {
799        cu(
800            handle.as_raw(),
801            op_a.raw(),
802            op_b.raw(),
803            alpha as *const T as *const c_void,
804            a.descr,
805            b.descr,
806            beta as *const T as *const c_void,
807            c.descr,
808            T::data_type(),
809            alg,
810            &mut size,
811        )
812    })?;
813    Ok(size)
814}
815
816/// `C = alpha * op(A) * op(B) + beta * C` with `A` sparse, `B`/`C` dense
817/// (`cusparseSpMM`).
818#[allow(clippy::too_many_arguments)]
819pub fn spmm<T: SparseScalar>(
820    handle: &Handle,
821    op_a: Op,
822    op_b: Op,
823    alpha: &T,
824    a: &SpMat<'_, T>,
825    b: &DnMat<'_, T>,
826    beta: &T,
827    c: &mut DnMat<'_, T>,
828    alg: SpMMAlg,
829    workspace: &mut DeviceBuffer<u8>,
830) -> Result<()> {
831    let c_api = cusparse()?;
832    let cu = c_api.cusparse_spmm()?;
833    check(unsafe {
834        cu(
835            handle.as_raw(),
836            op_a.raw(),
837            op_b.raw(),
838            alpha as *const T as *const c_void,
839            a.descr,
840            b.descr,
841            beta as *const T as *const c_void,
842            c.descr,
843            T::data_type(),
844            alg,
845            workspace.as_raw().0 as *mut c_void,
846        )
847    })
848}
849
850/// One-time preprocessing before [`spmm`]. Pre-computes algorithm-specific
851/// state into `workspace` so subsequent [`spmm`] calls (with the same A
852/// sparsity pattern + dimensions) are faster. Use this when the same
853/// matrix is multiplied many times.
854#[allow(clippy::too_many_arguments)]
855pub fn spmm_preprocess<T: SparseScalar>(
856    handle: &Handle,
857    op_a: Op,
858    op_b: Op,
859    alpha: &T,
860    a: &SpMat<'_, T>,
861    b: &DnMat<'_, T>,
862    beta: &T,
863    c: &mut DnMat<'_, T>,
864    alg: SpMMAlg,
865    workspace: &mut DeviceBuffer<u8>,
866) -> Result<()> {
867    let c_api = cusparse()?;
868    let cu = c_api.cusparse_spmm_preprocess()?;
869    check(unsafe {
870        cu(
871            handle.as_raw(),
872            op_a.raw(),
873            op_b.raw(),
874            alpha as *const T as *const c_void,
875            a.descr,
876            b.descr,
877            beta as *const T as *const c_void,
878            c.descr,
879            T::data_type(),
880            alg,
881            workspace.as_raw().0 as *mut c_void,
882        )
883    })
884}
885
886// ---- SpGEMM -------------------------------------------------------------
887
888/// Per-plan handle for a 3-phase SpGEMM computation.
889#[derive(Debug)]
890pub struct SpGEMMPlan {
891    raw: cusparseSpGEMMDescr_t,
892}
893
894impl SpGEMMPlan {
895    /// Create a new SpGEMM plan descriptor
896    /// (`cusparseSpGEMM_createDescr`).
897    pub fn new() -> Result<Self> {
898        let c = cusparse()?;
899        let cu = c.cusparse_spgemm_create_descr()?;
900        let mut d: cusparseSpGEMMDescr_t = core::ptr::null_mut();
901        check(unsafe { cu(&mut d) })?;
902        Ok(Self { raw: d })
903    }
904}
905
906impl Drop for SpGEMMPlan {
907    fn drop(&mut self) {
908        if let Ok(c) = cusparse()
909            && let Ok(cu) = c.cusparse_spgemm_destroy_descr()
910        {
911            let _ = unsafe { cu(self.raw) };
912        }
913    }
914}
915
916/// Phase 1: work-estimation. The caller provides `buffer1` whose size is
917/// returned in `size1`; pass `null` the first time, then allocate & re-call.
918///
919/// # Safety
920///
921/// `buffer1` must be either null (size-query mode) or a valid device
922/// pointer to at least `*size1` bytes of writable scratch memory that
923/// remains live for the duration of the underlying cuSPARSE call.
924#[allow(clippy::too_many_arguments)]
925pub unsafe fn spgemm_work_estimation<T: SparseScalar>(
926    handle: &Handle,
927    op_a: Op,
928    op_b: Op,
929    alpha: &T,
930    a: &SpMat<'_, T>,
931    b: &SpMat<'_, T>,
932    beta: &T,
933    c: &mut SpMat<'_, T>,
934    alg: SpGEMMAlg,
935    plan: &SpGEMMPlan,
936    size1: &mut usize,
937    buffer1: *mut c_void,
938) -> Result<()> {
939    unsafe {
940        let c_api = cusparse()?;
941        let cu = c_api.cusparse_spgemm_work_estimation()?;
942        check(cu(
943            handle.as_raw(),
944            op_a.raw(),
945            op_b.raw(),
946            alpha as *const T as *const c_void,
947            a.descr,
948            b.descr,
949            beta as *const T as *const c_void,
950            c.descr,
951            T::data_type(),
952            alg,
953            plan.raw,
954            size1,
955            buffer1,
956        ))
957    }
958}
959
960/// Phase 2: compute. Same two-step pattern for `buffer2`.
961///
962/// # Safety
963///
964/// `buffer2` must be either null (size-query mode) or a valid device
965/// pointer to at least `*size2` bytes of writable scratch memory that
966/// remains live for the duration of the underlying cuSPARSE call.
967#[allow(clippy::too_many_arguments)]
968pub unsafe fn spgemm_compute<T: SparseScalar>(
969    handle: &Handle,
970    op_a: Op,
971    op_b: Op,
972    alpha: &T,
973    a: &SpMat<'_, T>,
974    b: &SpMat<'_, T>,
975    beta: &T,
976    c: &mut SpMat<'_, T>,
977    alg: SpGEMMAlg,
978    plan: &SpGEMMPlan,
979    size2: &mut usize,
980    buffer2: *mut c_void,
981) -> Result<()> {
982    unsafe {
983        let c_api = cusparse()?;
984        let cu = c_api.cusparse_spgemm_compute()?;
985        check(cu(
986            handle.as_raw(),
987            op_a.raw(),
988            op_b.raw(),
989            alpha as *const T as *const c_void,
990            a.descr,
991            b.descr,
992            beta as *const T as *const c_void,
993            c.descr,
994            T::data_type(),
995            alg,
996            plan.raw,
997            size2,
998            buffer2,
999        ))
1000    }
1001}
1002
1003/// Phase 3: write output arrays into the pre-populated output `SpMat`.
1004#[allow(clippy::too_many_arguments)]
1005pub fn spgemm_copy<T: SparseScalar>(
1006    handle: &Handle,
1007    op_a: Op,
1008    op_b: Op,
1009    alpha: &T,
1010    a: &SpMat<'_, T>,
1011    b: &SpMat<'_, T>,
1012    beta: &T,
1013    c: &mut SpMat<'_, T>,
1014    alg: SpGEMMAlg,
1015    plan: &SpGEMMPlan,
1016) -> Result<()> {
1017    let c_api = cusparse()?;
1018    let cu = c_api.cusparse_spgemm_copy()?;
1019    check(unsafe {
1020        cu(
1021            handle.as_raw(),
1022            op_a.raw(),
1023            op_b.raw(),
1024            alpha as *const T as *const c_void,
1025            a.descr,
1026            b.descr,
1027            beta as *const T as *const c_void,
1028            c.descr,
1029            T::data_type(),
1030            alg,
1031            plan.raw,
1032        )
1033    })
1034}
1035
1036// ---- SpSV / SpSM --------------------------------------------------------
1037
1038/// Per-plan handle for an SpSV (sparse triangular solve, single RHS).
1039#[derive(Debug)]
1040pub struct SpSVPlan {
1041    raw: cusparseSpSVDescr_t,
1042}
1043
1044impl SpSVPlan {
1045    /// Create a new SpSV plan descriptor
1046    /// (`cusparseSpSV_createDescr`).
1047    pub fn new() -> Result<Self> {
1048        let c = cusparse()?;
1049        let cu = c.cusparse_spsv_create_descr()?;
1050        let mut d: cusparseSpSVDescr_t = core::ptr::null_mut();
1051        check(unsafe { cu(&mut d) })?;
1052        Ok(Self { raw: d })
1053    }
1054}
1055
1056impl Drop for SpSVPlan {
1057    fn drop(&mut self) {
1058        if let Ok(c) = cusparse()
1059            && let Ok(cu) = c.cusparse_spsv_destroy_descr()
1060        {
1061            let _ = unsafe { cu(self.raw) };
1062        }
1063    }
1064}
1065
1066/// Query buffer-size for a sparse triangular solve
1067/// (`cusparseSpSV_bufferSize`).
1068#[allow(clippy::too_many_arguments)]
1069pub fn spsv_buffer_size<T: SparseScalar>(
1070    handle: &Handle,
1071    op: Op,
1072    alpha: &T,
1073    a: &SpMat<'_, T>,
1074    x: &DnVec<'_, T>,
1075    y: &DnVec<'_, T>,
1076    alg: SpSVAlg,
1077    plan: &SpSVPlan,
1078) -> Result<usize> {
1079    let c = cusparse()?;
1080    let cu = c.cusparse_spsv_buffer_size()?;
1081    let mut size = 0usize;
1082    check(unsafe {
1083        cu(
1084            handle.as_raw(),
1085            op.raw(),
1086            alpha as *const T as *const c_void,
1087            a.descr,
1088            x.descr,
1089            y.descr,
1090            T::data_type(),
1091            alg,
1092            plan.raw,
1093            &mut size,
1094        )
1095    })?;
1096    Ok(size)
1097}
1098
1099/// Run the SpSV analysis phase
1100/// (`cusparseSpSV_analysis`) — must be called once before [`spsv_solve`].
1101#[allow(clippy::too_many_arguments)]
1102pub fn spsv_analysis<T: SparseScalar>(
1103    handle: &Handle,
1104    op: Op,
1105    alpha: &T,
1106    a: &SpMat<'_, T>,
1107    x: &DnVec<'_, T>,
1108    y: &DnVec<'_, T>,
1109    alg: SpSVAlg,
1110    plan: &SpSVPlan,
1111    workspace: &mut DeviceBuffer<u8>,
1112) -> Result<()> {
1113    let c = cusparse()?;
1114    let cu = c.cusparse_spsv_analysis()?;
1115    check(unsafe {
1116        cu(
1117            handle.as_raw(),
1118            op.raw(),
1119            alpha as *const T as *const c_void,
1120            a.descr,
1121            x.descr,
1122            y.descr,
1123            T::data_type(),
1124            alg,
1125            plan.raw,
1126            workspace.as_raw().0 as *mut c_void,
1127        )
1128    })
1129}
1130
1131/// Execute the sparse triangular solve `op(A) * y = alpha * x`
1132/// (`cusparseSpSV_solve`). Requires a prior [`spsv_analysis`] call.
1133#[allow(clippy::too_many_arguments)]
1134pub fn spsv_solve<T: SparseScalar>(
1135    handle: &Handle,
1136    op: Op,
1137    alpha: &T,
1138    a: &SpMat<'_, T>,
1139    x: &DnVec<'_, T>,
1140    y: &mut DnVec<'_, T>,
1141    alg: SpSVAlg,
1142    plan: &SpSVPlan,
1143) -> Result<()> {
1144    let c = cusparse()?;
1145    let cu = c.cusparse_spsv_solve()?;
1146    check(unsafe {
1147        cu(
1148            handle.as_raw(),
1149            op.raw(),
1150            alpha as *const T as *const c_void,
1151            a.descr,
1152            x.descr,
1153            y.descr,
1154            T::data_type(),
1155            alg,
1156            plan.raw,
1157        )
1158    })
1159}
1160
1161/// Per-plan handle for an SpSM (sparse triangular solve, multiple RHS).
1162#[derive(Debug)]
1163pub struct SpSMPlan {
1164    raw: cusparseSpSMDescr_t,
1165}
1166
1167impl SpSMPlan {
1168    /// Create a new SpSM plan descriptor
1169    /// (`cusparseSpSM_createDescr`).
1170    pub fn new() -> Result<Self> {
1171        let c = cusparse()?;
1172        let cu = c.cusparse_spsm_create_descr()?;
1173        let mut d: cusparseSpSMDescr_t = core::ptr::null_mut();
1174        check(unsafe { cu(&mut d) })?;
1175        Ok(Self { raw: d })
1176    }
1177}
1178
1179impl Drop for SpSMPlan {
1180    fn drop(&mut self) {
1181        if let Ok(c) = cusparse()
1182            && let Ok(cu) = c.cusparse_spsm_destroy_descr()
1183        {
1184            let _ = unsafe { cu(self.raw) };
1185        }
1186    }
1187}
1188
1189/// Query buffer-size for a multi-RHS sparse triangular solve
1190/// (`cusparseSpSM_bufferSize`).
1191#[allow(clippy::too_many_arguments)]
1192pub fn spsm_buffer_size<T: SparseScalar>(
1193    handle: &Handle,
1194    op_a: Op,
1195    op_b: Op,
1196    alpha: &T,
1197    a: &SpMat<'_, T>,
1198    b: &DnMat<'_, T>,
1199    c: &DnMat<'_, T>,
1200    alg: SpSMAlg,
1201    plan: &SpSMPlan,
1202) -> Result<usize> {
1203    let c_api = cusparse()?;
1204    let cu = c_api.cusparse_spsm_buffer_size()?;
1205    let mut size = 0usize;
1206    check(unsafe {
1207        cu(
1208            handle.as_raw(),
1209            op_a.raw(),
1210            op_b.raw(),
1211            alpha as *const T as *const c_void,
1212            a.descr,
1213            b.descr,
1214            c.descr,
1215            T::data_type(),
1216            alg,
1217            plan.raw,
1218            &mut size,
1219        )
1220    })?;
1221    Ok(size)
1222}
1223
1224/// Run the SpSM analysis phase
1225/// (`cusparseSpSM_analysis`) — must be called once before [`spsm_solve`].
1226#[allow(clippy::too_many_arguments)]
1227pub fn spsm_analysis<T: SparseScalar>(
1228    handle: &Handle,
1229    op_a: Op,
1230    op_b: Op,
1231    alpha: &T,
1232    a: &SpMat<'_, T>,
1233    b: &DnMat<'_, T>,
1234    c: &DnMat<'_, T>,
1235    alg: SpSMAlg,
1236    plan: &SpSMPlan,
1237    workspace: &mut DeviceBuffer<u8>,
1238) -> Result<()> {
1239    let c_api = cusparse()?;
1240    let cu = c_api.cusparse_spsm_analysis()?;
1241    check(unsafe {
1242        cu(
1243            handle.as_raw(),
1244            op_a.raw(),
1245            op_b.raw(),
1246            alpha as *const T as *const c_void,
1247            a.descr,
1248            b.descr,
1249            c.descr,
1250            T::data_type(),
1251            alg,
1252            plan.raw,
1253            workspace.as_raw().0 as *mut c_void,
1254        )
1255    })
1256}
1257
1258/// Execute the multi-RHS sparse triangular solve `op(A) * C = alpha * op(B)`
1259/// (`cusparseSpSM_solve`).
1260#[allow(clippy::too_many_arguments)]
1261pub fn spsm_solve<T: SparseScalar>(
1262    handle: &Handle,
1263    op_a: Op,
1264    op_b: Op,
1265    alpha: &T,
1266    a: &SpMat<'_, T>,
1267    b: &DnMat<'_, T>,
1268    c: &mut DnMat<'_, T>,
1269    alg: SpSMAlg,
1270    plan: &SpSMPlan,
1271) -> Result<()> {
1272    let c_api = cusparse()?;
1273    let cu = c_api.cusparse_spsm_solve()?;
1274    check(unsafe {
1275        cu(
1276            handle.as_raw(),
1277            op_a.raw(),
1278            op_b.raw(),
1279            alpha as *const T as *const c_void,
1280            a.descr,
1281            b.descr,
1282            c.descr,
1283            T::data_type(),
1284            alg,
1285            plan.raw,
1286        )
1287    })
1288}
1289
1290// ---- SDDMM -------------------------------------------------------------
1291
1292/// Query buffer-size for the sampled dense-dense matmul
1293/// (`cusparseSDDMM_bufferSize`).
1294#[allow(clippy::too_many_arguments)]
1295pub fn sddmm_buffer_size<T: SparseScalar>(
1296    handle: &Handle,
1297    op_a: Op,
1298    op_b: Op,
1299    alpha: &T,
1300    a: &DnMat<'_, T>,
1301    b: &DnMat<'_, T>,
1302    beta: &T,
1303    c: &SpMat<'_, T>,
1304    alg: SDDMMAlg,
1305) -> Result<usize> {
1306    let c_api = cusparse()?;
1307    let cu = c_api.cusparse_sddmm_buffer_size()?;
1308    let mut size = 0usize;
1309    check(unsafe {
1310        cu(
1311            handle.as_raw(),
1312            op_a.raw(),
1313            op_b.raw(),
1314            alpha as *const T as *const c_void,
1315            a.descr,
1316            b.descr,
1317            beta as *const T as *const c_void,
1318            c.descr,
1319            T::data_type(),
1320            alg,
1321            &mut size,
1322        )
1323    })?;
1324    Ok(size)
1325}
1326
1327/// Sampled dense-dense matrix multiply (`cusparseSDDMM`): computes
1328/// `C = alpha * op(A) * op(B) + beta * C` at the sparsity pattern of `C`.
1329#[allow(clippy::too_many_arguments)]
1330pub fn sddmm<T: SparseScalar>(
1331    handle: &Handle,
1332    op_a: Op,
1333    op_b: Op,
1334    alpha: &T,
1335    a: &DnMat<'_, T>,
1336    b: &DnMat<'_, T>,
1337    beta: &T,
1338    c: &mut SpMat<'_, T>,
1339    alg: SDDMMAlg,
1340    workspace: &mut DeviceBuffer<u8>,
1341) -> Result<()> {
1342    let c_api = cusparse()?;
1343    let cu = c_api.cusparse_sddmm()?;
1344    check(unsafe {
1345        cu(
1346            handle.as_raw(),
1347            op_a.raw(),
1348            op_b.raw(),
1349            alpha as *const T as *const c_void,
1350            a.descr,
1351            b.descr,
1352            beta as *const T as *const c_void,
1353            c.descr,
1354            T::data_type(),
1355            alg,
1356            workspace.as_raw().0 as *mut c_void,
1357        )
1358    })
1359}
1360
1361/// One-time preprocessing before [`sddmm`]. See [`spmm_preprocess`] for
1362/// the rationale.
1363#[allow(clippy::too_many_arguments)]
1364pub fn sddmm_preprocess<T: SparseScalar>(
1365    handle: &Handle,
1366    op_a: Op,
1367    op_b: Op,
1368    alpha: &T,
1369    a: &DnMat<'_, T>,
1370    b: &DnMat<'_, T>,
1371    beta: &T,
1372    c: &mut SpMat<'_, T>,
1373    alg: SDDMMAlg,
1374    workspace: &mut DeviceBuffer<u8>,
1375) -> Result<()> {
1376    let c_api = cusparse()?;
1377    let cu = c_api.cusparse_sddmm_preprocess()?;
1378    check(unsafe {
1379        cu(
1380            handle.as_raw(),
1381            op_a.raw(),
1382            op_b.raw(),
1383            alpha as *const T as *const c_void,
1384            a.descr,
1385            b.descr,
1386            beta as *const T as *const c_void,
1387            c.descr,
1388            T::data_type(),
1389            alg,
1390            workspace.as_raw().0 as *mut c_void,
1391        )
1392    })
1393}
1394
1395// ---- Sparse / dense conversions ----------------------------------------
1396
1397/// Query buffer-size for a sparse → dense conversion
1398/// (`cusparseSparseToDense_bufferSize`).
1399pub fn sparse_to_dense_buffer_size<T: SparseScalar>(
1400    handle: &Handle,
1401    sp: &SpMat<'_, T>,
1402    dn: &DnMat<'_, T>,
1403) -> Result<usize> {
1404    let c = cusparse()?;
1405    let cu = c.cusparse_sparse_to_dense_buffer_size()?;
1406    let mut size = 0usize;
1407    check(unsafe { cu(handle.as_raw(), sp.descr, dn.descr, 0, &mut size) })?;
1408    Ok(size)
1409}
1410
1411/// Convert a sparse matrix into a dense one
1412/// (`cusparseSparseToDense`).
1413pub fn sparse_to_dense<T: SparseScalar>(
1414    handle: &Handle,
1415    sp: &SpMat<'_, T>,
1416    dn: &mut DnMat<'_, T>,
1417    workspace: &mut DeviceBuffer<u8>,
1418) -> Result<()> {
1419    let c = cusparse()?;
1420    let cu = c.cusparse_sparse_to_dense()?;
1421    check(unsafe {
1422        cu(
1423            handle.as_raw(),
1424            sp.descr,
1425            dn.descr,
1426            0,
1427            workspace.as_raw().0 as *mut c_void,
1428        )
1429    })
1430}
1431
1432/// Query buffer-size for a dense → sparse conversion
1433/// (`cusparseDenseToSparse_bufferSize`).
1434pub fn dense_to_sparse_buffer_size<T: SparseScalar>(
1435    handle: &Handle,
1436    dn: &DnMat<'_, T>,
1437    sp: &SpMat<'_, T>,
1438) -> Result<usize> {
1439    let c = cusparse()?;
1440    let cu = c.cusparse_dense_to_sparse_buffer_size()?;
1441    let mut size = 0usize;
1442    check(unsafe { cu(handle.as_raw(), dn.descr, sp.descr, 0, &mut size) })?;
1443    Ok(size)
1444}
1445
1446/// Run the dense → sparse analysis phase
1447/// (`cusparseDenseToSparse_analysis`); fills in the sparsity pattern.
1448pub fn dense_to_sparse_analysis<T: SparseScalar>(
1449    handle: &Handle,
1450    dn: &DnMat<'_, T>,
1451    sp: &SpMat<'_, T>,
1452    workspace: &mut DeviceBuffer<u8>,
1453) -> Result<()> {
1454    let c = cusparse()?;
1455    let cu = c.cusparse_dense_to_sparse_analysis()?;
1456    check(unsafe {
1457        cu(
1458            handle.as_raw(),
1459            dn.descr,
1460            sp.descr,
1461            0,
1462            workspace.as_raw().0 as *mut c_void,
1463        )
1464    })
1465}
1466
1467/// Execute the dense → sparse conversion
1468/// (`cusparseDenseToSparse_convert`); fills in the sparse values.
1469pub fn dense_to_sparse_convert<T: SparseScalar>(
1470    handle: &Handle,
1471    dn: &DnMat<'_, T>,
1472    sp: &mut SpMat<'_, T>,
1473    workspace: &mut DeviceBuffer<u8>,
1474) -> Result<()> {
1475    let c = cusparse()?;
1476    let cu = c.cusparse_dense_to_sparse_convert()?;
1477    check(unsafe {
1478        cu(
1479            handle.as_raw(),
1480            dn.descr,
1481            sp.descr,
1482            0,
1483            workspace.as_raw().0 as *mut c_void,
1484        )
1485    })
1486}
1487
1488/// Workspace size in bytes for [`csr2csc_ex2`].
1489#[allow(clippy::too_many_arguments)]
1490pub fn csr2csc_ex2_buffer_size<T: SparseScalar + baracuda_types::DeviceRepr>(
1491    handle: &Handle,
1492    m: i32,
1493    n: i32,
1494    nnz: i32,
1495    csr_val: &DeviceBuffer<T>,
1496    csr_row_ptr: &DeviceBuffer<i32>,
1497    csr_col_ind: &DeviceBuffer<i32>,
1498    csc_val: &mut DeviceBuffer<T>,
1499    csc_col_ptr: &mut DeviceBuffer<i32>,
1500    csc_row_ind: &mut DeviceBuffer<i32>,
1501    copy_values: bool,
1502    idx_base: IndexBase,
1503    alg: Csr2CscAlg,
1504) -> Result<usize> {
1505    let c = cusparse()?;
1506    let cu = c.cusparse_csr2csc_ex2_buffer_size()?;
1507    let mut size = 0usize;
1508    check(unsafe {
1509        cu(
1510            handle.as_raw(),
1511            m,
1512            n,
1513            nnz,
1514            csr_val.as_raw().0 as *const c_void,
1515            csr_row_ptr.as_raw().0 as *const i32,
1516            csr_col_ind.as_raw().0 as *const i32,
1517            csc_val.as_raw().0 as *mut c_void,
1518            csc_col_ptr.as_raw().0 as *mut i32,
1519            csc_row_ind.as_raw().0 as *mut i32,
1520            T::data_type(),
1521            copy_values as i32,
1522            idx_base,
1523            alg,
1524            &mut size,
1525        )
1526    })?;
1527    Ok(size)
1528}
1529
1530/// Convert a CSR matrix to CSC format using the modern Ex2 entry point —
1531/// supports algorithm selection (`alg`) and arbitrary value types.
1532#[allow(clippy::too_many_arguments)]
1533pub fn csr2csc_ex2<T: SparseScalar + baracuda_types::DeviceRepr>(
1534    handle: &Handle,
1535    m: i32,
1536    n: i32,
1537    nnz: i32,
1538    csr_val: &DeviceBuffer<T>,
1539    csr_row_ptr: &DeviceBuffer<i32>,
1540    csr_col_ind: &DeviceBuffer<i32>,
1541    csc_val: &mut DeviceBuffer<T>,
1542    csc_col_ptr: &mut DeviceBuffer<i32>,
1543    csc_row_ind: &mut DeviceBuffer<i32>,
1544    copy_values: bool,
1545    idx_base: IndexBase,
1546    alg: Csr2CscAlg,
1547    workspace: &mut DeviceBuffer<u8>,
1548) -> Result<()> {
1549    let c = cusparse()?;
1550    let cu = c.cusparse_csr2csc_ex2()?;
1551    check(unsafe {
1552        cu(
1553            handle.as_raw(),
1554            m,
1555            n,
1556            nnz,
1557            csr_val.as_raw().0 as *const c_void,
1558            csr_row_ptr.as_raw().0 as *const i32,
1559            csr_col_ind.as_raw().0 as *const i32,
1560            csc_val.as_raw().0 as *mut c_void,
1561            csc_col_ptr.as_raw().0 as *mut i32,
1562            csc_row_ind.as_raw().0 as *mut i32,
1563            T::data_type(),
1564            copy_values as i32,
1565            idx_base,
1566            alg,
1567            workspace.as_raw().0 as *mut c_void,
1568        )
1569    })
1570}
1571
1572// ---- Sparse BLAS-1 helpers ---------------------------------------------
1573
1574/// Sparse AXPBY: `y = alpha * x + beta * y` over dense vectors
1575/// (`cusparseAxpby`).
1576pub fn axpby<T: SparseScalar>(
1577    handle: &Handle,
1578    alpha: &T,
1579    x: &DnVec<'_, T>,
1580    beta: &T,
1581    y: &mut DnVec<'_, T>,
1582) -> Result<()> {
1583    let c = cusparse()?;
1584    let cu = c.cusparse_axpby()?;
1585    check(unsafe {
1586        cu(
1587            handle.as_raw(),
1588            alpha as *const T as *const c_void,
1589            x.descr,
1590            beta as *const T as *const c_void,
1591            y.descr,
1592        )
1593    })
1594}
1595
1596/// Sparse gather: `x = y[indices]` (`cusparseGather`).
1597pub fn gather<T: SparseScalar>(
1598    handle: &Handle,
1599    y: &DnVec<'_, T>,
1600    x: &mut DnVec<'_, T>,
1601) -> Result<()> {
1602    let c = cusparse()?;
1603    let cu = c.cusparse_gather()?;
1604    check(unsafe { cu(handle.as_raw(), y.descr, x.descr) })
1605}
1606
1607/// Sparse scatter: `y[indices] = x` (`cusparseScatter`).
1608pub fn scatter<T: SparseScalar>(
1609    handle: &Handle,
1610    x: &DnVec<'_, T>,
1611    y: &mut DnVec<'_, T>,
1612) -> Result<()> {
1613    let c = cusparse()?;
1614    let cu = c.cusparse_scatter()?;
1615    check(unsafe { cu(handle.as_raw(), x.descr, y.descr) })
1616}
1617
1618/// Givens rotation applied to a pair of sparse/dense vectors
1619/// (`cusparseRot`).
1620pub fn rot<T: SparseScalar>(
1621    handle: &Handle,
1622    c_cos: &T,
1623    s_sin: &T,
1624    x: &mut DnVec<'_, T>,
1625    y: &mut DnVec<'_, T>,
1626) -> Result<()> {
1627    let c_api = cusparse()?;
1628    let cu = c_api.cusparse_rot()?;
1629    check(unsafe {
1630        cu(
1631            handle.as_raw(),
1632            c_cos as *const T as *const c_void,
1633            s_sin as *const T as *const c_void,
1634            x.descr,
1635            y.descr,
1636        )
1637    })
1638}
1639
1640// ---- Back-compat re-exports for existing users --------------------------
1641
1642/// Legacy alias kept for callers from v0.1 — prefer [`SpMat::csr`].
1643pub type CsrMatrix<'buf> = SpMat<'buf, f32>;
1644/// Legacy alias — prefer [`DnVec`].
1645pub type DenseVector<'buf, T> = DnVec<'buf, T>;
baracuda_cusparse/lib.rs

baracuda_cusparse/
lib.rs
