Add CUDA version flags for supporting TensorRT bindings that can work

with multiple versions of CUDA
mstallmo · Sep 14, 2020 · 42a38f7 · 42a38f7
1 parent 79d73ab
commit 42a38f7
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 5 deletions.
diff --git a/tensorrt-sys/Cargo.toml b/tensorrt-sys/Cargo.toml
@@ -11,10 +11,16 @@ description = "Low level wrapper around Nvidia's TensorRT library"
 [features]
 default = ["trt-515"]
 
-trt-515 = []
+trt-515 = ["cuda-101"]
 
 trt-713 = []
 
+cuda-101 = []
+
+cuda-102 = []
+
+cuda-110 = []
+
 [dependencies]
 libc = "0.2.62"
 

diff --git a/tensorrt-sys/build.rs b/tensorrt-sys/build.rs
@@ -43,8 +43,7 @@ fn configuration(full_library_path: &PathBuf) {
         println!("cargo:rustc-link-search=native={}", dst.display());
         println!("cargo:rustc-flags=-l dylib=nvinfer");
         println!("cargo:rustc-flags=-l dylib=nvparsers");
-        println!("cargo:rustc-flags=-L /usr/local/cuda-10.1/lib64");
-        println!("cargo:rustc-flags=-l dylib=cudart");
+        cuda_configuration();
     } else {
         panic!("Invalid nvinfer version found. Expected: libnvinfer.so.5.1.5, Found: {}", full_library_path.to_str().unwrap());
     }
@@ -58,13 +57,30 @@ fn configuration(full_library_path: &PathBuf) {
         println!("cargo:rustc-link-search=native={}", dst.display());
         println!("cargo:rustc-flags=-l dylib=nvinfer");
         println!("cargo:rustc-flags=-l dylib=nvparsers");
-        println!("cargo:rustc-flags=-L /usr/local/cuda-10.2/lib64");
-        println!("cargo:rustc-flags=-l dylib=cudart");
+        cuda_configuration();
     } else {
         panic!("Invalid nvinfer version found. Expected: libnvinfer.so.7.1.3, Found: {}", full_library_path.to_str().unwrap());
     }
 }
 
+#[cfg(feature = "cuda-101")]
+fn cuda_configuration() {
+    println!("cargo:rustc-flags=-L /usr/local/cuda-10.1/lib64");
+    println!("cargo:rustc-flags=-l dylib=cudart");
+}
+
+#[cfg(feature = "cuda-102")]
+fn cuda_configuration() {
+    println!("cargo:rustc-flags=-L /usr/local/cuda-10.2/lib64");
+    println!("cargo:rustc-flags=-l dylib=cudart");
+}
+
+#[cfg(feature = "cuda-110")]
+fn cuda_configuration() {
+    println!("cargo:rustc-flags=-L /usr/local/cuda-11.0/lib64");
+    println!("cargo:rustc-flags=-l dylib=cudart");
+}
+
 // Not sure if I love this solution but I think it's relatively robust enough for now on Unix systems.
 // Still have to thoroughly test what happens with a TRT library installed that's not done by the
 // dpkg. It's possible that we'll just have to fall back to only supporting one system library and assuming that