singe_cutensor/mp/
context.rs

1use std::{marker::PhantomData, 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    pub fn cuda_context(&self) -> &Arc<CudaContext> {
75        &self.handle.cuda_ctx
76    }
77
78    pub fn bind(&self) -> Result<()> {
79        Ok(self.cuda_context().bind()?)
80    }
81
82    pub fn local_device_id(&self) -> i32 {
83        self.handle.local_device_id
84    }
85
86    pub fn stream(&self) -> &StreamBinding {
87        &self.handle.stream
88    }
89
90    pub(crate) fn as_context_ref(&self) -> ContextRef<'comm> {
91        ContextRef {
92            handle: Arc::clone(&self.handle),
93        }
94    }
95
96    pub fn as_raw(&self) -> sys::cutensorMpHandle_t {
97        self.handle.raw
98    }
99}
100
101impl ContextRef<'_> {
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(crate) fn same_handle(&self, other: &Self) -> bool {
111        Arc::ptr_eq(&self.handle, &other.handle)
112    }
113
114    pub fn as_raw(&self) -> sys::cutensorMpHandle_t {
115        self.handle.raw
116    }
117}
118
119impl Drop for Handle<'_> {
120    fn drop(&mut self) {
121        unsafe {
122            if let Err(err) = try_ffi!(sys::cutensorMpDestroy(self.raw)) {
123                #[cfg(debug_assertions)]
124                eprintln!("failed to destroy cutensormp context: {err}");
125            }
126        }
127    }
128}
129
130pub(crate) fn validate_same_context(
131    expected: &ContextRef<'_>,
132    actual: &ContextRef<'_>,
133    name: &str,
134) -> Result<()> {
135    if !expected.same_handle(actual) {
136        return Err(Error::MpContextMismatch { name: name.into() });
137    }
138    Ok(())
139}
singe_cutensor/mp/context.rs

singe_cutensor/mp/
context.rs
