Remove direct dependency on libcuda

Pull Request resolved: #1926

Dynamically load  any cuda driver functions so that monarch library can be loaded a machine that doesn't have a gpu even if the library is built with one.

Differential Revision: [D87380631](https://our.internmc.facebook.com/intern/diff/D87380631/)
ghstack-source-id: 324224674

Author: zdevito

cuda-sys/build.rs

@@ -31,11 +31,11 @@ fn main() {
         .clang_arg("c++")
         .clang_arg("-std=gnu++20")
         .parse_callbacks(Box::new(bindgen::CargoCallbacks::new()))
-        // Allow the specified functions and types
-        .allowlist_function("cu.*")
-        .allowlist_function("CU.*")
-        .allowlist_type("cu.*")
-        .allowlist_type("CU.*")
+        // Allow the specified functions and types (CUDA Runtime API only)
+        .allowlist_function("cuda.*")
+        .allowlist_function("CUDA.*")
+        .allowlist_type("cuda.*")
+        .allowlist_type("CUDA.*")
         // Use newtype enum style
         .default_enum_style(bindgen::EnumVariation::NewType {
             is_bitfield: false,
@@ -78,7 +78,6 @@ fn main() {
         }
     };
     println!("cargo:rustc-link-search=native={}", cuda_lib_dir);
-    println!("cargo:rustc-link-lib=cuda");
     println!("cargo:rustc-link-lib=cudart");
 
     // Generate bindings - fail fast if this doesn't work

cuda-sys/src/lib.rs

@@ -34,20 +34,3 @@ mod inner {
 }
 
 pub use inner::*;
-
-#[cfg(test)]
-mod tests {
-    use std::mem::MaybeUninit;
-
-    use super::*;
-
-    #[test]
-    fn sanity() {
-        // SAFETY: testing bindings
-        unsafe {
-            let mut version = MaybeUninit::<i32>::uninit();
-            let result = cuDriverGetVersion(version.as_mut_ptr());
-            assert_eq!(result, cudaError_enum(0));
-        }
-    }
-}

cuda-sys/src/wrapper.h

@@ -8,5 +8,4 @@
 
 #pragma once
 
-#include <cuda.h>
 #include <cuda_runtime.h>

monarch_rdma/src/macros.rs

@@ -9,9 +9,9 @@
 #[macro_export]
 macro_rules! cu_check {
     ($result:expr) => {
-        if $result != cuda_sys::CUresult::CUDA_SUCCESS {
+        if $result != rdmaxcel_sys::CUDA_SUCCESS {
             let mut error_string: *const std::os::raw::c_char = std::ptr::null();
-            cuda_sys::cuGetErrorString($result, &mut error_string);
+            rdmaxcel_sys::rdmaxcel_cuGetErrorString($result, &mut error_string);
             panic!(
                 "cuda failure {}:{} {:?} '{}'",
                 file!(),

monarch_rdma/src/rdma_manager_actor.rs

@@ -366,13 +366,13 @@ impl RdmaManagerActor {
     ) -> Result<(RdmaMemoryRegionView, String), anyhow::Error> {
         unsafe {
             let mut mem_type: i32 = 0;
-            let ptr = addr as cuda_sys::CUdeviceptr;
-            let err = cuda_sys::cuPointerGetAttribute(
+            let ptr = addr as rdmaxcel_sys::CUdeviceptr;
+            let err = rdmaxcel_sys::rdmaxcel_cuPointerGetAttribute(
                 &mut mem_type as *mut _ as *mut std::ffi::c_void,
-                cuda_sys::CUpointer_attribute_enum::CU_POINTER_ATTRIBUTE_MEMORY_TYPE,
+                rdmaxcel_sys::CU_POINTER_ATTRIBUTE_MEMORY_TYPE,
                 ptr,
             );
-            let is_cuda = err == cuda_sys::CUresult::CUDA_SUCCESS;
+            let is_cuda = err == rdmaxcel_sys::CUDA_SUCCESS;
 
             let mut selected_rdma_device = None;
 
@@ -457,11 +457,11 @@ impl RdmaManagerActor {
                 mrv = maybe_mrv.unwrap();
             } else if is_cuda {
                 let mut fd: i32 = -1;
-                cuda_sys::cuMemGetHandleForAddressRange(
-                    &mut fd as *mut i32 as *mut std::ffi::c_void,
-                    addr as cuda_sys::CUdeviceptr,
+                rdmaxcel_sys::rdmaxcel_cuMemGetHandleForAddressRange(
+                    &mut fd,
+                    addr as rdmaxcel_sys::CUdeviceptr,
                     size,
-                    cuda_sys::CUmemRangeHandleType::CU_MEM_RANGE_HANDLE_TYPE_DMA_BUF_FD,
+                    rdmaxcel_sys::CU_MEM_RANGE_HANDLE_TYPE_DMA_BUF_FD,
                     0,
                 );
                 mr = rdmaxcel_sys::ibv_reg_dmabuf_mr(domain_pd, 0, size, 0, fd, access.0 as i32);

monarch_rdma/src/test_utils.rs

@@ -46,25 +46,25 @@ pub fn is_cuda_available() -> bool {
 fn check_cuda_available() -> bool {
     unsafe {
         // Try to initialize CUDA
-        let result = cuda_sys::cuInit(0);
+        let result = rdmaxcel_sys::rdmaxcel_cuInit(0);
 
-        if result != cuda_sys::CUresult::CUDA_SUCCESS {
+        if result != rdmaxcel_sys::CUDA_SUCCESS {
             return false;
         }
 
         // Check if there are any CUDA devices
         let mut device_count: i32 = 0;
-        let count_result = cuda_sys::cuDeviceGetCount(&mut device_count);
+        let count_result = rdmaxcel_sys::rdmaxcel_cuDeviceGetCount(&mut device_count);
 
-        if count_result != cuda_sys::CUresult::CUDA_SUCCESS || device_count <= 0 {
+        if count_result != rdmaxcel_sys::CUDA_SUCCESS || device_count <= 0 {
             return false;
         }
 
         // Try to get the first device to verify it's actually accessible
-        let mut device: cuda_sys::CUdevice = std::mem::zeroed();
-        let device_result = cuda_sys::cuDeviceGet(&mut device, 0);
+        let mut device: rdmaxcel_sys::CUdevice = std::mem::zeroed();
+        let device_result = rdmaxcel_sys::rdmaxcel_cuDeviceGet(&mut device, 0);
 
-        if device_result != cuda_sys::CUresult::CUDA_SUCCESS {
+        if device_result != rdmaxcel_sys::CUDA_SUCCESS {
             return false;
         }
 
@@ -270,8 +270,8 @@ pub mod test_utils {
         pub actor_2: ActorRef<RdmaManagerActor>,
         pub rdma_handle_1: RdmaBuffer,
         pub rdma_handle_2: RdmaBuffer,
-        cuda_context_1: Option<cuda_sys::CUcontext>,
-        cuda_context_2: Option<cuda_sys::CUcontext>,
+        cuda_context_1: Option<rdmaxcel_sys::CUcontext>,
+        cuda_context_2: Option<rdmaxcel_sys::CUcontext>,
     }
 
     #[derive(Debug, Clone)]
@@ -375,46 +375,53 @@ pub mod test_utils {
                 }
                 // CUDA case
                 unsafe {
-                    cu_check!(cuda_sys::cuInit(0));
+                    cu_check!(rdmaxcel_sys::rdmaxcel_cuInit(0));
 
-                    let mut dptr: cuda_sys::CUdeviceptr = std::mem::zeroed();
-                    let mut handle: cuda_sys::CUmemGenericAllocationHandle = std::mem::zeroed();
+                    let mut dptr: rdmaxcel_sys::CUdeviceptr = std::mem::zeroed();
+                    let mut handle: rdmaxcel_sys::CUmemGenericAllocationHandle = std::mem::zeroed();
 
-                    let mut device: cuda_sys::CUdevice = std::mem::zeroed();
-                    cu_check!(cuda_sys::cuDeviceGet(&mut device, accel.1 as i32));
+                    let mut device: rdmaxcel_sys::CUdevice = std::mem::zeroed();
+                    cu_check!(rdmaxcel_sys::rdmaxcel_cuDeviceGet(
+                        &mut device,
+                        accel.1 as i32
+                    ));
 
-                    let mut context: cuda_sys::CUcontext = std::mem::zeroed();
-                    cu_check!(cuda_sys::cuCtxCreate_v2(&mut context, 0, accel.1 as i32));
-                    cu_check!(cuda_sys::cuCtxSetCurrent(context));
+                    let mut context: rdmaxcel_sys::CUcontext = std::mem::zeroed();
+                    cu_check!(rdmaxcel_sys::rdmaxcel_cuCtxCreate_v2(
+                        &mut context,
+                        0,
+                        accel.1 as i32
+                    ));
+                    cu_check!(rdmaxcel_sys::rdmaxcel_cuCtxSetCurrent(context));
 
                     let mut granularity: usize = 0;
-                    let mut prop: cuda_sys::CUmemAllocationProp = std::mem::zeroed();
-                    prop.type_ = cuda_sys::CUmemAllocationType::CU_MEM_ALLOCATION_TYPE_PINNED;
-                    prop.location.type_ = cuda_sys::CUmemLocationType::CU_MEM_LOCATION_TYPE_DEVICE;
+                    let mut prop: rdmaxcel_sys::CUmemAllocationProp = std::mem::zeroed();
+                    prop.type_ = rdmaxcel_sys::CU_MEM_ALLOCATION_TYPE_PINNED;
+                    prop.location.type_ = rdmaxcel_sys::CU_MEM_LOCATION_TYPE_DEVICE;
                     prop.location.id = device;
                     prop.allocFlags.gpuDirectRDMACapable = 1;
                     prop.requestedHandleTypes =
-                        cuda_sys::CUmemAllocationHandleType::CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;
+                        rdmaxcel_sys::CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;
 
-                    cu_check!(cuda_sys::cuMemGetAllocationGranularity(
+                    cu_check!(rdmaxcel_sys::rdmaxcel_cuMemGetAllocationGranularity(
                         &mut granularity as *mut usize,
                         &prop,
-                        cuda_sys::CUmemAllocationGranularity_flags::CU_MEM_ALLOC_GRANULARITY_MINIMUM,
+                        rdmaxcel_sys::CU_MEM_ALLOC_GRANULARITY_MINIMUM,
                     ));
 
                     // ensure our size is aligned
                     let /*mut*/ padded_size: usize = ((buffer_size - 1) / granularity + 1) * granularity;
                     assert!(padded_size == buffer_size);
 
-                    cu_check!(cuda_sys::cuMemCreate(
-                        &mut handle as *mut cuda_sys::CUmemGenericAllocationHandle,
+                    cu_check!(rdmaxcel_sys::rdmaxcel_cuMemCreate(
+                        &mut handle as *mut rdmaxcel_sys::CUmemGenericAllocationHandle,
                         padded_size,
                         &prop,
                         0
                     ));
                     // reserve and map the memory
-                    cu_check!(cuda_sys::cuMemAddressReserve(
-                        &mut dptr as *mut cuda_sys::CUdeviceptr,
+                    cu_check!(rdmaxcel_sys::rdmaxcel_cuMemAddressReserve(
+                        &mut dptr as *mut rdmaxcel_sys::CUdeviceptr,
                         padded_size,
                         0,
                         0,
@@ -424,25 +431,28 @@ pub mod test_utils {
                     assert!(padded_size.is_multiple_of(granularity));
 
                     // fails if a add cu_check macro; but passes if we don't
-                    let err = cuda_sys::cuMemMap(
-                        dptr as cuda_sys::CUdeviceptr,
+                    let err = rdmaxcel_sys::rdmaxcel_cuMemMap(
+                        dptr as rdmaxcel_sys::CUdeviceptr,
                         padded_size,
                         0,
-                        handle as cuda_sys::CUmemGenericAllocationHandle,
+                        handle as rdmaxcel_sys::CUmemGenericAllocationHandle,
                         0,
                     );
-                    if err != cuda_sys::CUresult::CUDA_SUCCESS {
+                    if err != rdmaxcel_sys::CUDA_SUCCESS {
                         panic!("failed reserving and mapping memory {:?}", err);
                     }
 
                     // set access
-                    let mut access_desc: cuda_sys::CUmemAccessDesc = std::mem::zeroed();
-                    access_desc.location.type_ =
-                        cuda_sys::CUmemLocationType::CU_MEM_LOCATION_TYPE_DEVICE;
+                    let mut access_desc: rdmaxcel_sys::CUmemAccessDesc = std::mem::zeroed();
+                    access_desc.location.type_ = rdmaxcel_sys::CU_MEM_LOCATION_TYPE_DEVICE;
                     access_desc.location.id = device;
-                    access_desc.flags =
-                        cuda_sys::CUmemAccess_flags::CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
-                    cu_check!(cuda_sys::cuMemSetAccess(dptr, padded_size, &access_desc, 1));
+                    access_desc.flags = rdmaxcel_sys::CU_MEM_ACCESS_FLAGS_PROT_READWRITE;
+                    cu_check!(rdmaxcel_sys::rdmaxcel_cuMemSetAccess(
+                        dptr,
+                        padded_size,
+                        &access_desc,
+                        1
+                    ));
                     buf_vec.push(Buffer {
                         ptr: dptr,
                         len: padded_size,
@@ -460,11 +470,11 @@ pub mod test_utils {
                 }
                 unsafe {
                     // Use the CUDA context that was created for the first buffer
-                    cu_check!(cuda_sys::cuCtxSetCurrent(
+                    cu_check!(rdmaxcel_sys::rdmaxcel_cuCtxSetCurrent(
                         cuda_contexts[0].expect("No CUDA context found")
                     ));
 
-                    cu_check!(cuda_sys::cuMemcpyHtoD_v2(
+                    cu_check!(rdmaxcel_sys::rdmaxcel_cuMemcpyHtoD_v2(
                         buf_vec[0].ptr,
                         temp_buffer.as_ptr() as *const std::ffi::c_void,
                         temp_buffer.len()
@@ -514,30 +524,30 @@ pub mod test_utils {
                 .await?;
             if self.cuda_context_1.is_some() {
                 unsafe {
-                    cu_check!(cuda_sys::cuCtxSetCurrent(
+                    cu_check!(rdmaxcel_sys::rdmaxcel_cuCtxSetCurrent(
                         self.cuda_context_1.expect("No CUDA context found")
                     ));
-                    cu_check!(cuda_sys::cuMemUnmap(
-                        self.buffer_1.ptr as cuda_sys::CUdeviceptr,
+                    cu_check!(rdmaxcel_sys::rdmaxcel_cuMemUnmap(
+                        self.buffer_1.ptr as rdmaxcel_sys::CUdeviceptr,
                         self.buffer_1.len
                     ));
-                    cu_check!(cuda_sys::cuMemAddressFree(
-                        self.buffer_1.ptr as cuda_sys::CUdeviceptr,
+                    cu_check!(rdmaxcel_sys::rdmaxcel_cuMemAddressFree(
+                        self.buffer_1.ptr as rdmaxcel_sys::CUdeviceptr,
                         self.buffer_1.len
                     ));
                 }
             }
             if self.cuda_context_2.is_some() {
                 unsafe {
-                    cu_check!(cuda_sys::cuCtxSetCurrent(
+                    cu_check!(rdmaxcel_sys::rdmaxcel_cuCtxSetCurrent(
                         self.cuda_context_2.expect("No CUDA context found")
                     ));
-                    cu_check!(cuda_sys::cuMemUnmap(
-                        self.buffer_2.ptr as cuda_sys::CUdeviceptr,
+                    cu_check!(rdmaxcel_sys::rdmaxcel_cuMemUnmap(
+                        self.buffer_2.ptr as rdmaxcel_sys::CUdeviceptr,
                         self.buffer_2.len
                     ));
-                    cu_check!(cuda_sys::cuMemAddressFree(
-                        self.buffer_2.ptr as cuda_sys::CUdeviceptr,
+                    cu_check!(rdmaxcel_sys::rdmaxcel_cuMemAddressFree(
+                        self.buffer_2.ptr as rdmaxcel_sys::CUdeviceptr,
                         self.buffer_2.len
                     ));
                 }
@@ -579,12 +589,12 @@ pub mod test_utils {
                     let mut temp_buffer = vec![0u8; size].into_boxed_slice();
                     // SAFETY: The buffer is allocated with the correct size and the pointer is valid.
                     unsafe {
-                        cu_check!(cuda_sys::cuCtxSetCurrent(
+                        cu_check!(rdmaxcel_sys::rdmaxcel_cuCtxSetCurrent(
                             cuda_context.expect("No CUDA context found")
                         ));
-                        cu_check!(cuda_sys::cuMemcpyDtoH_v2(
+                        cu_check!(rdmaxcel_sys::rdmaxcel_cuMemcpyDtoH_v2(
                             temp_buffer.as_mut_ptr() as *mut std::ffi::c_void,
-                            virtual_addr as cuda_sys::CUdeviceptr,
+                            virtual_addr as rdmaxcel_sys::CUdeviceptr,
                             size
                         ));
                     }

rdmaxcel-sys/build.rs

@@ -21,10 +21,15 @@ fn main() {
     // Link against the mlx5 library
     println!("cargo:rustc-link-lib=mlx5");
 
+    // Link against dl for dynamic loading
+    println!("cargo:rustc-link-lib=dl");
+
     // Tell cargo to invalidate the built crate whenever the wrapper changes
     println!("cargo:rerun-if-changed=src/rdmaxcel.h");
     println!("cargo:rerun-if-changed=src/rdmaxcel.c");
     println!("cargo:rerun-if-changed=src/rdmaxcel.cpp");
+    println!("cargo:rerun-if-changed=src/driver_api.h");
+    println!("cargo:rerun-if-changed=src/driver_api.cpp");
 
     // Validate CUDA installation and get CUDA home path
     let cuda_home = match build_utils::validate_cuda_installation() {
@@ -88,6 +93,10 @@ fn main() {
         .allowlist_function("pt_cuda_allocator_compatibility")
         .allowlist_function("register_segments")
         .allowlist_function("deregister_segments")
+        .allowlist_function("rdmaxcel_cu.*")
+        .allowlist_function("get_cuda_pci_address_from_ptr")
+        .allowlist_function("rdmaxcel_print_device_info")
+        .allowlist_function("rdmaxcel_error_string")
         .allowlist_type("ibv_.*")
         .allowlist_type("mlx5dv_.*")
         .allowlist_type("mlx5_wqe_.*")
@@ -149,7 +158,8 @@ fn main() {
         }
     };
     println!("cargo:rustc-link-search=native={}", cuda_lib_dir);
-    println!("cargo:rustc-link-lib=cuda");
+    // Note: libcuda is now loaded dynamically via dlopen in driver_api.cpp
+    // Only link cudart (CUDA Runtime API)
     println!("cargo:rustc-link-lib=cudart");
 
     // Link PyTorch C++ libraries for c10 symbols
@@ -213,7 +223,8 @@ fn main() {
 
             // Compile the C++ source file for CUDA allocator compatibility
             let cpp_source_path = format!("{}/src/rdmaxcel.cpp", manifest_dir);
-            if Path::new(&cpp_source_path).exists() {
+            let driver_api_cpp_path = format!("{}/src/driver_api.cpp", manifest_dir);
+            if Path::new(&cpp_source_path).exists() && Path::new(&driver_api_cpp_path).exists() {
                 let mut libtorch_include_dirs: Vec<PathBuf> = vec![];
 
                 // Use the same approach as torch-sys: Python discovery first, env vars as fallback
@@ -249,6 +260,7 @@ fn main() {
                 let mut cpp_build = cc::Build::new();
                 cpp_build
                     .file(&cpp_source_path)
+                    .file(&driver_api_cpp_path)
                     .include(format!("{}/src", manifest_dir))
                     .flag("-fPIC")
                     .cpp(true)
@@ -270,7 +282,15 @@ fn main() {
 
                 cpp_build.compile("rdmaxcel_cpp");
             } else {
-                panic!("C++ source file not found at {}", cpp_source_path);
+                if !Path::new(&cpp_source_path).exists() {
+                    panic!("C++ source file not found at {}", cpp_source_path);
+                }
+                if !Path::new(&driver_api_cpp_path).exists() {
+                    panic!(
+                        "Driver API C++ source file not found at {}",
+                        driver_api_cpp_path
+                    );
+                }
             }
             // Compile the CUDA source file
             let cuda_source_path = format!("{}/src/rdmaxcel.cu", manifest_dir);