Prefer CUDA_HOME environment variable
This commit is contained in:
parent
6d3b8b6d7d
commit
1d01777a13
|
@ -121,6 +121,15 @@ fn locate_perf_libs() -> Option<PathBuf> {
|
||||||
}
|
}
|
||||||
|
|
||||||
fn find_cuda_home(perf_libs_path: &Path) -> Option<PathBuf> {
|
fn find_cuda_home(perf_libs_path: &Path) -> Option<PathBuf> {
|
||||||
|
if let Ok(cuda_home) = env::var("CUDA_HOME") {
|
||||||
|
let path = PathBuf::from(cuda_home);
|
||||||
|
if path.is_dir() {
|
||||||
|
info!("Using CUDA_HOME: {:?}", path);
|
||||||
|
return Some(path);
|
||||||
|
}
|
||||||
|
warn!("Ignoring CUDA_HOME, not a path: {:?}", path);
|
||||||
|
}
|
||||||
|
|
||||||
// Search /usr/local for a `cuda-` directory that matches a perf-libs subdirectory
|
// Search /usr/local for a `cuda-` directory that matches a perf-libs subdirectory
|
||||||
for entry in fs::read_dir(&perf_libs_path).unwrap() {
|
for entry in fs::read_dir(&perf_libs_path).unwrap() {
|
||||||
if let Ok(entry) = entry {
|
if let Ok(entry) = entry {
|
||||||
|
@ -138,6 +147,7 @@ fn find_cuda_home(perf_libs_path: &Path) -> Option<PathBuf> {
|
||||||
continue;
|
continue;
|
||||||
}
|
}
|
||||||
|
|
||||||
|
info!("CUDA installation found at {:?}", cuda_home);
|
||||||
return Some(cuda_home);
|
return Some(cuda_home);
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
@ -147,8 +157,6 @@ fn find_cuda_home(perf_libs_path: &Path) -> Option<PathBuf> {
|
||||||
pub fn init_cuda() {
|
pub fn init_cuda() {
|
||||||
if let Some(perf_libs_path) = locate_perf_libs() {
|
if let Some(perf_libs_path) = locate_perf_libs() {
|
||||||
if let Some(cuda_home) = find_cuda_home(&perf_libs_path) {
|
if let Some(cuda_home) = find_cuda_home(&perf_libs_path) {
|
||||||
info!("CUDA installation found at {:?}", cuda_home);
|
|
||||||
|
|
||||||
let cuda_lib64_dir = cuda_home.join("lib64");
|
let cuda_lib64_dir = cuda_home.join("lib64");
|
||||||
if cuda_lib64_dir.is_dir() {
|
if cuda_lib64_dir.is_dir() {
|
||||||
let ld_library_path = cuda_lib64_dir.to_str().unwrap_or("").to_string()
|
let ld_library_path = cuda_lib64_dir.to_str().unwrap_or("").to_string()
|
||||||
|
@ -160,7 +168,7 @@ pub fn init_cuda() {
|
||||||
// to ensure the correct CUDA version is used
|
// to ensure the correct CUDA version is used
|
||||||
env::set_var("LD_LIBRARY_PATH", ld_library_path)
|
env::set_var("LD_LIBRARY_PATH", ld_library_path)
|
||||||
} else {
|
} else {
|
||||||
warn!("{:?} does not exist", cuda_lib64_dir);
|
warn!("CUDA lib64 directory does not exist: {:?}", cuda_lib64_dir);
|
||||||
}
|
}
|
||||||
|
|
||||||
let libcuda_crypt = perf_libs_path
|
let libcuda_crypt = perf_libs_path
|
||||||
|
|
Loading…
Reference in New Issue