baracuda_runtime/
external.rs1use 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#[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 pub unsafe fn import(desc: &CUDA_EXTERNAL_MEMORY_HANDLE_DESC) -> Result<Self> {
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 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 #[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#[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 pub unsafe fn import(desc: &CUDA_EXTERNAL_SEMAPHORE_HANDLE_DESC) -> Result<Self> {
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 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, ¶ms, 1, stream.as_raw()) })
164 }
165
166 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, ¶ms, 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}