Add Default CUDA path list for Linux

.circleci/config.yml

@@ -4,7 +4,7 @@ version: 2
   name: Cargo build
   command: |
     export PATH=/root/.cargo/bin:$PATH
-    cargo build -vv
+    cargo build -v
 
 .job_apt_template: &job_apt
   steps:

Cargo.toml

@@ -13,3 +13,6 @@ keywords      = ["GPGPU", "CUDA", "ffi"]
 license       = "MIT"
 readme        = "README.md"
 categories    = []
+
+[build-dependencies]
+glob = "*"

build.rs

@@ -1,6 +1,9 @@
-use std::env;
+extern crate glob;
 
-fn find_library_paths() -> Vec<String> {
+use glob::glob;
+use std::{env, path::PathBuf};
+
+fn read_env() -> Vec<PathBuf> {
     if let Ok(path) = env::var("CUDA_LIBRARY_PATH") {
         // The location of the libcuda, libcudart, and libcublas can be hardcoded with the
         // CUDA_LIBRARY_PATH environment variable.
@@ -9,81 +12,113 @@ fn find_library_paths() -> Vec<String> {
         } else {
             ":"
         };
+        path.split(split_char).map(|s| PathBuf::from(s)).collect()
+    } else {
+        vec![]
+    }
+}
 
-        return path.split(split_char).map(|s| s.to_owned()).collect();
+fn find_cuda() -> Vec<PathBuf> {
+    let mut candidates = read_env();
+    candidates.push(PathBuf::from("/opt/cuda"));
+    candidates.push(PathBuf::from("/usr/local/cuda"));
+    for e in glob("/usr/local/cuda-*").unwrap() {
+        if let Ok(path) = e {
+            candidates.push(path)
+        }
     }
 
-    if cfg!(target_os = "windows") {
-        if let Ok(path) = env::var("CUDA_PATH") {
-            // If CUDA_LIBRARY_PATH is not found, then CUDA_PATH will be used when building for
-            // Windows to locate the Cuda installation. Cuda installs the full Cuda SDK for 64-bit,
-            // but only a limited set of libraries for 32-bit. Namely, it does not include cublas in
-            // 32-bit, which cuda-sys requires.
+    let mut valid_paths = vec![];
+    for base in &candidates {
+        let lib = PathBuf::from(base).join("lib64");
+        if lib.is_dir() {
+            valid_paths.push(lib.clone());
+            valid_paths.push(lib.join("stubs"));
+        }
+        let base = base.join("targets/x86_64-linux");
+        let header = base.join("include/cuda.h");
+        if header.is_file() {
+            valid_paths.push(base.join("lib"));
+            valid_paths.push(base.join("lib/stubs"));
+            continue;
+        }
+    }
+    eprintln!("Found CUDA paths: {:?}", valid_paths);
+    valid_paths
+}
 
-            // 'path' points to the base of the CUDA Installation. The lib directory is a
-            // sub-directory.
-            let path = std::path::Path::new(&path);
+fn find_cuda_windows() -> PathBuf {
+    let paths = read_env();
+    if !paths.is_empty() {
+        return paths[0].clone();
+    }
 
-            // To do this the right way, we check to see which target we're building for.
-            let target = env::var("TARGET")
-                .expect("cargo did not set the TARGET environment variable as required.");
+    if let Ok(path) = env::var("CUDA_PATH") {
+        // If CUDA_LIBRARY_PATH is not found, then CUDA_PATH will be used when building for
+        // Windows to locate the Cuda installation. Cuda installs the full Cuda SDK for 64-bit,
+        // but only a limited set of libraries for 32-bit. Namely, it does not include cublas in
+        // 32-bit, which cuda-sys requires.
 
-            // Targets use '-' separators. e.g. x86_64-pc-windows-msvc
-            let target_components: Vec<_> = target.as_str().split("-").collect();
+        // 'path' points to the base of the CUDA Installation. The lib directory is a
+        // sub-directory.
+        let path = PathBuf::from(path);
 
-            // We check that we're building for Windows. This code assumes that the layout in
-            // CUDA_PATH matches Windows.
-            if target_components[2] != "windows" {
-                println!(
-                    "INFO: The CUDA_PATH variable is only used by cuda-sys on Windows. Your target \
-                     is {}.",
-                    target
-                );
-                return vec![];
-            }
+        // To do this the right way, we check to see which target we're building for.
+        let target = env::var("TARGET")
+            .expect("cargo did not set the TARGET environment variable as required.");
 
-            // Sanity check that the second component of 'target' is "pc"
-            debug_assert_eq!(
-                "pc", target_components[1],
-                "Expected a Windows target to have the second component be 'pc'. Target: {}",
+        // Targets use '-' separators. e.g. x86_64-pc-windows-msvc
+        let target_components: Vec<_> = target.as_str().split("-").collect();
+
+        // We check that we're building for Windows. This code assumes that the layout in
+        // CUDA_PATH matches Windows.
+        if target_components[2] != "windows" {
+            panic!(
+                "The CUDA_PATH variable is only used by cuda-sys on Windows. Your target is {}.",
                 target
             );
+        }
 
-            // x86_64 should use the libs in the "lib/x64" directory. If we ever support i686 (which
-            // does not ship with cublas support), its libraries are in "lib/Win32".
-            let lib_path = match target_components[0] {
-                "x86_64" => "x64",
-                "i686" => {
-                    // lib path would be "Win32" if we support i686. "cublas" is not present in the
-                    // 32-bit install.
-                    println!("INFO: Rust cuda-sys does not currently support 32-bit Windows.");
-                    return vec![];
-                }
-                _ => {
-                    println!("INFO: Rust cuda-sys only supports the x86_64 Windows architecture.");
-                    return vec![];
-                }
-            };
+        // Sanity check that the second component of 'target' is "pc"
+        debug_assert_eq!(
+            "pc", target_components[1],
+            "Expected a Windows target to have the second component be 'pc'. Target: {}",
+            target
+        );
 
-            return vec![
-                // i.e. $CUDA_PATH/lib/x64
-                path.join("lib")
-                    .join(lib_path)
-                    .to_str()
-                    .unwrap()
-                    .to_string(),
-            ];
-        }
+        // x86_64 should use the libs in the "lib/x64" directory. If we ever support i686 (which
+        // does not ship with cublas support), its libraries are in "lib/Win32".
+        let lib_path = match target_components[0] {
+            "x86_64" => "x64",
+            "i686" => {
+                // lib path would be "Win32" if we support i686. "cublas" is not present in the
+                // 32-bit install.
+                panic!("Rust cuda-sys does not currently support 32-bit Windows.");
+            }
+            _ => {
+                panic!("Rust cuda-sys only supports the x86_64 Windows architecture.");
+            }
+        };
+
+        // i.e. $CUDA_PATH/lib/x64
+        return path.join("lib").join(lib_path);
     }
 
     // No idea where to look for CUDA
-    vec![]
+    panic!("CUDA cannot find");
 }
 
 fn main() {
-    for p in find_library_paths() {
-        println!("cargo:rustc-link-search=native={}", p);
-    }
+    if cfg!(target_os = "windows") {
+        println!(
+            "cargo:rustc-link-search=native={}",
+            find_cuda_windows().display()
+        );
+    } else {
+        for path in find_cuda() {
+            println!("cargo:rustc-link-search=native={}", path.display());
+        }
+    };
     println!("cargo:rustc-link-lib=dylib=cuda");
     println!("cargo:rustc-link-lib=dylib=cudart");
     println!("cargo:rustc-link-lib=dylib=cublas");