use std::path::{Path, PathBuf};
use std::{env, fs};
const PYTHON_PRINT_PYTORCH_DETAILS: &str = r"
import torch
from torch.utils import cpp_extension
print('LIBTORCH_VERSION:', torch.__version__.split('+')[0])
print('LIBTORCH_CXX11:', torch._C._GLIBCXX_USE_CXX11_ABI)
for include_path in cpp_extension.include_paths():
print('LIBTORCH_INCLUDE:', include_path)
for library_path in cpp_extension.library_paths():
print('LIBTORCH_LIB:', library_path)
";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum Os {
Linux,
Macos,
Windows,
}
#[allow(dead_code)]
#[derive(Debug, Clone)]
struct SystemInfo {
os: Os,
cxx11_abi: String,
libtorch_include_dirs: Vec<PathBuf>,
libtorch_lib_dir: PathBuf,
}
fn env_var_rerun(name: &str) -> Result<String, env::VarError> {
println!("cargo:rerun-if-env-changed={name}");
env::var(name)
}
impl SystemInfo {
fn new() -> Option<Self> {
let os = match env::var("CARGO_CFG_TARGET_OS")
.expect("Unable to get TARGET_OS")
.as_str()
{
"linux" => Os::Linux,
"windows" => Os::Windows,
"macos" => Os::Macos,
os => panic!("unsupported TARGET_OS '{os}'"),
};
let python_interpreter = match os {
Os::Windows => PathBuf::from("python.exe"),
Os::Linux | Os::Macos => {
if env::var_os("VIRTUAL_ENV").is_some() {
PathBuf::from("python")
} else {
PathBuf::from("python3")
}
}
};
let mut libtorch_include_dirs = vec![];
let mut libtorch_lib_dir = None;
let cxx11_abi = if env_var_rerun("LIBTORCH_USE_PYTORCH").is_ok() {
let output = std::process::Command::new(&python_interpreter)
.arg("-c")
.arg(PYTHON_PRINT_PYTORCH_DETAILS)
.output()
.expect("error running python interpreter");
let mut cxx11_abi = None;
for line in String::from_utf8_lossy(&output.stdout).lines() {
match line.strip_prefix("LIBTORCH_CXX11: ") {
Some("True") => cxx11_abi = Some("1".to_owned()),
Some("False") => cxx11_abi = Some("0".to_owned()),
_ => {}
}
if let Some(path) = line.strip_prefix("LIBTORCH_INCLUDE: ") {
libtorch_include_dirs.push(PathBuf::from(path))
}
if let Some(path) = line.strip_prefix("LIBTORCH_LIB: ") {
libtorch_lib_dir = Some(PathBuf::from(path))
}
}
match cxx11_abi {
Some(cxx11_abi) => cxx11_abi,
None => panic!("no cxx11 abi returned by python {output:?}"),
}
} else {
let libtorch = Self::prepare_libtorch_dir(os)?;
let includes = env_var_rerun("LIBTORCH_INCLUDE")
.map(PathBuf::from)
.unwrap_or_else(|_| libtorch.clone());
let lib = env_var_rerun("LIBTORCH_LIB")
.map(PathBuf::from)
.unwrap_or_else(|_| libtorch.clone());
libtorch_include_dirs.push(includes.join("include"));
libtorch_include_dirs.push(includes.join("include/torch/csrc/api/include"));
if lib.ends_with("lib") {
libtorch_lib_dir = Some(lib);
} else {
libtorch_lib_dir = Some(lib.join("lib"));
}
env_var_rerun("LIBTORCH_CXX11_ABI").unwrap_or_else(|_| "1".to_owned())
};
let libtorch_lib_dir = libtorch_lib_dir?;
Some(Self {
os,
cxx11_abi,
libtorch_include_dirs,
libtorch_lib_dir,
})
}
fn check_system_location(os: Os) -> Option<PathBuf> {
match os {
Os::Linux => Path::new("/usr/lib/libtorch.so")
.exists()
.then(|| PathBuf::from("/usr")),
_ => None,
}
}
fn prepare_libtorch_dir(os: Os) -> Option<PathBuf> {
if let Ok(libtorch) = env_var_rerun("DEP_TCH_LIBTORCH_LIB") {
Some(PathBuf::from(libtorch))
} else if let Ok(libtorch) = env_var_rerun("LIBTORCH") {
Some(PathBuf::from(libtorch))
} else if let Some(pathbuf) = Self::check_system_location(os) {
Some(pathbuf)
} else {
check_out_dir()
}
}
fn make(&self, use_cuda: bool, use_hip: bool) {
let cuda_dependency = if use_cuda || use_hip {
"src/cuda_hack/dummy_cuda_dependency.cpp"
} else {
"src/cuda_hack/fake_cuda_dependency.cpp"
};
println!("cargo:rerun-if-changed={cuda_dependency}");
match self.os {
Os::Linux | Os::Macos => {
cc::Build::new()
.cpp(true)
.pic(true)
.warnings(false)
.includes(&self.libtorch_include_dirs)
.flag(format!("-Wl,-rpath={}", self.libtorch_lib_dir.display()))
.flag("-std=c++17")
.flag(format!("-D_GLIBCXX_USE_CXX11_ABI={}", self.cxx11_abi))
.files(&[cuda_dependency])
.compile("burn-tch");
}
Os::Windows => {
cc::Build::new()
.cpp(true)
.pic(true)
.warnings(false)
.includes(&self.libtorch_include_dirs)
.flag("/std:c++17")
.files(&[cuda_dependency])
.compile("burn-tch");
}
};
}
fn make_cpu() {
let cuda_dependency = "src/cuda_hack/fake_cuda_dependency.cpp";
println!("cargo:rerun-if-changed={cuda_dependency}");
let os = env::var("CARGO_CFG_TARGET_OS").expect("Unable to get TARGET_OS");
match os.as_str() {
"windows" => {
cc::Build::new()
.cpp(true)
.pic(true)
.warnings(false)
.flag("/std:c++17")
.files(&[cuda_dependency])
.compile("burn-tch");
}
_ => {
cc::Build::new()
.cpp(true)
.pic(true)
.warnings(false)
.flag("-std=c++17")
.files(&[cuda_dependency])
.compile("tch");
}
};
}
}
fn check_out_dir() -> Option<PathBuf> {
let out_dir = env_var_rerun("OUT_DIR").ok()?;
let libtorch_dir = PathBuf::from(out_dir).join("libtorch");
libtorch_dir.exists().then_some(libtorch_dir)
}
fn main() {
let system_info = SystemInfo::new();
let out_dir = env_var_rerun("OUT_DIR").expect("Failed to get out dir");
let mut gpu_found = false;
let found_dir = system_info.is_some();
if let Some(system_info) = &system_info {
let si_lib = &system_info.libtorch_lib_dir;
let use_cuda =
si_lib.join("libtorch_cuda.so").exists() || si_lib.join("torch_cuda.dll").exists();
let use_hip =
si_lib.join("libtorch_hip.so").exists() || si_lib.join("torch_hip.dll").exists();
system_info.make(use_cuda, use_hip);
gpu_found = use_cuda || use_hip;
} else {
SystemInfo::make_cpu();
}
let check_file = PathBuf::from(out_dir).join("tch_gpu_check.rs");
if gpu_found {
fs::write(check_file, "#[allow(clippy::no_effect)]\n()").unwrap();
} else {
let message = if !found_dir {
r#"Could not find libtorch dir.
If you are trying to use the automatically downloaded version, the path is not directly available on Windows. Instead, try setting the `LIBTORCH` environment variable for the manual download instructions.
If the library has already been downloaded in the torch-sys OUT_DIR, you can point the variable to this path (or move the downloaded lib and point to it)."#
} else {
"No libtorch_cuda or libtorch_hip found. Download the GPU version of libtorch to use a GPU device"
};
fs::write(check_file, format!("panic!(\"{message}\")")).unwrap();
}
}