torch_build/
env.rs

1//! Environment variables and constatns.
2
3use crate::cuda::CudaArch;
4use cfg_if::cfg_if;
5use once_cell::sync::Lazy;
6use std::{
7    collections::HashSet,
8    ffi::{OsStr, OsString},
9    path::{Path, PathBuf},
10    process::Command,
11};
12
13/// The list of CUDA architectures given by `TORCH_CUDA_ARCH_LIST` environment variable.
14///
15/// If `TORCH_CUDA_ARCH_LIST` is not set, the default supported architectures are given.
16pub(crate) static TORCH_CUDA_ARCH_LIST: Lazy<HashSet<CudaArch>> = Lazy::new(|| {
17    if let Some(val) = rerun_env_string("TORCH_CUDA_ARCH_LIST") {
18        CudaArch::parse_list(&val)
19            .unwrap_or_else(|_| {
20                panic!(
21                    r#"unable to parse environment variable TORCH_CUDA_ARCH_LIST = "{}""#,
22                    val
23                )
24            })
25            .into_iter()
26            .collect()
27    } else {
28        let text = include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/TORCH_CUDA_ARCH_LIST"));
29        CudaArch::parse_list(text)
30            .expect("unable to load TORCH_CUDA_ARCH_LIST file")
31            .into_iter()
32            .collect()
33    }
34});
35
36pub(crate) static OUT_DIR: &str = env!("OUT_DIR");
37
38pub(crate) static TARGET: Lazy<Option<String>> = Lazy::new(|| rerun_env_string("TARGET"));
39
40/// The supported libtorch version.
41pub static TORCH_VERSION: &str =
42    include_str!(concat!(env!("CARGO_MANIFEST_DIR"), "/TORCH_VERSION"));
43
44/// The value of `LIBTORCH_CXX11_ABI` environment variable.
45pub static LIBTORCH_CXX11_ABI: Lazy<Option<bool>> = Lazy::new(|| {
46    rerun_env("LIBTORCH_CXX11_ABI").and_then(|val| {
47        cfg_if! {
48            if #[cfg(unix)] {
49                use std::os::unix::ffi::OsStrExt;
50                match val.as_bytes() {
51                    b"1" => Some(true),
52                    b"0" => Some(false),
53                    _ => {
54                        // warn
55                        None
56                    }
57                }
58            }
59            else {
60                match val.to_str() {
61                    Some("1") => Some(true),
62                    Some("0") => Some(false),
63                    _ => None,
64                }
65            }
66        }
67    })
68});
69
70/// The value of `LIBTORCH` environment variable.
71pub static LIBTORCH: Lazy<Option<PathBuf>> = Lazy::new(|| rerun_env_pathbuf("LIBTORCH"));
72
73/// The value of `TORCH_CUDA_VERSION` environment variable.
74pub static TORCH_CUDA_VERSION: Lazy<Option<String>> =
75    Lazy::new(|| rerun_env_string("TORCH_CUDA_VERSION"));
76
77/// The value of `CUDNN_HOME` environment variable, or `CUDNN_PATH` if `CUDNN_HOME` is not set.
78pub static CUDNN_HOME: Lazy<Option<PathBuf>> =
79    Lazy::new(|| rerun_env_pathbuf("CUDNN_HOME").or_else(|| rerun_env_pathbuf("CUDNN_PATH")));
80
81/// The value of `ROCM_HOME` environment variable, or `ROCM_PATH` if `ROCM_HOME` is not set.
82pub static ROCM_HOME: Lazy<Option<PathBuf>> = Lazy::new(|| {
83    let guess = rerun_env_pathbuf("ROCM_HOME")
84        .or_else(|| rerun_env_pathbuf("ROCM_PATH"))
85        .map(PathBuf::from);
86
87    #[cfg(unix)]
88    let guess = guess.or_else(|| {
89        Command::new("sh")
90            .arg("-c")
91            .arg("which hipcc | xargs readlink -f")
92            .output()
93            .ok()
94            .and_then(|output| output.status.success().then(|| output.stdout))
95            .and_then(|stdout| {
96                use std::os::unix::ffi::OsStrExt;
97
98                // strip trailing line breaks
99                let stdout = stdout.strip_suffix(b"\n")?;
100                let path = Path::new(OsStr::from_bytes(stdout));
101
102                let dir = path.parent()?.parent()?;
103                (dir.file_name()? == "hip")
104                    .then(|| dir.parent())
105                    .flatten()
106                    .map(PathBuf::from)
107            })
108    });
109
110    #[cfg(unix)]
111    let guess = guess.or_else(|| {
112        let dir = PathBuf::from("/opt/rocm");
113        dir.exists().then(|| dir)
114    });
115
116    guess
117});
118
119/// The CUDA installation directory on host system.
120///
121/// The value is determined in the following order.
122///
123/// 1. `CUDA_HOME` environment variable.
124/// 2. `CUDA_PATH` environment variable.
125/// 3. The path returned by `nvcc` if on Liunx or Mac.
126/// 4. `/usr/local/cuda` if on Debian or Ubuntu and the directory exists.
127pub static CUDA_HOME: Lazy<Option<PathBuf>> = Lazy::new(|| {
128    use os_info::Type::*;
129
130    let guess = rerun_env_pathbuf("CUDA_HOME")
131        .or_else(|| rerun_env_pathbuf("CUDA_PATH"))
132        .map(PathBuf::from);
133
134    #[cfg(any(target_os = "linux", target_os = "macos"))]
135    let guess = guess.or_else(|| {
136        Command::new("which")
137            .arg("nvcc")
138            .output()
139            .ok()
140            .and_then(|output| output.status.success().then(|| output.stdout))
141            .and_then(|stdout| {
142                use std::os::unix::ffi::OsStrExt;
143
144                // strip trailing line breaks
145                let stdout = stdout.strip_suffix(b"\n")?;
146                let path = Path::new(OsStr::from_bytes(stdout));
147                let dir = path.parent()?.parent()?.into();
148                Some(dir)
149            })
150    });
151
152    match os_info::get().os_type() {
153        Debian | Ubuntu => guess.or_else(|| {
154            let dir = PathBuf::from("/usr/local/cuda");
155            dir.exists().then(|| dir)
156        }),
157        _ => guess,
158    }
159});
160
161fn rerun_env(name: &str) -> Option<OsString> {
162    println!("cargo:rerun-if-env-changed={}", name);
163    std::env::var_os(name)
164}
165
166fn rerun_env_pathbuf(name: &str) -> Option<PathBuf> {
167    Some(rerun_env(name)?.into())
168}
169
170fn rerun_env_string(name: &str) -> Option<String> {
171    println!("cargo:rerun-if-env-changed={}", name);
172    std::env::var(name).ok()
173}
torch_build/env.rs

torch_build/
env.rs
