singe_cusparse/
vector.rs

1use std::{
2    marker::PhantomData,
3    mem::{ManuallyDrop, MaybeUninit},
4    ptr,
5    sync::Arc,
6};
7
8use singe_cuda::{
9    context::Context as CudaContext, data_type::DataTypeLike, memory::DeviceMemory,
10    types::DevicePtr,
11};
12
13use crate::{
14    context::Context,
15    error::{Error, Result},
16    sys, try_ffi,
17    types::{IndexBase, IndexType, IndexTypeLike},
18    utility::{to_i64, to_usize},
19};
20
21#[derive(Debug, Clone, Copy, PartialEq, Eq)]
22pub struct SparseVectorInfo {
23    pub size: usize,
24    pub nonzero_count: usize,
25    pub indices: DevicePtr,
26    pub values: DevicePtr,
27    pub index_type: IndexType,
28    pub index_base: IndexBase,
29}
30
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub struct DenseVectorInfo {
33    pub size: usize,
34    pub values: DevicePtr,
35}
36
37#[derive(Debug)]
38pub struct SparseVectorDescriptor<'a> {
39    handle: sys::cusparseSpVecDescr_t,
40    cuda_ctx: Arc<CudaContext>,
41    _buffers: PhantomData<&'a mut ()>,
42}
43
44#[derive(Debug)]
45pub struct DenseVectorDescriptor<'a> {
46    handle: sys::cusparseDnVecDescr_t,
47    cuda_ctx: Arc<CudaContext>,
48    _buffers: PhantomData<&'a mut ()>,
49}
50
51// Vector descriptors borrow mutable device buffers and can rebind their value
52// pointers, so they are movable but not shared concurrently.
53unsafe impl Send for SparseVectorDescriptor<'_> {}
54unsafe impl Send for DenseVectorDescriptor<'_> {}
55
56impl<'a> SparseVectorDescriptor<'a> {
57    /// Initializes a sparse vector descriptor.
58    ///
59    /// The descriptor borrows the device index and value buffers by pointer; the
60    /// buffers must outlive descriptor use or be replaced with [`SparseVectorDescriptor::set_values`].
61    pub fn create<I: IndexTypeLike, T: DataTypeLike>(
62        ctx: &Context,
63        size: usize,
64        nonzero_count: usize,
65        indices: &'a mut DeviceMemory<I>,
66        values: &'a mut DeviceMemory<T>,
67        index_base: IndexBase,
68    ) -> Result<Self> {
69        ctx.bind()?;
70        let size = to_i64(size, "size")?;
71        let nonzero_count = to_i64(nonzero_count, "nonzero_count")?;
72        let mut handle = ptr::null_mut();
73        unsafe {
74            try_ffi!(sys::cusparseCreateSpVec(
75                &raw mut handle,
76                size,
77                nonzero_count,
78                indices.as_mut_ptr().cast(),
79                values.as_mut_ptr().cast(),
80                I::index_type().into(),
81                index_base.into(),
82                T::data_type().into(),
83            ))?;
84        }
85
86        if handle.is_null() {
87            return Err(Error::NullHandle);
88        }
89
90        Ok(Self {
91            handle,
92            cuda_ctx: Arc::clone(ctx.cuda_context()),
93            _buffers: PhantomData,
94        })
95    }
96
97    /// Returns the index base of this sparse vector descriptor.
98    ///
99    /// # Errors
100    ///
101    /// Returns an error if cuSPARSE cannot report the sparse vector index base.
102    pub fn index_base(&self) -> Result<IndexBase> {
103        let mut value = sys::cusparseIndexBase_t::CUSPARSE_INDEX_BASE_ZERO;
104        unsafe {
105            try_ffi!(sys::cusparseSpVecGetIndexBase(
106                self.as_raw_const(),
107                &raw mut value
108            ))?;
109        }
110        Ok(value.into())
111    }
112
113    /// Returns the values pointer of this sparse vector descriptor.
114    ///
115    /// # Errors
116    ///
117    /// Returns an error if cuSPARSE cannot report the values pointer.
118    pub fn values(&self) -> Result<DevicePtr> {
119        let mut values = ptr::null_mut();
120        unsafe {
121            try_ffi!(sys::cusparseSpVecGetValues(self.as_raw(), &raw mut values))?;
122        }
123        Ok(unsafe { DevicePtr::from_raw(values.cast()) })
124    }
125
126    /// Sets the values pointer of this sparse vector descriptor.
127    ///
128    /// The pointed-to device buffer must remain valid while the descriptor uses it.
129    ///
130    /// # Errors
131    ///
132    /// Returns an error if cuSPARSE rejects the values pointer.
133    pub fn set_values<T: DataTypeLike>(&mut self, values: &'a mut DeviceMemory<T>) -> Result<()> {
134        unsafe {
135            try_ffi!(sys::cusparseSpVecSetValues(
136                self.as_raw(),
137                values.as_mut_ptr().cast()
138            ))?;
139        }
140        Ok(())
141    }
142
143    /// Returns the fields of this sparse vector descriptor.
144    ///
145    /// # Errors
146    ///
147    /// Returns an error if cuSPARSE cannot report the descriptor fields or if
148    /// reported sizes cannot be represented as `usize`.
149    pub fn info(&self) -> Result<SparseVectorInfo> {
150        let mut size = 0_i64;
151        let mut nnz = 0_i64;
152        let mut indices = ptr::null_mut();
153        let mut values = ptr::null_mut();
154        let mut index_type = sys::cusparseIndexType_t::CUSPARSE_INDEX_32I;
155        let mut index_base = sys::cusparseIndexBase_t::CUSPARSE_INDEX_BASE_ZERO;
156        let mut value_type = MaybeUninit::uninit();
157        unsafe {
158            try_ffi!(sys::cusparseSpVecGet(
159                self.as_raw(),
160                &raw mut size,
161                &raw mut nnz,
162                &raw mut indices,
163                &raw mut values,
164                &raw mut index_type,
165                &raw mut index_base,
166                value_type.as_mut_ptr(),
167            ))?;
168        }
169        Ok(SparseVectorInfo {
170            size: to_usize(size, "size")?,
171            nonzero_count: to_usize(nnz, "nonzero_count")?,
172            indices: unsafe { DevicePtr::from_raw(indices.cast()) },
173            values: unsafe { DevicePtr::from_raw(values.cast()) },
174            index_type: index_type.into(),
175            index_base: index_base.into(),
176        })
177    }
178
179    pub fn as_raw(&self) -> sys::cusparseSpVecDescr_t {
180        self.handle
181    }
182
183    pub fn as_raw_const(&self) -> sys::cusparseConstSpVecDescr_t {
184        self.handle.cast_const()
185    }
186
187    pub fn cuda_context(&self) -> &Arc<CudaContext> {
188        &self.cuda_ctx
189    }
190
191    pub(crate) fn ensure_context(&self, ctx: &Context) -> Result<()> {
192        if self.cuda_ctx.as_ref() != ctx.cuda_context().as_ref() {
193            return Err(Error::PlanContextMismatch);
194        }
195        Ok(())
196    }
197
198    /// Wraps an existing cuSPARSE sparse-vector descriptor and takes ownership
199    /// of it.
200    ///
201    /// # Safety
202    ///
203    /// `handle` must be a valid `cusparseSpVecDescr_t`. Any device buffers
204    /// referenced by the descriptor must remain valid for lifetime `'a`.
205    /// Ownership of `handle` is transferred to the returned descriptor, and the
206    /// handle must not be destroyed elsewhere after calling this function.
207    pub unsafe fn from_raw(handle: sys::cusparseSpVecDescr_t, ctx: &Context) -> Result<Self> {
208        if handle.is_null() {
209            return Err(Error::NullHandle);
210        }
211
212        Ok(Self {
213            handle,
214            cuda_ctx: Arc::clone(ctx.cuda_context()),
215            _buffers: PhantomData,
216        })
217    }
218
219    /// Consumes the descriptor and returns the raw cuSPARSE handle without
220    /// destroying it.
221    ///
222    /// The caller becomes responsible for eventually destroying the returned
223    /// handle with `cusparseDestroySpVec`.
224    pub fn into_raw(self) -> sys::cusparseSpVecDescr_t {
225        let descriptor = ManuallyDrop::new(self);
226        descriptor.handle
227    }
228}
229
230impl Drop for SparseVectorDescriptor<'_> {
231    fn drop(&mut self) {
232        unsafe {
233            if let Err(err) = try_ffi!(sys::cusparseDestroySpVec(self.as_raw_const())) {
234                #[cfg(debug_assertions)]
235                eprintln!("failed to destroy cusparse sparse vector descriptor: {err}");
236            }
237        }
238    }
239}
240
241impl<'a> DenseVectorDescriptor<'a> {
242    /// Initializes a dense vector descriptor.
243    ///
244    /// The descriptor borrows the device value buffer by pointer; the buffer must
245    /// outlive descriptor use or be replaced with [`DenseVectorDescriptor::set_values`].
246    ///
247    /// # Errors
248    ///
249    /// Returns an error if `size` cannot be represented by cuSPARSE, if cuSPARSE
250    /// cannot create the descriptor, or if it returns a null handle.
251    pub fn create<T: DataTypeLike>(
252        ctx: &Context,
253        size: usize,
254        values: &'a mut DeviceMemory<T>,
255    ) -> Result<Self> {
256        ctx.bind()?;
257        let size = to_i64(size, "size")?;
258        let mut handle = ptr::null_mut();
259        unsafe {
260            try_ffi!(sys::cusparseCreateDnVec(
261                &raw mut handle,
262                size,
263                values.as_mut_ptr().cast(),
264                T::data_type().into(),
265            ))?;
266        }
267
268        if handle.is_null() {
269            return Err(Error::NullHandle);
270        }
271
272        Ok(Self {
273            handle,
274            cuda_ctx: Arc::clone(ctx.cuda_context()),
275            _buffers: PhantomData,
276        })
277    }
278
279    /// Returns the values pointer of this dense vector descriptor.
280    ///
281    /// # Errors
282    ///
283    /// Returns an error if cuSPARSE cannot report the values pointer.
284    pub fn values(&self) -> Result<DevicePtr> {
285        let mut values = ptr::null_mut();
286        unsafe {
287            try_ffi!(sys::cusparseDnVecGetValues(self.as_raw(), &raw mut values))?;
288        }
289        Ok(unsafe { DevicePtr::from_raw(values.cast()) })
290    }
291
292    /// Sets the values pointer of this dense vector descriptor.
293    ///
294    /// The pointed-to device buffer must remain valid while the descriptor uses it.
295    ///
296    /// # Errors
297    ///
298    /// Returns an error if cuSPARSE rejects the values pointer.
299    pub fn set_values<T: DataTypeLike>(&mut self, values: &'a mut DeviceMemory<T>) -> Result<()> {
300        unsafe {
301            try_ffi!(sys::cusparseDnVecSetValues(
302                self.as_raw(),
303                values.as_mut_ptr().cast()
304            ))?;
305        }
306        Ok(())
307    }
308
309    /// Returns the fields of this dense vector descriptor.
310    ///
311    /// # Errors
312    ///
313    /// Returns an error if cuSPARSE cannot report the descriptor fields or if the
314    /// reported size cannot be represented as `usize`.
315    pub fn info(&self) -> Result<DenseVectorInfo> {
316        let mut size = 0_i64;
317        let mut values = ptr::null_mut();
318        let mut value_type = MaybeUninit::uninit();
319        unsafe {
320            try_ffi!(sys::cusparseDnVecGet(
321                self.as_raw(),
322                &raw mut size,
323                &raw mut values,
324                value_type.as_mut_ptr(),
325            ))?;
326        }
327        Ok(DenseVectorInfo {
328            size: to_usize(size, "size")?,
329            values: unsafe { DevicePtr::from_raw(values.cast()) },
330        })
331    }
332
333    pub fn as_raw(&self) -> sys::cusparseDnVecDescr_t {
334        self.handle
335    }
336
337    pub fn as_raw_const(&self) -> sys::cusparseConstDnVecDescr_t {
338        self.handle.cast_const()
339    }
340
341    pub fn cuda_context(&self) -> &Arc<CudaContext> {
342        &self.cuda_ctx
343    }
344
345    pub(crate) fn ensure_context(&self, ctx: &Context) -> Result<()> {
346        if self.cuda_ctx.as_ref() != ctx.cuda_context().as_ref() {
347            return Err(Error::PlanContextMismatch);
348        }
349        Ok(())
350    }
351
352    /// Wraps an existing cuSPARSE dense-vector descriptor and takes ownership
353    /// of it.
354    ///
355    /// # Safety
356    ///
357    /// `handle` must be a valid `cusparseDnVecDescr_t`. Any device buffer
358    /// referenced by the descriptor must remain valid for lifetime `'a`.
359    /// Ownership of `handle` is transferred to the returned descriptor, and the
360    /// handle must not be destroyed elsewhere after calling this function.
361    pub unsafe fn from_raw(handle: sys::cusparseDnVecDescr_t, ctx: &Context) -> Result<Self> {
362        if handle.is_null() {
363            return Err(Error::NullHandle);
364        }
365
366        Ok(Self {
367            handle,
368            cuda_ctx: Arc::clone(ctx.cuda_context()),
369            _buffers: PhantomData,
370        })
371    }
372
373    /// Consumes the descriptor and returns the raw cuSPARSE handle without
374    /// destroying it.
375    ///
376    /// The caller becomes responsible for eventually destroying the returned
377    /// handle with `cusparseDestroyDnVec`.
378    pub fn into_raw(self) -> sys::cusparseDnVecDescr_t {
379        let descriptor = ManuallyDrop::new(self);
380        descriptor.handle
381    }
382}
383
384impl Drop for DenseVectorDescriptor<'_> {
385    fn drop(&mut self) {
386        unsafe {
387            if let Err(err) = try_ffi!(sys::cusparseDestroyDnVec(self.as_raw_const())) {
388                #[cfg(debug_assertions)]
389                eprintln!("failed to destroy cusparse dense vector descriptor: {err}");
390            }
391        }
392    }
393}
394
395/// Gathers the elements of `dense_vector` into `sparse_vector`.
396///
397/// In other words,
398///
399/// ```text
400/// for i=0 to nnz-1
401/// sparse_values[i] = dense_values[sparse_indices[i]]
402/// ```
403///
404/// [`gather`] supports the following index type for representing `sparse_vector`:
405///
406/// * 32-bit indices ([`IndexType::I32`])
407/// * 64-bit indices ([`IndexType::I64`])
408///
409/// [`gather`] supports the following data types:
410///
411/// | `X`/`Y` | Notes |
412/// | --- | --- |
413/// | [`DataType::F16`](singe_cuda::data_type::DataType::F16) |  |
414/// | [`DataType::Bf16`](singe_cuda::data_type::DataType::Bf16) |  |
415/// | [`DataType::F32`](singe_cuda::data_type::DataType::F32) |  |
416/// | [`DataType::F64`](singe_cuda::data_type::DataType::F64) |  |
417/// | [`DataType::ComplexF16`](singe_cuda::data_type::DataType::ComplexF16) | Deprecated. |
418/// | [`DataType::ComplexBf16`](singe_cuda::data_type::DataType::ComplexBf16) | Deprecated. |
419/// | [`DataType::ComplexF32`](singe_cuda::data_type::DataType::ComplexF32) |  |
420/// | [`DataType::ComplexF64`](singe_cuda::data_type::DataType::ComplexF64) |  |
421///
422/// [`gather`] has the following constraints:
423///
424/// * The arrays representing `sparse_vector` must be aligned to 16 bytes.
425///
426/// [`gather`] has the following properties:
427///
428/// * Requires no extra storage.
429/// * Supports asynchronous execution.
430/// * Provides deterministic (bitwise) results for each run if `sparse_vector` indices are distinct.
431/// * Allows the indices of `sparse_vector` to be unsorted.
432///
433/// [`gather`] supports the following optimizations:
434///
435/// * CUDA graph capture.
436/// * Hardware Memory Compression.
437pub fn gather(
438    ctx: &Context,
439    sparse_vector: &mut SparseVectorDescriptor,
440    dense_vector: &DenseVectorDescriptor,
441) -> Result<()> {
442    sparse_vector.ensure_context(ctx)?;
443    dense_vector.ensure_context(ctx)?;
444    ctx.bind()?;
445    unsafe {
446        try_ffi!(sys::cusparseGather(
447            ctx.as_raw(),
448            dense_vector.as_raw_const(),
449            sparse_vector.as_raw(),
450        ))?;
451    }
452    Ok(())
453}
454
455/// Scatters the elements of `sparse_vector` into `dense_vector`.
456///
457/// In other words,
458///
459/// ```text
460/// for i=0 to nnz-1
461/// dense_values[sparse_indices[i]] = sparse_values[i]
462/// ```
463///
464/// [`scatter`] supports the following index type for representing `sparse_vector`:
465///
466/// * 32-bit indices ([`IndexType::I32`])
467/// * 64-bit indices ([`IndexType::I64`])
468///
469/// [`scatter`] supports the following data types:
470///
471/// | `X`/`Y` | Notes |
472/// | --- | --- |
473/// | [`DataType::I8`](singe_cuda::data_type::DataType::I8) |  |
474/// | [`DataType::F16`](singe_cuda::data_type::DataType::F16) |  |
475/// | [`DataType::Bf16`](singe_cuda::data_type::DataType::Bf16) |  |
476/// | [`DataType::F32`](singe_cuda::data_type::DataType::F32) |  |
477/// | [`DataType::F64`](singe_cuda::data_type::DataType::F64) |  |
478/// | [`DataType::ComplexF16`](singe_cuda::data_type::DataType::ComplexF16) | Deprecated. |
479/// | [`DataType::ComplexBf16`](singe_cuda::data_type::DataType::ComplexBf16) | Deprecated. |
480/// | [`DataType::ComplexF32`](singe_cuda::data_type::DataType::ComplexF32) |  |
481/// | [`DataType::ComplexF64`](singe_cuda::data_type::DataType::ComplexF64) |  |
482///
483/// [`scatter`] has the following constraints:
484///
485/// * The arrays representing `sparse_vector` must be aligned to 16 bytes.
486///
487/// [`scatter`] has the following properties:
488///
489/// * Requires no extra storage.
490/// * Supports asynchronous execution.
491/// * Provides deterministic (bitwise) results for each run if `sparse_vector` indices are distinct.
492/// * Allows the indices of `sparse_vector` to be unsorted.
493///
494/// [`scatter`] supports the following optimizations:
495///
496/// * CUDA graph capture.
497/// * Hardware Memory Compression.
498pub fn scatter(
499    ctx: &Context,
500    sparse_vector: &SparseVectorDescriptor,
501    dense_vector: &mut DenseVectorDescriptor,
502) -> Result<()> {
503    sparse_vector.ensure_context(ctx)?;
504    dense_vector.ensure_context(ctx)?;
505    ctx.bind()?;
506    unsafe {
507        try_ffi!(sys::cusparseScatter(
508            ctx.as_raw(),
509            sparse_vector.as_raw_const(),
510            dense_vector.as_raw(),
511        ))?;
512    }
513    Ok(())
514}
singe_cusparse/vector.rs

singe_cusparse/
vector.rs
