huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-16 10:38:54 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add compute cap env support to examples (#1400)

Browse Source
This commit is contained in:
emka
2023-12-03 17:00:24 +01:00
committed by GitHub
parent 16161145ae
commit 99b7273b03
1 changed files with 11 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
candle-examples/build.rs
View File

@ -168,8 +168,16 @@ fn set_cuda_include_dir() -> Result<()> {
#[allow(unused)]
fn compute_cap() -> Result<usize> {
// Grab compute code from nvidia-smi
let mut compute_cap = {
println!("cargo:rerun-if-env-changed=CUDA_COMPUTE_CAP");
// Try to parse compute cap from env
let mut compute_cap = if let Ok(compute_cap_str) = std::env::var("CUDA_COMPUTE_CAP") {
println!("cargo:rustc-env=CUDA_COMPUTE_CAP={compute_cap_str}");
compute_cap_str
.parse::<usize>()
.context("Could not parse code")?
} else {
// Grab compute cap from nvidia-smi
let out = std::process::Command::new("nvidia-smi")
.arg("--query-gpu=compute_cap")
.arg("--format=csv")
@ -185,6 +193,7 @@ fn compute_cap() -> Result<usize> {
.next()
.context("missing line in stdout")?
.replace('.', "");
println!("cargo:rustc-env=CUDA_COMPUTE_CAP={cap}");
cap.parse::<usize>()
.with_context(|| format!("cannot parse as int {cap}"))?
};