use libloading::Library;
use std::ffi::c_void;
use std::os::raw::{c_int, c_uint};
use std::sync::{Arc, OnceLock};
pub(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)]
pub(crate) struct CudaTable {
_lib: &'static Library,
pub graphics_gl_register_buffer:
unsafe extern "C" fn(*mut GraphicsResource, c_uint, c_uint) -> CudaError,
pub graphics_map_resources:
unsafe extern "C" fn(c_int, *mut GraphicsResource, *mut c_void) -> CudaError,
pub graphics_get_mapped_pointer:
unsafe extern "C" fn(*mut *mut c_void, *mut usize, GraphicsResource) -> CudaError,
pub graphics_unmap_resources:
unsafe extern "C" fn(c_int, *mut GraphicsResource, *mut c_void) -> CudaError,
pub graphics_unregister_resource: unsafe extern "C" fn(GraphicsResource) -> CudaError,
pub import_external_memory:
unsafe extern "C" fn(*mut ExternalMemory, *const c_void) -> CudaError,
pub external_memory_get_mapped_buffer:
unsafe extern "C" fn(*mut *mut c_void, ExternalMemory, *const c_void) -> CudaError,
pub destroy_external_memory: unsafe extern "C" fn(ExternalMemory) -> CudaError,
pub memcpy: unsafe extern "C" fn(*mut c_void, *const c_void, usize, c_int) -> CudaError,
pub free: unsafe extern "C" fn(*mut c_void) -> CudaError,
pub stream_create: unsafe extern "C" fn(*mut CudaStream) -> CudaError,
pub stream_synchronize: unsafe extern "C" fn(CudaStream) -> CudaError,
pub stream_destroy: unsafe extern "C" fn(CudaStream) -> CudaError,
}
static TABLE: OnceLock<Option<CudaTable>> = OnceLock::new();
fn load() -> Option<CudaTable> {
let lib = ["libcudart.so", "libcudart.so.12", "libcudart.so.11.0"]
.iter()
.find_map(|n| unsafe { Library::new(*n) }.ok())?;
let lib: &'static Library = Box::leak(Box::new(lib));
macro_rules! sym {
($n:literal) => {{
*unsafe { lib.get(concat!($n, "\0").as_bytes()) }.ok()?
}};
}
Some(CudaTable {
_lib: lib,
graphics_gl_register_buffer: sym!("cudaGraphicsGLRegisterBuffer"),
graphics_map_resources: sym!("cudaGraphicsMapResources"),
graphics_get_mapped_pointer: sym!("cudaGraphicsResourceGetMappedPointer"),
graphics_unmap_resources: sym!("cudaGraphicsUnmapResources"),
graphics_unregister_resource: sym!("cudaGraphicsUnregisterResource"),
import_external_memory: sym!("cudaImportExternalMemory"),
external_memory_get_mapped_buffer: sym!("cudaExternalMemoryGetMappedBuffer"),
destroy_external_memory: sym!("cudaDestroyExternalMemory"),
memcpy: sym!("cudaMemcpy"),
free: sym!("cudaFree"),
stream_create: sym!("cudaStreamCreate"),
stream_synchronize: sym!("cudaStreamSynchronize"),
stream_destroy: sym!("cudaStreamDestroy"),
})
}
pub(crate) fn table() -> Option<&'static CudaTable> {
TABLE.get_or_init(load).as_ref()
}
pub fn is_cuda_available() -> bool {
table().is_some()
}
pub const CUDA_MEMCPY_DEVICE_TO_HOST: c_int = 2;
pub unsafe fn memcpy_device_to_host(
host: *mut c_void,
device: *const c_void,
count: usize,
) -> bool {
match table() {
Some(t) => (t.memcpy)(host, device, count, CUDA_MEMCPY_DEVICE_TO_HOST) == 0,
None => false,
}
}
pub fn stream_create() -> Option<CudaStream> {
let t = table()?;
let mut stream: CudaStream = std::ptr::null_mut();
if unsafe { (t.stream_create)(&mut stream) } != 0 {
return None;
}
Some(stream)
}
pub unsafe fn stream_synchronize(stream: CudaStream) -> bool {
match table() {
Some(t) => (t.stream_synchronize)(stream) == 0,
None => false,
}
}
pub unsafe fn stream_destroy(stream: CudaStream) {
if let Some(t) = table() {
let _ = (t.stream_destroy)(stream);
}
}
pub fn gl_register_buffer(buffer_id: u32) -> Option<usize> {
let t = table()?;
let mut res: GraphicsResource = std::ptr::null_mut();
let rc = unsafe { (t.graphics_gl_register_buffer)(&mut res, buffer_id, 0) };
if rc != 0 {
log::debug!("cudaGraphicsGLRegisterBuffer(buffer={buffer_id}) failed: cudaError {rc}");
return None;
}
Some(res as usize)
}
pub fn gl_map_resource(resource: usize) -> Option<(usize, usize)> {
let t = table()?;
let mut res = resource as GraphicsResource;
if unsafe { (t.graphics_map_resources)(1, &mut res, std::ptr::null_mut()) } != 0 {
return None;
}
let (mut ptr, mut size) = (std::ptr::null_mut::<c_void>(), 0usize);
if unsafe { (t.graphics_get_mapped_pointer)(&mut ptr, &mut size, res) } != 0 {
unsafe {
(t.graphics_unmap_resources)(1, &mut res, std::ptr::null_mut());
}
return None;
}
Some((ptr as usize, size))
}
pub fn gl_unmap_resource(resource: usize) {
if let Some(t) = table() {
let mut r = resource as GraphicsResource;
unsafe {
(t.graphics_unmap_resources)(1, &mut r, std::ptr::null_mut());
}
}
}
pub fn gl_unregister_resource(resource: usize) {
if let Some(t) = table() {
unsafe {
(t.graphics_unregister_resource)(resource as GraphicsResource);
}
}
}
#[allow(dead_code)] pub(crate) const CUDA_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD: c_uint = 1;
#[allow(dead_code)] #[repr(C)]
pub(crate) struct CudaExternalMemoryHandleDesc {
pub type_: c_int,
pub _pad0: u32,
pub handle_fd: c_int,
pub _union_rest: [u32; 3],
pub size: u64,
pub flags: c_uint,
pub _tail: u32,
}
#[allow(dead_code)] #[repr(C)]
pub(crate) struct CudaExternalMemoryBufferDesc {
pub offset: u64,
pub size: u64,
pub flags: c_uint,
pub _tail: u32,
}
#[cfg(target_os = "linux")]
pub(crate) fn import_dma_fd(fd: i32, size: usize) -> Option<(ExternalMemory, *mut c_void)> {
use std::os::fd::{BorrowedFd, FromRawFd, IntoRawFd, OwnedFd};
let t = table()?;
let dup_fd = unsafe { BorrowedFd::borrow_raw(fd) }
.try_clone_to_owned()
.ok()?
.into_raw_fd();
let mut desc: CudaExternalMemoryHandleDesc = unsafe { std::mem::zeroed() };
desc.type_ = CUDA_EXTERNAL_MEMORY_HANDLE_TYPE_OPAQUE_FD as c_int;
desc.handle_fd = dup_fd;
desc.size = size as u64;
let mut ext: ExternalMemory = std::ptr::null_mut();
if unsafe { (t.import_external_memory)(&mut ext, &desc as *const _ as *const c_void) } != 0 {
drop(unsafe { OwnedFd::from_raw_fd(dup_fd) });
return None;
}
let bdesc = CudaExternalMemoryBufferDesc {
offset: 0,
size: size as u64,
flags: 0,
_tail: 0,
};
let mut dptr: *mut c_void = std::ptr::null_mut();
if unsafe {
(t.external_memory_get_mapped_buffer)(&mut dptr, ext, &bdesc as *const _ as *const c_void)
} != 0
{
unsafe { (t.destroy_external_memory)(ext) };
return None;
}
Some((ext, dptr))
}
pub trait CudaGlOps: Send + Sync {
fn map(&self, resource: GraphicsResource) -> Option<(*mut c_void, usize)>;
fn unmap(&self, resource: GraphicsResource);
fn unregister(&self, resource: GraphicsResource);
}
enum CudaBacking {
#[allow(dead_code)] GlBuffer {
resource: GraphicsResource,
ops: Arc<dyn CudaGlOps>,
},
#[allow(dead_code)] ExternalMem {
ext_mem: ExternalMemory,
dptr: *mut c_void,
},
}
unsafe impl Send for CudaBacking {}
unsafe impl Sync for CudaBacking {}
pub struct CudaHandle {
kind: CudaBacking,
size: usize,
}
impl std::fmt::Debug for CudaHandle {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
let kind = match &self.kind {
CudaBacking::GlBuffer { .. } => "GlBuffer",
CudaBacking::ExternalMem { .. } => "ExternalMem",
};
f.debug_struct("CudaHandle")
.field("kind", &kind)
.field("size", &self.size)
.finish()
}
}
impl CudaHandle {
pub fn new_gl(resource: GraphicsResource, size: usize, ops: Arc<dyn CudaGlOps>) -> Self {
Self {
kind: CudaBacking::GlBuffer { resource, ops },
size,
}
}
#[allow(dead_code)] pub(crate) fn new_external(ext_mem: ExternalMemory, dptr: *mut c_void, size: usize) -> Self {
Self {
kind: CudaBacking::ExternalMem { ext_mem, dptr },
size,
}
}
pub fn map(&self) -> Option<CudaMap<'_>> {
match &self.kind {
CudaBacking::GlBuffer { resource, ops } => {
let (ptr, len) = ops.map(*resource)?;
Some(CudaMap {
ptr,
len,
unmap: Some((ops.clone(), *resource)),
_marker: std::marker::PhantomData,
})
}
CudaBacking::ExternalMem { dptr, .. } => Some(CudaMap {
ptr: *dptr,
len: self.size,
unmap: None,
_marker: std::marker::PhantomData,
}),
}
}
}
impl Drop for CudaHandle {
fn drop(&mut self) {
match &self.kind {
CudaBacking::GlBuffer { resource, ops } => ops.unregister(*resource),
CudaBacking::ExternalMem { ext_mem, dptr: _ } => {
if let Some(t) = table() {
unsafe {
(t.destroy_external_memory)(*ext_mem);
}
}
}
}
}
}
pub struct CudaMap<'a> {
ptr: *mut c_void,
len: usize,
unmap: Option<(Arc<dyn CudaGlOps>, GraphicsResource)>,
_marker: std::marker::PhantomData<&'a ()>,
}
unsafe impl Send for CudaMap<'_> {}
unsafe impl Sync for CudaMap<'_> {}
impl CudaMap<'_> {
pub fn device_ptr(&self) -> *mut c_void {
self.ptr
}
pub fn len(&self) -> usize {
self.len
}
pub fn is_empty(&self) -> bool {
self.len == 0
}
}
impl Drop for CudaMap<'_> {
fn drop(&mut self) {
if let Some((ops, r)) = self.unmap.take() {
ops.unmap(r);
}
}
}
#[cfg(test)]
mod ext_mem_layout {
use super::*;
#[test]
fn external_memory_desc_abi() {
assert_eq!(std::mem::size_of::<CudaExternalMemoryHandleDesc>(), 40);
assert_eq!(std::mem::align_of::<CudaExternalMemoryHandleDesc>(), 8);
let d: CudaExternalMemoryHandleDesc = unsafe { std::mem::zeroed() };
let base = &d as *const _ as usize;
assert_eq!((&d.size as *const _ as usize) - base, 24);
assert_eq!((&d.flags as *const _ as usize) - base, 32);
assert_eq!(std::mem::size_of::<CudaExternalMemoryBufferDesc>(), 24);
let b: CudaExternalMemoryBufferDesc = unsafe { std::mem::zeroed() };
let bb = &b as *const _ as usize;
assert_eq!((&b.size as *const _ as usize) - bb, 8);
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn cuda_table_loads_when_libcudart_present() {
let avail = is_cuda_available();
if avail {
assert!(table().is_some(), "table present when available");
}
}
#[test]
fn pub_primitives_degrade_without_libcudart() {
if is_cuda_available() {
return;
}
assert!(gl_register_buffer(0).is_none());
assert!(gl_map_resource(0).is_none());
gl_unmap_resource(0); gl_unregister_resource(0); assert!(!unsafe { memcpy_device_to_host(std::ptr::null_mut(), std::ptr::null(), 0) });
}
}
#[cfg(test)]
mod handle_tests {
use super::*;
use std::sync::{
atomic::{AtomicUsize, Ordering},
Arc,
};
struct MockOps {
unmaps: Arc<AtomicUsize>,
unregisters: Arc<AtomicUsize>,
}
impl CudaGlOps for MockOps {
fn map(&self, _r: GraphicsResource) -> Option<(*mut std::ffi::c_void, usize)> {
Some((0x1000usize as *mut _, 4096))
}
fn unmap(&self, _r: GraphicsResource) {
self.unmaps.fetch_add(1, Ordering::SeqCst);
}
fn unregister(&self, _r: GraphicsResource) {
self.unregisters.fetch_add(1, Ordering::SeqCst);
}
}
#[test]
fn cudamap_guard_unmaps_on_drop_for_glbuffer() {
let unmaps = Arc::new(AtomicUsize::new(0));
let unregisters = Arc::new(AtomicUsize::new(0));
{
let h = CudaHandle::new_gl(
0x1usize as GraphicsResource,
4096,
Arc::new(MockOps {
unmaps: unmaps.clone(),
unregisters: unregisters.clone(),
}),
);
{
let m = h.map().expect("map");
assert_eq!(m.device_ptr() as usize, 0x1000);
assert_eq!(m.len(), 4096);
assert!(!m.is_empty());
}
assert_eq!(
unmaps.load(Ordering::SeqCst),
1,
"Drop must unmap a GlBuffer"
);
assert_eq!(unregisters.load(Ordering::SeqCst), 0);
}
assert_eq!(
unregisters.load(Ordering::SeqCst),
1,
"Dropping a GlBuffer handle must unregister"
);
}
struct NoneOps;
impl CudaGlOps for NoneOps {
fn map(&self, _r: GraphicsResource) -> Option<(*mut std::ffi::c_void, usize)> {
None
}
fn unmap(&self, _r: GraphicsResource) {}
fn unregister(&self, _r: GraphicsResource) {}
}
#[test]
fn glbuffer_map_returns_none_when_ops_map_fails() {
let h = CudaHandle::new_gl(0x9usize as GraphicsResource, 4096, Arc::new(NoneOps));
assert!(
h.map().is_none(),
"GlBuffer map propagates ops.map() failure"
);
}
#[test]
fn glbuffer_handle_debug_and_empty_map() {
let unmaps = Arc::new(AtomicUsize::new(0));
let unregisters = Arc::new(AtomicUsize::new(0));
let h = CudaHandle::new_gl(
0x2usize as GraphicsResource,
0,
Arc::new(MockOps {
unmaps: unmaps.clone(),
unregisters: unregisters.clone(),
}),
);
let dbg = format!("{h:?}");
assert!(
dbg.contains("GlBuffer"),
"debug names the backing kind: {dbg}"
);
assert!(dbg.contains("size"), "debug includes size: {dbg}");
}
#[test]
fn external_mem_map_is_persistent_and_debug_names_kind() {
let dptr = 0xCAFE_0000usize as *mut std::ffi::c_void;
let h = CudaHandle::new_external(std::ptr::null_mut(), dptr, 8192);
let dbg = format!("{h:?}");
assert!(dbg.contains("ExternalMem"), "debug names the kind: {dbg}");
{
let m = h.map().expect("ExternalMem map is always Some");
assert_eq!(m.device_ptr(), dptr, "persistent device ptr passthrough");
assert_eq!(m.len(), 8192);
assert!(!m.is_empty());
}
std::mem::forget(h);
}
#[test]
fn external_mem_zero_len_map_is_empty() {
let h = CudaHandle::new_external(std::ptr::null_mut(), std::ptr::null_mut(), 0);
{
let m = h.map().expect("map");
assert_eq!(m.len(), 0);
assert!(m.is_empty(), "zero-length mapping is empty");
assert!(m.device_ptr().is_null());
}
std::mem::forget(h); }
}