singe_cutensor/
tensor.rs

1use std::{mem::ManuallyDrop, ptr};
2
3use singe_cuda::data_type::{DataType, DataTypeLike};
4
5use crate::{
6    context::{Context, ContextRef},
7    error::{Error, Result},
8    sys, try_ffi,
9    utility::to_i64_vec,
10};
11
12#[derive(Debug)]
13pub struct TensorDescriptor {
14    handle: sys::cutensorTensorDescriptor_t,
15    context: ContextRef,
16    shape: Vec<u64>,
17    strides: Option<Vec<u64>>,
18    data_type: DataType,
19    alignment_requirement: u32,
20}
21
22impl TensorDescriptor {
23    pub fn create_for<T: DataTypeLike>(
24        ctx: &Context,
25        shape: &[u64],
26        alignment_requirement: u32,
27    ) -> Result<Self> {
28        Self::create(ctx, shape, T::data_type(), alignment_requirement)
29    }
30
31    pub fn create(
32        ctx: &Context,
33        shape: &[u64],
34        data_type: DataType,
35        alignment_requirement: u32,
36    ) -> Result<Self> {
37        Self::create_with_strides(ctx, shape, None, data_type, alignment_requirement)
38    }
39
40    /// Creates a tensor descriptor.
41    ///
42    /// This allocates a small amount of host memory.
43    ///
44    /// The descriptor frees the associated cuTENSOR resources when dropped.
45    ///
46    /// This non-blocking call is thread-safe but not reentrant.
47    ///
48    /// # Errors
49    ///
50    /// Returns an error if the context cannot be bound, `shape` and `strides`
51    /// have different ranks, an extent or stride cannot be represented for
52    /// cuTENSOR, the data type or descriptor layout is unsupported, or cuTENSOR
53    /// returns a null descriptor handle.
54    pub fn create_with_strides(
55        ctx: &Context,
56        shape: &[u64],
57        strides: Option<&[u64]>,
58        data_type: DataType,
59        alignment_requirement: u32,
60    ) -> Result<Self> {
61        ctx.bind()?;
62
63        let rank = shape.len() as u32;
64        if let Some(strides) = strides
65            && shape.len() != strides.len()
66        {
67            return Err(Error::TensorStrideMismatch {
68                rank,
69                stride_length: strides.len(),
70            });
71        }
72
73        let extent = to_i64_vec(shape, "shape")?;
74        let strides_vec = strides.map(<[u64]>::to_vec);
75        let stride_i64 = strides_vec
76            .as_ref()
77            .map(|values| to_i64_vec(values, "strides"))
78            .transpose()?;
79
80        let mut handle = ptr::null_mut();
81        unsafe {
82            try_ffi!(sys::cutensorCreateTensorDescriptor(
83                ctx.as_raw(),
84                &raw mut handle,
85                rank,
86                extent.as_ptr(),
87                stride_i64.as_ref().map_or(ptr::null(), Vec::as_ptr),
88                data_type.into(),
89                alignment_requirement,
90            ))?;
91        }
92
93        if handle.is_null() {
94            return Err(Error::NullHandle);
95        }
96
97        Ok(Self {
98            handle,
99            context: ctx.as_context_ref(),
100            shape: shape.to_vec(),
101            strides: strides_vec,
102            data_type,
103            alignment_requirement,
104        })
105    }
106
107    pub fn create_with_strides_for<T: DataTypeLike>(
108        ctx: &Context,
109        shape: &[u64],
110        strides: Option<&[u64]>,
111        alignment_requirement: u32,
112    ) -> Result<Self> {
113        Self::create_with_strides(ctx, shape, strides, T::data_type(), alignment_requirement)
114    }
115
116    pub fn shape(&self) -> &[u64] {
117        &self.shape
118    }
119
120    pub fn strides(&self) -> Option<&[u64]> {
121        self.strides.as_deref()
122    }
123
124    pub fn rank(&self) -> u32 {
125        self.shape.len() as u32
126    }
127
128    pub fn data_type(&self) -> DataType {
129        self.data_type
130    }
131
132    pub fn alignment_requirement(&self) -> u32 {
133        self.alignment_requirement
134    }
135
136    pub(crate) fn context(&self) -> &ContextRef {
137        &self.context
138    }
139
140    pub const fn as_raw(&self) -> sys::cutensorTensorDescriptor_t {
141        self.handle
142    }
143
144    /// Wraps an existing cuTENSOR tensor descriptor and takes ownership of it.
145    ///
146    /// # Safety
147    ///
148    /// `handle` must be a valid `cutensorTensorDescriptor_t` whose rank,
149    /// extents, strides, data type, and alignment requirement match the
150    /// metadata supplied to this function. Ownership of `handle` is transferred
151    /// to the returned descriptor, and the handle must not be destroyed
152    /// elsewhere after calling this function.
153    pub unsafe fn from_raw(
154        handle: sys::cutensorTensorDescriptor_t,
155        ctx: &Context,
156        shape: Vec<u64>,
157        strides: Option<Vec<u64>>,
158        data_type: DataType,
159        alignment_requirement: u32,
160    ) -> Self {
161        Self {
162            handle,
163            context: ctx.as_context_ref(),
164            shape,
165            strides,
166            data_type,
167            alignment_requirement,
168        }
169    }
170
171    /// Consumes the descriptor and returns the raw cuTENSOR tensor descriptor
172    /// handle without destroying it.
173    ///
174    /// The caller becomes responsible for eventually destroying the returned
175    /// handle with cuTENSOR.
176    pub fn into_raw(self) -> sys::cutensorTensorDescriptor_t {
177        let descriptor = ManuallyDrop::new(self);
178        descriptor.handle
179    }
180}
181
182impl Drop for TensorDescriptor {
183    fn drop(&mut self) {
184        if let Err(err) = self.context.bind() {
185            #[cfg(debug_assertions)]
186            eprintln!("failed to bind cutensor context before destroying tensor descriptor: {err}");
187        }
188
189        unsafe {
190            if let Err(err) = try_ffi!(sys::cutensorDestroyTensorDescriptor(self.handle)) {
191                #[cfg(debug_assertions)]
192                eprintln!("failed to destroy cutensor tensor descriptor: {err}");
193            }
194        }
195    }
196}
197
198#[derive(Debug)]
199pub struct BlockSparseTensorDescriptor {
200    handle: sys::cutensorBlockSparseTensorDescriptor_t,
201    context: ContextRef,
202    mode_count: u32,
203    non_zero_block_count: u64,
204    data_type: DataType,
205}
206
207impl BlockSparseTensorDescriptor {
208    /// Creates a block-sparse tensor descriptor.
209    ///
210    /// A block-sparse tensor descriptor fully specifies the data layout of a block-sparse tensor (currently limited to up to 8 modes).
211    ///
212    /// Consider an example for a block-sparse tensor of order 2, namely a block-sparse matrix `A`.
213    /// Its first mode (`rows`) is subdivided into 5 sections (with extents 4, 2, 3, 4, 5, respectively), and its second mode (`columns`) is subdivided into 3 sections (with extents 2, 3, 7).
214    /// The matrix has 8 non-zero blocks:
215    ///
216    /// | 2 columns | 3 columns | 7 columns |
217    /// | --- | --- | --- |
218    /// | 4 x 2 | | 4 x 7 |
219    /// | | 2 x 3 | |
220    /// | | 3 x 3 | 3 x 7 |
221    /// | 4 x 2 | | |
222    /// | | 5 x 3 | 5 x 7 |
223    ///
224    /// Each subsection has the same extent across its entire mode: every block
225    /// within the same section of a mode has the same extent.
226    /// For example, in the table above every block on the right has 7 columns,
227    /// and all blocks on the bottom have 5 rows.
228    ///
229    /// Only non-zero blocks are stored; zero blocks are left blank in the
230    /// representation above.
231    ///
232    /// For example:
233    ///
234    /// * strides of block #0: 5, 1, 10, 20.
235    ///   Sorted strides would be: 1, 5, 10, 20.
236    ///   The permutation to sort the strides in this case is to swap the first two elements.
237    /// * strides of block #1: 10, 1, 30, 60.
238    ///   Applying the permutation would result in: 1, 10, 30, 60.
239    ///   The result is sorted in ascending order, which is allowed.
240    /// * strides of block #2: 1, 5, 50, 100.
241    ///   Applying permutation would result in: 5, 1, 50, 100.
242    ///   The result is *not* sorted in ascending order, this is *not* allowed.
243    ///
244    /// This non-blocking call is thread-safe but not reentrant.
245    ///
246    /// # Errors
247    ///
248    /// Returns an error if the context cannot be bound, the section extent,
249    /// coordinate, or stride slices do not match the descriptor dimensions, a
250    /// value cannot be represented for cuTENSOR, cuTENSOR rejects the
251    /// descriptor, or cuTENSOR returns a null descriptor handle.
252    pub fn create(
253        ctx: &Context,
254        section_count_per_mode: &[u32],
255        section_extents: &[u64],
256        non_zero_block_count: u64,
257        non_zero_coordinates: &[i32],
258        block_strides: Option<&[u64]>,
259        data_type: DataType,
260    ) -> Result<Self> {
261        ctx.bind()?;
262
263        let mode_count = section_count_per_mode.len() as u32;
264        let expected_section_extents = section_count_per_mode
265            .iter()
266            .try_fold(0usize, |sum, &value| sum.checked_add(value as usize))
267            .ok_or(Error::OutOfRange {
268                name: "section_extents".into(),
269            })?;
270        if section_extents.len() != expected_section_extents {
271            return Err(Error::LengthMismatch {
272                name: "section_extents".into(),
273                expected: expected_section_extents,
274                actual: section_extents.len(),
275            });
276        }
277
278        let expected_coordinates =
279            mode_count
280                .checked_mul(non_zero_block_count as u32)
281                .ok_or(Error::OutOfRange {
282                    name: "non_zero_coordinates".into(),
283                })?;
284        if non_zero_coordinates.len() != expected_coordinates as usize {
285            return Err(Error::LengthMismatch {
286                name: "non_zero_coordinates".into(),
287                expected: expected_coordinates as usize,
288                actual: non_zero_coordinates.len(),
289            });
290        }
291
292        if let Some(block_strides) = block_strides
293            && block_strides.len() != expected_coordinates as usize
294        {
295            return Err(Error::TensorStrideMismatch {
296                rank: expected_coordinates,
297                stride_length: block_strides.len(),
298            });
299        }
300
301        let num_sections_per_mode = section_count_per_mode.to_vec();
302        let section_extents = to_i64_vec(section_extents, "section_extents")?;
303        let block_strides = block_strides
304            .map(|block_strides| to_i64_vec(block_strides, "block_strides"))
305            .transpose()?;
306
307        let mut handle = ptr::null_mut();
308        unsafe {
309            try_ffi!(sys::cutensorCreateBlockSparseTensorDescriptor(
310                ctx.as_raw(),
311                &raw mut handle,
312                mode_count,
313                non_zero_block_count,
314                num_sections_per_mode.as_ptr(),
315                section_extents.as_ptr(),
316                non_zero_coordinates.as_ptr(),
317                block_strides.as_ref().map_or(ptr::null(), Vec::as_ptr),
318                data_type.into(),
319            ))?;
320        }
321
322        if handle.is_null() {
323            return Err(Error::NullHandle);
324        }
325
326        Ok(Self {
327            handle,
328            context: ctx.as_context_ref(),
329            mode_count,
330            non_zero_block_count,
331            data_type,
332        })
333    }
334
335    pub fn mode_count(&self) -> u32 {
336        self.mode_count
337    }
338
339    pub fn non_zero_block_count(&self) -> u64 {
340        self.non_zero_block_count
341    }
342
343    pub fn data_type(&self) -> DataType {
344        self.data_type
345    }
346
347    pub(crate) fn context(&self) -> &ContextRef {
348        &self.context
349    }
350
351    pub const fn as_raw(&self) -> sys::cutensorBlockSparseTensorDescriptor_t {
352        self.handle
353    }
354
355    /// Wraps an existing cuTENSOR block-sparse tensor descriptor and takes
356    /// ownership of it.
357    ///
358    /// # Safety
359    ///
360    /// `handle` must be a valid `cutensorBlockSparseTensorDescriptor_t` whose
361    /// mode count, non-zero block count, and data type match the metadata
362    /// supplied to this function. Ownership of `handle` is transferred to the
363    /// returned descriptor, and the handle must not be destroyed elsewhere after
364    /// calling this function.
365    pub unsafe fn from_raw(
366        handle: sys::cutensorBlockSparseTensorDescriptor_t,
367        ctx: &Context,
368        mode_count: u32,
369        non_zero_block_count: u64,
370        data_type: DataType,
371    ) -> Self {
372        Self {
373            handle,
374            context: ctx.as_context_ref(),
375            mode_count,
376            non_zero_block_count,
377            data_type,
378        }
379    }
380
381    /// Consumes the descriptor and returns the raw cuTENSOR block-sparse tensor
382    /// descriptor handle without destroying it.
383    ///
384    /// The caller becomes responsible for eventually destroying the returned
385    /// handle with cuTENSOR.
386    pub fn into_raw(self) -> sys::cutensorBlockSparseTensorDescriptor_t {
387        let descriptor = ManuallyDrop::new(self);
388        descriptor.handle
389    }
390}
391
392impl Drop for BlockSparseTensorDescriptor {
393    fn drop(&mut self) {
394        if let Err(err) = self.context.bind() {
395            #[cfg(debug_assertions)]
396            eprintln!(
397                "failed to bind cutensor context before destroying block sparse tensor descriptor: {err}"
398            );
399        }
400
401        unsafe {
402            if let Err(err) = try_ffi!(sys::cutensorDestroyBlockSparseTensorDescriptor(self.handle))
403            {
404                #[cfg(debug_assertions)]
405                eprintln!("failed to destroy cutensor block sparse tensor descriptor: {err}");
406            }
407        }
408    }
409}
singe_cutensor/tensor.rs

singe_cutensor/
tensor.rs
