Skip to content

Commit

Permalink
fix: better CUDA detection (#61)
Browse files Browse the repository at this point in the history
# What ❔

This PR improves the detection of the CUDA toolkit under Linux, and the
detection of its version under all OSes.

## Why ❔

Under Linux, the CUDA detection is hardcoded to look into
`/usr/local/cuda`, which is not guaranteed to be the default
installation folder. The new method first looks for the `CUDA_PATH` env.
var., then default to `/usr/local/cuda` if it is empty.

The current way of finding out the installed CUDA version is to look for
a `version.json` file in the CUDA installation path. This file is not
guaranteed to exist. However, `nvcc` **is** guaranteed to be installed
with the toolkit; so the new method is to parse its output to find the
exact version there.

## Checklist
- [X] PR title corresponds to the body of PR (we generate changelog
entries from PRs).
- [X] Tests for the changes have been added / updated.
- [X] Documentation comments have been added / updated.
- [X] Code has been formatted via `zk fmt` and `zk lint`.
  • Loading branch information
delehef authored Jan 14, 2025
1 parent 0dda58b commit e902ca4
Show file tree
Hide file tree
Showing 2 changed files with 24 additions and 11 deletions.
4 changes: 2 additions & 2 deletions crates/cudart-sys/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ name = "era_cudart_sys"
description = "Raw CUDA bindings for ZKsync"

[dependencies]
serde_json = "1.0"
regex-lite = "0.1"

[build-dependencies]
serde_json = "1.0"
regex-lite = "0.1"
31 changes: 22 additions & 9 deletions crates/cudart-sys/src/utils.rs
Original file line number Diff line number Diff line change
Expand Up @@ -4,12 +4,16 @@ use std::path::{Path, PathBuf};
pub fn get_cuda_path() -> Option<&'static Path> {
#[cfg(target_os = "linux")]
{
let path = Path::new("/usr/local/cuda");
if path.exists() {
Some(path)
} else {
None
for path_name in [option_env!("CUDA_PATH"), Some("/usr/local/cuda")]
.iter()
.flatten()
{
let path = Path::new(path_name);
if path.exists() {
return Some(path);
}
}
None
}
#[cfg(target_os = "windows")]
{
Expand Down Expand Up @@ -44,10 +48,19 @@ pub fn get_cuda_version() -> Option<String> {
if let Some(version) = option_env!("CUDA_VERSION") {
Some(version.to_string())
} else if let Some(path) = get_cuda_path() {
let file = File::open(path.join("version.json")).expect("CUDA Toolkit should be installed");
let reader = std::io::BufReader::new(file);
let value: serde_json::Value = serde_json::from_reader(reader).unwrap();
Some(value["cuda"]["version"].as_str().unwrap().to_string())
let re = regex_lite::Regex::new(r"V(?<version>\d{2}\.\d+\.\d+)").unwrap();
let nvcc_out = std::process::Command::new("nvcc")
.arg("--version")
.output()
.expect("failed to start `nvcc`");
let nvcc_str = std::str::from_utf8(&nvcc_out.stdout).expect("`nvcc` output is not UTF8");
let captures = re.captures(&nvcc_str).unwrap();
let version = captures
.get(0)
.expect("unable to find nvcc version in the form VMM.mm.pp in the output of `nvcc --version`:\n{nvcc_str}")
.as_str()
.to_string();
Some(version)
} else {
None
}
Expand Down

0 comments on commit e902ca4

Please sign in to comment.