singe_cutensor/mp/
context.rs

1use std::{marker::PhantomData, mem::ManuallyDrop, ptr, sync::Arc};
2
3use singe_cuda::{context::Context as CudaContext, stream::StreamBinding};
4use singe_nccl::communicator::Communicator;
5
6use crate::{
7    error::{Error, Result},
8    sys, try_ffi,
9};
10
11#[derive(Debug, Clone)]
12pub struct Context<'comm> {
13    handle: Arc<Handle<'comm>>,
14}
15
16#[derive(Debug)]
17pub(crate) struct Handle<'comm> {
18    raw: sys::cutensorMpHandle_t,
19    cuda_ctx: Arc<CudaContext>,
20    local_device_id: i32,
21    stream: StreamBinding,
22    _communicator: PhantomData<&'comm Communicator>,
23}
24
25#[derive(Debug, Clone)]
26pub(crate) struct ContextRef<'comm> {
27    handle: Arc<Handle<'comm>>,
28}
29
30// cuTENSORMp contexts are created for a fixed communicator, CUDA context, and
31// stream binding. The Arc-backed handle can be shared immutably, while wrappers
32// remain movable.
33unsafe impl Send for Handle<'_> {}
34unsafe impl Sync for Handle<'_> {}
35unsafe impl Send for Context<'_> {}
36unsafe impl Send for ContextRef<'_> {}
37
38impl<'comm> Context<'comm> {
39    pub fn create(
40        communicator: &'comm Communicator,
41        local_device_id: i32,
42        stream: StreamBinding,
43    ) -> Result<Self> {
44        if communicator.cuda_context().as_ref() != stream.context() {
45            return Err(Error::StreamContextMismatch);
46        }
47        communicator.bind()?;
48
49        let mut handle = ptr::null_mut();
50        unsafe {
51            try_ffi!(sys::cutensorMpCreate(
52                &raw mut handle,
53                communicator.as_raw(),
54                local_device_id,
55                stream.as_raw(),
56            ))?;
57        }
58
59        if handle.is_null() {
60            return Err(Error::NullHandle);
61        }
62
63        Ok(Self {
64            handle: Arc::new(Handle {
65                raw: handle,
66                cuda_ctx: Arc::clone(communicator.cuda_context()),
67                local_device_id,
68                stream,
69                _communicator: PhantomData,
70            }),
71        })
72    }
73
74    /// Wraps an existing cuTENSORMp handle.
75    ///
76    /// # Safety
77    ///
78    /// `handle` must be a valid cuTENSORMp handle associated with `cuda_ctx`,
79    /// `local_device_id`, and `stream`. The returned context takes ownership
80    /// of `handle` and destroys it when the last clone is dropped.
81    pub unsafe fn from_raw(
82        handle: sys::cutensorMpHandle_t,
83        cuda_ctx: Arc<CudaContext>,
84        local_device_id: i32,
85        stream: StreamBinding,
86    ) -> Result<Self> {
87        if handle.is_null() {
88            return Err(Error::NullHandle);
89        }
90
91        Ok(Self {
92            handle: Arc::new(Handle {
93                raw: handle,
94                cuda_ctx,
95                local_device_id,
96                stream,
97                _communicator: PhantomData,
98            }),
99        })
100    }
101
102    pub fn cuda_context(&self) -> &Arc<CudaContext> {
103        &self.handle.cuda_ctx
104    }
105
106    pub fn bind(&self) -> Result<()> {
107        Ok(self.cuda_context().bind()?)
108    }
109
110    pub fn local_device_id(&self) -> i32 {
111        self.handle.local_device_id
112    }
113
114    pub fn stream(&self) -> &StreamBinding {
115        &self.handle.stream
116    }
117
118    pub(crate) fn as_context_ref(&self) -> ContextRef<'comm> {
119        ContextRef {
120            handle: Arc::clone(&self.handle),
121        }
122    }
123
124    pub fn as_raw(&self) -> sys::cutensorMpHandle_t {
125        self.handle.raw
126    }
127
128    /// Consumes this context and returns the owned raw cuTENSORMp handle.
129    ///
130    /// # Errors
131    ///
132    /// Returns [`Error::HandleShared`] if cloned contexts or context references
133    /// still point at the same handle.
134    pub fn into_raw(self) -> Result<sys::cutensorMpHandle_t> {
135        let handle = Arc::try_unwrap(self.handle).map_err(|_| Error::HandleShared)?;
136        let handle = ManuallyDrop::new(handle);
137        Ok(handle.raw)
138    }
139}
140
141impl ContextRef<'_> {
142    pub fn cuda_context(&self) -> &Arc<CudaContext> {
143        &self.handle.cuda_ctx
144    }
145
146    pub fn bind(&self) -> Result<()> {
147        Ok(self.cuda_context().bind()?)
148    }
149
150    pub(crate) fn same_handle(&self, other: &Self) -> bool {
151        Arc::ptr_eq(&self.handle, &other.handle)
152    }
153
154    pub fn as_raw(&self) -> sys::cutensorMpHandle_t {
155        self.handle.raw
156    }
157}
158
159impl Drop for Handle<'_> {
160    fn drop(&mut self) {
161        if let Err(err) = self.cuda_ctx.bind() {
162            #[cfg(debug_assertions)]
163            eprintln!("failed to bind cuda context before destroying cutensormp handle: {err}");
164        }
165
166        unsafe {
167            if let Err(err) = try_ffi!(sys::cutensorMpDestroy(self.raw)) {
168                #[cfg(debug_assertions)]
169                eprintln!("failed to destroy cutensormp context: {err}");
170            }
171        }
172    }
173}
174
175pub(crate) fn validate_same_context(
176    expected: &ContextRef<'_>,
177    actual: &ContextRef<'_>,
178    name: &str,
179) -> Result<()> {
180    if !expected.same_handle(actual) {
181        return Err(Error::MpContextMismatch { name: name.into() });
182    }
183    Ok(())
184}
singe_cutensor/mp/context.rs

singe_cutensor/mp/
context.rs
