Skip to main content

zenu_cuda_config/
lib.rs

1use glob::glob;
2use std::{env, path::PathBuf};
3
4#[must_use]
5pub fn read_env() -> Vec<PathBuf> {
6    if let Ok(path) = env::var("CUDA_LIBRARY_PATH") {
7        // The location of the libcuda, libcudart, and libcublas can be hardcoded with the
8        // CUDA_LIBRARY_PATH environment variable.
9        let split_char = if cfg!(target_os = "windows") {
10            ";"
11        } else {
12            ":"
13        };
14        path.split(split_char).map(PathBuf::from).collect()
15    } else {
16        vec![]
17    }
18}
19
20#[must_use]
21#[expect(clippy::missing_panics_doc)]
22pub fn find_cuda() -> Vec<PathBuf> {
23    let mut candidates = read_env();
24    candidates.push(PathBuf::from("/opt/cuda"));
25    candidates.push(PathBuf::from("/usr/local/cuda"));
26    candidates.push(PathBuf::from("/lib"));
27    for e in glob("/usr/local/cuda-*").unwrap().flatten() {
28        candidates.push(e);
29    }
30
31    let mut valid_paths = vec![];
32    for base in &candidates {
33        let lib = PathBuf::from(base).join("lib64");
34        if lib.is_dir() {
35            valid_paths.push(lib.clone());
36            valid_paths.push(lib.join("stubs"));
37        }
38        let base = base.join("targets/x86_64-linux");
39        let base = base.join("x86_64-linux");
40        let header = base.join("include/cuda.h");
41        if header.is_file() {
42            valid_paths.push(base.join("lib"));
43            valid_paths.push(base.join("lib/stubs"));
44            continue;
45        }
46    }
47    eprintln!("Found CUDA paths: {valid_paths:?}");
48    valid_paths
49}
50
51#[expect(clippy::missing_panics_doc)]
52#[expect(clippy::uninlined_format_args)]
53#[must_use]
54pub fn find_cuda_windows() -> PathBuf {
55    let paths = read_env();
56    if !paths.is_empty() {
57        return paths[0].clone();
58    }
59
60    if let Ok(path) = env::var("CUDA_PATH") {
61        // If CUDA_LIBRARY_PATH is not found, then CUDA_PATH will be used when building for
62        // Windows to locate the Cuda installation. Cuda installs the full Cuda SDK for 64-bit,
63        // but only a limited set of libraries for 32-bit. Namely, it does not include cublas in
64        // 32-bit, which cuda-sys requires.
65
66        // 'path' points to the base of the CUDA Installation. The lib directory is a
67        // sub-directory.
68        let path = PathBuf::from(path);
69
70        // To do this the right way, we check to see which target we're building for.
71        let target = env::var("TARGET")
72            .expect("cargo did not set the TARGET environment variable as required.");
73
74        // Targets use '-' separators. e.g. x86_64-pc-windows-msvc
75        let target_components: Vec<_> = target.as_str().split('-').collect();
76
77        // We check that we're building for Windows. This code assumes that the layout in
78        // CUDA_PATH matches Windows.
79        assert!(
80            target_components[2] == "windows",
81            "The CUDA_PATH variable is only used by cuda-sys on Windows. Your target is {}.",
82            target
83        );
84
85        // Sanity check that the second component of 'target' is "pc"
86        debug_assert_eq!(
87            "pc", target_components[1],
88            "Expected a Windows target to have the second component be 'pc'. Target: {}",
89            target
90        );
91
92        // x86_64 should use the libs in the "lib/x64" directory. If we ever support i686 (which
93        // does not ship with cublas support), its libraries are in "lib/Win32".
94        let lib_path = match *target_components.first().unwrap() {
95            "x86_64" => "x64",
96            "i686" => {
97                // lib path would be "Win32" if we support i686. "cublas" is not present in the
98                // 32-bit install.
99                panic!("Rust cuda-sys does not currently support 32-bit Windows.");
100            }
101            _ => {
102                panic!("Rust cuda-sys only supports the x86_64 Windows architecture.");
103            }
104        };
105
106        // i.e. $CUDA_PATH/lib/x64
107        return path.join("lib").join(lib_path);
108    }
109
110    // No idea where to look for CUDA
111    panic!("CUDA cannot find");
112}