singe_cutensor/mg/
contraction.rs

1use std::{mem::ManuallyDrop, ptr};
2
3use num_enum::{IntoPrimitive, TryFromPrimitive};
4use singe_core::impl_enum_conversion;
5use singe_cuda::stream::StreamBinding;
6use singe_cutensor_sys as sys;
7
8use crate::{
9    error::{Error, Result},
10    mg::{
11        context::{Context, ContextRef, validate_same_context},
12        copy::validate_count,
13        memory::{DistributedDeviceMemory, WorkspaceMemory},
14        tensor::TensorDescriptor,
15        types::Algorithm,
16        workspace::Workspace,
17    },
18    try_ffi,
19    types::{ComputeType, Mode, WorkspacePreference},
20    utility::modes_to_i32_vec,
21};
22
23#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, TryFromPrimitive, IntoPrimitive)]
24#[repr(u32)]
25#[non_exhaustive]
26pub enum ContractionFindAttribute {
27    Max = sys::cutensorMgContractionFindAttribute_t::CUTENSORMG_CONTRACTION_FIND_ATTRIBUTE_MAX as _,
28}
29
30impl_enum_conversion!(
31    sys::cutensorMgContractionFindAttribute_t,
32    ContractionFindAttribute
33);
34
35#[derive(Debug)]
36pub struct ContractionFind {
37    handle: sys::cutensorMgContractionFind_t,
38    context: ContextRef,
39    algorithm: Algorithm,
40}
41
42#[derive(Debug)]
43pub struct ContractionDescriptor {
44    handle: sys::cutensorMgContractionDescriptor_t,
45    context: ContextRef,
46    a_pointer_count: usize,
47    b_pointer_count: usize,
48    c_pointer_count: usize,
49    d_pointer_count: usize,
50}
51
52#[derive(Debug)]
53pub struct ContractionPlan {
54    handle: sys::cutensorMgContractionPlan_t,
55    context: ContextRef,
56    a_pointer_count: usize,
57    b_pointer_count: usize,
58    c_pointer_count: usize,
59    d_pointer_count: usize,
60    workspace: Workspace,
61}
62
63impl ContractionFind {
64    pub fn create(context: &Context, algorithm: Algorithm) -> Result<Self> {
65        let mut handle = ptr::null_mut();
66        unsafe {
67            try_ffi!(sys::cutensorMgCreateContractionFind(
68                context.as_raw(),
69                &raw mut handle,
70                algorithm.into(),
71            ))?;
72        }
73
74        if handle.is_null() {
75            return Err(Error::NullHandle);
76        }
77
78        Ok(Self {
79            handle,
80            context: context.as_context_ref(),
81            algorithm,
82        })
83    }
84
85    pub fn create_default(context: &Context) -> Result<Self> {
86        Self::create(context, Algorithm::Default)
87    }
88
89    /// Wraps an existing cuTENSORMg contraction find handle.
90    ///
91    /// # Safety
92    ///
93    /// `handle` must be a valid cuTENSORMg contraction find handle associated
94    /// with `context` and `algorithm`. The returned value takes ownership of
95    /// `handle` and destroys it on drop.
96    pub unsafe fn from_raw(
97        handle: sys::cutensorMgContractionFind_t,
98        context: &Context,
99        algorithm: Algorithm,
100    ) -> Result<Self> {
101        if handle.is_null() {
102            return Err(Error::NullHandle);
103        }
104
105        Ok(Self {
106            handle,
107            context: context.as_context_ref(),
108            algorithm,
109        })
110    }
111
112    pub fn algorithm(&self) -> Algorithm {
113        self.algorithm
114    }
115
116    pub fn set_attribute<T>(&mut self, attr: ContractionFindAttribute, value: &T) -> Result<()> {
117        unsafe {
118            try_ffi!(sys::cutensorMgContractionFindSetAttribute(
119                self.context.as_raw(),
120                self.handle,
121                attr.into(),
122                ptr::from_ref(value).cast(),
123                size_of::<T>() as i64,
124            ))?;
125        }
126        Ok(())
127    }
128
129    pub(crate) fn context(&self) -> &ContextRef {
130        &self.context
131    }
132
133    pub const fn as_raw(&self) -> sys::cutensorMgContractionFind_t {
134        self.handle
135    }
136
137    /// Consumes this find handle and returns the owned raw cuTENSORMg handle.
138    ///
139    /// The caller becomes responsible for destroying the handle.
140    pub fn into_raw(self) -> sys::cutensorMgContractionFind_t {
141        let this = ManuallyDrop::new(self);
142        this.handle
143    }
144}
145
146impl ContractionDescriptor {
147    pub fn create(
148        context: &Context,
149        a: &TensorDescriptor,
150        modes_a: &[Mode],
151        b: &TensorDescriptor,
152        modes_b: &[Mode],
153        c: &TensorDescriptor,
154        modes_c: &[Mode],
155        d: &TensorDescriptor,
156        modes_d: &[Mode],
157        compute: ComputeType,
158    ) -> Result<Self> {
159        let context_ref = context.as_context_ref();
160        validate_same_context(&context_ref, a.context(), "a")?;
161        validate_same_context(&context_ref, b.context(), "b")?;
162        validate_same_context(&context_ref, c.context(), "c")?;
163        validate_same_context(&context_ref, d.context(), "d")?;
164        validate_modes(a, modes_a)?;
165        validate_modes(b, modes_b)?;
166        validate_modes(c, modes_c)?;
167        validate_modes(d, modes_d)?;
168
169        let modes_a = modes_to_i32_vec(modes_a);
170        let modes_b = modes_to_i32_vec(modes_b);
171        let modes_c = modes_to_i32_vec(modes_c);
172        let modes_d = modes_to_i32_vec(modes_d);
173        let mut handle = ptr::null_mut();
174        unsafe {
175            try_ffi!(sys::cutensorMgCreateContractionDescriptor(
176                context.as_raw(),
177                &raw mut handle,
178                a.as_raw(),
179                modes_a.as_ptr(),
180                b.as_raw(),
181                modes_b.as_ptr(),
182                c.as_raw(),
183                modes_c.as_ptr(),
184                d.as_raw(),
185                modes_d.as_ptr(),
186                compute.into(),
187            ))?;
188        }
189
190        if handle.is_null() {
191            return Err(Error::NullHandle);
192        }
193
194        Ok(Self {
195            handle,
196            context: context.as_context_ref(),
197            a_pointer_count: a.devices().len(),
198            b_pointer_count: b.devices().len(),
199            c_pointer_count: c.devices().len(),
200            d_pointer_count: d.devices().len(),
201        })
202    }
203
204    /// Wraps an existing cuTENSORMg contraction descriptor.
205    ///
206    /// # Safety
207    ///
208    /// `handle` must be a valid cuTENSORMg contraction descriptor associated
209    /// with `context`. Pointer counts must match the tensors used to create the
210    /// descriptor. The returned value takes ownership of `handle` and destroys
211    /// it on drop.
212    pub unsafe fn from_raw(
213        handle: sys::cutensorMgContractionDescriptor_t,
214        context: &Context,
215        a_pointer_count: usize,
216        b_pointer_count: usize,
217        c_pointer_count: usize,
218        d_pointer_count: usize,
219    ) -> Result<Self> {
220        if handle.is_null() {
221            return Err(Error::NullHandle);
222        }
223
224        Ok(Self {
225            handle,
226            context: context.as_context_ref(),
227            a_pointer_count,
228            b_pointer_count,
229            c_pointer_count,
230            d_pointer_count,
231        })
232    }
233
234    pub fn workspace(
235        &self,
236        find: &ContractionFind,
237        preference: WorkspacePreference,
238    ) -> Result<Workspace> {
239        validate_same_context(&self.context, find.context(), "find")?;
240        let mut device_workspace_size = vec![0; self.context.device_count()];
241        let mut host_workspace_size = 0;
242        unsafe {
243            try_ffi!(sys::cutensorMgContractionGetWorkspace(
244                self.context.as_raw(),
245                self.handle,
246                find.as_raw(),
247                preference.into(),
248                device_workspace_size.as_mut_ptr(),
249                &raw mut host_workspace_size,
250            ))?;
251        }
252        Workspace::from_raw(device_workspace_size, host_workspace_size)
253    }
254
255    pub(crate) fn context(&self) -> &ContextRef {
256        &self.context
257    }
258
259    pub const fn as_raw(&self) -> sys::cutensorMgContractionDescriptor_t {
260        self.handle
261    }
262
263    /// Consumes this descriptor and returns the owned raw cuTENSORMg descriptor.
264    ///
265    /// The caller becomes responsible for destroying the descriptor.
266    pub fn into_raw(self) -> sys::cutensorMgContractionDescriptor_t {
267        let this = ManuallyDrop::new(self);
268        this.handle
269    }
270}
271
272impl ContractionPlan {
273    pub fn create(
274        context: &Context,
275        desc: &ContractionDescriptor,
276        find: &ContractionFind,
277        workspace: Workspace,
278    ) -> Result<Self> {
279        validate_same_context(&context.as_context_ref(), desc.context(), "desc")?;
280        validate_same_context(desc.context(), find.context(), "find")?;
281        workspace.validate_for_context(desc.context())?;
282
283        let mut handle = ptr::null_mut();
284        unsafe {
285            try_ffi!(sys::cutensorMgCreateContractionPlan(
286                context.as_raw(),
287                &raw mut handle,
288                desc.as_raw(),
289                find.as_raw(),
290                workspace.device_sizes_ptr(),
291                workspace.host_size(),
292            ))?;
293        }
294
295        if handle.is_null() {
296            return Err(Error::NullHandle);
297        }
298
299        Ok(Self {
300            handle,
301            context: context.as_context_ref(),
302            a_pointer_count: desc.a_pointer_count,
303            b_pointer_count: desc.b_pointer_count,
304            c_pointer_count: desc.c_pointer_count,
305            d_pointer_count: desc.d_pointer_count,
306            workspace,
307        })
308    }
309
310    /// Wraps an existing cuTENSORMg contraction plan.
311    ///
312    /// # Safety
313    ///
314    /// `handle` must be a valid cuTENSORMg contraction plan associated with
315    /// `context` and `workspace`. Pointer counts must match the descriptors
316    /// used to create the plan. The returned value takes ownership of `handle`
317    /// and destroys it on drop.
318    pub unsafe fn from_raw(
319        handle: sys::cutensorMgContractionPlan_t,
320        context: &Context,
321        a_pointer_count: usize,
322        b_pointer_count: usize,
323        c_pointer_count: usize,
324        d_pointer_count: usize,
325        workspace: Workspace,
326    ) -> Result<Self> {
327        if handle.is_null() {
328            return Err(Error::NullHandle);
329        }
330
331        Ok(Self {
332            handle,
333            context: context.as_context_ref(),
334            a_pointer_count,
335            b_pointer_count,
336            c_pointer_count,
337            d_pointer_count,
338            workspace,
339        })
340    }
341
342    pub fn workspace(&self) -> &Workspace {
343        &self.workspace
344    }
345
346    pub fn create_workspace_memory(&self) -> Result<WorkspaceMemory> {
347        WorkspaceMemory::create(&self.context, &self.workspace)
348    }
349
350    pub fn contract<TA, TB, TC, TD, TScalar>(
351        &self,
352        alpha: &TScalar,
353        a: &DistributedDeviceMemory<TA>,
354        b: &DistributedDeviceMemory<TB>,
355        beta: &TScalar,
356        c: &DistributedDeviceMemory<TC>,
357        d: &mut DistributedDeviceMemory<TD>,
358        workspace: &mut WorkspaceMemory,
359        streams: &[StreamBinding],
360    ) -> Result<()> {
361        self.validate_execution_slices(
362            a.pointer_count(),
363            b.pointer_count(),
364            c.pointer_count(),
365            d.pointer_count(),
366            workspace.device_pointer_count(),
367            streams.len(),
368        )?;
369
370        let a = a.const_ptrs();
371        let b = b.const_ptrs();
372        let c = c.const_ptrs();
373        let mut d = d.mut_ptrs();
374        let mut device_workspace = workspace.device_mut_ptrs();
375
376        unsafe {
377            self.contract_raw(
378                ptr::from_ref(alpha).cast(),
379                &a,
380                &b,
381                ptr::from_ref(beta).cast(),
382                &c,
383                &mut d,
384                &mut device_workspace,
385                workspace.host_mut_ptr(),
386                streams,
387            )
388        }
389    }
390
391    pub fn contract_in_place<TA, TB, TC, TScalar>(
392        &self,
393        alpha: &TScalar,
394        a: &DistributedDeviceMemory<TA>,
395        b: &DistributedDeviceMemory<TB>,
396        beta: &TScalar,
397        c_and_d: &mut DistributedDeviceMemory<TC>,
398        workspace: &mut WorkspaceMemory,
399        streams: &[StreamBinding],
400    ) -> Result<()> {
401        self.validate_execution_slices(
402            a.pointer_count(),
403            b.pointer_count(),
404            c_and_d.pointer_count(),
405            c_and_d.pointer_count(),
406            workspace.device_pointer_count(),
407            streams.len(),
408        )?;
409
410        let a = a.const_ptrs();
411        let b = b.const_ptrs();
412        let c = c_and_d.const_ptrs();
413        let mut d = c_and_d.mut_ptrs();
414        let mut device_workspace = workspace.device_mut_ptrs();
415
416        unsafe {
417            self.contract_raw(
418                ptr::from_ref(alpha).cast(),
419                &a,
420                &b,
421                ptr::from_ref(beta).cast(),
422                &c,
423                &mut d,
424                &mut device_workspace,
425                workspace.host_mut_ptr(),
426                streams,
427            )
428        }
429    }
430
431    pub fn contract_in_place_on_default_streams<TA, TB, TC, TScalar>(
432        &self,
433        alpha: &TScalar,
434        a: &DistributedDeviceMemory<TA>,
435        b: &DistributedDeviceMemory<TB>,
436        beta: &TScalar,
437        c_and_d: &mut DistributedDeviceMemory<TC>,
438        workspace: &mut WorkspaceMemory,
439    ) -> Result<()> {
440        self.validate_execution_slices(
441            a.pointer_count(),
442            b.pointer_count(),
443            c_and_d.pointer_count(),
444            c_and_d.pointer_count(),
445            self.context.device_count(),
446            self.context.device_count(),
447        )?;
448
449        let a = a.const_ptrs();
450        let b = b.const_ptrs();
451        let c = c_and_d.const_ptrs();
452        let mut d = c_and_d.mut_ptrs();
453        let mut device_workspace = workspace.device_mut_ptrs();
454
455        unsafe {
456            self.contract_raw_on_default_streams(
457                ptr::from_ref(alpha).cast(),
458                &a,
459                &b,
460                ptr::from_ref(beta).cast(),
461                &c,
462                &mut d,
463                &mut device_workspace,
464                workspace.host_mut_ptr(),
465            )
466        }
467    }
468
469    /// Executes this contraction plan with raw distributed tensor pointers.
470    ///
471    /// # Safety
472    ///
473    /// Pointer arrays must match the descriptors used for plan creation. Scalars must
474    /// have the type expected by the cuTENSORMg data-type combination.
475    pub unsafe fn contract_raw(
476        &self,
477        alpha: *const (),
478        a: &[*const ()],
479        b: &[*const ()],
480        beta: *const (),
481        c: &[*const ()],
482        d: &mut [*mut ()],
483        device_workspace: &mut [*mut ()],
484        host_workspace: *mut (),
485        streams: &[StreamBinding],
486    ) -> Result<()> {
487        self.validate_execution_slices(
488            a.len(),
489            b.len(),
490            c.len(),
491            d.len(),
492            device_workspace.len(),
493            streams.len(),
494        )?;
495        let mut a = a.to_vec();
496        let mut b = b.to_vec();
497        let mut c = c.to_vec();
498        let mut streams: Vec<_> = streams.iter().map(StreamBinding::as_raw).collect();
499        unsafe {
500            try_ffi!(sys::cutensorMgContraction(
501                self.context.as_raw(),
502                self.handle,
503                alpha as _,
504                a.as_mut_ptr() as _,
505                b.as_mut_ptr() as _,
506                beta as _,
507                c.as_mut_ptr() as _,
508                d.as_mut_ptr() as _,
509                device_workspace.as_mut_ptr() as _,
510                host_workspace as _,
511                streams.as_mut_ptr(),
512            ))?;
513        }
514        Ok(())
515    }
516
517    /// Executes this contraction plan on each device's default stream.
518    ///
519    /// # Safety
520    ///
521    /// Pointer arrays must match the descriptors used for plan creation. Scalars must
522    /// have the type expected by the cuTENSORMg data-type combination.
523    pub unsafe fn contract_raw_on_default_streams(
524        &self,
525        alpha: *const (),
526        a: &[*const ()],
527        b: &[*const ()],
528        beta: *const (),
529        c: &[*const ()],
530        d: &mut [*mut ()],
531        device_workspace: &mut [*mut ()],
532        host_workspace: *mut (),
533    ) -> Result<()> {
534        self.validate_execution_slices(
535            a.len(),
536            b.len(),
537            c.len(),
538            d.len(),
539            device_workspace.len(),
540            self.context.device_count(),
541        )?;
542        let mut a = a.to_vec();
543        let mut b = b.to_vec();
544        let mut c = c.to_vec();
545        let mut streams = vec![ptr::null_mut(); self.context.device_count()];
546        unsafe {
547            try_ffi!(sys::cutensorMgContraction(
548                self.context.as_raw(),
549                self.handle,
550                alpha as _,
551                a.as_mut_ptr() as _,
552                b.as_mut_ptr() as _,
553                beta as _,
554                c.as_mut_ptr() as _,
555                d.as_mut_ptr() as _,
556                device_workspace.as_mut_ptr() as _,
557                host_workspace as _,
558                streams.as_mut_ptr(),
559            ))?;
560        }
561        Ok(())
562    }
563
564    fn validate_execution_slices(
565        &self,
566        a_count: usize,
567        b_count: usize,
568        c_count: usize,
569        d_count: usize,
570        device_workspace_count: usize,
571        stream_count: usize,
572    ) -> Result<()> {
573        validate_count("a", self.a_pointer_count, a_count)?;
574        validate_count("b", self.b_pointer_count, b_count)?;
575        validate_count("c", self.c_pointer_count, c_count)?;
576        validate_count("d", self.d_pointer_count, d_count)?;
577        validate_count(
578            "device_workspace",
579            self.context.device_count(),
580            device_workspace_count,
581        )?;
582        validate_count("streams", self.context.device_count(), stream_count)
583    }
584
585    pub const fn as_raw(&self) -> sys::cutensorMgContractionPlan_t {
586        self.handle
587    }
588
589    /// Consumes this plan and returns the owned raw cuTENSORMg contraction plan.
590    ///
591    /// The caller becomes responsible for destroying the plan.
592    pub fn into_raw(self) -> sys::cutensorMgContractionPlan_t {
593        let this = ManuallyDrop::new(self);
594        this.handle
595    }
596}
597
598impl Drop for ContractionFind {
599    fn drop(&mut self) {
600        unsafe {
601            if let Err(err) = try_ffi!(sys::cutensorMgDestroyContractionFind(self.handle)) {
602                #[cfg(debug_assertions)]
603                eprintln!("failed to destroy cutensormg contraction find: {err}");
604            }
605        }
606    }
607}
608
609impl Drop for ContractionDescriptor {
610    fn drop(&mut self) {
611        unsafe {
612            if let Err(err) = try_ffi!(sys::cutensorMgDestroyContractionDescriptor(self.handle)) {
613                #[cfg(debug_assertions)]
614                eprintln!("failed to destroy cutensormg contraction descriptor: {err}");
615            }
616        }
617    }
618}
619
620impl Drop for ContractionPlan {
621    fn drop(&mut self) {
622        unsafe {
623            if let Err(err) = try_ffi!(sys::cutensorMgDestroyContractionPlan(self.handle)) {
624                #[cfg(debug_assertions)]
625                eprintln!("failed to destroy cutensormg contraction plan: {err}");
626            }
627        }
628    }
629}
630
631fn validate_modes(desc: &TensorDescriptor, modes: &[Mode]) -> Result<()> {
632    if desc.rank() as usize != modes.len() {
633        return Err(Error::TensorModeMismatch {
634            rank: desc.rank(),
635            mode_length: modes.len(),
636        });
637    }
638    Ok(())
639}
singe_cutensor/mg/contraction.rs

singe_cutensor/mg/
contraction.rs
