Skip to main content

mnemosyne_backend/backends/cuda/
unified.rs

1use core::ffi::c_void;
2use core::sync::atomic::Ordering;
3
4use super::registry::CUDA_ALLOCATIONS;
5use super::{CudaAllocOps, CudaAllocationRegistry, loader, managed_raw_alloc, managed_raw_free};
6
7/// A zero-copy memory backend mapping memory blocks directly using CUDA
8/// managed memory.
9///
10/// `allocate` returns null when the NVIDIA driver is not loaded, when the
11/// driver allocation fails, or when the bounded CUDA allocation registry is
12/// full (the fresh allocation is released first). There is no host fallback;
13/// callers must select another backend on null.
14pub struct CudaUnifiedBackend;
15
16impl CudaAllocOps for CudaUnifiedBackend {
17    #[inline]
18    fn registry() -> &'static CudaAllocationRegistry {
19        &CUDA_ALLOCATIONS
20    }
21
22    #[inline]
23    fn alloc_sym() -> *mut c_void {
24        loader::CU_MEM_ALLOC_MANAGED.load(Ordering::Acquire)
25    }
26
27    #[inline]
28    fn free_sym() -> *mut c_void {
29        loader::CU_MEM_FREE.load(Ordering::Acquire)
30    }
31
32    #[inline]
33    unsafe fn raw_alloc(alloc_sym: *mut c_void, size: usize) -> *mut u8 {
34        // SAFETY: forwarded caller contract (resolved `cuMemAllocManaged`).
35        unsafe { managed_raw_alloc(alloc_sym, size) }
36    }
37
38    #[inline]
39    unsafe fn raw_free(free_sym: *mut c_void, ptr: *mut u8) -> core::ffi::c_int {
40        // SAFETY: forwarded caller contract (resolved `cuMemFree`, live ptr).
41        unsafe { managed_raw_free(free_sym, ptr) }
42    }
43}
44
45impl_cuda_memory_backend!(CudaUnifiedBackend);