singe_cusparse/
operation.rs

1use std::{marker::PhantomData, mem::ManuallyDrop, ptr, sync::Arc};
2
3use singe_cuda::context::Context as CudaContext;
4use singe_cuda::data_type::DataType;
5use singe_cuda::types::DevicePtr;
6
7use crate::{
8    context::Context,
9    error::{Error, Result},
10    matrix::{DenseMatrixDescriptor, SparseMatrixDescriptor},
11    sys, try_ffi,
12    types::{Operation, SpMmOpAlgorithm},
13};
14
15#[derive(Debug)]
16pub struct SpGemmDescriptor {
17    handle: sys::cusparseSpGEMMDescr_t,
18    cuda_ctx: Arc<CudaContext>,
19}
20
21#[derive(Debug)]
22pub struct SpSvDescriptor {
23    handle: sys::cusparseSpSVDescr_t,
24    cuda_ctx: Arc<CudaContext>,
25}
26
27#[derive(Debug)]
28pub struct SpSmDescriptor {
29    handle: sys::cusparseSpSMDescr_t,
30    cuda_ctx: Arc<CudaContext>,
31}
32
33#[derive(Debug)]
34pub struct SpMmOpPlan<'a> {
35    context: &'a Context,
36    handle: sys::cusparseSpMMOpPlan_t,
37    _matrix_a: PhantomData<&'a SparseMatrixDescriptor<'a>>,
38    _matrix_b: PhantomData<&'a DenseMatrixDescriptor<'a>>,
39    _matrix_c: PhantomData<&'a DenseMatrixDescriptor<'a>>,
40}
41
42// Operation descriptors and plans own cuSPARSE analysis state but do not expose
43// pointer rebinding through shared references, so immutable sharing is allowed.
44unsafe impl Send for SpGemmDescriptor {}
45unsafe impl Sync for SpGemmDescriptor {}
46unsafe impl Send for SpSvDescriptor {}
47unsafe impl Sync for SpSvDescriptor {}
48unsafe impl Send for SpSmDescriptor {}
49unsafe impl Sync for SpSmDescriptor {}
50unsafe impl Send for SpMmOpPlan<'_> {}
51unsafe impl Sync for SpMmOpPlan<'_> {}
52
53impl SpGemmDescriptor {
54    pub fn create(ctx: &Context) -> Result<Self> {
55        ctx.bind()?;
56
57        let mut handle = ptr::null_mut();
58        unsafe {
59            try_ffi!(sys::cusparseSpGEMM_createDescr(&raw mut handle))?;
60        }
61
62        if handle.is_null() {
63            return Err(Error::NullHandle);
64        }
65
66        Ok(Self {
67            handle,
68            cuda_ctx: Arc::clone(ctx.cuda_context()),
69        })
70    }
71
72    pub fn num_products(&self) -> Result<usize> {
73        let mut value = 0_i64;
74        unsafe {
75            try_ffi!(sys::cusparseSpGEMM_getNumProducts(
76                self.as_raw(),
77                &raw mut value,
78            ))?;
79        }
80
81        usize::try_from(value).map_err(|_| Error::OutOfRange {
82            name: "num_products".into(),
83        })
84    }
85
86    pub fn as_raw(&self) -> sys::cusparseSpGEMMDescr_t {
87        self.handle
88    }
89
90    pub(crate) fn ensure_context(&self, ctx: &Context) -> Result<()> {
91        if self.cuda_ctx.as_ref() != ctx.cuda_context().as_ref() {
92            return Err(Error::PlanContextMismatch);
93        }
94        Ok(())
95    }
96
97    /// Wraps an existing cuSPARSE SpGEMM descriptor and takes ownership of it.
98    ///
99    /// # Safety
100    ///
101    /// `handle` must be a valid `cusparseSpGEMMDescr_t` associated with `ctx`.
102    /// Ownership of `handle` is transferred to the returned descriptor, and the
103    /// handle must not be destroyed elsewhere after calling this function.
104    pub unsafe fn from_raw(handle: sys::cusparseSpGEMMDescr_t, ctx: &Context) -> Result<Self> {
105        if handle.is_null() {
106            return Err(Error::NullHandle);
107        }
108
109        Ok(Self {
110            handle,
111            cuda_ctx: Arc::clone(ctx.cuda_context()),
112        })
113    }
114
115    /// Consumes the descriptor and returns the raw cuSPARSE handle without
116    /// destroying it.
117    ///
118    /// The caller becomes responsible for eventually destroying the returned
119    /// handle with `cusparseSpGEMM_destroyDescr`.
120    pub fn into_raw(self) -> sys::cusparseSpGEMMDescr_t {
121        let descriptor = ManuallyDrop::new(self);
122        descriptor.handle
123    }
124}
125
126impl Drop for SpGemmDescriptor {
127    fn drop(&mut self) {
128        unsafe {
129            if let Err(err) = try_ffi!(sys::cusparseSpGEMM_destroyDescr(self.handle)) {
130                #[cfg(debug_assertions)]
131                eprintln!("failed to destroy cusparse spgemm descriptor: {err}");
132            }
133        }
134    }
135}
136
137impl SpSvDescriptor {
138    pub fn create(ctx: &Context) -> Result<Self> {
139        ctx.bind()?;
140
141        let mut handle = ptr::null_mut();
142        unsafe {
143            try_ffi!(sys::cusparseSpSV_createDescr(&raw mut handle))?;
144        }
145
146        if handle.is_null() {
147            return Err(Error::NullHandle);
148        }
149
150        Ok(Self {
151            handle,
152            cuda_ctx: Arc::clone(ctx.cuda_context()),
153        })
154    }
155
156    pub fn as_raw(&self) -> sys::cusparseSpSVDescr_t {
157        self.handle
158    }
159
160    pub(crate) fn ensure_context(&self, ctx: &Context) -> Result<()> {
161        if self.cuda_ctx.as_ref() != ctx.cuda_context().as_ref() {
162            return Err(Error::PlanContextMismatch);
163        }
164        Ok(())
165    }
166
167    /// Wraps an existing cuSPARSE SpSV descriptor and takes ownership of it.
168    ///
169    /// # Safety
170    ///
171    /// `handle` must be a valid `cusparseSpSVDescr_t` associated with `ctx`.
172    /// Ownership of `handle` is transferred to the returned descriptor, and the
173    /// handle must not be destroyed elsewhere after calling this function.
174    pub unsafe fn from_raw(handle: sys::cusparseSpSVDescr_t, ctx: &Context) -> Result<Self> {
175        if handle.is_null() {
176            return Err(Error::NullHandle);
177        }
178
179        Ok(Self {
180            handle,
181            cuda_ctx: Arc::clone(ctx.cuda_context()),
182        })
183    }
184
185    /// Consumes the descriptor and returns the raw cuSPARSE handle without
186    /// destroying it.
187    ///
188    /// The caller becomes responsible for eventually destroying the returned
189    /// handle with `cusparseSpSV_destroyDescr`.
190    pub fn into_raw(self) -> sys::cusparseSpSVDescr_t {
191        let descriptor = ManuallyDrop::new(self);
192        descriptor.handle
193    }
194}
195
196impl Drop for SpSvDescriptor {
197    fn drop(&mut self) {
198        unsafe {
199            if let Err(err) = try_ffi!(sys::cusparseSpSV_destroyDescr(self.handle)) {
200                #[cfg(debug_assertions)]
201                eprintln!("failed to destroy cusparse spsv descriptor: {err}");
202            }
203        }
204    }
205}
206
207impl SpSmDescriptor {
208    pub fn create(ctx: &Context) -> Result<Self> {
209        ctx.bind()?;
210
211        let mut handle = ptr::null_mut();
212        unsafe {
213            try_ffi!(sys::cusparseSpSM_createDescr(&raw mut handle))?;
214        }
215
216        if handle.is_null() {
217            return Err(Error::NullHandle);
218        }
219
220        Ok(Self {
221            handle,
222            cuda_ctx: Arc::clone(ctx.cuda_context()),
223        })
224    }
225
226    pub fn as_raw(&self) -> sys::cusparseSpSMDescr_t {
227        self.handle
228    }
229
230    pub(crate) fn ensure_context(&self, ctx: &Context) -> Result<()> {
231        if self.cuda_ctx.as_ref() != ctx.cuda_context().as_ref() {
232            return Err(Error::PlanContextMismatch);
233        }
234        Ok(())
235    }
236
237    /// Wraps an existing cuSPARSE SpSM descriptor and takes ownership of it.
238    ///
239    /// # Safety
240    ///
241    /// `handle` must be a valid `cusparseSpSMDescr_t` associated with `ctx`.
242    /// Ownership of `handle` is transferred to the returned descriptor, and the
243    /// handle must not be destroyed elsewhere after calling this function.
244    pub unsafe fn from_raw(handle: sys::cusparseSpSMDescr_t, ctx: &Context) -> Result<Self> {
245        if handle.is_null() {
246            return Err(Error::NullHandle);
247        }
248
249        Ok(Self {
250            handle,
251            cuda_ctx: Arc::clone(ctx.cuda_context()),
252        })
253    }
254
255    /// Consumes the descriptor and returns the raw cuSPARSE handle without
256    /// destroying it.
257    ///
258    /// The caller becomes responsible for eventually destroying the returned
259    /// handle with `cusparseSpSM_destroyDescr`.
260    pub fn into_raw(self) -> sys::cusparseSpSMDescr_t {
261        let descriptor = ManuallyDrop::new(self);
262        descriptor.handle
263    }
264}
265
266impl<'a> SpMmOpPlan<'a> {
267    pub fn create(
268        ctx: &'a Context,
269        op_a: Operation,
270        op_b: Operation,
271        matrix_a: &'a SparseMatrixDescriptor<'a>,
272        matrix_b: &'a DenseMatrixDescriptor<'a>,
273        matrix_c: &'a mut DenseMatrixDescriptor<'a>,
274        compute_type: DataType,
275        algorithm: SpMmOpAlgorithm,
276        add_operation_ltoir: Option<&[u8]>,
277        mul_operation_ltoir: Option<&[u8]>,
278        epilogue_ltoir: Option<&[u8]>,
279    ) -> Result<(Self, usize)> {
280        matrix_a.ensure_context(ctx)?;
281        matrix_b.ensure_context(ctx)?;
282        matrix_c.ensure_context(ctx)?;
283        ctx.bind()?;
284
285        let mut handle = ptr::null_mut();
286        let mut workspace_size = 0;
287        let (add_ptr, add_len) = ltoir_parts(add_operation_ltoir);
288        let (mul_ptr, mul_len) = ltoir_parts(mul_operation_ltoir);
289        let (epilogue_ptr, epilogue_len) = ltoir_parts(epilogue_ltoir);
290        unsafe {
291            try_ffi!(sys::cusparseSpMMOp_createPlan(
292                ctx.as_raw(),
293                &raw mut handle,
294                op_a.into(),
295                op_b.into(),
296                matrix_a.as_raw_const(),
297                matrix_b.as_raw_const(),
298                matrix_c.as_raw(),
299                compute_type.into(),
300                algorithm.into(),
301                add_ptr.cast(),
302                add_len as _,
303                mul_ptr.cast(),
304                mul_len as _,
305                epilogue_ptr.cast(),
306                epilogue_len as _,
307                &raw mut workspace_size,
308            ))?;
309        }
310
311        if handle.is_null() {
312            return Err(Error::NullHandle);
313        }
314
315        Ok((
316            Self {
317                context: ctx,
318                handle,
319                _matrix_a: PhantomData,
320                _matrix_b: PhantomData,
321                _matrix_c: PhantomData,
322            },
323            usize::try_from(workspace_size).map_err(|_| Error::OutOfRange {
324                name: "spmm op workspace size".into(),
325            })?,
326        ))
327    }
328
329    pub fn context(&self) -> &Context {
330        self.context
331    }
332
333    /// NVRTC and nvJitLink are not currently available on Arm64 Android platforms.
334    ///
335    /// This operation does not support Android and Tegra platforms except Judy (sm87).
336    ///
337    /// Experimental: multiplies `matrix_a` and `matrix_b` with custom operators.
338    ///
339    /// where
340    ///
341    /// * `op(A)` is a sparse matrix of size $m \times k$.
342    /// * `op(B)` is a dense matrix of size $k \times n$.
343    /// * `C` is a dense matrix of size $m \times n$.
344    /// * $\oplus$, $\otimes$, and $\text{epilogue}$ are custom **add**, **mul**, and **epilogue** operators respectively.
345    ///
346    /// `op(A)` is selected by `op_a` and may be `A`, `A^T`, or `A^H`.
347    /// `op(B)` is selected by `op_b` and may be `B`, `B^T`, or `B^H`.
348    ///
349    /// Only `op_a == Operation::NonTranspose` is currently supported.
350    ///
351    /// [`SpMmOpPlan::create`] returns the workspace size together with the compiled kernel state needed by [`SpMmOpPlan::execute`].
352    ///
353    /// The custom add, multiply, and epilogue operators must accept and return the selected compute type.
354    /// The compute type may be `float`, `double`, `cuComplex`, `cuDoubleComplex`, or `int`.
355    ///
356    /// [`SpMmOpPlan::execute`] supports the following sparse matrix formats:
357    ///
358    /// * [`Format::Csr`](crate::types::Format::Csr)
359    ///
360    /// [`SpMmOpPlan::execute`] supports the following index type for representing `matrix_a`:
361    ///
362    /// * 32-bit indices ([`IndexType::I32`](crate::types::IndexType::I32))
363    /// * 64-bit indices ([`IndexType::I64`](crate::types::IndexType::I64))
364    ///
365    /// [`SpMmOpPlan::execute`] supports the following data types:
366    ///
367    /// Uniform-precision computation:
368    ///
369    /// | `A`/`B`/`C`/`compute_type` |
370    /// | --- |
371    /// | [`DataType::F32`] |
372    /// | [`DataType::F64`] |
373    /// | [`DataType::ComplexF32`] |
374    /// | [`DataType::ComplexF64`] |
375    ///
376    /// Mixed-precision computation:
377    ///
378    /// | `A`/`B` | `C` | `compute_type` |
379    /// | --- | --- | --- |
380    /// | [`DataType::I8`] | [`DataType::I32`] | [`DataType::I32`] |
381    /// | [`DataType::I8`] | [`DataType::F32`] | [`DataType::F32`] |
382    /// | [`DataType::F16`] | | |
383    /// | [`DataType::Bf16`] | | |
384    /// | [`DataType::F16`] | [`DataType::F16`] | |
385    /// | [`DataType::Bf16`] | [`DataType::Bf16`] | |
386    ///
387    /// [`SpMmOpPlan::execute`] supports the following algorithms:
388    ///
389    /// | Algorithm | Notes |
390    /// | --- | --- |
391    /// | [`SpMmOpAlgorithm::Default`] | Default algorithm for any sparse matrix format. |
392    ///
393    /// **Performance notes:**
394    ///
395    /// * Row-major layout provides higher performance than column-major.
396    ///
397    /// [`SpMmOpPlan::execute`] has the following properties:
398    ///
399    /// * Requires extra storage.
400    /// * Supports asynchronous execution.
401    /// * Provides deterministic (bitwise) results for each run.
402    /// * Allows the indices of `matrix_a` to be unsorted.
403    ///
404    /// [`SpMmOpPlan::execute`] supports the following optimizations:
405    ///
406    /// * CUDA graph capture.
407    /// * Hardware Memory Compression.
408    ///
409    /// # Errors
410    ///
411    /// Returns an error if the CUDA context cannot be bound or if cuSPARSE
412    /// rejects the prepared SpMM operation.
413    #[allow(deprecated)]
414    pub fn execute(&self, external_buffer: Option<DevicePtr>) -> Result<()> {
415        self.context.bind()?;
416
417        unsafe {
418            try_ffi!(sys::cusparseSpMMOp(
419                self.as_raw(),
420                external_buffer.unwrap_or(DevicePtr::null()).as_ptr() as _,
421            ))?;
422        }
423        Ok(())
424    }
425
426    pub fn as_raw(&self) -> sys::cusparseSpMMOpPlan_t {
427        self.handle
428    }
429
430    /// Wraps an existing cuSPARSE SpMM operation plan and takes ownership of it.
431    ///
432    /// # Safety
433    ///
434    /// `handle` must be a valid `cusparseSpMMOpPlan_t` associated with `ctx`
435    /// and with matrix descriptors that remain valid for lifetime `'a`.
436    /// Ownership of `handle` is transferred to the returned plan, and the
437    /// handle must not be destroyed elsewhere after calling this function.
438    pub unsafe fn from_raw(ctx: &'a Context, handle: sys::cusparseSpMMOpPlan_t) -> Result<Self> {
439        if handle.is_null() {
440            return Err(Error::NullHandle);
441        }
442
443        Ok(Self {
444            context: ctx,
445            handle,
446            _matrix_a: PhantomData,
447            _matrix_b: PhantomData,
448            _matrix_c: PhantomData,
449        })
450    }
451
452    /// Consumes the plan and returns the raw cuSPARSE handle without destroying
453    /// it.
454    ///
455    /// The caller becomes responsible for eventually destroying the returned
456    /// handle with `cusparseSpMMOp_destroyPlan`.
457    pub fn into_raw(self) -> sys::cusparseSpMMOpPlan_t {
458        let plan = ManuallyDrop::new(self);
459        plan.handle
460    }
461}
462
463impl Drop for SpMmOpPlan<'_> {
464    fn drop(&mut self) {
465        if let Err(err) = self.context.bind() {
466            #[cfg(debug_assertions)]
467            eprintln!("failed to bind cusparse spmm op plan context during drop: {err}");
468            return;
469        }
470
471        unsafe {
472            if let Err(err) = try_ffi!(sys::cusparseSpMMOp_destroyPlan(self.handle)) {
473                #[cfg(debug_assertions)]
474                eprintln!("failed to destroy cusparse spmm op plan: {err}");
475            }
476        }
477    }
478}
479
480fn ltoir_parts(buffer: Option<&[u8]>) -> (*const u8, usize) {
481    match buffer {
482        Some(buffer) => (buffer.as_ptr(), buffer.len()),
483        None => (ptr::null(), 0),
484    }
485}
486
487impl Drop for SpSmDescriptor {
488    fn drop(&mut self) {
489        unsafe {
490            if let Err(err) = try_ffi!(sys::cusparseSpSM_destroyDescr(self.handle)) {
491                #[cfg(debug_assertions)]
492                eprintln!("failed to destroy cusparse spsm descriptor: {err}");
493            }
494        }
495    }
496}
singe_cusparse/operation.rs

singe_cusparse/
operation.rs
