Skip to main content

baracuda_runtime/
external.rs

1//! External memory / semaphore interop via the Runtime API.
2//!
3//! Mirrors [`baracuda_driver::external`]. Because `cudaExternalMemory_t`
4//! and `CUexternalMemory` are typedef-compatible (same underlying C
5//! pointer), the two wrappers are interchangeable at the handle level —
6//! this module exists so Runtime-API users don't have to pull in the
7//! Driver crate just for external-resource import.
8//!
9//! Struct layouts (`CUDA_EXTERNAL_MEMORY_HANDLE_DESC` etc.) are shared
10//! with the Driver API — populate using the same typed builders in
11//! [`baracuda_cuda_sys::types`].
12
13use std::sync::Arc;
14
15use baracuda_cuda_sys::runtime::{cudaExternalMemory_t, cudaExternalSemaphore_t, runtime};
16use baracuda_cuda_sys::types::{
17    CUDA_EXTERNAL_MEMORY_BUFFER_DESC, CUDA_EXTERNAL_MEMORY_HANDLE_DESC,
18    CUDA_EXTERNAL_SEMAPHORE_HANDLE_DESC, CUDA_EXTERNAL_SEMAPHORE_SIGNAL_PARAMS,
19    CUDA_EXTERNAL_SEMAPHORE_WAIT_PARAMS,
20};
21
22use crate::error::{check, Result};
23use crate::stream::Stream;
24
25/// An imported external-memory handle (Vulkan `VkDeviceMemory`, D3D12
26/// heap / resource, NvSciBuf, DMA-buf FD, ...). Destroyed on drop.
27#[derive(Clone)]
28pub struct ExternalMemory {
29    inner: Arc<ExternalMemoryInner>,
30}
31
32struct ExternalMemoryInner {
33    handle: cudaExternalMemory_t,
34}
35
36unsafe impl Send for ExternalMemoryInner {}
37unsafe impl Sync for ExternalMemoryInner {}
38
39impl core::fmt::Debug for ExternalMemoryInner {
40    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
41        f.debug_struct("ExternalMemory")
42            .field("handle", &self.handle)
43            .finish_non_exhaustive()
44    }
45}
46
47impl core::fmt::Debug for ExternalMemory {
48    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
49        self.inner.fmt(f)
50    }
51}
52
53impl ExternalMemory {
54    /// Import an external-memory handle described by `desc`.
55    ///
56    /// # Safety
57    ///
58    /// `desc.handle` must describe a live OS object that the process can
59    /// access (file descriptor, NT HANDLE, NvSciBufObj, ...). CUDA
60    /// retains a reference to the underlying memory until this
61    /// `ExternalMemory` drops.
62    pub unsafe fn import(desc: &CUDA_EXTERNAL_MEMORY_HANDLE_DESC) -> Result<Self> { unsafe {
63        let r = runtime()?;
64        let cu = r.cuda_import_external_memory()?;
65        let mut handle: cudaExternalMemory_t = core::ptr::null_mut();
66        check(cu(&mut handle, desc))?;
67        Ok(Self {
68            inner: Arc::new(ExternalMemoryInner { handle }),
69        })
70    }}
71
72    /// Expose a subregion of the imported memory as a device pointer
73    /// valid in the *current* CUDA context.
74    pub fn mapped_buffer(
75        &self,
76        offset: u64,
77        size: u64,
78        flags: u32,
79    ) -> Result<*mut core::ffi::c_void> {
80        let r = runtime()?;
81        let cu = r.cuda_external_memory_get_mapped_buffer()?;
82        let desc = CUDA_EXTERNAL_MEMORY_BUFFER_DESC {
83            offset,
84            size,
85            flags,
86            reserved: [0; 16],
87        };
88        let mut ptr: *mut core::ffi::c_void = core::ptr::null_mut();
89        check(unsafe { cu(&mut ptr, self.inner.handle, &desc) })?;
90        Ok(ptr)
91    }
92
93    /// Raw handle. Interchangeable with `CUexternalMemory` at the ABI
94    /// level — cast via `as baracuda_cuda_sys::CUexternalMemory` if you
95    /// need to interop with the Driver-side wrapper.
96    #[inline]
97    pub fn as_raw(&self) -> cudaExternalMemory_t {
98        self.inner.handle
99    }
100}
101
102impl Drop for ExternalMemoryInner {
103    fn drop(&mut self) {
104        if self.handle.is_null() {
105            return;
106        }
107        if let Ok(r) = runtime() {
108            if let Ok(cu) = r.cuda_destroy_external_memory() {
109                let _ = unsafe { cu(self.handle) };
110            }
111        }
112    }
113}
114
115/// An imported external-semaphore handle. Destroyed on drop.
116#[derive(Clone)]
117pub struct ExternalSemaphore {
118    inner: Arc<ExternalSemaphoreInner>,
119}
120
121struct ExternalSemaphoreInner {
122    handle: cudaExternalSemaphore_t,
123}
124
125unsafe impl Send for ExternalSemaphoreInner {}
126unsafe impl Sync for ExternalSemaphoreInner {}
127
128impl core::fmt::Debug for ExternalSemaphoreInner {
129    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
130        f.debug_struct("ExternalSemaphore")
131            .field("handle", &self.handle)
132            .finish_non_exhaustive()
133    }
134}
135
136impl core::fmt::Debug for ExternalSemaphore {
137    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
138        self.inner.fmt(f)
139    }
140}
141
142impl ExternalSemaphore {
143    /// Import an external-semaphore handle.
144    ///
145    /// # Safety
146    ///
147    /// Same discipline as [`ExternalMemory::import`].
148    pub unsafe fn import(desc: &CUDA_EXTERNAL_SEMAPHORE_HANDLE_DESC) -> Result<Self> { unsafe {
149        let r = runtime()?;
150        let cu = r.cuda_import_external_semaphore()?;
151        let mut handle: cudaExternalSemaphore_t = core::ptr::null_mut();
152        check(cu(&mut handle, desc))?;
153        Ok(Self {
154            inner: Arc::new(ExternalSemaphoreInner { handle }),
155        })
156    }}
157
158    /// Enqueue a signal of fence value `value` on `stream`.
159    pub fn signal_fence_async(&self, value: u64, stream: &Stream) -> Result<()> {
160        let r = runtime()?;
161        let cu = r.cuda_signal_external_semaphores_async()?;
162        let params = CUDA_EXTERNAL_SEMAPHORE_SIGNAL_PARAMS::fence_value(value);
163        check(unsafe { cu(&self.inner.handle, &params, 1, stream.as_raw()) })
164    }
165
166    /// Enqueue a wait for fence value `value` on `stream`.
167    pub fn wait_fence_async(&self, value: u64, stream: &Stream) -> Result<()> {
168        let r = runtime()?;
169        let cu = r.cuda_wait_external_semaphores_async()?;
170        let params = CUDA_EXTERNAL_SEMAPHORE_WAIT_PARAMS::fence_value(value);
171        check(unsafe { cu(&self.inner.handle, &params, 1, stream.as_raw()) })
172    }
173
174    #[inline]
175    pub fn as_raw(&self) -> cudaExternalSemaphore_t {
176        self.inner.handle
177    }
178}
179
180impl Drop for ExternalSemaphoreInner {
181    fn drop(&mut self) {
182        if self.handle.is_null() {
183            return;
184        }
185        if let Ok(r) = runtime() {
186            if let Ok(cu) = r.cuda_destroy_external_semaphore() {
187                let _ = unsafe { cu(self.handle) };
188            }
189        }
190    }
191}