singe_cutensor/
plan.rs

1#[allow(unused_imports)]
2use crate::error::Status;
3
4use std::{mem::MaybeUninit, ptr};
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    fn set_attribute<T>(&mut self, attr: PlanPreferenceAttribute, value: &T) -> Result<()> {
61        validate_plan_preference_attribute_size(attr, size_of::<T>())?;
62
63        self.context.bind()?;
64        unsafe {
65            try_ffi!(sys::cutensorPlanPreferenceSetAttribute(
66                self.context.as_raw(),
67                self.handle,
68                attr.into(),
69                ptr::from_ref(value).cast(),
70                size_of::<T>() as _,
71            ))?;
72        }
73        Ok(())
74    }
75
76    fn attribute<T>(&self, attr: PlanPreferenceAttribute) -> Result<T> {
77        validate_plan_preference_attribute_size(attr, size_of::<T>())?;
78
79        self.context.bind()?;
80        let mut value = MaybeUninit::<T>::uninit();
81        unsafe {
82            try_ffi!(sys::cutensorPlanPreferenceGetAttribute(
83                self.context.as_raw(),
84                self.handle,
85                attr.into(),
86                value.as_mut_ptr().cast(),
87                size_of::<T>() as _,
88            ))?;
89            Ok(value.assume_init())
90        }
91    }
92
93    pub fn set_autotune_mode(&mut self, mode: AutotuneMode) -> Result<()> {
94        self.set_attribute(PlanPreferenceAttribute::AutotuneMode, &mode)
95    }
96
97    pub fn set_cache_mode(&mut self, mode: CacheMode) -> Result<()> {
98        self.set_attribute(PlanPreferenceAttribute::CacheMode, &mode)
99    }
100
101    pub fn set_incremental_count(&mut self, count: i32) -> Result<()> {
102        self.set_attribute(PlanPreferenceAttribute::IncrementalCount, &count)
103    }
104
105    pub fn set_algorithm(&mut self, algorithm: Algorithm) -> Result<()> {
106        self.set_attribute(PlanPreferenceAttribute::Algorithm, &algorithm)
107    }
108
109    pub fn set_kernel_rank(&mut self, rank: i32) -> Result<()> {
110        self.set_attribute(PlanPreferenceAttribute::KernelRank, &rank)
111    }
112
113    pub fn set_jit_mode(&mut self, mode: JitMode) -> Result<()> {
114        self.set_attribute(PlanPreferenceAttribute::JitMode, &mode)
115    }
116
117    pub fn autotune_mode(&self) -> Result<AutotuneMode> {
118        self.attribute(PlanPreferenceAttribute::AutotuneMode)
119    }
120
121    pub fn cache_mode(&self) -> Result<CacheMode> {
122        self.attribute(PlanPreferenceAttribute::CacheMode)
123    }
124
125    pub fn incremental_count(&self) -> Result<i32> {
126        let value = self.attribute(PlanPreferenceAttribute::IncrementalCount)?;
127        Ok(value)
128    }
129
130    pub fn algorithm(&self) -> Result<Algorithm> {
131        self.attribute(PlanPreferenceAttribute::Algorithm)
132    }
133
134    pub fn kernel_rank(&self) -> Result<i32> {
135        let value = self.attribute(PlanPreferenceAttribute::KernelRank)?;
136        Ok(value)
137    }
138
139    pub fn jit_mode(&self) -> Result<JitMode> {
140        self.attribute(PlanPreferenceAttribute::JitMode)
141    }
142
143    pub fn gpu_arch(&self) -> Result<u32> {
144        let value: u32 = self.attribute(PlanPreferenceAttribute::GpuArch)?;
145        Ok(value)
146    }
147
148    /// Plans for a specific GPU architecture instead of the architecture associated with the context.
149    /// The value encodes the SM version as `10 * SM.major + SM.minor`.
150    /// Currently only SM versions 80, 90 and 100 are supported.
151    ///
152    /// # Errors
153    ///
154    /// Returns an error if `gpu_arch` cannot be represented for cuTENSOR or if
155    /// cuTENSOR rejects the architecture value.
156    pub fn set_gpu_arch(&mut self, gpu_arch: u32) -> Result<()> {
157        let gpu_arch = to_i32(gpu_arch, "gpu_arch")?;
158        self.set_attribute(PlanPreferenceAttribute::GpuArch, &gpu_arch)
159    }
160
161    pub const fn as_raw(&self) -> sys::cutensorPlanPreference_t {
162        self.handle
163    }
164}
165
166impl Drop for PlanPreference {
167    fn drop(&mut self) {
168        if let Err(err) = self.context.bind() {
169            #[cfg(debug_assertions)]
170            eprintln!("failed to bind cutensor context before destroying plan preference: {err}");
171        }
172
173        unsafe {
174            if let Err(err) = try_ffi!(sys::cutensorDestroyPlanPreference(self.handle)) {
175                #[cfg(debug_assertions)]
176                eprintln!("failed to destroy cutensor plan preference: {err}");
177            }
178        }
179    }
180}
181
182#[derive(Debug)]
183pub struct Plan {
184    handle: sys::cutensorPlan_t,
185    context: ContextRef,
186    required_workspace_size: u64,
187    signature: OperationSignature,
188}
189
190impl Plan {
191    /// Determines the required workspace size for the given operation encoded by `operation`.
192    ///
193    /// # Errors
194    ///
195    /// Returns an error if the context cannot be bound, `operation` and
196    /// `preference` do not belong to a compatible context, cuTENSOR rejects the
197    /// workspace preference, or cuTENSOR cannot estimate the workspace size.
198    pub fn estimate_workspace_size(
199        ctx: &Context,
200        operation: &OperationDescriptor,
201        preference: &PlanPreference,
202        workspace_preference: WorkspacePreference,
203    ) -> Result<u64> {
204        ctx.bind()?;
205
206        let mut workspace_size = 0;
207        unsafe {
208            try_ffi!(sys::cutensorEstimateWorkspaceSize(
209                ctx.as_raw(),
210                operation.as_raw(),
211                preference.as_raw(),
212                workspace_preference.into(),
213                &raw mut workspace_size,
214            ))?;
215        }
216        Ok(workspace_size)
217    }
218
219    pub fn create(
220        ctx: &Context,
221        operation: &OperationDescriptor,
222        preference: &PlanPreference,
223        workspace_size_limit: u64,
224    ) -> Result<Self> {
225        ctx.bind()?;
226
227        let mut handle = ptr::null_mut();
228        unsafe {
229            try_ffi!(sys::cutensorCreatePlan(
230                ctx.as_raw(),
231                &raw mut handle,
232                operation.as_raw(),
233                preference.as_raw(),
234                workspace_size_limit,
235            ))?;
236        }
237
238        if handle.is_null() {
239            return Err(Error::NullHandle);
240        }
241
242        let mut required_workspace_size = 0_u64;
243        unsafe {
244            try_ffi!(sys::cutensorPlanGetAttribute(
245                ctx.as_raw(),
246                handle,
247                PlanAttribute::RequiredWorkspace.into(),
248                (&raw mut required_workspace_size).cast(),
249                size_of::<u64>() as _,
250            ))?;
251        }
252
253        Ok(Self {
254            handle,
255            context: ctx.as_context_ref(),
256            required_workspace_size,
257            signature: operation.signature(),
258        })
259    }
260
261    pub fn required_workspace_size(&self) -> u64 {
262        self.required_workspace_size
263    }
264
265    pub fn required_workspace_size_bytes(&self) -> Result<usize> {
266        to_usize(self.required_workspace_size, "required workspace size")
267    }
268
269    pub(crate) fn context(&self) -> &ContextRef {
270        &self.context
271    }
272
273    pub(crate) fn signature(&self) -> OperationSignature {
274        self.signature
275    }
276
277    pub const fn as_raw(&self) -> sys::cutensorPlan_t {
278        self.handle
279    }
280}
281
282impl Drop for Plan {
283    fn drop(&mut self) {
284        if let Err(err) = self.context.bind() {
285            #[cfg(debug_assertions)]
286            eprintln!("failed to bind cutensor context before destroying plan: {err}");
287        }
288
289        unsafe {
290            if let Err(err) = try_ffi!(sys::cutensorDestroyPlan(self.handle)) {
291                #[cfg(debug_assertions)]
292                eprintln!("failed to destroy cutensor plan: {err}");
293            }
294        }
295    }
296}
297
298fn validate_plan_preference_attribute_size(
299    attr: PlanPreferenceAttribute,
300    actual: usize,
301) -> Result<()> {
302    let expected = match attr {
303        PlanPreferenceAttribute::AutotuneMode => size_of::<sys::cutensorAutotuneMode_t>(),
304        PlanPreferenceAttribute::CacheMode => size_of::<sys::cutensorCacheMode_t>(),
305        PlanPreferenceAttribute::IncrementalCount
306        | PlanPreferenceAttribute::KernelRank
307        | PlanPreferenceAttribute::GpuArch => size_of::<i32>(),
308        PlanPreferenceAttribute::Algorithm => size_of::<sys::cutensorAlgo_t>(),
309        PlanPreferenceAttribute::JitMode => size_of::<sys::cutensorJitMode_t>(),
310    };
311
312    if actual != expected {
313        return Err(Error::PlanPreferenceInvalidAttributeSize {
314            attr,
315            expected,
316            actual,
317        });
318    }
319
320    Ok(())
321}
singe_cutensor/plan.rs

singe_cutensor/
plan.rs
