Add bindings for cudnn and nccl modules

coreylowman · Jun 27, 2024 · 2ccb3fe · 2ccb3fe
1 parent 35fa946
commit 2ccb3fe
Show file tree

Hide file tree

Showing 7 changed files with 9,695 additions and 4 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -25,7 +25,6 @@ features = ["cuda-12050", "f16", "cudnn"]
 
 [features]
 default = ["std", "cublas", "cublaslt", "cudnn", "curand", "driver", "nccl", "nvrtc"]
-no-cudnn = ["std", "cublas", "cublaslt", "curand", "driver", "nvrtc"]
 
 cuda-version-from-build-system = []
 cuda-11040 = []

diff --git a/src/cudnn/result.rs b/src/cudnn/result.rs
@@ -46,7 +46,7 @@ pub fn get_cudart_version() -> usize {
 
 /// Runs all *VersionCheck functions.
 pub fn version_check() -> Result<(), CudnnError> {
-    #[cfg(not(any(feature = "cuda-12030", feature = "cuda-12040", feature = "cuda-12050")))]
+    #[cfg(not(any(feature = "cuda-11040", feature = "cuda-12030", feature = "cuda-12040", feature = "cuda-12050")))]
     unsafe {
         lib().cudnnAdvInferVersionCheck().result()?;
         lib().cudnnAdvTrainVersionCheck().result()?;
@@ -55,6 +55,12 @@ pub fn version_check() -> Result<(), CudnnError> {
         lib().cudnnOpsInferVersionCheck().result()?;
         lib().cudnnOpsTrainVersionCheck().result()?;
     }
+    #[cfg(feature = "cuda-11040")]
+    unsafe {
+        lib().cudnnAdvVersionCheck().result()?;
+        lib().cudnnCnnVersionCheck().result()?;
+        lib().cudnnOpsVersionCheck().result()?;
+    }
     #[cfg(feature = "cuda-12030")]
     unsafe {
         lib().cudnnAdvVersionCheck().result()?;

diff --git a/src/cudnn/sys/mod.rs b/src/cudnn/sys/mod.rs
@@ -1,3 +1,8 @@
+#[cfg(feature = "cuda-11040")]
+mod sys_11040;
+#[cfg(feature = "cuda-11040")]
+pub use sys_11040::*;
+
 #[cfg(feature = "cuda-11050")]
 mod sys_11050;
 #[cfg(feature = "cuda-11050")]