feat: no_cuda via env

matter-labs · Sep 23, 2024 · a680433 · a680433
1 parent 9bba6d2
commit a680433
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 11 deletions.
diff --git a/crates/boojum-cuda/build/main.rs b/crates/boojum-cuda/build/main.rs
@@ -10,12 +10,10 @@ fn main() {
     gates::generate();
     poseidon_constants::generate();
     println!("cargo::rustc-check-cfg=cfg(no_cuda)");
-    #[cfg(no_cuda)]
-    {
+    if era_cudart_sys::is_no_cuda() {
         println!("cargo::warning={}", era_cudart_sys::no_cuda_message!());
-    }
-    #[cfg(not(no_cuda))]
-    {
+        println!("cargo::rustc-cfg=no_cuda");
+    } else {
         use era_cudart_sys::{get_cuda_lib_path, get_cuda_version};
         use std::env::var;
         let cuda_version =

diff --git a/crates/cudart-sys/build.rs b/crates/cudart-sys/build.rs
@@ -4,12 +4,10 @@ include!("src/utils.rs");
 
 fn main() {
     println!("cargo::rustc-check-cfg=cfg(no_cuda)");
-    #[cfg(no_cuda)]
-    {
+    if is_no_cuda() {
         println!("cargo::warning={}", no_cuda_message!());
-    }
-    #[cfg(not(no_cuda))]
-    {
+        println!("cargo::rustc-cfg=no_cuda");
+    } else {
         let cuda_version =
             get_cuda_version().expect("Failed to determine the CUDA Toolkit version.");
         if !cuda_version.starts_with("12.") {

diff --git a/crates/cudart-sys/src/utils.rs b/crates/cudart-sys/src/utils.rs
@@ -53,7 +53,15 @@ pub fn get_cuda_version() -> Option<String> {
     }
 }
 
-#[cfg(no_cuda)]
+pub fn is_no_cuda() -> bool {
+    if cfg!(no_cuda) {
+        true
+    } else {
+        let no_cuda = option_env!("ZKSYNC_USE_CUDA_STUBS").unwrap_or("");
+        no_cuda == "1" || no_cuda.to_lowercase() == "true" || no_cuda.to_lowercase() == "yes"
+    }
+}
+
 #[macro_export]
 macro_rules! no_cuda_message {
     () => {