singe_cutensor/
execution.rs

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

singe_cutensor/
execution.rs
