singe_cutensor/
plan.rs

1use std::{
2    mem::{ManuallyDrop, MaybeUninit},
3    ptr,
4};
5
6use crate::{
7    context::{Context, ContextRef},
8    error::{Error, Result},
9    operation::{OperationDescriptor, OperationSignature},
10    sys, try_ffi,
11    types::{
12        Algorithm, AutotuneMode, CacheMode, JitMode, PlanAttribute, PlanPreferenceAttribute,
13        WorkspacePreference,
14    },
15    utility::{to_i32, to_usize},
16};
17
18#[derive(Debug)]
19pub struct PlanPreference {
20    handle: sys::cutensorPlanPreference_t,
21    context: ContextRef,
22}
23
24impl PlanPreference {
25    pub fn create_default(ctx: &Context) -> Result<Self> {
26        Self::create(ctx, Algorithm::Default, JitMode::Default)
27    }
28
29    /// Creates a plan preference that limits the applicable kernels for a plan or operation.
30    ///
31    /// The preference is tied to `ctx` and frees its cuTENSOR handle when dropped.
32    ///
33    /// # Errors
34    ///
35    /// Returns an error if the context is not initialized, cuTENSOR rejects the
36    /// requested algorithm or JIT mode, or cuTENSOR does not return a valid handle.
37    pub fn create(ctx: &Context, algorithm: Algorithm, jit_mode: JitMode) -> Result<Self> {
38        ctx.bind()?;
39
40        let mut handle = ptr::null_mut();
41        unsafe {
42            try_ffi!(sys::cutensorCreatePlanPreference(
43                ctx.as_raw(),
44                &raw mut handle,
45                algorithm.into(),
46                jit_mode.into(),
47            ))?;
48        }
49
50        if handle.is_null() {
51            return Err(Error::NullHandle);
52        }
53
54        Ok(Self {
55            handle,
56            context: ctx.as_context_ref(),
57        })
58    }
59
60    /// Wraps an existing cuTENSOR plan-preference handle and takes ownership of it.
61    ///
62    /// # Safety
63    ///
64    /// `handle` must be a valid cuTENSOR plan-preference handle associated with
65    /// `ctx`. Ownership of `handle` is transferred to the returned preference,
66    /// and the handle must not be destroyed elsewhere after calling this
67    /// function.
68    pub unsafe fn from_raw(handle: sys::cutensorPlanPreference_t, ctx: &Context) -> Result<Self> {
69        if handle.is_null() {
70            return Err(Error::NullHandle);
71        }
72
73        Ok(Self {
74            handle,
75            context: ctx.as_context_ref(),
76        })
77    }
78
79    fn set_attribute<T>(&mut self, attr: PlanPreferenceAttribute, value: &T) -> Result<()> {
80        validate_plan_preference_attribute_size(attr, size_of::<T>())?;
81
82        self.context.bind()?;
83        unsafe {
84            try_ffi!(sys::cutensorPlanPreferenceSetAttribute(
85                self.context.as_raw(),
86                self.handle,
87                attr.into(),
88                ptr::from_ref(value).cast(),
89                size_of::<T>() as _,
90            ))?;
91        }
92        Ok(())
93    }
94
95    fn attribute<T>(&self, attr: PlanPreferenceAttribute) -> Result<T> {
96        validate_plan_preference_attribute_size(attr, size_of::<T>())?;
97
98        self.context.bind()?;
99        let mut value = MaybeUninit::<T>::uninit();
100        unsafe {
101            try_ffi!(sys::cutensorPlanPreferenceGetAttribute(
102                self.context.as_raw(),
103                self.handle,
104                attr.into(),
105                value.as_mut_ptr().cast(),
106                size_of::<T>() as _,
107            ))?;
108            Ok(value.assume_init())
109        }
110    }
111
112    pub fn set_autotune_mode(&mut self, mode: AutotuneMode) -> Result<()> {
113        self.set_attribute(PlanPreferenceAttribute::AutotuneMode, &mode)
114    }
115
116    pub fn set_cache_mode(&mut self, mode: CacheMode) -> Result<()> {
117        self.set_attribute(PlanPreferenceAttribute::CacheMode, &mode)
118    }
119
120    pub fn set_incremental_count(&mut self, count: i32) -> Result<()> {
121        self.set_attribute(PlanPreferenceAttribute::IncrementalCount, &count)
122    }
123
124    pub fn set_algorithm(&mut self, algorithm: Algorithm) -> Result<()> {
125        self.set_attribute(PlanPreferenceAttribute::Algorithm, &algorithm)
126    }
127
128    pub fn set_kernel_rank(&mut self, rank: i32) -> Result<()> {
129        self.set_attribute(PlanPreferenceAttribute::KernelRank, &rank)
130    }
131
132    pub fn set_jit_mode(&mut self, mode: JitMode) -> Result<()> {
133        self.set_attribute(PlanPreferenceAttribute::JitMode, &mode)
134    }
135
136    pub fn autotune_mode(&self) -> Result<AutotuneMode> {
137        self.attribute(PlanPreferenceAttribute::AutotuneMode)
138    }
139
140    pub fn cache_mode(&self) -> Result<CacheMode> {
141        self.attribute(PlanPreferenceAttribute::CacheMode)
142    }
143
144    pub fn incremental_count(&self) -> Result<i32> {
145        let value = self.attribute(PlanPreferenceAttribute::IncrementalCount)?;
146        Ok(value)
147    }
148
149    pub fn algorithm(&self) -> Result<Algorithm> {
150        self.attribute(PlanPreferenceAttribute::Algorithm)
151    }
152
153    pub fn kernel_rank(&self) -> Result<i32> {
154        let value = self.attribute(PlanPreferenceAttribute::KernelRank)?;
155        Ok(value)
156    }
157
158    pub fn jit_mode(&self) -> Result<JitMode> {
159        self.attribute(PlanPreferenceAttribute::JitMode)
160    }
161
162    pub fn gpu_arch(&self) -> Result<u32> {
163        let value: u32 = self.attribute(PlanPreferenceAttribute::GpuArch)?;
164        Ok(value)
165    }
166
167    /// Plans for a specific GPU architecture instead of the architecture associated with the context.
168    /// The value encodes the SM version as `10 * SM.major + SM.minor`.
169    /// Currently only SM versions 80, 90 and 100 are supported.
170    ///
171    /// # Errors
172    ///
173    /// Returns an error if `gpu_arch` cannot be represented for cuTENSOR or if
174    /// cuTENSOR rejects the architecture value.
175    pub fn set_gpu_arch(&mut self, gpu_arch: u32) -> Result<()> {
176        let gpu_arch = to_i32(gpu_arch, "gpu_arch")?;
177        self.set_attribute(PlanPreferenceAttribute::GpuArch, &gpu_arch)
178    }
179
180    pub const fn as_raw(&self) -> sys::cutensorPlanPreference_t {
181        self.handle
182    }
183
184    /// Consumes the preference and returns the raw cuTENSOR handle without
185    /// destroying it.
186    ///
187    /// The caller becomes responsible for eventually destroying the returned
188    /// handle with cuTENSOR.
189    pub fn into_raw(self) -> sys::cutensorPlanPreference_t {
190        let preference = ManuallyDrop::new(self);
191        preference.handle
192    }
193}
194
195impl Drop for PlanPreference {
196    fn drop(&mut self) {
197        if let Err(err) = self.context.bind() {
198            #[cfg(debug_assertions)]
199            eprintln!("failed to bind cutensor context before destroying plan preference: {err}");
200        }
201
202        unsafe {
203            if let Err(err) = try_ffi!(sys::cutensorDestroyPlanPreference(self.handle)) {
204                #[cfg(debug_assertions)]
205                eprintln!("failed to destroy cutensor plan preference: {err}");
206            }
207        }
208    }
209}
210
211#[derive(Debug)]
212pub struct Plan {
213    handle: sys::cutensorPlan_t,
214    context: ContextRef,
215    required_workspace_size: u64,
216    signature: OperationSignature,
217}
218
219unsafe impl Send for Plan {}
220
221impl Plan {
222    /// Determines the required workspace size for the given operation encoded by `operation`.
223    ///
224    /// # Errors
225    ///
226    /// Returns an error if the context cannot be bound, `operation` and
227    /// `preference` do not belong to a compatible context, cuTENSOR rejects the
228    /// workspace preference, or cuTENSOR cannot estimate the workspace size.
229    pub fn estimate_workspace_size(
230        ctx: &Context,
231        operation: &OperationDescriptor,
232        preference: &PlanPreference,
233        workspace_preference: WorkspacePreference,
234    ) -> Result<u64> {
235        ctx.bind()?;
236
237        let mut workspace_size = 0;
238        unsafe {
239            try_ffi!(sys::cutensorEstimateWorkspaceSize(
240                ctx.as_raw(),
241                operation.as_raw(),
242                preference.as_raw(),
243                workspace_preference.into(),
244                &raw mut workspace_size,
245            ))?;
246        }
247        Ok(workspace_size)
248    }
249
250    pub fn create(
251        ctx: &Context,
252        operation: &OperationDescriptor,
253        preference: &PlanPreference,
254        workspace_size_limit: u64,
255    ) -> Result<Self> {
256        ctx.bind()?;
257
258        let mut handle = ptr::null_mut();
259        unsafe {
260            try_ffi!(sys::cutensorCreatePlan(
261                ctx.as_raw(),
262                &raw mut handle,
263                operation.as_raw(),
264                preference.as_raw(),
265                workspace_size_limit,
266            ))?;
267        }
268
269        if handle.is_null() {
270            return Err(Error::NullHandle);
271        }
272
273        let mut required_workspace_size = 0_u64;
274        unsafe {
275            try_ffi!(sys::cutensorPlanGetAttribute(
276                ctx.as_raw(),
277                handle,
278                PlanAttribute::RequiredWorkspace.into(),
279                (&raw mut required_workspace_size).cast(),
280                size_of::<u64>() as _,
281            ))?;
282        }
283
284        Ok(Self {
285            handle,
286            context: ctx.as_context_ref(),
287            required_workspace_size,
288            signature: operation.signature(),
289        })
290    }
291
292    /// Wraps an existing cuTENSOR plan handle and takes ownership of it.
293    ///
294    /// # Safety
295    ///
296    /// `handle` must be a valid cuTENSOR plan handle associated with `ctx` and
297    /// compatible with `signature`. `required_workspace_size` must match the
298    /// plan's required workspace size. Ownership of `handle` is transferred to
299    /// the returned plan, and the handle must not be destroyed elsewhere after
300    /// calling this function.
301    pub unsafe fn from_raw(
302        handle: sys::cutensorPlan_t,
303        ctx: &Context,
304        required_workspace_size: u64,
305        signature: OperationSignature,
306    ) -> Self {
307        Self {
308            handle,
309            context: ctx.as_context_ref(),
310            required_workspace_size,
311            signature,
312        }
313    }
314
315    pub fn required_workspace_size(&self) -> u64 {
316        self.required_workspace_size
317    }
318
319    pub fn required_workspace_size_bytes(&self) -> Result<usize> {
320        to_usize(self.required_workspace_size, "required workspace size")
321    }
322
323    pub(crate) fn context(&self) -> &ContextRef {
324        &self.context
325    }
326
327    pub(crate) fn signature(&self) -> OperationSignature {
328        self.signature
329    }
330
331    pub const fn as_raw(&self) -> sys::cutensorPlan_t {
332        self.handle
333    }
334
335    /// Consumes the plan and returns the raw cuTENSOR plan handle without
336    /// destroying it.
337    ///
338    /// The caller becomes responsible for eventually destroying the returned
339    /// handle with cuTENSOR.
340    pub fn into_raw(self) -> sys::cutensorPlan_t {
341        let plan = ManuallyDrop::new(self);
342        plan.handle
343    }
344}
345
346impl Drop for Plan {
347    fn drop(&mut self) {
348        if let Err(err) = self.context.bind() {
349            #[cfg(debug_assertions)]
350            eprintln!("failed to bind cutensor context before destroying plan: {err}");
351        }
352
353        unsafe {
354            if let Err(err) = try_ffi!(sys::cutensorDestroyPlan(self.handle)) {
355                #[cfg(debug_assertions)]
356                eprintln!("failed to destroy cutensor plan: {err}");
357            }
358        }
359    }
360}
361
362fn validate_plan_preference_attribute_size(
363    attr: PlanPreferenceAttribute,
364    actual: usize,
365) -> Result<()> {
366    let expected = match attr {
367        PlanPreferenceAttribute::AutotuneMode => size_of::<sys::cutensorAutotuneMode_t>(),
368        PlanPreferenceAttribute::CacheMode => size_of::<sys::cutensorCacheMode_t>(),
369        PlanPreferenceAttribute::IncrementalCount
370        | PlanPreferenceAttribute::KernelRank
371        | PlanPreferenceAttribute::GpuArch => size_of::<i32>(),
372        PlanPreferenceAttribute::Algorithm => size_of::<sys::cutensorAlgo_t>(),
373        PlanPreferenceAttribute::JitMode => size_of::<sys::cutensorJitMode_t>(),
374    };
375
376    if actual != expected {
377        return Err(Error::PlanPreferenceInvalidAttributeSize {
378            attr,
379            expected,
380            actual,
381        });
382    }
383
384    Ok(())
385}
singe_cutensor/plan.rs

singe_cutensor/
plan.rs
