chore(cuda): bump up backend to 1.2.1-rc.4

Cargo.toml

@@ -113,11 +113,11 @@ lto = "thin"
 
 [workspace.dependencies]
 # Stark Backend
-openvm-stark-backend = { git = "https://github.com/openvm-org/stark-backend.git", tag = "v1.2.1-rc.3", default-features = false }
-openvm-stark-sdk = { git = "https://github.com/openvm-org/stark-backend.git", tag = "v1.2.1-rc.3", default-features = false }
-openvm-cuda-backend = { git = "https://github.com/openvm-org/stark-backend.git", tag = "v1.2.1-rc.3", default-features = false }
-openvm-cuda-builder = { git = "https://github.com/openvm-org/stark-backend.git", tag = "v1.2.1-rc.3", default-features = false }
-openvm-cuda-common = { git = "https://github.com/openvm-org/stark-backend.git", tag = "v1.2.1-rc.3", default-features = false }
+openvm-stark-backend = { git = "https://github.com/openvm-org/stark-backend.git", branch="fix/async-errors", default-features = false }
+openvm-stark-sdk = { git = "https://github.com/openvm-org/stark-backend.git", branch="fix/async-errors", default-features = false }
+openvm-cuda-backend = { git = "https://github.com/openvm-org/stark-backend.git", branch="fix/async-errors", default-features = false }
+openvm-cuda-builder = { git = "https://github.com/openvm-org/stark-backend.git", branch="fix/async-errors", default-features = false }
+openvm-cuda-common = { git = "https://github.com/openvm-org/stark-backend.git", branch="fix/async-errors", default-features = false }
 
 # OpenVM
 openvm-sdk = { path = "crates/sdk", default-features = false }

benchmarks/prove/Cargo.toml

@@ -21,7 +21,7 @@ openvm-native-recursion = { workspace = true, features = ["test-utils"] }
 
 clap = { version = "4.5.9", features = ["derive", "env"] }
 eyre.workspace = true
-tokio = { version = "1.43.1", features = ["rt", "rt-multi-thread", "macros"] }
+tokio = { version = "1.43.1", features = ["rt", "rt-multi-thread", "macros", "time"] }
 rand_chacha = { version = "0.3", default-features = false }
 k256 = { workspace = true, features = ["ecdsa"] }
 tiny-keccak.workspace = true

benchmarks/prove/src/bin/async_regex.rs

@@ -0,0 +1,92 @@
+use std::{
+    env::var,
+    sync::{
+        atomic::{AtomicUsize, Ordering},
+        Arc,
+    },
+    time::Duration,
+};
+
+use clap::Parser;
+use openvm_benchmarks_prove::util::BenchmarkCli;
+use openvm_benchmarks_utils::get_programs_dir;
+use openvm_sdk::{
+    config::{SdkVmBuilder, SdkVmConfig},
+    prover::AppProver,
+    DefaultStarkEngine, Sdk, StdIn, F,
+};
+use openvm_stark_sdk::config::setup_tracing;
+use tokio::{spawn, task::spawn_blocking, time::sleep};
+
+#[tokio::main]
+async fn main() -> eyre::Result<()> {
+    setup_tracing();
+    let args = BenchmarkCli::parse();
+    let mut config = SdkVmConfig::from_toml(include_str!("../../../guest/regex/openvm.toml"))?;
+    if let Some(max_height) = args.max_segment_length {
+        config
+            .app_vm_config
+            .as_mut()
+            .segmentation_limits
+            .max_trace_height = max_height;
+    }
+    if let Some(max_cells) = args.segment_max_cells {
+        config.app_vm_config.as_mut().segmentation_limits.max_cells = max_cells;
+    }
+
+    let sdk = Sdk::new(config)?;
+
+    let manifest_dir = get_programs_dir().join("regex");
+    let elf = sdk.build(Default::default(), manifest_dir, &None, None)?;
+    let app_exe = sdk.convert_to_exe(elf)?;
+
+    let data = include_str!("../../../guest/regex/regex_email.txt");
+    let fe_bytes = data.to_owned().into_bytes();
+    let input = StdIn::<F>::from_bytes(&fe_bytes);
+
+    let (app_pk, _app_vk) = sdk.app_keygen();
+
+    let max_par_jobs: usize = var("MAX_PAR_JOBS").map(|m| m.parse()).unwrap_or(Ok(1))?;
+    let num_jobs: usize = var("NUM_JOBS").map(|m| m.parse()).unwrap_or(Ok(20))?;
+    let cur_num_jobs = Arc::new(AtomicUsize::new(0));
+
+    let mut tasks = Vec::with_capacity(num_jobs);
+    for idx in 0..num_jobs {
+        let cur_num_jobs = cur_num_jobs.clone();
+        let app_exe = app_exe.clone();
+        let app_pk = app_pk.clone();
+        let input = input.clone();
+        let task = spawn(async move {
+            loop {
+                let c = cur_num_jobs.fetch_add(1, Ordering::SeqCst);
+                if c < max_par_jobs {
+                    tracing::info!("Acquired job {}, cur num jobs {}", idx, c + 1);
+                    break;
+                }
+                cur_num_jobs.fetch_sub(1, Ordering::SeqCst);
+                sleep(Duration::from_millis(100)).await;
+            }
+            let res = spawn_blocking(move || -> eyre::Result<()> {
+                let mut prover = AppProver::<DefaultStarkEngine, _>::new(
+                    SdkVmBuilder,
+                    &app_pk.app_vm_pk,
+                    app_exe,
+                    app_pk.leaf_verifier_program_commit(),
+                )?;
+                let _proof = prover.prove(input)?;
+                Ok(())
+            })
+            .await?;
+            let prev_num = cur_num_jobs.fetch_sub(1, Ordering::SeqCst);
+            tracing::info!("Decrement cur_num_jobs {} to {}", prev_num, prev_num - 1);
+            res
+        });
+        tasks.push(task);
+        sleep(Duration::from_millis(1000)).await;
+    }
+    for task in tasks {
+        task.await??;
+    }
+
+    Ok(())
+}

extensions/keccak256/circuit/Cargo.toml

@@ -35,6 +35,7 @@ cfg-if.workspace = true
 openvm-stark-sdk = { workspace = true }
 openvm-circuit = { workspace = true, features = ["test-utils"] }
 hex.workspace = true
+tokio = { version = "1.47.1", features = ["full"] }
 
 [build-dependencies]
 openvm-cuda-builder = { workspace = true, optional = true }

extensions/keccak256/circuit/src/tests.rs

@@ -406,3 +406,103 @@ fn test_keccak256_cuda_tracegen() {
         .simple_test()
         .unwrap();
 }
+
+#[cfg(feature = "cuda")]
+#[test]
+fn test_keccak256_cuda_tracegen_multi() {
+    let num_threads: usize = std::env::var("NUM_THREADS")
+        .ok()
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(2);
+
+    let num_tasks: usize = std::env::var("NUM_TASKS")
+        .ok()
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(num_threads * 4);
+
+    let runtime = tokio::runtime::Builder::new_multi_thread()
+        .max_blocking_threads(num_threads)
+        .enable_all()
+        .build()
+        .unwrap();
+
+    runtime.block_on(async {
+        let tasks_per_thread = num_tasks.div_ceil(num_threads);
+        let mut worker_handles = Vec::new();
+
+        for worker_idx in 0..num_threads {
+            let start_task = worker_idx * tasks_per_thread;
+            let end_task = std::cmp::min(start_task + tasks_per_thread, num_tasks);
+
+            let worker_handle = tokio::task::spawn(async move {
+                for task_id in start_task..end_task {
+                    tokio::task::spawn_blocking(move || {
+                        println!("[worker {}, task {}] Starting test", worker_idx, task_id);
+
+                        let mut rng = create_seeded_rng();
+                        let mut tester = GpuChipTestBuilder::default()
+                            .with_bitwise_op_lookup(default_bitwise_lookup_bus());
+
+                        let mut harness = create_cuda_harness(&tester);
+
+                        let num_ops: usize = 10;
+                        for _ in 0..num_ops {
+                            set_and_execute(
+                                &mut tester,
+                                &mut harness.executor,
+                                &mut harness.dense_arena,
+                                &mut rng,
+                                KECCAK256,
+                                None,
+                                None,
+                                None,
+                            );
+                        }
+
+                        for len in [0, 135, 136, 137, 2000] {
+                            set_and_execute(
+                                &mut tester,
+                                &mut harness.executor,
+                                &mut harness.dense_arena,
+                                &mut rng,
+                                KECCAK256,
+                                None,
+                                Some(len),
+                                None,
+                            );
+                        }
+
+                        harness
+                            .dense_arena
+                            .get_record_seeker::<KeccakVmRecordMut, _>()
+                            .transfer_to_matrix_arena(&mut harness.matrix_arena);
+
+                        tester
+                            .build()
+                            .load_gpu_harness(harness)
+                            .finalize()
+                            .simple_test()
+                            .unwrap();
+
+                        println!(
+                            "[worker {}, task {}] Test completed ✅",
+                            worker_idx, task_id
+                        );
+                    })
+                    .await
+                    .expect("task failed");
+                }
+            });
+            worker_handles.push(worker_handle);
+        }
+
+        for handle in worker_handles {
+            handle.await.expect("worker failed");
+        }
+
+        println!(
+            "\nAll {} tasks completed on {} workers.",
+            num_tasks, num_threads
+        );
+    });
+}