singe_cutensor/mp/
plan.rs

1use std::{
2    mem::{ManuallyDrop, MaybeUninit},
3    ptr,
4};
5
6use singe_cuda::memory::DeviceMemory;
7use singe_cutensor_sys as sys;
8
9use crate::{
10    error::{Error, Result},
11    mp::{
12        context::{Context, ContextRef, validate_same_context},
13        memory::WorkspaceMemory,
14        tensor::TensorDescriptor,
15        types::{Algorithm, PlanAttribute},
16    },
17    operation::ComputeDescriptor,
18    try_ffi,
19    types::{Mode, Operator},
20    utility::modes_to_i32_vec,
21};
22
23#[derive(Debug)]
24pub struct OperationDescriptor<'comm> {
25    handle: sys::cutensorMpOperationDescriptor_t,
26    context: ContextRef<'comm>,
27}
28
29#[derive(Debug)]
30pub struct PlanPreference<'comm> {
31    handle: sys::cutensorMpPlanPreference_t,
32    context: ContextRef<'comm>,
33    algorithm: Algorithm,
34    device_workspace_limit: u64,
35    host_workspace_limit: u64,
36}
37
38#[derive(Debug)]
39pub struct Plan<'comm> {
40    handle: sys::cutensorMpPlan_t,
41    context: ContextRef<'comm>,
42}
43
44#[derive(Debug, Clone, Copy, PartialEq, Eq)]
45pub struct Workspace {
46    device_size: u64,
47    host_size: u64,
48}
49
50impl<'comm> OperationDescriptor<'comm> {
51    pub fn create_contraction(
52        context: &Context<'comm>,
53        a: &TensorDescriptor<'comm>,
54        modes_a: &[Mode],
55        op_a: Operator,
56        b: &TensorDescriptor<'comm>,
57        modes_b: &[Mode],
58        op_b: Operator,
59        c: &TensorDescriptor<'comm>,
60        modes_c: &[Mode],
61        op_c: Operator,
62        d: &TensorDescriptor<'comm>,
63        modes_d: &[Mode],
64        compute: ComputeDescriptor,
65    ) -> Result<Self> {
66        let context_ref = context.as_context_ref();
67        validate_same_context(&context_ref, a.context(), "a")?;
68        validate_same_context(&context_ref, b.context(), "b")?;
69        validate_same_context(&context_ref, c.context(), "c")?;
70        validate_same_context(&context_ref, d.context(), "d")?;
71        validate_modes(a, modes_a)?;
72        validate_modes(b, modes_b)?;
73        validate_modes(c, modes_c)?;
74        validate_modes(d, modes_d)?;
75
76        let modes_a = modes_to_i32_vec(modes_a);
77        let modes_b = modes_to_i32_vec(modes_b);
78        let modes_c = modes_to_i32_vec(modes_c);
79        let modes_d = modes_to_i32_vec(modes_d);
80        let mut handle = ptr::null_mut();
81        context.bind()?;
82        unsafe {
83            try_ffi!(sys::cutensorMpCreateContraction(
84                context.as_raw(),
85                &raw mut handle,
86                a.as_raw(),
87                modes_a.as_ptr(),
88                op_a.into(),
89                b.as_raw(),
90                modes_b.as_ptr(),
91                op_b.into(),
92                c.as_raw(),
93                modes_c.as_ptr(),
94                op_c.into(),
95                d.as_raw(),
96                modes_d.as_ptr(),
97                compute.as_raw(),
98            ))?;
99        }
100
101        if handle.is_null() {
102            return Err(Error::NullHandle);
103        }
104
105        Ok(Self {
106            handle,
107            context: context.as_context_ref(),
108        })
109    }
110
111    /// Wraps an existing cuTENSORMp operation descriptor.
112    ///
113    /// # Safety
114    ///
115    /// `handle` must be a valid cuTENSORMp operation descriptor associated with
116    /// `context`. The returned value takes ownership of `handle` and destroys
117    /// it on drop.
118    pub unsafe fn from_raw(
119        handle: sys::cutensorMpOperationDescriptor_t,
120        context: &Context<'comm>,
121    ) -> Result<Self> {
122        if handle.is_null() {
123            return Err(Error::NullHandle);
124        }
125
126        Ok(Self {
127            handle,
128            context: context.as_context_ref(),
129        })
130    }
131
132    pub(crate) fn context(&self) -> &ContextRef<'comm> {
133        &self.context
134    }
135
136    pub const fn as_raw(&self) -> sys::cutensorMpOperationDescriptor_t {
137        self.handle
138    }
139
140    /// Consumes this descriptor and returns the owned raw cuTENSORMp operation descriptor.
141    ///
142    /// The caller becomes responsible for destroying the descriptor.
143    pub fn into_raw(self) -> sys::cutensorMpOperationDescriptor_t {
144        let this = ManuallyDrop::new(self);
145        this.handle
146    }
147}
148
149impl<'comm> PlanPreference<'comm> {
150    pub fn create(
151        context: &Context<'comm>,
152        algorithm: Algorithm,
153        device_workspace_limit: u64,
154        host_workspace_limit: u64,
155    ) -> Result<Self> {
156        let mut handle = ptr::null_mut();
157        context.bind()?;
158        unsafe {
159            try_ffi!(sys::cutensorMpCreatePlanPreference(
160                context.as_raw(),
161                &raw mut handle,
162                algorithm.into(),
163                device_workspace_limit,
164                host_workspace_limit,
165            ))?;
166        }
167
168        if handle.is_null() {
169            return Err(Error::NullHandle);
170        }
171
172        Ok(Self {
173            handle,
174            context: context.as_context_ref(),
175            algorithm,
176            device_workspace_limit,
177            host_workspace_limit,
178        })
179    }
180
181    /// Wraps an existing cuTENSORMp plan preference.
182    ///
183    /// # Safety
184    ///
185    /// `handle` must be a valid cuTENSORMp plan preference associated with
186    /// `context`. Metadata arguments must describe the raw preference. The
187    /// returned value takes ownership of `handle` and destroys it on drop.
188    pub unsafe fn from_raw(
189        handle: sys::cutensorMpPlanPreference_t,
190        context: &Context<'comm>,
191        algorithm: Algorithm,
192        device_workspace_limit: u64,
193        host_workspace_limit: u64,
194    ) -> Result<Self> {
195        if handle.is_null() {
196            return Err(Error::NullHandle);
197        }
198
199        Ok(Self {
200            handle,
201            context: context.as_context_ref(),
202            algorithm,
203            device_workspace_limit,
204            host_workspace_limit,
205        })
206    }
207
208    pub fn algorithm(&self) -> Algorithm {
209        self.algorithm
210    }
211
212    pub fn device_workspace_limit(&self) -> u64 {
213        self.device_workspace_limit
214    }
215
216    pub fn host_workspace_limit(&self) -> u64 {
217        self.host_workspace_limit
218    }
219
220    pub(crate) fn context(&self) -> &ContextRef<'comm> {
221        &self.context
222    }
223
224    pub const fn as_raw(&self) -> sys::cutensorMpPlanPreference_t {
225        self.handle
226    }
227
228    /// Consumes this preference and returns the owned raw cuTENSORMp plan preference.
229    ///
230    /// The caller becomes responsible for destroying the preference.
231    pub fn into_raw(self) -> sys::cutensorMpPlanPreference_t {
232        let this = ManuallyDrop::new(self);
233        this.handle
234    }
235}
236
237impl<'comm> Plan<'comm> {
238    pub fn create(
239        context: &Context<'comm>,
240        desc: &OperationDescriptor<'comm>,
241        preference: Option<&PlanPreference<'comm>>,
242    ) -> Result<Self> {
243        let context_ref = context.as_context_ref();
244        validate_same_context(&context_ref, desc.context(), "desc")?;
245        if let Some(preference) = preference {
246            validate_same_context(&context_ref, preference.context(), "preference")?;
247        }
248
249        let mut handle = ptr::null_mut();
250        context.bind()?;
251        unsafe {
252            try_ffi!(sys::cutensorMpCreatePlan(
253                context.as_raw(),
254                &raw mut handle,
255                desc.as_raw(),
256                preference.map_or(ptr::null_mut(), |preference| preference.as_raw()),
257            ))?;
258        }
259
260        if handle.is_null() {
261            return Err(Error::NullHandle);
262        }
263
264        Ok(Self {
265            handle,
266            context: context.as_context_ref(),
267        })
268    }
269
270    /// Wraps an existing cuTENSORMp plan.
271    ///
272    /// # Safety
273    ///
274    /// `handle` must be a valid cuTENSORMp plan associated with `context`. The
275    /// returned value takes ownership of `handle` and destroys it on drop.
276    pub unsafe fn from_raw(
277        handle: sys::cutensorMpPlan_t,
278        context: &Context<'comm>,
279    ) -> Result<Self> {
280        if handle.is_null() {
281            return Err(Error::NullHandle);
282        }
283
284        Ok(Self {
285            handle,
286            context: context.as_context_ref(),
287        })
288    }
289
290    pub fn workspace(&self) -> Result<Workspace> {
291        Ok(Workspace {
292            device_size: self.attribute(PlanAttribute::RequiredWorkspaceDevice)?,
293            host_size: self.attribute(PlanAttribute::RequiredWorkspaceHost)?,
294        })
295    }
296
297    pub fn create_workspace_memory(&self) -> Result<WorkspaceMemory> {
298        WorkspaceMemory::create(self.workspace()?)
299    }
300
301    pub fn contract<TA, TB, TC, TD, TScalar>(
302        &self,
303        alpha: &TScalar,
304        a: &DeviceMemory<TA>,
305        b: &DeviceMemory<TB>,
306        beta: &TScalar,
307        c: &DeviceMemory<TC>,
308        d: &mut DeviceMemory<TD>,
309        workspace: &mut WorkspaceMemory,
310    ) -> Result<()> {
311        unsafe {
312            self.contract_raw(
313                ptr::from_ref(alpha).cast(),
314                a.as_ptr().cast(),
315                b.as_ptr().cast(),
316                ptr::from_ref(beta).cast(),
317                c.as_ptr().cast(),
318                d.as_mut_ptr().cast(),
319                workspace.device_mut_ptr(),
320                workspace.host_mut_ptr(),
321            )
322        }
323    }
324
325    pub fn contract_in_place<TA, TB, TC, TScalar>(
326        &self,
327        alpha: &TScalar,
328        a: &DeviceMemory<TA>,
329        b: &DeviceMemory<TB>,
330        beta: &TScalar,
331        c_and_d: &mut DeviceMemory<TC>,
332        workspace: &mut WorkspaceMemory,
333    ) -> Result<()> {
334        unsafe {
335            self.contract_raw(
336                ptr::from_ref(alpha).cast(),
337                a.as_ptr().cast(),
338                b.as_ptr().cast(),
339                ptr::from_ref(beta).cast(),
340                c_and_d.as_ptr().cast(),
341                c_and_d.as_mut_ptr().cast(),
342                workspace.device_mut_ptr(),
343                workspace.host_mut_ptr(),
344            )
345        }
346    }
347
348    /// Executes this distributed contraction with raw local tensor pointers.
349    ///
350    /// # Safety
351    ///
352    /// Pointers must refer to the local rank's tensor shards and workspaces
353    /// matching the descriptors and plan. Scalars must have the type expected by
354    /// cuTENSORMp for the data-type combination. All pointers must remain valid
355    /// for the duration of the launched operation, and `d`, `device_workspace`,
356    /// and `host_workspace` must be writable according to the plan's workspace
357    /// requirements.
358    pub unsafe fn contract_raw(
359        &self,
360        alpha: *const (),
361        a: *const (),
362        b: *const (),
363        beta: *const (),
364        c: *const (),
365        d: *mut (),
366        device_workspace: *mut (),
367        host_workspace: *mut (),
368    ) -> Result<()> {
369        self.context.bind()?;
370        unsafe {
371            try_ffi!(sys::cutensorMpContract(
372                self.context.as_raw(),
373                self.handle,
374                alpha.cast(),
375                a.cast(),
376                b.cast(),
377                beta.cast(),
378                c.cast(),
379                d.cast(),
380                device_workspace.cast(),
381                host_workspace.cast(),
382            ))?;
383        }
384        Ok(())
385    }
386
387    pub fn as_raw(&self) -> sys::cutensorMpPlan_t {
388        self.handle
389    }
390
391    /// Consumes this plan and returns the owned raw cuTENSORMp plan.
392    ///
393    /// The caller becomes responsible for destroying the plan.
394    pub fn into_raw(self) -> sys::cutensorMpPlan_t {
395        let this = ManuallyDrop::new(self);
396        this.handle
397    }
398
399    fn attribute<T: Copy>(&self, attr: PlanAttribute) -> Result<T> {
400        let mut value = MaybeUninit::<T>::uninit();
401        self.context.bind()?;
402        unsafe {
403            try_ffi!(sys::cutensorMpPlanGetAttribute(
404                self.context.as_raw(),
405                self.handle,
406                attr.into(),
407                value.as_mut_ptr().cast(),
408                size_of::<T>() as u64,
409            ))?;
410            Ok(value.assume_init())
411        }
412    }
413}
414
415impl Workspace {
416    pub fn device_size(&self) -> u64 {
417        self.device_size
418    }
419
420    pub fn host_size(&self) -> u64 {
421        self.host_size
422    }
423}
424
425impl Drop for OperationDescriptor<'_> {
426    fn drop(&mut self) {
427        if let Err(err) = self.context.bind() {
428            #[cfg(debug_assertions)]
429            eprintln!(
430                "failed to bind cutensormp context before destroying operation descriptor: {err}"
431            );
432        }
433
434        unsafe {
435            if let Err(err) = try_ffi!(sys::cutensorMpDestroyOperationDescriptor(self.handle)) {
436                #[cfg(debug_assertions)]
437                eprintln!("failed to destroy cutensormp operation descriptor: {err}");
438            }
439        }
440    }
441}
442
443impl Drop for PlanPreference<'_> {
444    fn drop(&mut self) {
445        if let Err(err) = self.context.bind() {
446            #[cfg(debug_assertions)]
447            eprintln!("failed to bind cutensormp context before destroying plan preference: {err}");
448        }
449
450        unsafe {
451            if let Err(err) = try_ffi!(sys::cutensorMpDestroyPlanPreference(self.handle)) {
452                #[cfg(debug_assertions)]
453                eprintln!("failed to destroy cutensormp plan preference: {err}");
454            }
455        }
456    }
457}
458
459impl Drop for Plan<'_> {
460    fn drop(&mut self) {
461        if let Err(err) = self.context.bind() {
462            #[cfg(debug_assertions)]
463            eprintln!("failed to bind cutensormp context before destroying plan: {err}");
464        }
465
466        unsafe {
467            if let Err(err) = try_ffi!(sys::cutensorMpDestroyPlan(self.handle)) {
468                #[cfg(debug_assertions)]
469                eprintln!("failed to destroy cutensormp plan: {err}");
470            }
471        }
472    }
473}
474
475fn validate_modes(desc: &TensorDescriptor<'_>, modes: &[Mode]) -> Result<()> {
476    if desc.rank() as usize != modes.len() {
477        return Err(Error::TensorModeMismatch {
478            rank: desc.rank(),
479            mode_length: modes.len(),
480        });
481    }
482    Ok(())
483}
singe_cutensor/mp/plan.rs

singe_cutensor/mp/
plan.rs
