1use libloading::Library;
13use std::ffi::c_void;
14use std::os::raw::{c_int, c_uint};
15use std::sync::{Arc, OnceLock};
16
17pub(crate) type CudaError = c_int; pub type GraphicsResource = *mut c_void; pub(crate) type ExternalMemory = *mut c_void; pub type CudaStream = *mut c_void; #[allow(non_snake_case, dead_code)]
23pub(crate) struct CudaTable {
24 _lib: &'static Library,
25 pub graphics_gl_register_buffer:
26 unsafe extern "C" fn(*mut GraphicsResource, c_uint, c_uint) -> CudaError,
27 pub graphics_map_resources:
28 unsafe extern "C" fn(c_int, *mut GraphicsResource, *mut c_void) -> CudaError,
29 pub graphics_get_mapped_pointer:
30 unsafe extern "C" fn(*mut *mut c_void, *mut usize, GraphicsResource) -> CudaError,
31 pub graphics_unmap_resources:
32 unsafe extern "C" fn(c_int, *mut GraphicsResource, *mut c_void) -> CudaError,
33 pub graphics_unregister_resource: unsafe extern "C" fn(GraphicsResource) -> CudaError,
34 pub import_external_memory:
35 unsafe extern "C" fn(*mut ExternalMemory, *const c_void) -> CudaError,
36 pub external_memory_get_mapped_buffer:
37 unsafe extern "C" fn(*mut *mut c_void, ExternalMemory, *const c_void) -> CudaError,
38 pub destroy_external_memory: unsafe extern "C" fn(ExternalMemory) -> CudaError,
39 pub memcpy: unsafe extern "C" fn(*mut c_void, *const c_void, usize, c_int) -> CudaError,
40 pub free: unsafe extern "C" fn(*mut c_void) -> CudaError,
41 pub stream_create: unsafe extern "C" fn(*mut CudaStream) -> CudaError,
42 pub stream_synchronize: unsafe extern "C" fn(CudaStream) -> CudaError,
43 pub stream_destroy: unsafe extern "C" fn(CudaStream) -> CudaError,
44}
45
46static TABLE: OnceLock<Option<CudaTable>> = OnceLock::new();
47
48fn load() -> Option<CudaTable> {
49 let lib = ["libcudart.so", "libcudart.so.12", "libcudart.so.11.0"]
50 .iter()
51 .find_map(|n| unsafe { Library::new(*n) }.ok())?;
52 let lib: &'static Library = Box::leak(Box::new(lib));
53 macro_rules! sym {
54 ($n:literal) => {{
55 *unsafe { lib.get(concat!($n, "\0").as_bytes()) }.ok()?
56 }};
57 }
58 Some(CudaTable {
59 _lib: lib,
60 graphics_gl_register_buffer: sym!("cudaGraphicsGLRegisterBuffer"),
61 graphics_map_resources: sym!("cudaGraphicsMapResources"),
62 graphics_get_mapped_pointer: sym!("cudaGraphicsResourceGetMappedPointer"),
63 graphics_unmap_resources: sym!("cudaGraphicsUnmapResources"),
64 graphics_unregister_resource: sym!("cudaGraphicsUnregisterResource"),
65 import_external_memory: sym!("cudaImportExternalMemory"),
66 external_memory_get_mapped_buffer: sym!("cudaExternalMemoryGetMappedBuffer"),
67 destroy_external_memory: sym!("cudaDestroyExternalMemory"),
68 memcpy: sym!("cudaMemcpy"),
69 free: sym!("cudaFree"),
70 stream_create: sym!("cudaStreamCreate"),
71 stream_synchronize: sym!("cudaStreamSynchronize"),
72 stream_destroy: sym!("cudaStreamDestroy"),
73 })
74}
75
76pub(crate) fn table() -> Option<&'static CudaTable> {
77 TABLE.get_or_init(load).as_ref()
78}
79
80pub fn is_cuda_available() -> bool {
82 table().is_some()
83}
84
85pub const CUDA_MEMCPY_DEVICE_TO_HOST: c_int = 2;
87
88pub unsafe fn memcpy_device_to_host(
98 host: *mut c_void,
99 device: *const c_void,
100 count: usize,
101) -> bool {
102 match table() {
103 Some(t) => (t.memcpy)(host, device, count, CUDA_MEMCPY_DEVICE_TO_HOST) == 0,
104 None => false,
105 }
106}
107
108pub fn stream_create() -> Option<CudaStream> {
113 let t = table()?;
114 let mut stream: CudaStream = std::ptr::null_mut();
115 if unsafe { (t.stream_create)(&mut stream) } != 0 {
116 return None;
117 }
118 Some(stream)
119}
120
121pub unsafe fn stream_synchronize(stream: CudaStream) -> bool {
129 match table() {
130 Some(t) => (t.stream_synchronize)(stream) == 0,
131 None => false,
132 }
133}
134
135pub unsafe fn stream_destroy(stream: CudaStream) {
142 if let Some(t) = table() {
143 let _ = (t.stream_destroy)(stream);
144 }
145}
146
147pub fn gl_register_buffer(buffer_id: u32) -> Option<usize> {
151 let t = table()?;
152 let mut res: GraphicsResource = std::ptr::null_mut();
153 let rc = unsafe { (t.graphics_gl_register_buffer)(&mut res, buffer_id, 0) };
155 if rc != 0 {
156 log::debug!("cudaGraphicsGLRegisterBuffer(buffer={buffer_id}) failed: cudaError {rc}");
157 return None;
158 }
159 Some(res as usize)
160}
161
162pub fn gl_map_resource(resource: usize) -> Option<(usize, usize)> {
164 let t = table()?;
165 let mut res = resource as GraphicsResource;
166 if unsafe { (t.graphics_map_resources)(1, &mut res, std::ptr::null_mut()) } != 0 {
167 return None;
168 }
169 let (mut ptr, mut size) = (std::ptr::null_mut::<c_void>(), 0usize);
170 if unsafe { (t.graphics_get_mapped_pointer)(&mut ptr, &mut size, res) } != 0 {
171 unsafe {
172 (t.graphics_unmap_resources)(1, &mut res, std::ptr::null_mut());
173 }
174 return None;
175 }
176 Some((ptr as usize, size))
177}
178
179pub fn gl_unmap_resource(resource: usize) {
181 if let Some(t) = table() {
182 let mut r = resource as GraphicsResource;
183 unsafe {
184 (t.graphics_unmap_resources)(1, &mut r, std::ptr::null_mut());
185 }
186 }
187}
188
189pub fn gl_unregister_resource(resource: usize) {
191 if let Some(t) = table() {
192 unsafe {
193 (t.graphics_unregister_resource)(resource as GraphicsResource);
194 }
195 }
196}
197
198#[allow(dead_code)] pub(crate) const CUDA_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD: c_uint = 1;
212
213#[allow(dead_code)] #[repr(C)]
225pub(crate) struct CudaExternalMemoryHandleDesc {
226 pub type_: c_int,
227 pub _pad0: u32,
228 pub handle_fd: c_int,
229 pub _union_rest: [u32; 3],
230 pub size: u64,
231 pub flags: c_uint,
232 pub _tail: u32,
233}
234
235#[allow(dead_code)] #[repr(C)]
244pub(crate) struct CudaExternalMemoryBufferDesc {
245 pub offset: u64,
246 pub size: u64,
247 pub flags: c_uint,
248 pub _tail: u32,
249}
250
251#[cfg(target_os = "linux")]
267pub(crate) fn import_dma_fd(fd: i32, size: usize) -> Option<(ExternalMemory, *mut c_void)> {
268 use std::os::fd::{BorrowedFd, FromRawFd, IntoRawFd, OwnedFd};
269 let t = table()?;
270 let dup_fd = unsafe { BorrowedFd::borrow_raw(fd) }
275 .try_clone_to_owned()
276 .ok()?
277 .into_raw_fd();
278 let mut desc: CudaExternalMemoryHandleDesc = unsafe { std::mem::zeroed() };
279 desc.type_ = CUDA_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD as c_int;
280 desc.handle_fd = dup_fd;
281 desc.size = size as u64;
282 let mut ext: ExternalMemory = std::ptr::null_mut();
283 if unsafe { (t.import_external_memory)(&mut ext, &desc as *const _ as *const c_void) } != 0 {
284 drop(unsafe { OwnedFd::from_raw_fd(dup_fd) });
286 return None;
287 }
288 let bdesc = CudaExternalMemoryBufferDesc {
290 offset: 0,
291 size: size as u64,
292 flags: 0,
293 _tail: 0,
294 };
295 let mut dptr: *mut c_void = std::ptr::null_mut();
296 if unsafe {
297 (t.external_memory_get_mapped_buffer)(&mut dptr, ext, &bdesc as *const _ as *const c_void)
298 } != 0
299 {
300 unsafe { (t.destroy_external_memory)(ext) };
301 return None;
302 }
303 Some((ext, dptr))
304}
305
306pub trait CudaGlOps: Send + Sync {
309 fn map(&self, resource: GraphicsResource) -> Option<(*mut c_void, usize)>;
310 fn unmap(&self, resource: GraphicsResource);
311 fn unregister(&self, resource: GraphicsResource);
312}
313
314enum CudaBacking {
315 #[allow(dead_code)] GlBuffer {
317 resource: GraphicsResource,
318 ops: Arc<dyn CudaGlOps>,
319 },
320 #[allow(dead_code)] ExternalMem {
322 ext_mem: ExternalMemory,
323 dptr: *mut c_void,
324 },
325}
326
327unsafe impl Send for CudaBacking {}
330unsafe impl Sync for CudaBacking {}
331
332pub struct CudaHandle {
334 kind: CudaBacking,
335 size: usize,
336}
337
338impl std::fmt::Debug for CudaHandle {
339 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
340 let kind = match &self.kind {
341 CudaBacking::GlBuffer { .. } => "GlBuffer",
342 CudaBacking::ExternalMem { .. } => "ExternalMem",
343 };
344 f.debug_struct("CudaHandle")
345 .field("kind", &kind)
346 .field("size", &self.size)
347 .finish()
348 }
349}
350
351impl CudaHandle {
352 pub fn new_gl(resource: GraphicsResource, size: usize, ops: Arc<dyn CudaGlOps>) -> Self {
356 Self {
357 kind: CudaBacking::GlBuffer { resource, ops },
358 size,
359 }
360 }
361
362 #[allow(dead_code)] pub(crate) fn new_external(ext_mem: ExternalMemory, dptr: *mut c_void, size: usize) -> Self {
364 Self {
365 kind: CudaBacking::ExternalMem { ext_mem, dptr },
366 size,
367 }
368 }
369
370 pub fn map(&self) -> Option<CudaMap<'_>> {
373 match &self.kind {
374 CudaBacking::GlBuffer { resource, ops } => {
375 let (ptr, len) = ops.map(*resource)?;
376 Some(CudaMap {
377 ptr,
378 len,
379 unmap: Some((ops.clone(), *resource)),
380 _marker: std::marker::PhantomData,
381 })
382 }
383 CudaBacking::ExternalMem { dptr, .. } => Some(CudaMap {
384 ptr: *dptr,
385 len: self.size,
386 unmap: None,
387 _marker: std::marker::PhantomData,
388 }),
389 }
390 }
391}
392
393impl Drop for CudaHandle {
394 fn drop(&mut self) {
395 match &self.kind {
396 CudaBacking::GlBuffer { resource, ops } => ops.unregister(*resource),
397 CudaBacking::ExternalMem { ext_mem, dptr: _ } => {
398 if let Some(t) = table() {
404 unsafe {
405 (t.destroy_external_memory)(*ext_mem);
406 }
407 }
408 }
409 }
410 }
411}
412
413pub struct CudaMap<'a> {
417 ptr: *mut c_void,
418 len: usize,
419 unmap: Option<(Arc<dyn CudaGlOps>, GraphicsResource)>,
420 _marker: std::marker::PhantomData<&'a ()>,
421}
422
423unsafe impl Send for CudaMap<'_> {}
427unsafe impl Sync for CudaMap<'_> {}
428
429impl CudaMap<'_> {
430 pub fn device_ptr(&self) -> *mut c_void {
432 self.ptr
433 }
434
435 pub fn len(&self) -> usize {
437 self.len
438 }
439
440 pub fn is_empty(&self) -> bool {
442 self.len == 0
443 }
444}
445
446impl Drop for CudaMap<'_> {
447 fn drop(&mut self) {
448 if let Some((ops, r)) = self.unmap.take() {
449 ops.unmap(r);
450 }
451 }
452}
453
454#[cfg(test)]
455mod ext_mem_layout {
456 use super::*;
457 #[test]
458 fn external_memory_desc_abi() {
459 assert_eq!(std::mem::size_of::<CudaExternalMemoryHandleDesc>(), 40);
460 assert_eq!(std::mem::align_of::<CudaExternalMemoryHandleDesc>(), 8);
461 let d: CudaExternalMemoryHandleDesc = unsafe { std::mem::zeroed() };
463 let base = &d as *const _ as usize;
464 assert_eq!((&d.size as *const _ as usize) - base, 24);
465 assert_eq!((&d.flags as *const _ as usize) - base, 32);
466 assert_eq!(std::mem::size_of::<CudaExternalMemoryBufferDesc>(), 24);
467 let b: CudaExternalMemoryBufferDesc = unsafe { std::mem::zeroed() };
468 let bb = &b as *const _ as usize;
469 assert_eq!((&b.size as *const _ as usize) - bb, 8);
470 }
471}
472
473#[cfg(test)]
474mod tests {
475 use super::*;
476 #[test]
477 fn cuda_table_loads_when_libcudart_present() {
478 let avail = is_cuda_available();
479 if avail {
480 assert!(table().is_some(), "table present when available");
481 }
482 }
484
485 #[test]
486 fn pub_primitives_degrade_without_libcudart() {
487 if is_cuda_available() {
492 return;
493 }
494 assert!(gl_register_buffer(0).is_none());
495 assert!(gl_map_resource(0).is_none());
496 gl_unmap_resource(0); gl_unregister_resource(0); assert!(!unsafe { memcpy_device_to_host(std::ptr::null_mut(), std::ptr::null(), 0) });
501 }
502}
503
504#[cfg(test)]
505mod handle_tests {
506 use super::*;
507 use std::sync::{
508 atomic::{AtomicUsize, Ordering},
509 Arc,
510 };
511 struct MockOps {
512 unmaps: Arc<AtomicUsize>,
513 unregisters: Arc<AtomicUsize>,
514 }
515 impl CudaGlOps for MockOps {
516 fn map(&self, _r: GraphicsResource) -> Option<(*mut std::ffi::c_void, usize)> {
517 Some((0x1000usize as *mut _, 4096))
518 }
519 fn unmap(&self, _r: GraphicsResource) {
520 self.unmaps.fetch_add(1, Ordering::SeqCst);
521 }
522 fn unregister(&self, _r: GraphicsResource) {
523 self.unregisters.fetch_add(1, Ordering::SeqCst);
524 }
525 }
526 #[test]
527 fn cudamap_guard_unmaps_on_drop_for_glbuffer() {
528 let unmaps = Arc::new(AtomicUsize::new(0));
529 let unregisters = Arc::new(AtomicUsize::new(0));
530 {
531 let h = CudaHandle::new_gl(
532 0x1usize as GraphicsResource,
533 4096,
534 Arc::new(MockOps {
535 unmaps: unmaps.clone(),
536 unregisters: unregisters.clone(),
537 }),
538 );
539 {
540 let m = h.map().expect("map");
541 assert_eq!(m.device_ptr() as usize, 0x1000);
542 assert_eq!(m.len(), 4096);
543 assert!(!m.is_empty());
544 }
545 assert_eq!(
547 unmaps.load(Ordering::SeqCst),
548 1,
549 "Drop must unmap a GlBuffer"
550 );
551 assert_eq!(unregisters.load(Ordering::SeqCst), 0);
552 }
553 assert_eq!(
555 unregisters.load(Ordering::SeqCst),
556 1,
557 "Dropping a GlBuffer handle must unregister"
558 );
559 }
560
561 struct NoneOps;
563 impl CudaGlOps for NoneOps {
564 fn map(&self, _r: GraphicsResource) -> Option<(*mut std::ffi::c_void, usize)> {
565 None
566 }
567 fn unmap(&self, _r: GraphicsResource) {}
568 fn unregister(&self, _r: GraphicsResource) {}
569 }
570 #[test]
571 fn glbuffer_map_returns_none_when_ops_map_fails() {
572 let h = CudaHandle::new_gl(0x9usize as GraphicsResource, 4096, Arc::new(NoneOps));
573 assert!(
574 h.map().is_none(),
575 "GlBuffer map propagates ops.map() failure"
576 );
577 }
578
579 #[test]
580 fn glbuffer_handle_debug_and_empty_map() {
581 let unmaps = Arc::new(AtomicUsize::new(0));
582 let unregisters = Arc::new(AtomicUsize::new(0));
583 let h = CudaHandle::new_gl(
584 0x2usize as GraphicsResource,
585 0,
586 Arc::new(MockOps {
587 unmaps: unmaps.clone(),
588 unregisters: unregisters.clone(),
589 }),
590 );
591 let dbg = format!("{h:?}");
592 assert!(
593 dbg.contains("GlBuffer"),
594 "debug names the backing kind: {dbg}"
595 );
596 assert!(dbg.contains("size"), "debug includes size: {dbg}");
597 }
598
599 #[test]
600 fn external_mem_map_is_persistent_and_debug_names_kind() {
601 let dptr = 0xCAFE_0000usize as *mut std::ffi::c_void;
604 let h = CudaHandle::new_external(std::ptr::null_mut(), dptr, 8192);
605 let dbg = format!("{h:?}");
606 assert!(dbg.contains("ExternalMem"), "debug names the kind: {dbg}");
607 {
608 let m = h.map().expect("ExternalMem map is always Some");
609 assert_eq!(m.device_ptr(), dptr, "persistent device ptr passthrough");
610 assert_eq!(m.len(), 8192);
611 assert!(!m.is_empty());
612 }
614 std::mem::forget(h);
617 }
618
619 #[test]
620 fn external_mem_zero_len_map_is_empty() {
621 let h = CudaHandle::new_external(std::ptr::null_mut(), std::ptr::null_mut(), 0);
622 {
623 let m = h.map().expect("map");
624 assert_eq!(m.len(), 0);
625 assert!(m.is_empty(), "zero-length mapping is empty");
626 assert!(m.device_ptr().is_null());
627 }
628 std::mem::forget(h); }
630}