singe_cutensor/
execution.rs

1#[allow(unused_imports)]
2use crate::error::Status;
3
4use std::ptr;
5
6use singe_cuda::{data_type::DataTypeLike, memory::DeviceMemory, stream::Stream, types::DevicePtr};
7
8use crate::{
9    error::{Error, Result},
10    operation::OperationKind,
11    plan::Plan,
12    sys, try_ffi,
13};
14
15impl Plan {
16    /// Performs the three-input element-wise tensor operation encoded by this plan.
17    ///
18    /// The operation has the form:
19    ///
20    /// $$ D\_{\Pi^C(i\_0,i\_1,...,i\_n)} = \Phi\_{ABC}(\Phi\_{AB}(\alpha op\_A(A\_{\Pi^A(i\_0,i\_1,...,i\_n)}), \beta op\_B(B\_{\Pi^B(i\_0,i\_1,...,i\_n)})), \gamma op\_C(C\_{\Pi^C(i\_0,i\_1,...,i\_n)})) $$
21    ///
22    /// See [`OperationDescriptor::elementwise_trinary`](crate::operation::OperationDescriptor::elementwise_trinary) for details.
23    ///
24    /// The call enqueues asynchronous CUDA work on `stream`. cuTENSOR treats this
25    /// operation as thread-safe but not reentrant.
26    ///
27    /// # Errors
28    ///
29    /// Returns an error if `stream` belongs to a different CUDA context, scalar
30    /// or tensor types do not match the plan, tensor dimensions or modes are
31    /// invalid, a data type or operation combination is unsupported, or
32    /// cuTENSOR rejects the operation.
33    pub fn elementwise_trinary<TA, TB, TC, TD, TScalar>(
34        &self,
35        alpha: &TScalar,
36        a: &DeviceMemory<TA>,
37        beta: &TScalar,
38        b: &DeviceMemory<TB>,
39        gamma: &TScalar,
40        c: &DeviceMemory<TC>,
41        d: &mut DeviceMemory<TD>,
42        stream: &Stream,
43    ) -> Result<()>
44    where
45        TA: DataTypeLike,
46        TB: DataTypeLike,
47        TC: DataTypeLike,
48        TD: DataTypeLike,
49        TScalar: DataTypeLike,
50    {
51        self.context().ensure_stream(stream)?;
52        validate_elementwise_trinary_types::<TA, TB, TC, TD, TScalar>(self)?;
53
54        unsafe {
55            try_ffi!(sys::cutensorElementwiseTrinaryExecute(
56                self.context().as_raw(),
57                self.as_raw(),
58                ptr::from_ref(alpha) as _,
59                a.as_ptr() as _,
60                ptr::from_ref(beta) as _,
61                b.as_ptr() as _,
62                ptr::from_ref(gamma) as _,
63                c.as_ptr() as _,
64                d.as_mut_ptr() as _,
65                stream.as_raw(),
66            ))?;
67        }
68
69        Ok(())
70    }
71
72    /// Performs the two-input element-wise tensor operation encoded by this plan.
73    ///
74    /// The operation has the form:
75    ///
76    /// $$ D\_{\Pi^C(i\_0,i\_1,...,i\_n)} = \Phi\_{AC}(\alpha op\_A(A\_{\Pi^A(i\_0,i\_1,...,i\_n)}), \gamma op\_C(C\_{\Pi^C(i\_0,i\_1,...,i\_n)})) $$
77    ///
78    /// See [`OperationDescriptor::elementwise_binary`](crate::operation::OperationDescriptor::elementwise_binary) for details.
79    ///
80    /// The call enqueues asynchronous CUDA work on `stream`. cuTENSOR treats this
81    /// operation as thread-safe but not reentrant.
82    ///
83    /// # Errors
84    ///
85    /// Returns an error if `stream` belongs to a different CUDA context, scalar
86    /// or tensor types do not match the plan, tensor dimensions or modes are
87    /// invalid, a data type or operation combination is unsupported, or
88    /// cuTENSOR rejects the operation.
89    pub fn elementwise_binary<TA, TC, TD, TScalar>(
90        &self,
91        alpha: &TScalar,
92        a: &DeviceMemory<TA>,
93        gamma: &TScalar,
94        c: &DeviceMemory<TC>,
95        d: &mut DeviceMemory<TD>,
96        stream: &Stream,
97    ) -> Result<()>
98    where
99        TA: DataTypeLike,
100        TC: DataTypeLike,
101        TD: DataTypeLike,
102        TScalar: DataTypeLike,
103    {
104        self.context().ensure_stream(stream)?;
105        validate_elementwise_binary_types::<TA, TC, TD, TScalar>(self)?;
106
107        unsafe {
108            try_ffi!(sys::cutensorElementwiseBinaryExecute(
109                self.context().as_raw(),
110                self.as_raw(),
111                ptr::from_ref(alpha) as _,
112                a.as_ptr() as _,
113                ptr::from_ref(gamma) as _,
114                c.as_ptr() as _,
115                d.as_mut_ptr() as _,
116                stream.as_raw(),
117            ))?;
118        }
119
120        Ok(())
121    }
122
123    /// Performs the tensor permutation encoded by this plan.
124    ///
125    /// The operation has the form:
126    ///
127    /// $$ B\_{\Pi^B(i\_0,i\_1,...,i\_n)} = \alpha op\_A(A\_{\Pi^A(i\_0,i\_1,...,i\_n)}) $$
128    ///
129    /// Consequently, this performs an out-of-place tensor permutation.
130    ///
131    /// Where:
132    ///
133    /// * `A` and `B` are multi-mode tensors of arbitrary data types.
134    /// * $\Pi^A$ and $\Pi^B$ permute the modes of `A` and `B`, respectively.
135    /// * $op\_A$ is the unary element-wise operator specified when creating the
136    ///   operation descriptor with [`OperationDescriptor::permutation`](crate::operation::OperationDescriptor::permutation).
137    ///
138    /// The call enqueues asynchronous CUDA work on `stream`. cuTENSOR treats this
139    /// operation as thread-safe but not reentrant.
140    ///
141    /// # Errors
142    ///
143    /// Returns an error if `stream` belongs to a different CUDA context, scalar
144    /// or tensor types do not match the plan, tensor dimensions or modes are
145    /// invalid, a data type or operation combination is unsupported, or
146    /// cuTENSOR rejects the operation.
147    pub fn permute<TA, TB, TScalar>(
148        &self,
149        alpha: &TScalar,
150        a: &DeviceMemory<TA>,
151        b: &mut DeviceMemory<TB>,
152        stream: &Stream,
153    ) -> Result<()>
154    where
155        TA: DataTypeLike,
156        TB: DataTypeLike,
157        TScalar: DataTypeLike,
158    {
159        self.context().ensure_stream(stream)?;
160        validate_permute_types::<TA, TB, TScalar>(self)?;
161
162        unsafe {
163            try_ffi!(sys::cutensorPermute(
164                self.context().as_raw(),
165                self.as_raw(),
166                ptr::from_ref(alpha) as _,
167                a.as_ptr() as _,
168                b.as_mut_ptr() as _,
169                stream.as_raw(),
170            ))?;
171        }
172
173        Ok(())
174    }
175
176    /// Computes the tensor contraction $D = alpha \cdot A \cdot B + beta \cdot C$.
177    ///
178    /// $$ \mathcal{D}\_{{modes}\_\mathcal{D}} \gets \alpha \cdot \mathcal{A}\_{{modes}\_\mathcal{A}} B\_{{modes}\_\mathcal{B}} + \beta \mathcal{C}\_{{modes}\_\mathcal{C}} $$
179    ///
180    /// The active CUDA device must match the CUDA device that was active when
181    /// the plan was created.
182    ///
183    /// See [NVIDIA/CUDALibrarySamples](https://github.com/NVIDIA/CUDALibrarySamples/tree/master/cuTENSOR/contraction.cu) for a concrete example.
184    ///
185    /// # Errors
186    ///
187    /// Returns an error if `stream` belongs to a different CUDA context, scalar
188    /// or tensor types do not match the plan, the active device does not match
189    /// the plan, the provided workspace is smaller than
190    /// [`Plan::required_workspace_size`], the CUDA driver is insufficient, or
191    /// cuTENSOR rejects or fails the operation.
192    pub fn contract<TA, TB, TC, TD, TScalar>(
193        &self,
194        alpha: &TScalar,
195        a: &DeviceMemory<TA>,
196        b: &DeviceMemory<TB>,
197        beta: &TScalar,
198        c: &DeviceMemory<TC>,
199        d: &mut DeviceMemory<TD>,
200        workspace: Option<&mut DeviceMemory<u8>>,
201        stream: &Stream,
202    ) -> Result<()>
203    where
204        TA: DataTypeLike,
205        TB: DataTypeLike,
206        TC: DataTypeLike,
207        TD: DataTypeLike,
208        TScalar: DataTypeLike,
209    {
210        self.context().ensure_stream(stream)?;
211        validate_contraction_types::<TA, TB, TC, TD, TScalar>(self)?;
212        let (workspace_ptr, workspace_size) =
213            workspace_ptr_and_size(workspace, self.required_workspace_size())?;
214
215        unsafe {
216            try_ffi!(sys::cutensorContract(
217                self.context().as_raw(),
218                self.as_raw(),
219                ptr::from_ref(alpha) as _,
220                a.as_ptr() as _,
221                b.as_ptr() as _,
222                ptr::from_ref(beta) as _,
223                c.as_ptr() as _,
224                d.as_mut_ptr() as _,
225                workspace_ptr as _,
226                workspace_size,
227                stream.as_raw(),
228            ))?;
229        }
230
231        Ok(())
232    }
233
234    /// Performs the tensor reduction that is encoded by this plan.
235    ///
236    /// See [`OperationDescriptor::reduction`](crate::operation::OperationDescriptor::reduction) for details.
237    /// The call enqueues asynchronous CUDA work on `stream`. The provided
238    /// `workspace`, when required by the plan, must remain valid until the
239    /// queued work has completed.
240    ///
241    /// # Errors
242    ///
243    /// Returns an error if `stream` belongs to a different CUDA context, the
244    /// scalar and tensor types do not match the plan, the provided workspace is
245    /// smaller than [`Plan::required_workspace_size`], or cuTENSOR rejects the
246    /// operation.
247    pub fn reduce<TA, TC, TD, TScalar>(
248        &self,
249        alpha: &TScalar,
250        a: &DeviceMemory<TA>,
251        beta: &TScalar,
252        c: &DeviceMemory<TC>,
253        d: &mut DeviceMemory<TD>,
254        workspace: Option<&mut DeviceMemory<u8>>,
255        stream: &Stream,
256    ) -> Result<()>
257    where
258        TA: DataTypeLike,
259        TC: DataTypeLike,
260        TD: DataTypeLike,
261        TScalar: DataTypeLike,
262    {
263        self.context().ensure_stream(stream)?;
264        validate_reduction_types::<TA, TC, TD, TScalar>(self)?;
265        let (workspace_ptr, workspace_size) =
266            workspace_ptr_and_size(workspace, self.required_workspace_size())?;
267
268        unsafe {
269            try_ffi!(sys::cutensorReduce(
270                self.context().as_raw(),
271                self.as_raw(),
272                ptr::from_ref(alpha) as _,
273                a.as_ptr() as _,
274                ptr::from_ref(beta) as _,
275                c.as_ptr() as _,
276                d.as_mut_ptr() as _,
277                workspace_ptr as _,
278                workspace_size,
279                stream.as_raw(),
280            ))?;
281        }
282
283        Ok(())
284    }
285
286    /// Computes the tensor contraction $E = alpha \cdot A \cdot B \cdot C + beta \cdot D$.
287    ///
288    /// $$ \mathcal{E}\_{{modes}\_\mathcal{E}} \gets \alpha \cdot \mathcal{A}\_{{modes}\_\mathcal{A}} \mathcal{B}\_{{modes}\_\mathcal{B}} \mathcal{C}\_{{modes}\_\mathcal{C}} + \beta \mathcal{D}\_{{modes}\_\mathcal{D}} $$
289    ///
290    /// The active CUDA device must match the CUDA device that was active when
291    /// the plan was created.
292    ///
293    /// See [NVIDIA/CUDALibrarySamples](https://github.com/NVIDIA/CUDALibrarySamples/tree/master/cuTENSOR/contraction_trinary.cu) for a concrete example.
294    ///
295    /// # Errors
296    ///
297    /// Returns an error if `stream` belongs to a different CUDA context, scalar
298    /// or tensor types do not match the plan, the active device does not match
299    /// the plan, the provided workspace is smaller than
300    /// [`Plan::required_workspace_size`], the CUDA driver is insufficient, or
301    /// cuTENSOR rejects or fails the operation.
302    pub fn contract_trinary<TA, TB, TC, TD, TE, TScalar>(
303        &self,
304        alpha: &TScalar,
305        a: &DeviceMemory<TA>,
306        b: &DeviceMemory<TB>,
307        c: &DeviceMemory<TC>,
308        beta: &TScalar,
309        d: &DeviceMemory<TD>,
310        e: &mut DeviceMemory<TE>,
311        workspace: Option<&mut DeviceMemory<u8>>,
312        stream: &Stream,
313    ) -> Result<()>
314    where
315        TA: DataTypeLike,
316        TB: DataTypeLike,
317        TC: DataTypeLike,
318        TD: DataTypeLike,
319        TE: DataTypeLike,
320        TScalar: DataTypeLike,
321    {
322        self.context().ensure_stream(stream)?;
323        validate_contraction_trinary_types::<TA, TB, TC, TD, TE, TScalar>(self)?;
324        let (workspace_ptr, workspace_size) =
325            workspace_ptr_and_size(workspace, self.required_workspace_size())?;
326
327        unsafe {
328            try_ffi!(sys::cutensorContractTrinary(
329                self.context().as_raw(),
330                self.as_raw(),
331                ptr::from_ref(alpha) as _,
332                a.as_ptr() as _,
333                b.as_ptr() as _,
334                c.as_ptr() as _,
335                ptr::from_ref(beta) as _,
336                d.as_ptr() as _,
337                e.as_mut_ptr() as _,
338                workspace_ptr as _,
339                workspace_size,
340                stream.as_raw(),
341            ))?;
342        }
343
344        Ok(())
345    }
346
347    /// Computes the block-sparse tensor contraction $D = alpha \cdot A \cdot B + beta \cdot C$.
348    ///
349    /// $$ \mathcal{D}\_{{modes}\_\mathcal{D}} \gets \alpha \cdot \mathcal{A}\_{{modes}\_\mathcal{A}} B\_{{modes}\_\mathcal{B}} + \beta \mathcal{C}\_{{modes}\_\mathcal{C}} $$
350    ///
351    /// The active CUDA device must match the CUDA device that was active when
352    /// the plan was created.
353    ///
354    /// The array parameters `A`, `B`, `C`, and `D` are host arrays containing pointers to GPU-accessible memory.
355    /// For example, `A` is a host array whose size equals the number of non-zero blocks in tensor $\mathcal{A}$.
356    /// `A\[i\]` is a pointer to the GPU-accessible memory location of block number `i` of $\mathcal{A}$.
357    /// The blocks are numbered in the same way as in the construction of $\mathcal{A}$’s block-sparse tensor descriptor.
358    /// The same applies to the other array parameters `B`, `C`, and `D`.
359    ///
360    /// # Errors
361    ///
362    /// Returns an error if `stream` belongs to a different CUDA context, the
363    /// pointer slices do not match the block counts encoded in the plan, the
364    /// active device does not match the plan, the provided workspace is smaller
365    /// than [`Plan::required_workspace_size`], the CUDA driver is insufficient,
366    /// or cuTENSOR rejects or fails the operation.
367    pub fn block_sparse_contract<TScalar>(
368        &self,
369        alpha: &TScalar,
370        a: &[DevicePtr],
371        b: &[DevicePtr],
372        beta: &TScalar,
373        c: &[DevicePtr],
374        d: &mut [DevicePtr],
375        workspace: Option<&mut DeviceMemory<u8>>,
376        stream: &Stream,
377    ) -> Result<()>
378    where
379        TScalar: DataTypeLike,
380    {
381        self.context().ensure_stream(stream)?;
382        validate_block_sparse_contract_arguments::<TScalar>(self, a, b, c, d)?;
383        let (workspace_ptr, workspace_size) =
384            workspace_ptr_and_size(workspace, self.required_workspace_size())?;
385
386        let a = a
387            .iter()
388            .map(|ptr| ptr.as_ptr().cast_const())
389            .collect::<Vec<_>>();
390        let b = b
391            .iter()
392            .map(|ptr| ptr.as_ptr().cast_const())
393            .collect::<Vec<_>>();
394        let c = c
395            .iter()
396            .map(|ptr| ptr.as_ptr().cast_const())
397            .collect::<Vec<_>>();
398        let mut d = d.iter().map(|ptr| ptr.as_ptr()).collect::<Vec<_>>();
399
400        unsafe {
401            try_ffi!(sys::cutensorBlockSparseContract(
402                self.context().as_raw(),
403                self.as_raw(),
404                ptr::from_ref(alpha) as _,
405                a.as_ptr() as _,
406                b.as_ptr() as _,
407                ptr::from_ref(beta) as _,
408                c.as_ptr() as _,
409                d.as_mut_ptr() as _,
410                workspace_ptr as _,
411                workspace_size,
412                stream.as_raw(),
413            ))?;
414        }
415
416        Ok(())
417    }
418}
419
420fn validate_elementwise_trinary_types<TA, TB, TC, TD, TScalar>(plan: &Plan) -> Result<()>
421where
422    TA: DataTypeLike,
423    TB: DataTypeLike,
424    TC: DataTypeLike,
425    TD: DataTypeLike,
426    TScalar: DataTypeLike,
427{
428    let signature = plan.signature();
429    let OperationKind::ElementwiseTrinary { a, b, c, d } = signature.kind else {
430        return Err(Error::PlanOperationMismatch {
431            expected: "elementwise_trinary".into(),
432            actual: signature.kind.name().into(),
433        });
434    };
435    validate_scalar_type(signature.scalar_type, TScalar::data_type())?;
436    validate_tensor_type(a, TA::data_type())?;
437    validate_tensor_type(b, TB::data_type())?;
438    validate_tensor_type(c, TC::data_type())?;
439    validate_tensor_type(d, TD::data_type())?;
440    Ok(())
441}
442
443fn validate_elementwise_binary_types<TA, TC, TD, TScalar>(plan: &Plan) -> Result<()>
444where
445    TA: DataTypeLike,
446    TC: DataTypeLike,
447    TD: DataTypeLike,
448    TScalar: DataTypeLike,
449{
450    let signature = plan.signature();
451    let OperationKind::ElementwiseBinary { a, c, d } = signature.kind else {
452        return Err(Error::PlanOperationMismatch {
453            expected: "elementwise_binary".into(),
454            actual: signature.kind.name().into(),
455        });
456    };
457    validate_scalar_type(signature.scalar_type, TScalar::data_type())?;
458    validate_tensor_type(a, TA::data_type())?;
459    validate_tensor_type(c, TC::data_type())?;
460    validate_tensor_type(d, TD::data_type())?;
461    Ok(())
462}
463
464fn validate_permute_types<TA, TB, TScalar>(plan: &Plan) -> Result<()>
465where
466    TA: DataTypeLike,
467    TB: DataTypeLike,
468    TScalar: DataTypeLike,
469{
470    let signature = plan.signature();
471    let OperationKind::Permutation { a, b } = signature.kind else {
472        return Err(Error::PlanOperationMismatch {
473            expected: "permute".into(),
474            actual: signature.kind.name().into(),
475        });
476    };
477    validate_scalar_type(signature.scalar_type, TScalar::data_type())?;
478    validate_tensor_type(a, TA::data_type())?;
479    validate_tensor_type(b, TB::data_type())?;
480    Ok(())
481}
482
483fn validate_contraction_types<TA, TB, TC, TD, TScalar>(plan: &Plan) -> Result<()>
484where
485    TA: DataTypeLike,
486    TB: DataTypeLike,
487    TC: DataTypeLike,
488    TD: DataTypeLike,
489    TScalar: DataTypeLike,
490{
491    let signature = plan.signature();
492    let OperationKind::Contraction { a, b, c, d } = signature.kind else {
493        return Err(Error::PlanOperationMismatch {
494            expected: "contract".into(),
495            actual: signature.kind.name().into(),
496        });
497    };
498    validate_scalar_type(signature.scalar_type, TScalar::data_type())?;
499    validate_tensor_type(a, TA::data_type())?;
500    validate_tensor_type(b, TB::data_type())?;
501    validate_tensor_type(c, TC::data_type())?;
502    validate_tensor_type(d, TD::data_type())?;
503    Ok(())
504}
505
506fn validate_reduction_types<TA, TC, TD, TScalar>(plan: &Plan) -> Result<()>
507where
508    TA: DataTypeLike,
509    TC: DataTypeLike,
510    TD: DataTypeLike,
511    TScalar: DataTypeLike,
512{
513    let signature = plan.signature();
514    let OperationKind::Reduction { a, c, d } = signature.kind else {
515        return Err(Error::PlanOperationMismatch {
516            expected: "reduce".into(),
517            actual: signature.kind.name().into(),
518        });
519    };
520    validate_scalar_type(signature.scalar_type, TScalar::data_type())?;
521    validate_tensor_type(a, TA::data_type())?;
522    validate_tensor_type(c, TC::data_type())?;
523    validate_tensor_type(d, TD::data_type())?;
524    Ok(())
525}
526
527fn validate_contraction_trinary_types<TA, TB, TC, TD, TE, TScalar>(plan: &Plan) -> Result<()>
528where
529    TA: DataTypeLike,
530    TB: DataTypeLike,
531    TC: DataTypeLike,
532    TD: DataTypeLike,
533    TE: DataTypeLike,
534    TScalar: DataTypeLike,
535{
536    let signature = plan.signature();
537    let OperationKind::ContractionTrinary { a, b, c, d, e } = signature.kind else {
538        return Err(Error::PlanOperationMismatch {
539            expected: "contract_trinary".into(),
540            actual: signature.kind.name().into(),
541        });
542    };
543    validate_scalar_type(signature.scalar_type, TScalar::data_type())?;
544    validate_tensor_type(a, TA::data_type())?;
545    validate_tensor_type(b, TB::data_type())?;
546    validate_tensor_type(c, TC::data_type())?;
547    validate_tensor_type(d, TD::data_type())?;
548    validate_tensor_type(e, TE::data_type())?;
549    Ok(())
550}
551
552fn validate_block_sparse_contract_arguments<TScalar>(
553    plan: &Plan,
554    a: &[DevicePtr],
555    b: &[DevicePtr],
556    c: &[DevicePtr],
557    d: &[DevicePtr],
558) -> Result<()>
559where
560    TScalar: DataTypeLike,
561{
562    let signature = plan.signature();
563    let OperationKind::BlockSparseContraction {
564        a_non_zero_blocks,
565        b_non_zero_blocks,
566        c_non_zero_blocks,
567        d_non_zero_blocks,
568    } = signature.kind
569    else {
570        return Err(Error::PlanOperationMismatch {
571            expected: "block_sparse_contract".into(),
572            actual: signature.kind.name().into(),
573        });
574    };
575    validate_scalar_type(signature.scalar_type, TScalar::data_type())?;
576    validate_pointer_count("a", a_non_zero_blocks as usize, a.len())?;
577    validate_pointer_count("b", b_non_zero_blocks as usize, b.len())?;
578    validate_pointer_count("c", c_non_zero_blocks as usize, c.len())?;
579    validate_pointer_count("d", d_non_zero_blocks as usize, d.len())?;
580    Ok(())
581}
582
583fn validate_tensor_type(
584    expected: singe_cuda::data_type::DataType,
585    actual: singe_cuda::data_type::DataType,
586) -> Result<()> {
587    if expected != actual {
588        return Err(Error::TensorMemoryDataTypeMismatch {
589            descriptor: expected,
590            memory: actual,
591        });
592    }
593    Ok(())
594}
595
596fn validate_scalar_type(
597    expected: singe_cuda::data_type::DataType,
598    actual: singe_cuda::data_type::DataType,
599) -> Result<()> {
600    if expected != actual {
601        return Err(Error::ScalarDataTypeMismatch {
602            descriptor: expected,
603            memory: actual,
604        });
605    }
606    Ok(())
607}
608
609fn validate_pointer_count(name: &str, expected: usize, actual: usize) -> Result<()> {
610    if expected != actual {
611        return Err(Error::PointerCountMismatch {
612            name: name.into(),
613            expected,
614            actual,
615        });
616    }
617    Ok(())
618}
619
620fn workspace_ptr_and_size(
621    workspace: Option<&mut DeviceMemory<u8>>,
622    required: u64,
623) -> Result<(*mut (), u64)> {
624    let Some(workspace) = workspace else {
625        if required == 0 {
626            return Ok((ptr::null_mut(), 0));
627        }
628
629        return Err(Error::InsufficientWorkspaceSize {
630            required,
631            actual: 0,
632        });
633    };
634
635    let actual = workspace.byte_len() as u64;
636    if actual < required {
637        return Err(Error::InsufficientWorkspaceSize { required, actual });
638    }
639    if required != 0
640        && !workspace.as_ptr().is_null()
641        && !(workspace.as_ptr() as usize).is_multiple_of(256)
642    {
643        return Err(Error::WorkspaceMisaligned {
644            required_alignment: 256,
645            size: workspace.byte_len(),
646        });
647    }
648
649    Ok((workspace.as_mut_ptr() as _, required))
650}
651
652#[cfg(all(test, feature = "testing"))]
653mod tests {
654    use super::*;
655    use crate::{
656        error::Error,
657        operation::{ComputeDescriptor, OperationDescriptor, TensorOperand},
658        plan::{Plan, PlanPreference},
659        tensor::TensorDescriptor,
660        testing::setup_context,
661        types::{Operator, WorkspacePreference},
662    };
663    use singe_core::assert_close;
664    use singe_cuda::data_type::DataType;
665
666    fn packed_offset(indices: &[u64], extents: &[u64]) -> usize {
667        let mut stride = 1_usize;
668        let mut offset = 0_usize;
669
670        for (&index, &extent) in indices.iter().zip(extents) {
671            offset += index as usize * stride;
672            stride *= extent as usize;
673        }
674
675        offset
676    }
677
678    #[test]
679    fn test_permute_matches_cpu_reference() -> Result<()> {
680        let context = setup_context()?;
681        let stream = context.cuda_context().create_stream()?;
682
683        let extent_a = vec![2, 3];
684        let extent_b = vec![3, 2];
685        let mode_a = vec!['i'.into(), 'j'.into()];
686        let mode_b = vec!['j'.into(), 'i'.into()];
687
688        let host_a = vec![1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0];
689        let device_a = DeviceMemory::from_slice(&host_a)?;
690        let mut device_b = DeviceMemory::<f32>::create(host_a.len())?;
691
692        const ALIGNMENT: u32 = 128;
693
694        let descriptor_a = TensorDescriptor::create(&context, &extent_a, DataType::F32, ALIGNMENT)?;
695        let descriptor_b = TensorDescriptor::create(&context, &extent_b, DataType::F32, ALIGNMENT)?;
696
697        let permutation = OperationDescriptor::permutation(
698            &context,
699            TensorOperand::identity(&descriptor_a, &mode_a),
700            TensorOperand::identity(&descriptor_b, &mode_b),
701            ComputeDescriptor::f32(),
702        )?;
703        let preference = PlanPreference::create_default(&context)?;
704        let workspace_estimate = Plan::estimate_workspace_size(
705            &context,
706            &permutation,
707            &preference,
708            WorkspacePreference::Default,
709        )?;
710        let plan = Plan::create(&context, &permutation, &preference, workspace_estimate)?;
711
712        let alpha = 2.5_f32;
713        plan.permute(&alpha, &device_a, &mut device_b, &stream)?;
714
715        stream.synchronize()?;
716
717        let result = device_b.copy_to_host_vec()?;
718        let mut expected = vec![0.0_f32; host_a.len()];
719
720        for i in 0..extent_a[0] {
721            for j in 0..extent_a[1] {
722                let a_offset = packed_offset(&[i, j], &extent_a);
723                let b_offset = packed_offset(&[j, i], &extent_b);
724                expected[b_offset] = alpha * host_a[a_offset];
725            }
726        }
727
728        assert_close!(&result, &expected);
729
730        Ok(())
731    }
732
733    #[test]
734    fn test_reduce_matches_cpu_reference() -> Result<()> {
735        let context = setup_context()?;
736        let stream = context.cuda_context().create_stream()?;
737
738        let extent_a = vec![2, 3];
739        let extent_cd = vec![2];
740        let mode_a = vec!['i'.into(), 'j'.into()];
741        let mode_cd = vec!['i'.into()];
742
743        let host_a = vec![1.0_f32, 2.0, 3.0, 4.0, 5.0, 6.0];
744        let host_c = vec![0.5_f32, -1.0];
745
746        let device_a = DeviceMemory::from_slice(&host_a)?;
747        let device_c = DeviceMemory::from_slice(&host_c)?;
748        let mut device_d = DeviceMemory::from_slice(&host_c)?;
749
750        const ALIGNMENT: u32 = 128;
751
752        let descriptor_a = TensorDescriptor::create(&context, &extent_a, DataType::F32, ALIGNMENT)?;
753        let descriptor_cd =
754            TensorDescriptor::create(&context, &extent_cd, DataType::F32, ALIGNMENT)?;
755
756        let reduction = OperationDescriptor::reduction(
757            &context,
758            TensorOperand::identity(&descriptor_a, &mode_a),
759            TensorOperand::identity(&descriptor_cd, &mode_cd),
760            TensorOperand::identity(&descriptor_cd, &mode_cd),
761            Operator::Add,
762            ComputeDescriptor::f32(),
763        )?;
764        let preference = PlanPreference::create_default(&context)?;
765        let workspace_estimate = Plan::estimate_workspace_size(
766            &context,
767            &reduction,
768            &preference,
769            WorkspacePreference::Default,
770        )?;
771        let plan = Plan::create(&context, &reduction, &preference, workspace_estimate)?;
772
773        let alpha = 1.25_f32;
774        let beta = -0.5_f32;
775
776        if plan.required_workspace_size() == 0 {
777            plan.reduce(
778                &alpha,
779                &device_a,
780                &beta,
781                &device_c,
782                &mut device_d,
783                None,
784                &stream,
785            )?;
786        } else {
787            let mut workspace =
788                DeviceMemory::<u8>::create(plan.required_workspace_size() as usize)?;
789            plan.reduce(
790                &alpha,
791                &device_a,
792                &beta,
793                &device_c,
794                &mut device_d,
795                Some(&mut workspace),
796                &stream,
797            )?;
798        }
799
800        stream.synchronize()?;
801
802        let result = device_d.copy_to_host_vec()?;
803        let mut expected = vec![0.0_f32; extent_cd[0] as usize];
804
805        for i in 0..extent_cd[0] as usize {
806            let mut sum = 0.0_f32;
807            for j in 0..extent_a[1] {
808                let a_offset = packed_offset(&[i as u64, j], &extent_a);
809                sum += host_a[a_offset];
810            }
811            expected[i] = alpha * sum + beta * host_c[i];
812        }
813
814        assert_close!(&result, &expected);
815
816        Ok(())
817    }
818
819    #[test]
820    fn test_padding_attributes_validate_inputs() -> Result<()> {
821        let context = setup_context()?;
822
823        let extent_a = vec![2, 3];
824        let extent_b = vec![3, 2];
825        let mode_a = vec!['i'.into(), 'j'.into()];
826        let mode_b = vec!['j'.into(), 'i'.into()];
827        const ALIGNMENT: u32 = 128;
828
829        let descriptor_a = TensorDescriptor::create(&context, &extent_a, DataType::F32, ALIGNMENT)?;
830        let descriptor_b = TensorDescriptor::create(&context, &extent_b, DataType::F32, ALIGNMENT)?;
831
832        let mut permutation = OperationDescriptor::permutation(
833            &context,
834            TensorOperand::identity(&descriptor_a, &mode_a),
835            TensorOperand::identity(&descriptor_b, &mode_b),
836            ComputeDescriptor::f32(),
837        )?;
838
839        let err = permutation.set_padding_left(&[1]).unwrap_err();
840        assert!(matches!(
841            err,
842            Error::LengthMismatch {
843                expected: 2,
844                actual: 1,
845                ..
846            }
847        ));
848
849        let err = permutation.set_padding_right(&[1]).unwrap_err();
850        assert!(matches!(
851            err,
852            Error::LengthMismatch {
853                expected: 2,
854                actual: 1,
855                ..
856            }
857        ));
858
859        let err = permutation.set_padding_left(&[1, u32::MAX]).unwrap_err();
860        assert!(matches!(err, Error::OutOfRange { .. }));
861
862        let err = permutation.set_padding_right(&[1, u32::MAX]).unwrap_err();
863        assert!(matches!(err, Error::OutOfRange { .. }));
864
865        let err = permutation.set_padding_value(7_i32).unwrap_err();
866        assert!(matches!(
867            err,
868            Error::ScalarDataTypeMismatch {
869                descriptor: DataType::F32,
870                memory: DataType::I32,
871            }
872        ));
873
874        Ok(())
875    }
876
877    #[test]
878    fn test_operation_creation_validates_matching_output_descriptors_and_modes() -> Result<()> {
879        let context = setup_context()?;
880
881        const ALIGNMENT: u32 = 128;
882
883        let extent_a = vec![2, 3];
884        let extent_cd = vec![2, 3];
885        let extent_e = vec![3, 2];
886
887        let mode_a = vec!['i'.into(), 'j'.into()];
888        let mode_cd = vec!['i'.into(), 'j'.into()];
889        let mode_e = vec!['j'.into(), 'i'.into()];
890
891        let descriptor_a = TensorDescriptor::create(&context, &extent_a, DataType::F32, ALIGNMENT)?;
892        let descriptor_cd =
893            TensorDescriptor::create(&context, &extent_cd, DataType::F32, ALIGNMENT)?;
894        let descriptor_e = TensorDescriptor::create(&context, &extent_e, DataType::F32, ALIGNMENT)?;
895
896        let err = OperationDescriptor::elementwise_binary(
897            &context,
898            TensorOperand::identity(&descriptor_a, &mode_a),
899            TensorOperand::identity(&descriptor_cd, &mode_cd),
900            TensorOperand::identity(&descriptor_e, &mode_e),
901            Operator::Add,
902            ComputeDescriptor::f32(),
903        )
904        .unwrap_err();
905        assert!(matches!(err, Error::DescriptorMismatch { .. }));
906
907        let err = OperationDescriptor::reduction(
908            &context,
909            TensorOperand::identity(&descriptor_a, &mode_a),
910            TensorOperand::identity(&descriptor_cd, &mode_cd),
911            TensorOperand::identity(&descriptor_cd, &mode_e),
912            Operator::Add,
913            ComputeDescriptor::f32(),
914        )
915        .unwrap_err();
916        assert!(matches!(err, Error::ModeMismatch { .. }));
917
918        let err = OperationDescriptor::contraction(
919            &context,
920            TensorOperand::identity(&descriptor_a, &mode_a),
921            TensorOperand::identity(&descriptor_a, &mode_a),
922            TensorOperand::identity(&descriptor_cd, &mode_cd),
923            TensorOperand::identity(&descriptor_e, &mode_e),
924            ComputeDescriptor::f32(),
925        )
926        .unwrap_err();
927        assert!(matches!(err, Error::DescriptorMismatch { .. }));
928
929        let err = OperationDescriptor::contraction_trinary(
930            &context,
931            TensorOperand::identity(&descriptor_a, &mode_a),
932            TensorOperand::identity(&descriptor_a, &mode_a),
933            TensorOperand::identity(&descriptor_a, &mode_a),
934            TensorOperand::identity(&descriptor_cd, &mode_cd),
935            TensorOperand::identity(&descriptor_cd, &mode_e),
936            ComputeDescriptor::f32(),
937        )
938        .unwrap_err();
939        assert!(matches!(err, Error::ModeMismatch { .. }));
940
941        Ok(())
942    }
943
944    #[test]
945    fn test_contraction_matches_cpu_reference() -> Result<()> {
946        let context = setup_context()?;
947        let stream = context.cuda_context().create_stream()?;
948
949        let mode_c = vec!['m'.into(), 'u'.into(), 'n'.into(), 'v'.into()];
950        let mode_a = vec!['m'.into(), 'h'.into(), 'k'.into(), 'n'.into()];
951        let mode_b = vec!['u'.into(), 'k'.into(), 'v'.into(), 'h'.into()];
952
953        let extent_m = 2;
954        let extent_u = 2;
955        let extent_n = 3;
956        let extent_v = 2;
957        let extent_h = 2;
958        let extent_k = 2;
959
960        let extent_c = vec![extent_m, extent_u, extent_n, extent_v];
961        let extent_a = vec![extent_m, extent_h, extent_k, extent_n];
962        let extent_b = vec![extent_u, extent_k, extent_v, extent_h];
963
964        let elements_a = extent_a.iter().product::<u64>();
965        let elements_b = extent_b.iter().product::<u64>();
966        let elements_c = extent_c.iter().product::<u64>();
967
968        let host_a = (0..elements_a)
969            .map(|index| (index as f32 + 1.0) * 0.25)
970            .collect::<Vec<_>>();
971        let host_b = (0..elements_b)
972            .map(|index| ((index % 7) as f32 - 3.0) * 0.5)
973            .collect::<Vec<_>>();
974        let host_c = (0..elements_c)
975            .map(|index| (index as f32 - 4.0) * 0.125)
976            .collect::<Vec<_>>();
977
978        let device_a = DeviceMemory::from_slice(&host_a)?;
979        let device_b = DeviceMemory::from_slice(&host_b)?;
980        let device_c_input = DeviceMemory::from_slice(&host_c)?;
981        let mut device_c = DeviceMemory::from_slice(&host_c)?;
982
983        const ALIGNMENT: u32 = 128;
984
985        let descriptor_a = TensorDescriptor::create(&context, &extent_a, DataType::F32, ALIGNMENT)?;
986        let descriptor_b = TensorDescriptor::create(&context, &extent_b, DataType::F32, ALIGNMENT)?;
987        let descriptor_c = TensorDescriptor::create(&context, &extent_c, DataType::F32, ALIGNMENT)?;
988
989        let contraction = OperationDescriptor::contraction(
990            &context,
991            TensorOperand::identity(&descriptor_a, &mode_a),
992            TensorOperand::identity(&descriptor_b, &mode_b),
993            TensorOperand::identity(&descriptor_c, &mode_c),
994            TensorOperand::identity(&descriptor_c, &mode_c),
995            ComputeDescriptor::f32(),
996        )?;
997
998        assert_eq!(contraction.scalar_type()?, DataType::F32);
999
1000        let preference = PlanPreference::create_default(&context)?;
1001        let workspace_estimate = Plan::estimate_workspace_size(
1002            &context,
1003            &contraction,
1004            &preference,
1005            WorkspacePreference::Default,
1006        )?;
1007        let plan = Plan::create(&context, &contraction, &preference, workspace_estimate)?;
1008
1009        let mut workspace = DeviceMemory::<u8>::create(plan.required_workspace_size() as usize)?;
1010
1011        let alpha = 1.1_f32;
1012        let beta = 0.0_f32;
1013
1014        plan.contract(
1015            &alpha,
1016            &device_a,
1017            &device_b,
1018            &beta,
1019            &device_c_input,
1020            &mut device_c,
1021            Some(&mut workspace),
1022            &stream,
1023        )?;
1024
1025        stream.synchronize()?;
1026
1027        let result = device_c.copy_to_host_vec()?;
1028        let mut expected = vec![0.0_f32; elements_c as usize];
1029
1030        for m in 0..extent_m {
1031            for u in 0..extent_u {
1032                for n in 0..extent_n {
1033                    for v in 0..extent_v {
1034                        let mut sum = 0.0_f32;
1035                        for h in 0..extent_h {
1036                            for k in 0..extent_k {
1037                                let a_offset = packed_offset(&[m, h, k, n], &extent_a);
1038                                let b_offset = packed_offset(&[u, k, v, h], &extent_b);
1039                                sum += host_a[a_offset] * host_b[b_offset];
1040                            }
1041                        }
1042
1043                        let c_offset = packed_offset(&[m, u, n, v], &extent_c);
1044                        expected[c_offset] = alpha * sum + beta * host_c[c_offset];
1045                    }
1046                }
1047            }
1048        }
1049
1050        assert_close!(&result, &expected);
1051
1052        Ok(())
1053    }
1054}
singe_cutensor/execution.rs

singe_cutensor/
execution.rs
