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