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    pub fn bsr(
399        brows: i64,
400        bcols: i64,
401        bnnz: i64,
402        row_block_dim: i64,
403        col_block_dim: i64,
404        order: Order,
405        row_offsets: &'buf mut DeviceBuffer<i32>,
406        col_indices: &'buf mut DeviceBuffer<i32>,
407        values: &'buf mut DeviceBuffer<T>,
408    ) -> Result<Self> {
409        let c = cusparse()?;
410        let cu = c.cusparse_create_bsr()?;
411        let mut descr: cusparseSpMatDescr_t = core::ptr::null_mut();
412        check(unsafe {
413            cu(
414                &mut descr,
415                brows,
416                bcols,
417                bnnz,
418                row_block_dim,
419                col_block_dim,
420                row_offsets.as_raw().0 as *mut c_void,
421                col_indices.as_raw().0 as *mut c_void,
422                values.as_raw().0 as *mut c_void,
423                cusparseIndexType_t::I32I,
424                cusparseIndexType_t::I32I,
425                cusparseIndexBase_t::Zero,
426                T::data_type(),
427                order.raw(),
428            )
429        })?;
430        Ok(Self {
431            descr,
432            _markers: PhantomData,
433        })
434    }
435
436    /// Build a COO (coordinate) descriptor.
437    pub fn coo(
438        rows: i64,
439        cols: i64,
440        nnz: i64,
441        row_indices: &'buf mut DeviceBuffer<i32>,
442        col_indices: &'buf mut DeviceBuffer<i32>,
443        values: &'buf mut DeviceBuffer<T>,
444    ) -> Result<Self> {
445        let c = cusparse()?;
446        let cu = c.cusparse_create_coo()?;
447        let mut descr: cusparseSpMatDescr_t = core::ptr::null_mut();
448        check(unsafe {
449            cu(
450                &mut descr,
451                rows,
452                cols,
453                nnz,
454                row_indices.as_raw().0 as *mut c_void,
455                col_indices.as_raw().0 as *mut c_void,
456                values.as_raw().0 as *mut c_void,
457                cusparseIndexType_t::I32I,
458                cusparseIndexBase_t::Zero,
459                T::data_type(),
460            )
461        })?;
462        Ok(Self {
463            descr,
464            _markers: PhantomData,
465        })
466    }
467}
468
469impl<T> SpMat<'_, T> {
470    /// Sparse matrix dimensions: `(rows, cols, nnz)`.
471    pub fn shape(&self) -> Result<(i64, i64, i64)> {
472        let c = cusparse()?;
473        let cu = c.cusparse_sp_mat_get_size()?;
474        let (mut r, mut col, mut nz) = (0i64, 0i64, 0i64);
475        check(unsafe { cu(self.descr, &mut r, &mut col, &mut nz) })?;
476        Ok((r, col, nz))
477    }
478
479    /// Rebind a CSR descriptor's underlying device pointers without
480    /// rebuilding it. Saves descriptor recreation when the same shape
481    /// is reused with new data.
482    ///
483    /// # Safety
484    ///
485    /// All three pointers must be live device allocations matching the
486    /// original `(rows + 1, nnz, nnz)` element counts and the original
487    /// element types. They must stay valid until the next operation
488    /// on this descriptor completes.
489    pub unsafe fn set_csr_pointers(
490        &self,
491        row_offsets: *mut c_void,
492        col_indices: *mut c_void,
493        values: *mut c_void,
494    ) -> Result<()> {
495        unsafe {
496            let c = cusparse()?;
497            let cu = c.cusparse_csr_set_pointers()?;
498            check(cu(self.descr, row_offsets, col_indices, values))
499        }
500    }
501
502    /// Rebind a CSC descriptor's underlying device pointers.
503    ///
504    /// # Safety
505    ///
506    /// See [`Self::set_csr_pointers`].
507    pub unsafe fn set_csc_pointers(
508        &self,
509        col_offsets: *mut c_void,
510        row_indices: *mut c_void,
511        values: *mut c_void,
512    ) -> Result<()> {
513        unsafe {
514            let c = cusparse()?;
515            let cu = c.cusparse_csc_set_pointers()?;
516            check(cu(self.descr, col_offsets, row_indices, values))
517        }
518    }
519
520    /// Rebind a COO descriptor's underlying device pointers.
521    ///
522    /// # Safety
523    ///
524    /// See [`Self::set_csr_pointers`].
525    pub unsafe fn set_coo_pointers(
526        &self,
527        row_indices: *mut c_void,
528        col_indices: *mut c_void,
529        values: *mut c_void,
530    ) -> Result<()> {
531        unsafe {
532            let c = cusparse()?;
533            let cu = c.cusparse_coo_set_pointers()?;
534            check(cu(self.descr, row_indices, col_indices, values))
535        }
536    }
537
538    /// Set the fill-triangle attribute (for triangular solves).
539    pub fn set_fill(&self, fill: Fill) -> Result<()> {
540        let c = cusparse()?;
541        let cu = c.cusparse_sp_mat_set_attribute()?;
542        let raw = fill.raw();
543        check(unsafe {
544            cu(
545                self.descr,
546                cusparseSpMatAttribute_t::FillMode,
547                &raw as *const _ as *const c_void,
548                core::mem::size_of::<cusparseFillMode_t>(),
549            )
550        })
551    }
552
553    /// Set the diagonal-type attribute (unit vs non-unit, for triangular solves).
554    pub fn set_diag(&self, diag: Diag) -> Result<()> {
555        let c = cusparse()?;
556        let cu = c.cusparse_sp_mat_set_attribute()?;
557        let raw = diag.raw();
558        check(unsafe {
559            cu(
560                self.descr,
561                cusparseSpMatAttribute_t::DiagType,
562                &raw as *const _ as *const c_void,
563                core::mem::size_of::<cusparseDiagType_t>(),
564            )
565        })
566    }
567
568    /// Raw `cusparseSpMatDescr_t`. Use with care.
569    #[inline]
570    pub fn as_raw(&self) -> cusparseSpMatDescr_t {
571        self.descr
572    }
573}
574
575impl<T> Drop for SpMat<'_, T> {
576    fn drop(&mut self) {
577        if let Ok(c) = cusparse()
578            && let Ok(cu) = c.cusparse_destroy_sp_mat()
579        {
580            let _ = unsafe { cu(self.descr) };
581        }
582    }
583}
584
585// ---- Dense vector / matrix -----------------------------------------------
586
587/// Dense vector descriptor for cuSPARSE's generic API.
588///
589/// Borrows the underlying [`DeviceBuffer`]; the buffer cannot be freed while
590/// the descriptor is alive.
591pub struct DnVec<'buf, T> {
592    descr: cusparseDnVecDescr_t,
593    _marker: PhantomData<&'buf mut T>,
594}
595
596unsafe impl<T> Send for DnVec<'_, T> {}
597
598impl<T> core::fmt::Debug for DnVec<'_, T> {
599    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
600        f.debug_struct("DnVec")
601            .field("descr", &self.descr)
602            .finish_non_exhaustive()
603    }
604}
605
606impl<'buf, T: SparseScalar + baracuda_types::DeviceRepr> DnVec<'buf, T> {
607    /// Build a dense-vector descriptor (`cusparseCreateDnVec`)
608    /// over the buffer `values`.
609    pub fn new(values: &'buf mut DeviceBuffer<T>) -> Result<Self> {
610        let c = cusparse()?;
611        let cu = c.cusparse_create_dn_vec()?;
612        let mut descr: cusparseDnVecDescr_t = core::ptr::null_mut();
613        check(unsafe {
614            cu(
615                &mut descr,
616                values.len() as i64,
617                values.as_raw().0 as *mut c_void,
618                T::data_type(),
619            )
620        })?;
621        Ok(Self {
622            descr,
623            _marker: PhantomData,
624        })
625    }
626}
627
628impl<T> DnVec<'_, T> {
629    /// Raw `cusparseDnVecDescr_t`. Use with care.
630    #[inline]
631    pub fn as_raw(&self) -> cusparseDnVecDescr_t {
632        self.descr
633    }
634}
635
636impl<T> Drop for DnVec<'_, T> {
637    fn drop(&mut self) {
638        if let Ok(c) = cusparse()
639            && let Ok(cu) = c.cusparse_destroy_dn_vec()
640        {
641            let _ = unsafe { cu(self.descr) };
642        }
643    }
644}
645
646/// Dense matrix descriptor for cuSPARSE's generic API.
647///
648/// Borrows the underlying [`DeviceBuffer`]; the buffer cannot be freed while
649/// the descriptor is alive.
650pub struct DnMat<'buf, T> {
651    descr: cusparseDnMatDescr_t,
652    _marker: PhantomData<&'buf mut T>,
653}
654
655unsafe impl<T> Send for DnMat<'_, T> {}
656
657impl<T> core::fmt::Debug for DnMat<'_, T> {
658    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
659        f.debug_struct("DnMat")
660            .field("descr", &self.descr)
661            .finish_non_exhaustive()
662    }
663}
664
665impl<'buf, T: SparseScalar + baracuda_types::DeviceRepr> DnMat<'buf, T> {
666    /// Build a dense-matrix descriptor (`cusparseCreateDnMat`) for the
667    /// given shape, leading dimension, and storage order.
668    pub fn new(
669        rows: i64,
670        cols: i64,
671        ld: i64,
672        order: Order,
673        values: &'buf mut DeviceBuffer<T>,
674    ) -> Result<Self> {
675        let c = cusparse()?;
676        let cu = c.cusparse_create_dn_mat()?;
677        let mut descr: cusparseDnMatDescr_t = core::ptr::null_mut();
678        check(unsafe {
679            cu(
680                &mut descr,
681                rows,
682                cols,
683                ld,
684                values.as_raw().0 as *mut c_void,
685                T::data_type(),
686                order.raw(),
687            )
688        })?;
689        Ok(Self {
690            descr,
691            _marker: PhantomData,
692        })
693    }
694}
695
696impl<T> DnMat<'_, T> {
697    /// Raw `cusparseDnMatDescr_t`. Use with care.
698    #[inline]
699    pub fn as_raw(&self) -> cusparseDnMatDescr_t {
700        self.descr
701    }
702}
703
704impl<T> Drop for DnMat<'_, T> {
705    fn drop(&mut self) {
706        if let Ok(c) = cusparse()
707            && let Ok(cu) = c.cusparse_destroy_dn_mat()
708        {
709            let _ = unsafe { cu(self.descr) };
710        }
711    }
712}
713
714// ---- SpMV ---------------------------------------------------------------
715
716/// Query buffer-size for `y = alpha * op(A) * x + beta * y`.
717pub fn spmv_buffer_size<T: SparseScalar>(
718    handle: &Handle,
719    op: Op,
720    alpha: &T,
721    a: &SpMat<'_, T>,
722    x: &DnVec<'_, T>,
723    beta: &T,
724    y: &DnVec<'_, T>,
725    alg: SpMVAlg,
726) -> Result<usize> {
727    let c = cusparse()?;
728    let cu = c.cusparse_spmv_buffer_size()?;
729    let mut size: usize = 0;
730    check(unsafe {
731        cu(
732            handle.as_raw(),
733            op.raw(),
734            alpha as *const T as *const c_void,
735            a.descr,
736            x.descr,
737            beta as *const T as *const c_void,
738            y.descr,
739            T::data_type(),
740            alg,
741            &mut size,
742        )
743    })?;
744    Ok(size)
745}
746
747/// `y = alpha * op(A) * x + beta * y`.
748pub fn spmv<T: SparseScalar>(
749    handle: &Handle,
750    op: Op,
751    alpha: &T,
752    a: &SpMat<'_, T>,
753    x: &DnVec<'_, T>,
754    beta: &T,
755    y: &mut DnVec<'_, T>,
756    alg: SpMVAlg,
757    workspace: &mut DeviceBuffer<u8>,
758) -> Result<()> {
759    let c = cusparse()?;
760    let cu = c.cusparse_spmv()?;
761    check(unsafe {
762        cu(
763            handle.as_raw(),
764            op.raw(),
765            alpha as *const T as *const c_void,
766            a.descr,
767            x.descr,
768            beta as *const T as *const c_void,
769            y.descr,
770            T::data_type(),
771            alg,
772            workspace.as_raw().0 as *mut c_void,
773        )
774    })
775}
776
777// ---- SpMM ---------------------------------------------------------------
778
779/// Query buffer-size for `C = alpha * op(A) * op(B) + beta * C`, `A` sparse.
780pub fn spmm_buffer_size<T: SparseScalar>(
781    handle: &Handle,
782    op_a: Op,
783    op_b: Op,
784    alpha: &T,
785    a: &SpMat<'_, T>,
786    b: &DnMat<'_, T>,
787    beta: &T,
788    c: &DnMat<'_, T>,
789    alg: SpMMAlg,
790) -> Result<usize> {
791    let c_api = cusparse()?;
792    let cu = c_api.cusparse_spmm_buffer_size()?;
793    let mut size = 0usize;
794    check(unsafe {
795        cu(
796            handle.as_raw(),
797            op_a.raw(),
798            op_b.raw(),
799            alpha as *const T as *const c_void,
800            a.descr,
801            b.descr,
802            beta as *const T as *const c_void,
803            c.descr,
804            T::data_type(),
805            alg,
806            &mut size,
807        )
808    })?;
809    Ok(size)
810}
811
812/// `C = alpha * op(A) * op(B) + beta * C` with `A` sparse, `B`/`C` dense
813/// (`cusparseSpMM`).
814pub fn spmm<T: SparseScalar>(
815    handle: &Handle,
816    op_a: Op,
817    op_b: Op,
818    alpha: &T,
819    a: &SpMat<'_, T>,
820    b: &DnMat<'_, T>,
821    beta: &T,
822    c: &mut DnMat<'_, T>,
823    alg: SpMMAlg,
824    workspace: &mut DeviceBuffer<u8>,
825) -> Result<()> {
826    let c_api = cusparse()?;
827    let cu = c_api.cusparse_spmm()?;
828    check(unsafe {
829        cu(
830            handle.as_raw(),
831            op_a.raw(),
832            op_b.raw(),
833            alpha as *const T as *const c_void,
834            a.descr,
835            b.descr,
836            beta as *const T as *const c_void,
837            c.descr,
838            T::data_type(),
839            alg,
840            workspace.as_raw().0 as *mut c_void,
841        )
842    })
843}
844
845/// One-time preprocessing before [`spmm`]. Pre-computes algorithm-specific
846/// state into `workspace` so subsequent [`spmm`] calls (with the same A
847/// sparsity pattern + dimensions) are faster. Use this when the same
848/// matrix is multiplied many times.
849pub fn spmm_preprocess<T: SparseScalar>(
850    handle: &Handle,
851    op_a: Op,
852    op_b: Op,
853    alpha: &T,
854    a: &SpMat<'_, T>,
855    b: &DnMat<'_, T>,
856    beta: &T,
857    c: &mut DnMat<'_, T>,
858    alg: SpMMAlg,
859    workspace: &mut DeviceBuffer<u8>,
860) -> Result<()> {
861    let c_api = cusparse()?;
862    let cu = c_api.cusparse_spmm_preprocess()?;
863    check(unsafe {
864        cu(
865            handle.as_raw(),
866            op_a.raw(),
867            op_b.raw(),
868            alpha as *const T as *const c_void,
869            a.descr,
870            b.descr,
871            beta as *const T as *const c_void,
872            c.descr,
873            T::data_type(),
874            alg,
875            workspace.as_raw().0 as *mut c_void,
876        )
877    })
878}
879
880// ---- SpGEMM -------------------------------------------------------------
881
882/// Per-plan handle for a 3-phase SpGEMM computation.
883#[derive(Debug)]
884pub struct SpGEMMPlan {
885    raw: cusparseSpGEMMDescr_t,
886}
887
888impl SpGEMMPlan {
889    /// Create a new SpGEMM plan descriptor
890    /// (`cusparseSpGEMM_createDescr`).
891    pub fn new() -> Result<Self> {
892        let c = cusparse()?;
893        let cu = c.cusparse_spgemm_create_descr()?;
894        let mut d: cusparseSpGEMMDescr_t = core::ptr::null_mut();
895        check(unsafe { cu(&mut d) })?;
896        Ok(Self { raw: d })
897    }
898}
899
900impl Drop for SpGEMMPlan {
901    fn drop(&mut self) {
902        if let Ok(c) = cusparse()
903            && let Ok(cu) = c.cusparse_spgemm_destroy_descr()
904        {
905            let _ = unsafe { cu(self.raw) };
906        }
907    }
908}
909
910/// Phase 1: work-estimation. The caller provides `buffer1` whose size is
911/// returned in `size1`; pass `null` the first time, then allocate & re-call.
912///
913/// # Safety
914///
915/// `buffer1` must be either null (size-query mode) or a valid device
916/// pointer to at least `*size1` bytes of writable scratch memory that
917/// remains live for the duration of the underlying cuSPARSE call.
918pub unsafe fn spgemm_work_estimation<T: SparseScalar>(
919    handle: &Handle,
920    op_a: Op,
921    op_b: Op,
922    alpha: &T,
923    a: &SpMat<'_, T>,
924    b: &SpMat<'_, T>,
925    beta: &T,
926    c: &mut SpMat<'_, T>,
927    alg: SpGEMMAlg,
928    plan: &SpGEMMPlan,
929    size1: &mut usize,
930    buffer1: *mut c_void,
931) -> Result<()> {
932    unsafe {
933        let c_api = cusparse()?;
934        let cu = c_api.cusparse_spgemm_work_estimation()?;
935        check(cu(
936            handle.as_raw(),
937            op_a.raw(),
938            op_b.raw(),
939            alpha as *const T as *const c_void,
940            a.descr,
941            b.descr,
942            beta as *const T as *const c_void,
943            c.descr,
944            T::data_type(),
945            alg,
946            plan.raw,
947            size1,
948            buffer1,
949        ))
950    }
951}
952
953/// Phase 2: compute. Same two-step pattern for `buffer2`.
954///
955/// # Safety
956///
957/// `buffer2` must be either null (size-query mode) or a valid device
958/// pointer to at least `*size2` bytes of writable scratch memory that
959/// remains live for the duration of the underlying cuSPARSE call.
960pub unsafe fn spgemm_compute<T: SparseScalar>(
961    handle: &Handle,
962    op_a: Op,
963    op_b: Op,
964    alpha: &T,
965    a: &SpMat<'_, T>,
966    b: &SpMat<'_, T>,
967    beta: &T,
968    c: &mut SpMat<'_, T>,
969    alg: SpGEMMAlg,
970    plan: &SpGEMMPlan,
971    size2: &mut usize,
972    buffer2: *mut c_void,
973) -> Result<()> {
974    unsafe {
975        let c_api = cusparse()?;
976        let cu = c_api.cusparse_spgemm_compute()?;
977        check(cu(
978            handle.as_raw(),
979            op_a.raw(),
980            op_b.raw(),
981            alpha as *const T as *const c_void,
982            a.descr,
983            b.descr,
984            beta as *const T as *const c_void,
985            c.descr,
986            T::data_type(),
987            alg,
988            plan.raw,
989            size2,
990            buffer2,
991        ))
992    }
993}
994
995/// Phase 3: write output arrays into the pre-populated output `SpMat`.
996pub fn spgemm_copy<T: SparseScalar>(
997    handle: &Handle,
998    op_a: Op,
999    op_b: Op,
1000    alpha: &T,
1001    a: &SpMat<'_, T>,
1002    b: &SpMat<'_, T>,
1003    beta: &T,
1004    c: &mut SpMat<'_, T>,
1005    alg: SpGEMMAlg,
1006    plan: &SpGEMMPlan,
1007) -> Result<()> {
1008    let c_api = cusparse()?;
1009    let cu = c_api.cusparse_spgemm_copy()?;
1010    check(unsafe {
1011        cu(
1012            handle.as_raw(),
1013            op_a.raw(),
1014            op_b.raw(),
1015            alpha as *const T as *const c_void,
1016            a.descr,
1017            b.descr,
1018            beta as *const T as *const c_void,
1019            c.descr,
1020            T::data_type(),
1021            alg,
1022            plan.raw,
1023        )
1024    })
1025}
1026
1027// ---- SpSV / SpSM --------------------------------------------------------
1028
1029/// Per-plan handle for an SpSV (sparse triangular solve, single RHS).
1030#[derive(Debug)]
1031pub struct SpSVPlan {
1032    raw: cusparseSpSVDescr_t,
1033}
1034
1035impl SpSVPlan {
1036    /// Create a new SpSV plan descriptor
1037    /// (`cusparseSpSV_createDescr`).
1038    pub fn new() -> Result<Self> {
1039        let c = cusparse()?;
1040        let cu = c.cusparse_spsv_create_descr()?;
1041        let mut d: cusparseSpSVDescr_t = core::ptr::null_mut();
1042        check(unsafe { cu(&mut d) })?;
1043        Ok(Self { raw: d })
1044    }
1045}
1046
1047impl Drop for SpSVPlan {
1048    fn drop(&mut self) {
1049        if let Ok(c) = cusparse()
1050            && let Ok(cu) = c.cusparse_spsv_destroy_descr()
1051        {
1052            let _ = unsafe { cu(self.raw) };
1053        }
1054    }
1055}
1056
1057/// Query buffer-size for a sparse triangular solve
1058/// (`cusparseSpSV_bufferSize`).
1059pub fn spsv_buffer_size<T: SparseScalar>(
1060    handle: &Handle,
1061    op: Op,
1062    alpha: &T,
1063    a: &SpMat<'_, T>,
1064    x: &DnVec<'_, T>,
1065    y: &DnVec<'_, T>,
1066    alg: SpSVAlg,
1067    plan: &SpSVPlan,
1068) -> Result<usize> {
1069    let c = cusparse()?;
1070    let cu = c.cusparse_spsv_buffer_size()?;
1071    let mut size = 0usize;
1072    check(unsafe {
1073        cu(
1074            handle.as_raw(),
1075            op.raw(),
1076            alpha as *const T as *const c_void,
1077            a.descr,
1078            x.descr,
1079            y.descr,
1080            T::data_type(),
1081            alg,
1082            plan.raw,
1083            &mut size,
1084        )
1085    })?;
1086    Ok(size)
1087}
1088
1089/// Run the SpSV analysis phase
1090/// (`cusparseSpSV_analysis`) — must be called once before [`spsv_solve`].
1091pub fn spsv_analysis<T: SparseScalar>(
1092    handle: &Handle,
1093    op: Op,
1094    alpha: &T,
1095    a: &SpMat<'_, T>,
1096    x: &DnVec<'_, T>,
1097    y: &DnVec<'_, T>,
1098    alg: SpSVAlg,
1099    plan: &SpSVPlan,
1100    workspace: &mut DeviceBuffer<u8>,
1101) -> Result<()> {
1102    let c = cusparse()?;
1103    let cu = c.cusparse_spsv_analysis()?;
1104    check(unsafe {
1105        cu(
1106            handle.as_raw(),
1107            op.raw(),
1108            alpha as *const T as *const c_void,
1109            a.descr,
1110            x.descr,
1111            y.descr,
1112            T::data_type(),
1113            alg,
1114            plan.raw,
1115            workspace.as_raw().0 as *mut c_void,
1116        )
1117    })
1118}
1119
1120/// Execute the sparse triangular solve `op(A) * y = alpha * x`
1121/// (`cusparseSpSV_solve`). Requires a prior [`spsv_analysis`] call.
1122pub fn spsv_solve<T: SparseScalar>(
1123    handle: &Handle,
1124    op: Op,
1125    alpha: &T,
1126    a: &SpMat<'_, T>,
1127    x: &DnVec<'_, T>,
1128    y: &mut DnVec<'_, T>,
1129    alg: SpSVAlg,
1130    plan: &SpSVPlan,
1131) -> Result<()> {
1132    let c = cusparse()?;
1133    let cu = c.cusparse_spsv_solve()?;
1134    check(unsafe {
1135        cu(
1136            handle.as_raw(),
1137            op.raw(),
1138            alpha as *const T as *const c_void,
1139            a.descr,
1140            x.descr,
1141            y.descr,
1142            T::data_type(),
1143            alg,
1144            plan.raw,
1145        )
1146    })
1147}
1148
1149/// Per-plan handle for an SpSM (sparse triangular solve, multiple RHS).
1150#[derive(Debug)]
1151pub struct SpSMPlan {
1152    raw: cusparseSpSMDescr_t,
1153}
1154
1155impl SpSMPlan {
1156    /// Create a new SpSM plan descriptor
1157    /// (`cusparseSpSM_createDescr`).
1158    pub fn new() -> Result<Self> {
1159        let c = cusparse()?;
1160        let cu = c.cusparse_spsm_create_descr()?;
1161        let mut d: cusparseSpSMDescr_t = core::ptr::null_mut();
1162        check(unsafe { cu(&mut d) })?;
1163        Ok(Self { raw: d })
1164    }
1165}
1166
1167impl Drop for SpSMPlan {
1168    fn drop(&mut self) {
1169        if let Ok(c) = cusparse()
1170            && let Ok(cu) = c.cusparse_spsm_destroy_descr()
1171        {
1172            let _ = unsafe { cu(self.raw) };
1173        }
1174    }
1175}
1176
1177/// Query buffer-size for a multi-RHS sparse triangular solve
1178/// (`cusparseSpSM_bufferSize`).
1179pub fn spsm_buffer_size<T: SparseScalar>(
1180    handle: &Handle,
1181    op_a: Op,
1182    op_b: Op,
1183    alpha: &T,
1184    a: &SpMat<'_, T>,
1185    b: &DnMat<'_, T>,
1186    c: &DnMat<'_, T>,
1187    alg: SpSMAlg,
1188    plan: &SpSMPlan,
1189) -> Result<usize> {
1190    let c_api = cusparse()?;
1191    let cu = c_api.cusparse_spsm_buffer_size()?;
1192    let mut size = 0usize;
1193    check(unsafe {
1194        cu(
1195            handle.as_raw(),
1196            op_a.raw(),
1197            op_b.raw(),
1198            alpha as *const T as *const c_void,
1199            a.descr,
1200            b.descr,
1201            c.descr,
1202            T::data_type(),
1203            alg,
1204            plan.raw,
1205            &mut size,
1206        )
1207    })?;
1208    Ok(size)
1209}
1210
1211/// Run the SpSM analysis phase
1212/// (`cusparseSpSM_analysis`) — must be called once before [`spsm_solve`].
1213pub fn spsm_analysis<T: SparseScalar>(
1214    handle: &Handle,
1215    op_a: Op,
1216    op_b: Op,
1217    alpha: &T,
1218    a: &SpMat<'_, T>,
1219    b: &DnMat<'_, T>,
1220    c: &DnMat<'_, T>,
1221    alg: SpSMAlg,
1222    plan: &SpSMPlan,
1223    workspace: &mut DeviceBuffer<u8>,
1224) -> Result<()> {
1225    let c_api = cusparse()?;
1226    let cu = c_api.cusparse_spsm_analysis()?;
1227    check(unsafe {
1228        cu(
1229            handle.as_raw(),
1230            op_a.raw(),
1231            op_b.raw(),
1232            alpha as *const T as *const c_void,
1233            a.descr,
1234            b.descr,
1235            c.descr,
1236            T::data_type(),
1237            alg,
1238            plan.raw,
1239            workspace.as_raw().0 as *mut c_void,
1240        )
1241    })
1242}
1243
1244/// Execute the multi-RHS sparse triangular solve `op(A) * C = alpha * op(B)`
1245/// (`cusparseSpSM_solve`).
1246pub fn spsm_solve<T: SparseScalar>(
1247    handle: &Handle,
1248    op_a: Op,
1249    op_b: Op,
1250    alpha: &T,
1251    a: &SpMat<'_, T>,
1252    b: &DnMat<'_, T>,
1253    c: &mut DnMat<'_, T>,
1254    alg: SpSMAlg,
1255    plan: &SpSMPlan,
1256) -> Result<()> {
1257    let c_api = cusparse()?;
1258    let cu = c_api.cusparse_spsm_solve()?;
1259    check(unsafe {
1260        cu(
1261            handle.as_raw(),
1262            op_a.raw(),
1263            op_b.raw(),
1264            alpha as *const T as *const c_void,
1265            a.descr,
1266            b.descr,
1267            c.descr,
1268            T::data_type(),
1269            alg,
1270            plan.raw,
1271        )
1272    })
1273}
1274
1275// ---- SDDMM -------------------------------------------------------------
1276
1277/// Query buffer-size for the sampled dense-dense matmul
1278/// (`cusparseSDDMM_bufferSize`).
1279pub fn sddmm_buffer_size<T: SparseScalar>(
1280    handle: &Handle,
1281    op_a: Op,
1282    op_b: Op,
1283    alpha: &T,
1284    a: &DnMat<'_, T>,
1285    b: &DnMat<'_, T>,
1286    beta: &T,
1287    c: &SpMat<'_, T>,
1288    alg: SDDMMAlg,
1289) -> Result<usize> {
1290    let c_api = cusparse()?;
1291    let cu = c_api.cusparse_sddmm_buffer_size()?;
1292    let mut size = 0usize;
1293    check(unsafe {
1294        cu(
1295            handle.as_raw(),
1296            op_a.raw(),
1297            op_b.raw(),
1298            alpha as *const T as *const c_void,
1299            a.descr,
1300            b.descr,
1301            beta as *const T as *const c_void,
1302            c.descr,
1303            T::data_type(),
1304            alg,
1305            &mut size,
1306        )
1307    })?;
1308    Ok(size)
1309}
1310
1311/// Sampled dense-dense matrix multiply (`cusparseSDDMM`): computes
1312/// `C = alpha * op(A) * op(B) + beta * C` at the sparsity pattern of `C`.
1313pub fn sddmm<T: SparseScalar>(
1314    handle: &Handle,
1315    op_a: Op,
1316    op_b: Op,
1317    alpha: &T,
1318    a: &DnMat<'_, T>,
1319    b: &DnMat<'_, T>,
1320    beta: &T,
1321    c: &mut SpMat<'_, T>,
1322    alg: SDDMMAlg,
1323    workspace: &mut DeviceBuffer<u8>,
1324) -> Result<()> {
1325    let c_api = cusparse()?;
1326    let cu = c_api.cusparse_sddmm()?;
1327    check(unsafe {
1328        cu(
1329            handle.as_raw(),
1330            op_a.raw(),
1331            op_b.raw(),
1332            alpha as *const T as *const c_void,
1333            a.descr,
1334            b.descr,
1335            beta as *const T as *const c_void,
1336            c.descr,
1337            T::data_type(),
1338            alg,
1339            workspace.as_raw().0 as *mut c_void,
1340        )
1341    })
1342}
1343
1344/// One-time preprocessing before [`sddmm`]. See [`spmm_preprocess`] for
1345/// the rationale.
1346pub fn sddmm_preprocess<T: SparseScalar>(
1347    handle: &Handle,
1348    op_a: Op,
1349    op_b: Op,
1350    alpha: &T,
1351    a: &DnMat<'_, T>,
1352    b: &DnMat<'_, T>,
1353    beta: &T,
1354    c: &mut SpMat<'_, T>,
1355    alg: SDDMMAlg,
1356    workspace: &mut DeviceBuffer<u8>,
1357) -> Result<()> {
1358    let c_api = cusparse()?;
1359    let cu = c_api.cusparse_sddmm_preprocess()?;
1360    check(unsafe {
1361        cu(
1362            handle.as_raw(),
1363            op_a.raw(),
1364            op_b.raw(),
1365            alpha as *const T as *const c_void,
1366            a.descr,
1367            b.descr,
1368            beta as *const T as *const c_void,
1369            c.descr,
1370            T::data_type(),
1371            alg,
1372            workspace.as_raw().0 as *mut c_void,
1373        )
1374    })
1375}
1376
1377// ---- Sparse / dense conversions ----------------------------------------
1378
1379/// Query buffer-size for a sparse → dense conversion
1380/// (`cusparseSparseToDense_bufferSize`).
1381pub fn sparse_to_dense_buffer_size<T: SparseScalar>(
1382    handle: &Handle,
1383    sp: &SpMat<'_, T>,
1384    dn: &DnMat<'_, T>,
1385) -> Result<usize> {
1386    let c = cusparse()?;
1387    let cu = c.cusparse_sparse_to_dense_buffer_size()?;
1388    let mut size = 0usize;
1389    check(unsafe { cu(handle.as_raw(), sp.descr, dn.descr, 0, &mut size) })?;
1390    Ok(size)
1391}
1392
1393/// Convert a sparse matrix into a dense one
1394/// (`cusparseSparseToDense`).
1395pub fn sparse_to_dense<T: SparseScalar>(
1396    handle: &Handle,
1397    sp: &SpMat<'_, T>,
1398    dn: &mut DnMat<'_, T>,
1399    workspace: &mut DeviceBuffer<u8>,
1400) -> Result<()> {
1401    let c = cusparse()?;
1402    let cu = c.cusparse_sparse_to_dense()?;
1403    check(unsafe {
1404        cu(
1405            handle.as_raw(),
1406            sp.descr,
1407            dn.descr,
1408            0,
1409            workspace.as_raw().0 as *mut c_void,
1410        )
1411    })
1412}
1413
1414/// Query buffer-size for a dense → sparse conversion
1415/// (`cusparseDenseToSparse_bufferSize`).
1416pub fn dense_to_sparse_buffer_size<T: SparseScalar>(
1417    handle: &Handle,
1418    dn: &DnMat<'_, T>,
1419    sp: &SpMat<'_, T>,
1420) -> Result<usize> {
1421    let c = cusparse()?;
1422    let cu = c.cusparse_dense_to_sparse_buffer_size()?;
1423    let mut size = 0usize;
1424    check(unsafe { cu(handle.as_raw(), dn.descr, sp.descr, 0, &mut size) })?;
1425    Ok(size)
1426}
1427
1428/// Run the dense → sparse analysis phase
1429/// (`cusparseDenseToSparse_analysis`); fills in the sparsity pattern.
1430pub fn dense_to_sparse_analysis<T: SparseScalar>(
1431    handle: &Handle,
1432    dn: &DnMat<'_, T>,
1433    sp: &SpMat<'_, T>,
1434    workspace: &mut DeviceBuffer<u8>,
1435) -> Result<()> {
1436    let c = cusparse()?;
1437    let cu = c.cusparse_dense_to_sparse_analysis()?;
1438    check(unsafe {
1439        cu(
1440            handle.as_raw(),
1441            dn.descr,
1442            sp.descr,
1443            0,
1444            workspace.as_raw().0 as *mut c_void,
1445        )
1446    })
1447}
1448
1449/// Execute the dense → sparse conversion
1450/// (`cusparseDenseToSparse_convert`); fills in the sparse values.
1451pub fn dense_to_sparse_convert<T: SparseScalar>(
1452    handle: &Handle,
1453    dn: &DnMat<'_, T>,
1454    sp: &mut SpMat<'_, T>,
1455    workspace: &mut DeviceBuffer<u8>,
1456) -> Result<()> {
1457    let c = cusparse()?;
1458    let cu = c.cusparse_dense_to_sparse_convert()?;
1459    check(unsafe {
1460        cu(
1461            handle.as_raw(),
1462            dn.descr,
1463            sp.descr,
1464            0,
1465            workspace.as_raw().0 as *mut c_void,
1466        )
1467    })
1468}
1469
1470/// Workspace size in bytes for [`csr2csc_ex2`].
1471pub fn csr2csc_ex2_buffer_size<T: SparseScalar + baracuda_types::DeviceRepr>(
1472    handle: &Handle,
1473    m: i32,
1474    n: i32,
1475    nnz: i32,
1476    csr_val: &DeviceBuffer<T>,
1477    csr_row_ptr: &DeviceBuffer<i32>,
1478    csr_col_ind: &DeviceBuffer<i32>,
1479    csc_val: &mut DeviceBuffer<T>,
1480    csc_col_ptr: &mut DeviceBuffer<i32>,
1481    csc_row_ind: &mut DeviceBuffer<i32>,
1482    copy_values: bool,
1483    idx_base: IndexBase,
1484    alg: Csr2CscAlg,
1485) -> Result<usize> {
1486    let c = cusparse()?;
1487    let cu = c.cusparse_csr2csc_ex2_buffer_size()?;
1488    let mut size = 0usize;
1489    check(unsafe {
1490        cu(
1491            handle.as_raw(),
1492            m,
1493            n,
1494            nnz,
1495            csr_val.as_raw().0 as *const c_void,
1496            csr_row_ptr.as_raw().0 as *const i32,
1497            csr_col_ind.as_raw().0 as *const i32,
1498            csc_val.as_raw().0 as *mut c_void,
1499            csc_col_ptr.as_raw().0 as *mut i32,
1500            csc_row_ind.as_raw().0 as *mut i32,
1501            T::data_type(),
1502            copy_values as i32,
1503            idx_base,
1504            alg,
1505            &mut size,
1506        )
1507    })?;
1508    Ok(size)
1509}
1510
1511/// Convert a CSR matrix to CSC format using the modern Ex2 entry point —
1512/// supports algorithm selection (`alg`) and arbitrary value types.
1513pub fn csr2csc_ex2<T: SparseScalar + baracuda_types::DeviceRepr>(
1514    handle: &Handle,
1515    m: i32,
1516    n: i32,
1517    nnz: i32,
1518    csr_val: &DeviceBuffer<T>,
1519    csr_row_ptr: &DeviceBuffer<i32>,
1520    csr_col_ind: &DeviceBuffer<i32>,
1521    csc_val: &mut DeviceBuffer<T>,
1522    csc_col_ptr: &mut DeviceBuffer<i32>,
1523    csc_row_ind: &mut DeviceBuffer<i32>,
1524    copy_values: bool,
1525    idx_base: IndexBase,
1526    alg: Csr2CscAlg,
1527    workspace: &mut DeviceBuffer<u8>,
1528) -> Result<()> {
1529    let c = cusparse()?;
1530    let cu = c.cusparse_csr2csc_ex2()?;
1531    check(unsafe {
1532        cu(
1533            handle.as_raw(),
1534            m,
1535            n,
1536            nnz,
1537            csr_val.as_raw().0 as *const c_void,
1538            csr_row_ptr.as_raw().0 as *const i32,
1539            csr_col_ind.as_raw().0 as *const i32,
1540            csc_val.as_raw().0 as *mut c_void,
1541            csc_col_ptr.as_raw().0 as *mut i32,
1542            csc_row_ind.as_raw().0 as *mut i32,
1543            T::data_type(),
1544            copy_values as i32,
1545            idx_base,
1546            alg,
1547            workspace.as_raw().0 as *mut c_void,
1548        )
1549    })
1550}
1551
1552// ---- Sparse BLAS-1 helpers ---------------------------------------------
1553
1554/// Sparse AXPBY: `y = alpha * x + beta * y` over dense vectors
1555/// (`cusparseAxpby`).
1556pub fn axpby<T: SparseScalar>(
1557    handle: &Handle,
1558    alpha: &T,
1559    x: &DnVec<'_, T>,
1560    beta: &T,
1561    y: &mut DnVec<'_, T>,
1562) -> Result<()> {
1563    let c = cusparse()?;
1564    let cu = c.cusparse_axpby()?;
1565    check(unsafe {
1566        cu(
1567            handle.as_raw(),
1568            alpha as *const T as *const c_void,
1569            x.descr,
1570            beta as *const T as *const c_void,
1571            y.descr,
1572        )
1573    })
1574}
1575
1576/// Sparse gather: `x = y[indices]` (`cusparseGather`).
1577pub fn gather<T: SparseScalar>(
1578    handle: &Handle,
1579    y: &DnVec<'_, T>,
1580    x: &mut DnVec<'_, T>,
1581) -> Result<()> {
1582    let c = cusparse()?;
1583    let cu = c.cusparse_gather()?;
1584    check(unsafe { cu(handle.as_raw(), y.descr, x.descr) })
1585}
1586
1587/// Sparse scatter: `y[indices] = x` (`cusparseScatter`).
1588pub fn scatter<T: SparseScalar>(
1589    handle: &Handle,
1590    x: &DnVec<'_, T>,
1591    y: &mut DnVec<'_, T>,
1592) -> Result<()> {
1593    let c = cusparse()?;
1594    let cu = c.cusparse_scatter()?;
1595    check(unsafe { cu(handle.as_raw(), x.descr, y.descr) })
1596}
1597
1598/// Givens rotation applied to a pair of sparse/dense vectors
1599/// (`cusparseRot`).
1600pub fn rot<T: SparseScalar>(
1601    handle: &Handle,
1602    c_cos: &T,
1603    s_sin: &T,
1604    x: &mut DnVec<'_, T>,
1605    y: &mut DnVec<'_, T>,
1606) -> Result<()> {
1607    let c_api = cusparse()?;
1608    let cu = c_api.cusparse_rot()?;
1609    check(unsafe {
1610        cu(
1611            handle.as_raw(),
1612            c_cos as *const T as *const c_void,
1613            s_sin as *const T as *const c_void,
1614            x.descr,
1615            y.descr,
1616        )
1617    })
1618}
1619
1620// ---- Back-compat re-exports for existing users --------------------------
1621
1622/// Legacy alias kept for callers from v0.1 — prefer [`SpMat::csr`].
1623pub type CsrMatrix<'buf> = SpMat<'buf, f32>;
1624/// Legacy alias — prefer [`DnVec`].
1625pub type DenseVector<'buf, T> = DnVec<'buf, T>;
baracuda_cusparse/lib.rs

baracuda_cusparse/
lib.rs
