singe_cutensor/mp/
plan.rs

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

singe_cutensor/mp/
plan.rs
