From f77ff80ee4bbe6d83ffab9b9915fc922ef87c1ad Mon Sep 17 00:00:00 2001
From: Robert Remen <robik75@gmail.com>
Date: Tue, 24 Sep 2024 12:43:46 +0200
Subject: [PATCH] feat: enable compilation with CUDA stubs via the
 `ZKSYNC_USE_CUDA_STUBS` environment variable (#29)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

# What ❔

This PR is a follow-up to

https://github.com/matter-labs/era-cuda/pull/9
https://github.com/matter-labs/era-boojum-cuda/pull/36

making it possible to enable the compilation with CUDA stubs by setting
the environment variable `ZKSYNC_USE_CUDA_STUBS` to `1` or `true` or
`yes` in any capitalization in addition to the `no_cuda` configuration
flag.

## Why ❔

Other way to enable the compilation with stubs besides a configuration
flag is desirable.

## Checklist

- [x] PR title corresponds to the body of PR (we generate changelog
entries from PRs).
- [x] Code has been formatted via `cargo fmt` and linted via `cargo
clippy`.
---
 README.md                        |  4 ++++
 crates/boojum-cuda/build/main.rs |  8 +++-----
 crates/cudart-sys/build.rs       |  8 +++-----
 crates/cudart-sys/src/utils.rs   | 10 +++++++++-
 4 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/README.md b/README.md
index ab8809d..d96e085 100644
--- a/README.md
+++ b/README.md
@@ -25,6 +25,10 @@ Then add the following variable to your config (`.bashrc`/`.zshrc`):
 export BELLMAN_CUDA_DIR=<PATH_TO>/era-bellman-cuda
 ```
 
+Alternatively, if you can't or don't want to install the CUDA toolkit, you can compile the crates `cudart-sys` and `boojum-cuda` and crates that depend on them without CUDA toolkit.
+Doing so will result in stubs replacing calls to CUDA API and any GPU device code. The code will compile but any call to one of the stubs will result in an error during runtime.
+To compile in this mode, either include the rucstc `cfg` flag named `no_cuda`, for example by setting the `RUSTFLAGS` environment variable to  `--cfg no_cuda`, or by setting the environment variable `ZKSYNC_USE_CUDA_STUBS` to `1` or `true` or `yes` in any capitalization.
+
 ## Crates
 
 - [bindings-generator](./crates/bindings-generator/)
diff --git a/crates/boojum-cuda/build/main.rs b/crates/boojum-cuda/build/main.rs
index 5850091..a141f0f 100644
--- 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
index 9b96a2b..9089e2a 100644
--- 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
index bb305be..025caf9 100644
--- 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 {
     () => {