singe_cutensor/mg/
contraction.rs

1use std::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)]
25pub enum ContractionFindAttribute {
26    Max = sys::cutensorMgContractionFindAttribute_t::CUTENSORMG_CONTRACTION_FIND_ATTRIBUTE_MAX as _,
27}
28
29impl_enum_conversion!(
30    sys::cutensorMgContractionFindAttribute_t,
31    ContractionFindAttribute
32);
33
34#[derive(Debug)]
35pub struct ContractionFind {
36    handle: sys::cutensorMgContractionFind_t,
37    context: ContextRef,
38    algorithm: Algorithm,
39}
40
41#[derive(Debug)]
42pub struct ContractionDescriptor {
43    handle: sys::cutensorMgContractionDescriptor_t,
44    context: ContextRef,
45    a_pointer_count: usize,
46    b_pointer_count: usize,
47    c_pointer_count: usize,
48    d_pointer_count: usize,
49}
50
51#[derive(Debug)]
52pub struct ContractionPlan {
53    handle: sys::cutensorMgContractionPlan_t,
54    context: ContextRef,
55    a_pointer_count: usize,
56    b_pointer_count: usize,
57    c_pointer_count: usize,
58    d_pointer_count: usize,
59    workspace: Workspace,
60}
61
62impl ContractionFind {
63    pub fn create(context: &Context, algorithm: Algorithm) -> Result<Self> {
64        let mut handle = ptr::null_mut();
65        unsafe {
66            try_ffi!(sys::cutensorMgCreateContractionFind(
67                context.as_raw(),
68                &raw mut handle,
69                algorithm.into(),
70            ))?;
71        }
72
73        if handle.is_null() {
74            return Err(Error::NullHandle);
75        }
76
77        Ok(Self {
78            handle,
79            context: context.as_context_ref(),
80            algorithm,
81        })
82    }
83
84    pub fn create_default(context: &Context) -> Result<Self> {
85        Self::create(context, Algorithm::Default)
86    }
87
88    pub fn algorithm(&self) -> Algorithm {
89        self.algorithm
90    }
91
92    pub fn set_attribute<T>(&mut self, attr: ContractionFindAttribute, value: &T) -> Result<()> {
93        unsafe {
94            try_ffi!(sys::cutensorMgContractionFindSetAttribute(
95                self.context.as_raw(),
96                self.handle,
97                attr.into(),
98                ptr::from_ref(value).cast(),
99                size_of::<T>() as i64,
100            ))?;
101        }
102        Ok(())
103    }
104
105    pub(crate) fn context(&self) -> &ContextRef {
106        &self.context
107    }
108
109    pub const fn as_raw(&self) -> sys::cutensorMgContractionFind_t {
110        self.handle
111    }
112}
113
114impl ContractionDescriptor {
115    #[allow(clippy::too_many_arguments)]
116    pub fn create(
117        context: &Context,
118        a: &TensorDescriptor,
119        modes_a: &[Mode],
120        b: &TensorDescriptor,
121        modes_b: &[Mode],
122        c: &TensorDescriptor,
123        modes_c: &[Mode],
124        d: &TensorDescriptor,
125        modes_d: &[Mode],
126        compute: ComputeType,
127    ) -> Result<Self> {
128        let context_ref = context.as_context_ref();
129        validate_same_context(&context_ref, a.context(), "a")?;
130        validate_same_context(&context_ref, b.context(), "b")?;
131        validate_same_context(&context_ref, c.context(), "c")?;
132        validate_same_context(&context_ref, d.context(), "d")?;
133        validate_modes(a, modes_a)?;
134        validate_modes(b, modes_b)?;
135        validate_modes(c, modes_c)?;
136        validate_modes(d, modes_d)?;
137
138        let modes_a = modes_to_i32_vec(modes_a);
139        let modes_b = modes_to_i32_vec(modes_b);
140        let modes_c = modes_to_i32_vec(modes_c);
141        let modes_d = modes_to_i32_vec(modes_d);
142        let mut handle = ptr::null_mut();
143        unsafe {
144            try_ffi!(sys::cutensorMgCreateContractionDescriptor(
145                context.as_raw(),
146                &raw mut handle,
147                a.as_raw(),
148                modes_a.as_ptr(),
149                b.as_raw(),
150                modes_b.as_ptr(),
151                c.as_raw(),
152                modes_c.as_ptr(),
153                d.as_raw(),
154                modes_d.as_ptr(),
155                compute.into(),
156            ))?;
157        }
158
159        if handle.is_null() {
160            return Err(Error::NullHandle);
161        }
162
163        Ok(Self {
164            handle,
165            context: context.as_context_ref(),
166            a_pointer_count: a.devices().len(),
167            b_pointer_count: b.devices().len(),
168            c_pointer_count: c.devices().len(),
169            d_pointer_count: d.devices().len(),
170        })
171    }
172
173    pub fn workspace(
174        &self,
175        find: &ContractionFind,
176        preference: WorkspacePreference,
177    ) -> Result<Workspace> {
178        validate_same_context(&self.context, find.context(), "find")?;
179        let mut device_workspace_size = vec![0; self.context.device_count()];
180        let mut host_workspace_size = 0;
181        unsafe {
182            try_ffi!(sys::cutensorMgContractionGetWorkspace(
183                self.context.as_raw(),
184                self.handle,
185                find.as_raw(),
186                preference.into(),
187                device_workspace_size.as_mut_ptr(),
188                &raw mut host_workspace_size,
189            ))?;
190        }
191        Workspace::from_raw(device_workspace_size, host_workspace_size)
192    }
193
194    pub(crate) fn context(&self) -> &ContextRef {
195        &self.context
196    }
197
198    pub const fn as_raw(&self) -> sys::cutensorMgContractionDescriptor_t {
199        self.handle
200    }
201}
202
203impl ContractionPlan {
204    pub fn create(
205        context: &Context,
206        desc: &ContractionDescriptor,
207        find: &ContractionFind,
208        workspace: Workspace,
209    ) -> Result<Self> {
210        validate_same_context(&context.as_context_ref(), desc.context(), "desc")?;
211        validate_same_context(desc.context(), find.context(), "find")?;
212        workspace.validate_for_context(desc.context())?;
213
214        let mut handle = ptr::null_mut();
215        unsafe {
216            try_ffi!(sys::cutensorMgCreateContractionPlan(
217                context.as_raw(),
218                &raw mut handle,
219                desc.as_raw(),
220                find.as_raw(),
221                workspace.device_sizes_ptr(),
222                workspace.host_size(),
223            ))?;
224        }
225
226        if handle.is_null() {
227            return Err(Error::NullHandle);
228        }
229
230        Ok(Self {
231            handle,
232            context: context.as_context_ref(),
233            a_pointer_count: desc.a_pointer_count,
234            b_pointer_count: desc.b_pointer_count,
235            c_pointer_count: desc.c_pointer_count,
236            d_pointer_count: desc.d_pointer_count,
237            workspace,
238        })
239    }
240
241    pub fn workspace(&self) -> &Workspace {
242        &self.workspace
243    }
244
245    pub fn create_workspace_memory(&self) -> Result<WorkspaceMemory> {
246        WorkspaceMemory::create(&self.context, &self.workspace)
247    }
248
249    pub fn contract<TA, TB, TC, TD, TScalar>(
250        &self,
251        alpha: &TScalar,
252        a: &DistributedDeviceMemory<TA>,
253        b: &DistributedDeviceMemory<TB>,
254        beta: &TScalar,
255        c: &DistributedDeviceMemory<TC>,
256        d: &mut DistributedDeviceMemory<TD>,
257        workspace: &mut WorkspaceMemory,
258        streams: &[StreamBinding],
259    ) -> Result<()> {
260        self.validate_execution_slices(
261            a.pointer_count(),
262            b.pointer_count(),
263            c.pointer_count(),
264            d.pointer_count(),
265            workspace.device_pointer_count(),
266            streams.len(),
267        )?;
268
269        let a = a.const_ptrs();
270        let b = b.const_ptrs();
271        let c = c.const_ptrs();
272        let mut d = d.mut_ptrs();
273        let mut device_workspace = workspace.device_mut_ptrs();
274
275        unsafe {
276            self.contract_raw(
277                ptr::from_ref(alpha).cast(),
278                &a,
279                &b,
280                ptr::from_ref(beta).cast(),
281                &c,
282                &mut d,
283                &mut device_workspace,
284                workspace.host_mut_ptr(),
285                streams,
286            )
287        }
288    }
289
290    pub fn contract_in_place<TA, TB, TC, TScalar>(
291        &self,
292        alpha: &TScalar,
293        a: &DistributedDeviceMemory<TA>,
294        b: &DistributedDeviceMemory<TB>,
295        beta: &TScalar,
296        c_and_d: &mut DistributedDeviceMemory<TC>,
297        workspace: &mut WorkspaceMemory,
298        streams: &[StreamBinding],
299    ) -> Result<()> {
300        self.validate_execution_slices(
301            a.pointer_count(),
302            b.pointer_count(),
303            c_and_d.pointer_count(),
304            c_and_d.pointer_count(),
305            workspace.device_pointer_count(),
306            streams.len(),
307        )?;
308
309        let a = a.const_ptrs();
310        let b = b.const_ptrs();
311        let c = c_and_d.const_ptrs();
312        let mut d = c_and_d.mut_ptrs();
313        let mut device_workspace = workspace.device_mut_ptrs();
314
315        unsafe {
316            self.contract_raw(
317                ptr::from_ref(alpha).cast(),
318                &a,
319                &b,
320                ptr::from_ref(beta).cast(),
321                &c,
322                &mut d,
323                &mut device_workspace,
324                workspace.host_mut_ptr(),
325                streams,
326            )
327        }
328    }
329
330    pub fn contract_in_place_on_default_streams<TA, TB, TC, TScalar>(
331        &self,
332        alpha: &TScalar,
333        a: &DistributedDeviceMemory<TA>,
334        b: &DistributedDeviceMemory<TB>,
335        beta: &TScalar,
336        c_and_d: &mut DistributedDeviceMemory<TC>,
337        workspace: &mut WorkspaceMemory,
338    ) -> Result<()> {
339        self.validate_execution_slices(
340            a.pointer_count(),
341            b.pointer_count(),
342            c_and_d.pointer_count(),
343            c_and_d.pointer_count(),
344            self.context.device_count(),
345            self.context.device_count(),
346        )?;
347
348        let a = a.const_ptrs();
349        let b = b.const_ptrs();
350        let c = c_and_d.const_ptrs();
351        let mut d = c_and_d.mut_ptrs();
352        let mut device_workspace = workspace.device_mut_ptrs();
353
354        unsafe {
355            self.contract_raw_on_default_streams(
356                ptr::from_ref(alpha).cast(),
357                &a,
358                &b,
359                ptr::from_ref(beta).cast(),
360                &c,
361                &mut d,
362                &mut device_workspace,
363                workspace.host_mut_ptr(),
364            )
365        }
366    }
367
368    /// Executes this contraction plan with raw distributed tensor pointers.
369    ///
370    /// # Safety
371    ///
372    /// Pointer arrays must match the descriptors used for plan creation. Scalars must
373    /// have the type expected by the cuTENSORMg data-type combination.
374    pub unsafe fn contract_raw(
375        &self,
376        alpha: *const (),
377        a: &[*const ()],
378        b: &[*const ()],
379        beta: *const (),
380        c: &[*const ()],
381        d: &mut [*mut ()],
382        device_workspace: &mut [*mut ()],
383        host_workspace: *mut (),
384        streams: &[StreamBinding],
385    ) -> Result<()> {
386        self.validate_execution_slices(
387            a.len(),
388            b.len(),
389            c.len(),
390            d.len(),
391            device_workspace.len(),
392            streams.len(),
393        )?;
394        let mut a = a.to_vec();
395        let mut b = b.to_vec();
396        let mut c = c.to_vec();
397        let mut streams: Vec<_> = streams.iter().map(StreamBinding::as_raw).collect();
398        unsafe {
399            try_ffi!(sys::cutensorMgContraction(
400                self.context.as_raw(),
401                self.handle,
402                alpha as _,
403                a.as_mut_ptr() as _,
404                b.as_mut_ptr() as _,
405                beta as _,
406                c.as_mut_ptr() as _,
407                d.as_mut_ptr() as _,
408                device_workspace.as_mut_ptr() as _,
409                host_workspace as _,
410                streams.as_mut_ptr(),
411            ))?;
412        }
413        Ok(())
414    }
415
416    /// Executes this contraction plan on each device's default stream.
417    ///
418    /// # Safety
419    ///
420    /// Pointer arrays must match the descriptors used for plan creation. Scalars must
421    /// have the type expected by the cuTENSORMg data-type combination.
422    pub unsafe fn contract_raw_on_default_streams(
423        &self,
424        alpha: *const (),
425        a: &[*const ()],
426        b: &[*const ()],
427        beta: *const (),
428        c: &[*const ()],
429        d: &mut [*mut ()],
430        device_workspace: &mut [*mut ()],
431        host_workspace: *mut (),
432    ) -> Result<()> {
433        self.validate_execution_slices(
434            a.len(),
435            b.len(),
436            c.len(),
437            d.len(),
438            device_workspace.len(),
439            self.context.device_count(),
440        )?;
441        let mut a = a.to_vec();
442        let mut b = b.to_vec();
443        let mut c = c.to_vec();
444        let mut streams = vec![ptr::null_mut(); self.context.device_count()];
445        unsafe {
446            try_ffi!(sys::cutensorMgContraction(
447                self.context.as_raw(),
448                self.handle,
449                alpha as _,
450                a.as_mut_ptr() as _,
451                b.as_mut_ptr() as _,
452                beta as _,
453                c.as_mut_ptr() as _,
454                d.as_mut_ptr() as _,
455                device_workspace.as_mut_ptr() as _,
456                host_workspace as _,
457                streams.as_mut_ptr(),
458            ))?;
459        }
460        Ok(())
461    }
462
463    fn validate_execution_slices(
464        &self,
465        a_count: usize,
466        b_count: usize,
467        c_count: usize,
468        d_count: usize,
469        device_workspace_count: usize,
470        stream_count: usize,
471    ) -> Result<()> {
472        validate_count("a", self.a_pointer_count, a_count)?;
473        validate_count("b", self.b_pointer_count, b_count)?;
474        validate_count("c", self.c_pointer_count, c_count)?;
475        validate_count("d", self.d_pointer_count, d_count)?;
476        validate_count(
477            "device_workspace",
478            self.context.device_count(),
479            device_workspace_count,
480        )?;
481        validate_count("streams", self.context.device_count(), stream_count)
482    }
483
484    pub const fn as_raw(&self) -> sys::cutensorMgContractionPlan_t {
485        self.handle
486    }
487}
488
489impl Drop for ContractionFind {
490    fn drop(&mut self) {
491        unsafe {
492            if let Err(err) = try_ffi!(sys::cutensorMgDestroyContractionFind(self.handle)) {
493                #[cfg(debug_assertions)]
494                eprintln!("failed to destroy cutensormg contraction find: {err}");
495            }
496        }
497    }
498}
499
500impl Drop for ContractionDescriptor {
501    fn drop(&mut self) {
502        unsafe {
503            if let Err(err) = try_ffi!(sys::cutensorMgDestroyContractionDescriptor(self.handle)) {
504                #[cfg(debug_assertions)]
505                eprintln!("failed to destroy cutensormg contraction descriptor: {err}");
506            }
507        }
508    }
509}
510
511impl Drop for ContractionPlan {
512    fn drop(&mut self) {
513        unsafe {
514            if let Err(err) = try_ffi!(sys::cutensorMgDestroyContractionPlan(self.handle)) {
515                #[cfg(debug_assertions)]
516                eprintln!("failed to destroy cutensormg contraction plan: {err}");
517            }
518        }
519    }
520}
521
522fn validate_modes(desc: &TensorDescriptor, modes: &[Mode]) -> Result<()> {
523    if desc.rank() as usize != modes.len() {
524        return Err(Error::TensorModeMismatch {
525            rank: desc.rank(),
526            mode_length: modes.len(),
527        });
528    }
529    Ok(())
530}
singe_cutensor/mg/contraction.rs

singe_cutensor/mg/
contraction.rs
