Load nccl dynamically

Pull Request resolved: #2088

nccl can be a big library, but we do not want to force our build to try to load it on a machine that maybe doesn't have cuda. So we do the same as we do with the libcuda api and dynamically load it.
ghstack-source-id: 328272743
@exported-using-ghexport

Differential Revision: [D88672908](https://our.internmc.facebook.com/intern/diff/D88672908/)

Author: zdevito

.gitignore

@@ -33,3 +33,4 @@ docs/_build/**
 docs/build/**
 docs/**/generated/**
 */sg_execution_times.rst
+nccl/**

nccl-sys/Cargo.toml

@@ -13,3 +13,4 @@ serde = { version = "1.0.185", features = ["derive", "rc"] }
 [build-dependencies]
 bindgen = "0.70.1"
 build_utils = { path = "../build_utils" }
+cc = "1.0"

nccl-sys/build.rs

@@ -13,21 +13,32 @@ fn main() {}
 
 #[cfg(not(target_os = "macos"))]
 fn main() {
+    // Compile the bridge.cpp file
+    let mut cc_builder = cc::Build::new();
+    cc_builder
+        .cpp(true)
+        .file("src/bridge.cpp")
+        .flag("-std=c++14");
+
+    // Include CUDA headers
+    if let Some(cuda_home) = build_utils::find_cuda_home() {
+        cc_builder.include(format!("{}/include", cuda_home));
+    }
+
+    cc_builder.compile("nccl_bridge");
+
     let mut builder = bindgen::Builder::default()
-        .header("src/nccl.h")
+        .header("src/bridge.h")
         .clang_arg("-x")
         .clang_arg("c++")
         .clang_arg("-std=c++14")
-        .clang_arg(format!(
-            "-I{}/include",
-            build_utils::find_cuda_home().unwrap()
-        ))
         .parse_callbacks(Box::new(bindgen::CargoCallbacks::new()))
-        // Communicator creation and management
-        .allowlist_function("ncclGetLastError")
-        .allowlist_function("ncclGetErrorString")
+        // Version and error handling
         .allowlist_function("ncclGetVersion")
         .allowlist_function("ncclGetUniqueId")
+        .allowlist_function("ncclGetErrorString")
+        .allowlist_function("ncclGetLastError")
+        // Communicator creation and management
         .allowlist_function("ncclCommInitRank")
         .allowlist_function("ncclCommInitAll")
         .allowlist_function("ncclCommInitRankConfig")
@@ -60,15 +71,20 @@ fn main() {
         // User-defined reduction operators
         .allowlist_function("ncclRedOpCreatePreMulSum")
         .allowlist_function("ncclRedOpDestroy")
-        // Random nccl stuff we want
-        .allowlist_function("cudaStream.*")
+        // CUDA runtime functions
         .allowlist_function("cudaSetDevice")
+        .allowlist_function("cudaStreamSynchronize")
+        // Types
         .allowlist_type("ncclComm_t")
         .allowlist_type("ncclResult_t")
         .allowlist_type("ncclDataType_t")
         .allowlist_type("ncclRedOp_t")
         .allowlist_type("ncclScalarResidence_t")
         .allowlist_type("ncclSimInfo_t")
+        .allowlist_type("ncclConfig_t")
+        .allowlist_type("cudaError_t")
+        .allowlist_type("cudaStream_t")
+        // Constants
         .allowlist_var("NCCL_SPLIT_NOCOLOR")
         .allowlist_var("NCCL_MAJOR")
         .allowlist_var("NCCL_MINOR")
@@ -79,6 +95,11 @@ fn main() {
             is_global: false,
         });
 
+    // Include CUDA headers
+    if let Some(cuda_home) = build_utils::find_cuda_home() {
+        builder = builder.clang_arg(format!("-I{}/include", cuda_home));
+    }
+
     // Include headers and libs from the active environment
     let python_config = match build_utils::python_env_dirs() {
         Ok(config) => config,
@@ -103,13 +124,17 @@ fn main() {
 
     // Write the bindings to the $OUT_DIR/bindings.rs file.
     let out_path = PathBuf::from(std::env::var("OUT_DIR").unwrap());
+
+    // Generate bindings (NCCL + CUDA runtime combined)
     builder
         .generate()
         .expect("Unable to generate bindings")
         .write_to_file(out_path.join("bindings.rs"))
         .expect("Couldn't write bindings!");
 
-    println!("cargo::rustc-link-lib=nccl");
+    // We no longer link against nccl directly since we dlopen it
+    // But we do link against CUDA runtime
+    println!("cargo::rustc-link-lib=cudart");
     println!("cargo::rustc-cfg=cargo");
     println!("cargo::rustc-check-cfg=cfg(cargo)");
 }