Skip to content

Commit

Permalink
respect CUDA_PATH and friends in build.rs (with windows support)
Browse files Browse the repository at this point in the history
  • Loading branch information
gerwin3 committed Jan 16, 2024
1 parent 9ac0141 commit 9d59010
Showing 1 changed file with 22 additions and 23 deletions.
45 changes: 22 additions & 23 deletions build.rs
Original file line number Diff line number Diff line change
@@ -1,40 +1,39 @@
#[cfg(unix)]
fn main() {
let cuda_path = std::env::var("CUDA_PATH").unwrap_or_else(|_| "/usr/local/cuda".to_string());
let cuda_path = std::env::var("CUDA_PATH").map(std::path::PathBuf::from);

let cuda_include_path =
std::env::var("CUDA_INCLUDE_PATH").unwrap_or_else(|_| format!("{cuda_path}/include"));
#[cfg(not(windows))]
let cuda_path = cuda_path.unwrap_or_else(|_| std::path::PathBuf::from("/usr/local/cuda"));
#[cfg(windows)]
let cuda_path = cuda_path.expect("Missing environment variable `CUDA_PATH`.");

let cuda_lib_path =
std::env::var("CUDA_LIB_PATH").unwrap_or_else(|_| format!("{cuda_path}/lib64"));
let cuda_include_path = std::env::var("CUDA_INCLUDE_PATH")
.map(std::path::PathBuf::from)
.unwrap_or_else(|_| cuda_path.join("include"));

let cuda_lib_path = std::env::var("CUDA_LIB_PATH")
.map(std::path::PathBuf::from)
.unwrap_or_else(|_| {
#[cfg(not(windows))]
{
cuda_path.join("lib64")
}
#[cfg(windows)]
{
cuda_path.join("lib").join("x64")
}
});

cpp_build::Config::new()
.include(cuda_include_path)
.build("src/lib.rs");

println!("cargo:rustc-link-search={cuda_lib_path}");
println!("cargo:rustc-link-search={}", cuda_lib_path.display());
println!("cargo:rustc-link-lib=cudart");

#[cfg(feature = "npp")]
link_npp_libraries();
}

#[cfg(windows)]
fn main() {
let cuda_path = std::env::var("CUDA_PATH").expect("Missing environment variable `CUDA_PATH`.");
let cuda_path = std::path::Path::new(&cuda_path);
cpp_build::Config::new()
.include(cuda_path.join("include"))
.build("src/lib.rs");
println!(
"cargo:rustc-link-search={}",
cuda_path.join("lib").join("x64").display()
);
println!("cargo:rustc-link-lib=cudart");
#[cfg(feature = "npp")]
link_npp_libraries();
}

#[cfg(feature = "npp")]
fn link_npp_libraries() {
println!("cargo:rustc-link-lib=nppc");
Expand Down

0 comments on commit 9d59010

Please sign in to comment.