Risc0 CUDA support

bin/citrea/Cargo.toml

@@ -92,7 +92,7 @@ citrea-e2e = { git = "https://github.com/chainwayxyz/citrea-e2e", rev = "0a6492e
 
 [features]
 default = [] # Deviate from convention by making the "native" feature active by default. This aligns with how this package is meant to be used (as a binary first, library second).
-
+cuda = ["citrea-risc0-bonsai-adapter/cuda"]
 bench = ["hex"] # "sov-risc0-adapter/bench", "risc0/bench"]
 
 [[bin]]

crates/risc0-bonsai/Cargo.toml

@@ -36,3 +36,4 @@ native = [
   "sov-rollup-interface/native",
 ]
 bench = ["native"]
+cuda = ["risc0-zkvm/cuda"]

crates/risc0-bonsai/src/host.rs

@@ -1,4 +1,5 @@
 //! This module implements the [`ZkvmHost`] trait for the RISC0 VM.
+use core::str;
 use std::thread;
 use std::time::Duration;
 
@@ -9,11 +10,11 @@ use bonsai_sdk::blocking::Client;
 use borsh::{BorshDeserialize, BorshSerialize};
 use risc0_zkvm::sha::Digest;
 use risc0_zkvm::{
-    compute_image_id, ExecutorEnvBuilder, ExecutorImpl, Journal, LocalProver, Prover, Receipt,
+    compute_image_id, ExecutorEnvBuilder, ExecutorImpl, ExternalProver, Journal, Prover,
+    ProverOpts, Receipt,
 };
 use sov_db::ledger_db::{LedgerDB, ProvingServiceLedgerOps};
 use sov_risc0_adapter::guest::Risc0Guest;
-use sov_risc0_adapter::host::add_benchmarking_callbacks;
 use sov_rollup_interface::zk::{Proof, Zkvm, ZkvmHost};
 use tracing::{error, info, warn};
 
@@ -299,35 +300,54 @@ impl<'a> ZkvmHost for Risc0BonsaiHost<'a> {
         let proof = match (self.client.as_ref(), with_proof) {
             // Local execution. If mode is Execute, we always do local execution.
             (_, false) => {
-                let env = add_benchmarking_callbacks(ExecutorEnvBuilder::default())
+                let env = ExecutorEnvBuilder::default()
                     .write_slice(&self.env)
                     .build()
-                    .unwrap();
+                    .expect("Proving environment should be built");
                 let mut executor = ExecutorImpl::from_elf(env, self.elf)?;
 
                 let session = executor.run()?;
+
+                tracing::info!(
+                    "Local execution completed.\nProving stats:\nSegments: {}\nTotal Cycles: {}\nUser Cycles: {}",
+                    session.segments.len(),
+                    session.total_cycles,
+                    session.user_cycles
+                );
                 let data =
                     bincode::serialize(&session.journal.expect("Journal shouldn't be empty"))?;
 
                 Ok(Proof::PublicInput(data))
             }
             // Local proving
             (None, true) => {
-                let env = add_benchmarking_callbacks(ExecutorEnvBuilder::default())
+                let env = ExecutorEnvBuilder::default()
+                    .segment_limit_po2(21)
                     .write_slice(&self.env)
                     .build()
-                    .unwrap();
+                    .expect("Proving environment should be built");
 
-                let prover = LocalProver::new("citrea");
-                let receipt = prover.prove(env, self.elf)?.receipt;
+                // RISC0_SERVER_PATH is used to specify the path to the r0vm binary
+                // As per https://github.com/chainwayxyz/citrea/issues/1381, it's a good idea to
+                // start using ric0 env var names
+                let prover = ExternalProver::new("citrea", std::env::var("RISC0_SERVER_PATH")?);
+
+                let start = std::time::Instant::now();
+                let proving_session =
+                    prover.prove_with_opts(env, self.elf, &ProverOpts::groth16())?;
+                let end = std::time::Instant::now();
 
                 tracing::info!("Local proving completed");
 
-                receipt.verify(self.image_id)?;
+                proving_session.receipt.verify(self.image_id)?;
 
-                tracing::info!("Verified the receipt");
+                tracing::info!(
+                    "Proving stats:\n{:?}\nTime: {}",
+                    proving_session.stats,
+                    (end - start).as_secs_f64()
+                );
 
-                let serialized_receipt = bincode::serialize(&receipt)?;
+                let serialized_receipt = bincode::serialize(&proving_session.receipt)?;
 
                 Ok(Proof::Full(serialized_receipt))
             }

crates/sovereign-sdk/adapters/risc0/src/lib.rs

@@ -8,9 +8,6 @@ use serde::{Deserialize, Serialize};
 use sov_rollup_interface::zk::Matches;
 
 pub mod guest;
-#[cfg(feature = "native")]
-pub mod host;
-
 #[cfg(feature = "bench")]
 pub mod metrics;