#![allow(clippy::expect_used)]
#![allow(clippy::manual_assert)]
#![allow(clippy::similar_names)]
#![allow(clippy::needless_pass_by_value)]
#![allow(clippy::unwrap_used)]
use std::env;
use std::path::{Path, PathBuf};
use std::process::Command;
fn main() {
println!("cargo:rerun-if-changed=build.rs");
if cfg!(feature = "cuda") {
cuda_build();
}
}
#[cfg_attr(not(feature = "cuda"), allow(dead_code))]
fn cuda_build() {
let nvcc = locate_nvcc().unwrap_or_else(|| {
panic!(
"nvcc not found. Set CUDA_HOME, or install the CUDA toolkit in one of: \
/opt/cuda, /usr/local/cuda. Or build without --features cuda."
);
});
let cuda_root = nvcc
.parent()
.and_then(Path::parent)
.expect("nvcc has a parent directory")
.to_path_buf();
let cuda_include = cuda_root.join("include");
let cuda_lib64 = cuda_root.join("lib64");
let manifest_dir = PathBuf::from(env::var("CARGO_MANIFEST_DIR").unwrap());
let cuda_src_dir = manifest_dir
.parent()
.unwrap()
.parent()
.unwrap()
.join("cuda");
let kernels_cu = cuda_src_dir.join("kernels.cu");
let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap());
let obj_path = out_dir.join("kernels.o");
let archive_path = out_dir.join("libdsfb_gpu_kernels.a");
println!("cargo:rerun-if-changed={}", kernels_cu.display());
println!(
"cargo:rerun-if-changed={}",
cuda_src_dir.join("common.cuh").display()
);
println!(
"cargo:rerun-if-changed={}",
cuda_src_dir.join("layout.cuh").display()
);
println!("cargo:rerun-if-env-changed=CUDA_HOME");
let nvcc_flags: &[&str] = &[
"--std=c++17",
"-gencode=arch=compute_75,code=sm_75",
"-gencode=arch=compute_80,code=sm_80",
"-gencode=arch=compute_89,code=sm_89",
"-gencode=arch=compute_89,code=compute_89",
"--fmad=false",
"-O2",
"-Xcompiler",
"-fPIC",
"-Xptxas",
"-O2",
"-DDSFB_GPU_FIXED_POINT",
];
let mut cmd = Command::new(&nvcc);
cmd.args(nvcc_flags)
.arg("-c")
.arg(&kernels_cu)
.arg("-o")
.arg(&obj_path)
.arg("-I")
.arg(&cuda_src_dir)
.arg("-I")
.arg(&cuda_include);
let status = cmd.status().expect("failed to invoke nvcc");
if !status.success() {
panic!(
"nvcc compilation failed for {}: {status}",
kernels_cu.display()
);
}
let ar = env::var("AR").unwrap_or_else(|_| "ar".to_string());
let status = Command::new(&ar)
.arg("rcs")
.arg(&archive_path)
.arg(&obj_path)
.status()
.expect("failed to invoke ar");
if !status.success() {
panic!("ar failed building {}: {status}", archive_path.display());
}
println!("cargo:rustc-link-search=native={}", out_dir.display());
println!("cargo:rustc-link-lib=static=dsfb_gpu_kernels");
println!("cargo:rustc-link-search=native={}", cuda_lib64.display());
println!("cargo:rustc-link-lib=dylib=cudart");
println!("cargo:rustc-link-lib=dylib=stdc++");
}
#[cfg_attr(not(feature = "cuda"), allow(dead_code))]
fn locate_nvcc() -> Option<PathBuf> {
if let Ok(home) = env::var("CUDA_HOME") {
let p = Path::new(&home).join("bin").join("nvcc");
if p.is_file() {
return Some(p);
}
}
if let Some(path) = env::var_os("PATH") {
for dir in env::split_paths(&path) {
let p = dir.join("nvcc");
if p.is_file() {
return Some(p);
}
}
}
for prefix in ["/opt/cuda", "/usr/local/cuda"] {
let p = Path::new(prefix).join("bin").join("nvcc");
if p.is_file() {
return Some(p);
}
}
None
}