singe_cutensor/mg/
context.rs

1use std::{mem::ManuallyDrop, ptr, sync::Arc};
2
3use crate::{
4    error::{Error, Result},
5    sys, try_ffi,
6    utility::to_u32,
7};
8
9#[derive(Debug, Clone)]
10pub struct Context {
11    handle: Arc<Handle>,
12}
13
14#[derive(Debug)]
15struct Handle {
16    raw: sys::cutensorMgHandle_t,
17    devices: Vec<i32>,
18}
19
20#[derive(Debug, Clone)]
21pub struct ContextRef {
22    handle: Arc<Handle>,
23}
24
25// cuTENSORMg contexts are immutable after creation for a fixed device set. The
26// Arc-backed handle can be shared, while wrappers remain freely movable.
27unsafe impl Send for Handle {}
28unsafe impl Sync for Handle {}
29unsafe impl Send for Context {}
30unsafe impl Send for ContextRef {}
31
32impl Context {
33    /// Creates a cuTENSORMg handle for the devices that may participate in later operations.
34    ///
35    /// # Errors
36    ///
37    /// Returns an error if the device count cannot be represented by cuTENSORMg,
38    /// if cuTENSORMg cannot create a handle, or if it returns a null handle.
39    pub fn create(devices: &[i32]) -> Result<Self> {
40        let num_devices = to_u32(devices.len(), "devices")?;
41        let mut handle = ptr::null_mut();
42        unsafe {
43            try_ffi!(sys::cutensorMgCreate(
44                &raw mut handle,
45                num_devices,
46                devices.as_ptr(),
47            ))?;
48        }
49
50        if handle.is_null() {
51            return Err(Error::NullHandle);
52        }
53
54        Ok(Self {
55            handle: Arc::new(Handle {
56                raw: handle,
57                devices: devices.to_vec(),
58            }),
59        })
60    }
61
62    /// Wraps an existing cuTENSORMg handle.
63    ///
64    /// # Safety
65    ///
66    /// `handle` must be a valid cuTENSORMg handle for `devices`. The returned
67    /// context takes ownership of `handle` and destroys it when the last clone
68    /// is dropped.
69    pub unsafe fn from_raw(handle: sys::cutensorMgHandle_t, devices: Vec<i32>) -> Result<Self> {
70        if handle.is_null() {
71            return Err(Error::NullHandle);
72        }
73
74        Ok(Self {
75            handle: Arc::new(Handle {
76                raw: handle,
77                devices,
78            }),
79        })
80    }
81
82    pub fn devices(&self) -> &[i32] {
83        &self.handle.devices
84    }
85
86    pub fn device_count(&self) -> usize {
87        self.handle.devices.len()
88    }
89
90    pub(crate) fn as_context_ref(&self) -> ContextRef {
91        ContextRef {
92            handle: Arc::clone(&self.handle),
93        }
94    }
95
96    pub fn as_raw(&self) -> sys::cutensorMgHandle_t {
97        self.handle.raw
98    }
99
100    /// Consumes this context and returns the owned raw cuTENSORMg handle.
101    ///
102    /// # Errors
103    ///
104    /// Returns [`Error::HandleShared`] if cloned contexts or context references
105    /// still point at the same handle.
106    pub fn into_raw(self) -> Result<sys::cutensorMgHandle_t> {
107        let handle = Arc::try_unwrap(self.handle).map_err(|_| Error::HandleShared)?;
108        let handle = ManuallyDrop::new(handle);
109        Ok(handle.raw)
110    }
111}
112
113impl ContextRef {
114    pub fn devices(&self) -> &[i32] {
115        &self.handle.devices
116    }
117
118    pub fn device_count(&self) -> usize {
119        self.handle.devices.len()
120    }
121
122    pub(crate) fn same_handle(&self, other: &Self) -> bool {
123        Arc::ptr_eq(&self.handle, &other.handle)
124    }
125
126    pub fn as_raw(&self) -> sys::cutensorMgHandle_t {
127        self.handle.raw
128    }
129}
130
131impl Drop for Handle {
132    fn drop(&mut self) {
133        unsafe {
134            if let Err(err) = try_ffi!(sys::cutensorMgDestroy(self.raw)) {
135                #[cfg(debug_assertions)]
136                eprintln!("failed to destroy cutensormg context: {err}");
137            }
138        }
139    }
140}
141
142pub(crate) fn validate_same_context(
143    expected: &ContextRef,
144    actual: &ContextRef,
145    name: &str,
146) -> Result<()> {
147    if !expected.same_handle(actual) {
148        return Err(Error::MgContextMismatch { name: name.into() });
149    }
150    Ok(())
151}
singe_cutensor/mg/context.rs

singe_cutensor/mg/
context.rs
