From dda015f2aa462656487f81e00a7721c2f71cc22f Mon Sep 17 00:00:00 2001
From: Morgan Funtowicz <morgan@huggingface.co>
Date: Thu, 25 Jul 2024 10:29:51 +0000
Subject: [PATCH] add some custom stuff for nccl linkage

---
 backends/trtllm/build.rs | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/backends/trtllm/build.rs b/backends/trtllm/build.rs
index db5ca5f8..c8890394 100644
--- a/backends/trtllm/build.rs
+++ b/backends/trtllm/build.rs
@@ -10,10 +10,11 @@ const CUDA_REQUIRED_VERSION: &str = "12.5";
 const MPI_REQUIRED_VERSION: &str = "4.1";
 const INSTALL_PREFIX: Option<&str> = option_env!("CMAKE_INSTALL_PREFIX");
 const TENSORRT_ROOT_DIR: Option<&str> = option_env!("TENSORRT_ROOT_DIR");
+const NCCL_ROOT_DIR: Option<&str> = option_env!("NCCL_ROOT_DIR");
 
 // Dependencies
 const BACKEND_DEPS: [&str; 2] = ["tgi_trtllm_backend_impl", "tgi_trtllm_backend"];
-const CUDA_TRANSITIVE_DEPS: [&str; 5] = ["cuda", "cudart", "cublas", "nvidia-ml", "nccl"];
+const CUDA_TRANSITIVE_DEPS: [&str; 4] = ["cuda", "cudart", "cublas", "nvidia-ml"];
 const TENSORRT_LLM_TRANSITIVE_DEPS: [(&str, &str); 5] = [
     ("dylib", "tensorrt-llm"),
     ("static", "tensorrt_llm_executor_static"),
@@ -124,6 +125,14 @@ fn main() {
         probe!(name, CUDA_REQUIRED_VERSION);
     });
 
+    // NCCL is slightly trickier because it might not have a pkgconfig installed
+    let nccl_library_path = NCCL_ROOT_DIR.unwrap_or(&format!(
+        "/usr/local/{}-linux-gnu",
+        env!("CARGO_CFG_TARGET_ARCH")
+    ));
+    println!(r"cargo:rustc-link-search=native={}", nccl_library_path);
+    println!("cargo:rustc-link-lib=dylib=nccl");
+
     // TensorRT
     let tensort_library_path = TENSORRT_ROOT_DIR.unwrap_or("/usr/local/tensorrt/lib");
     println!(r"cargo:rustc-link-search=native={}", tensort_library_path);