build(cudnn-sys): Add CUDNN_INCLUDE_DIR (#213)

* build(cudnn-sys): Add CUDNN_INCLUDE_DIR

Enables users to specify a non-standard cuDNN install path. This seems
to be needed for the newer editions of the CUDA toolkit, as cuDNN isn't
included by default (at least in the Fedora repo's, you have to install
from a tarball)

Author: giantcow

crates/cudnn-sys/build/cudnn_sdk.rs

@@ -1,6 +1,8 @@
+use std::env;
 use std::error;
 use std::fs;
 use std::path;
+use std::path::Path;
 
 /// Represents the cuDNN SDK installation.
 #[derive(Debug, Clone)]
@@ -57,17 +59,25 @@ impl CudnnSdk {
     }
 
     fn find_cudnn_include_dir() -> Result<path::PathBuf, Box<dyn error::Error>> {
+        let cudnn_include_dir = env::var_os("CUDNN_INCLUDE_DIR");
+
         #[cfg(not(target_os = "windows"))]
         const CUDNN_DEFAULT_PATHS: &[&str] = &["/usr/include", "/usr/local/include"];
         #[cfg(target_os = "windows")]
         const CUDNN_DEFAULT_PATHS: &[&str] = &[
             "C:/Program Files/NVIDIA/CUDNN/v9.x/include",
             "C:/Program Files/NVIDIA/CUDNN/v8.x/include",
         ];
-        CUDNN_DEFAULT_PATHS
+
+        let mut cudnn_paths: Vec<&Path> = CUDNN_DEFAULT_PATHS.iter().map(Path::new).collect();
+        if let Some(override_path) = &cudnn_include_dir {
+            cudnn_paths.push(Path::new(override_path));
+        }
+
+        cudnn_paths
             .iter()
-            .find(|s| Self::is_cudnn_include_path(s))
-            .map(path::PathBuf::from)
+            .find(|p| Self::is_cudnn_include_path(p))
+            .map(|p| p.to_path_buf())
             .ok_or("Cannot find cuDNN include directory.".into())
     }