feat(host): GPU prover setup cache wired into GpuProverBuilder

Cargo.toml

@@ -62,8 +62,8 @@ insta = "1"
 tera = { version = "1", default-features = false }
 
 # Dependencies for airbender-crypto
-common_constants = { git = "https://github.com/matter-labs/zksync-airbender", rev = "73d69b5", default-features = false }
-blake2s_u32 = { git = "https://github.com/matter-labs/zksync-airbender", rev = "73d69b5", default-features = false }
+common_constants = { git = "https://github.com/matter-labs/zksync-airbender", rev = "f1e26aa5", default-features = false }
+blake2s_u32 = { git = "https://github.com/matter-labs/zksync-airbender", rev = "f1e26aa5", default-features = false }
 blake2 = { version = "0.10", default-features = false }
 sha2 = { version = "0.10", default-features = false }
 k256 = { version = "0.13", default-features = false }
@@ -91,11 +91,11 @@ ark-test-curves = { version = "0.5", default-features = false }
 rand_core = { version = "0.6.4", default-features = false }
 
 # Airbender dependencies
-riscv_common = { git = "https://github.com/matter-labs/zksync-airbender", rev = "73d69b5" }
-execution_utils = { git = "https://github.com/matter-labs/zksync-airbender", rev = "73d69b5" }
-gpu_prover = { git = "https://github.com/matter-labs/zksync-airbender", rev = "73d69b5" }
-riscv_transpiler = { git = "https://github.com/matter-labs/zksync-airbender", rev = "73d69b5" }
-verifier_common = { git = "https://github.com/matter-labs/zksync-airbender", rev = "73d69b5" }
+riscv_common = { git = "https://github.com/matter-labs/zksync-airbender", rev = "f1e26aa5" }
+execution_utils = { git = "https://github.com/matter-labs/zksync-airbender", rev = "f1e26aa5" }
+gpu_prover = { git = "https://github.com/matter-labs/zksync-airbender", rev = "f1e26aa5" }
+riscv_transpiler = { git = "https://github.com/matter-labs/zksync-airbender", rev = "f1e26aa5" }
+verifier_common = { git = "https://github.com/matter-labs/zksync-airbender", rev = "f1e26aa5" }
 
 # These packages can require too much stack space to compile,
 # which can be increased with `RUST_MIN_STACK=16777216` environment variable.

crates/airbender-host/Cargo.toml

@@ -19,6 +19,7 @@ gpu-prover = ["dep:gpu_prover", "execution_utils/gpu_prover"]
 airbender-core = { path = "../airbender-core" }
 airbender-codec = { path = "../airbender-codec" }
 serde = { workspace = true, features = ["derive"] }
+bincode = { workspace = true, features = ["serde", "std"] }
 execution_utils = { workspace = true }
 gpu_prover = { workspace = true, optional = true }
 riscv_transpiler = { workspace = true, features = ["jit", "flamegraph"] }

crates/airbender-host/src/error.rs

@@ -16,6 +16,8 @@ pub enum HostError {
     Prover(String),
     #[error("verification error: {0}")]
     Verification(String),
+    #[error("setup cache not found in path: {0}")]
+    SetupCacheNotFound(String),
 }
 
 pub type Result<T> = std::result::Result<T, HostError>;

crates/airbender-host/src/prover/gpu_prover.rs

@@ -4,11 +4,11 @@ use super::{
 use crate::error::{HostError, Result};
 use crate::proof::{Proof, RealProof};
 use crate::security::SecurityLevel;
-use execution_utils::unrolled_gpu::UnrolledProver;
+use execution_utils::unrolled_gpu::{UnrolledProver, UnrolledProverCache};
 use gpu_prover::execution::prover::ExecutionProverConfiguration;
 use riscv_transpiler::abstractions::non_determinism::QuasiUARTSource;
 use std::any::Any;
-use std::path::{Path, PathBuf};
+use std::path::{self, Path, PathBuf};
 use std::sync::atomic::{AtomicBool, Ordering};
 use std::sync::{mpsc, Mutex};
 use std::thread::JoinHandle;
@@ -19,6 +19,7 @@ pub struct GpuProverBuilder {
     worker_threads: Option<usize>,
     security: SecurityLevel,
     level: ProverLevel,
+    setup_cache_path: Option<PathBuf>,
 }
 
 impl GpuProverBuilder {
@@ -28,6 +29,7 @@ impl GpuProverBuilder {
             worker_threads: None,
             security: SecurityLevel::default(),
             level: ProverLevel::RecursionUnified,
+            setup_cache_path: None,
         }
     }
 
@@ -53,12 +55,33 @@ impl GpuProverBuilder {
         self
     }
 
+    /// Load pre-computed per-level setup data from `path` instead of running
+    /// the multi-minute compute. The file must exist (built earlier with a
+    /// dedicated generator that calls `UnrolledProver::dump_cache`); a missing
+    /// file at `build` time produces `HostError::SetupCacheNotFound`.
+    ///
+    /// The caller is responsible for keying the path on whatever distinguishes
+    /// the underlying binaries (e.g. an app-binary hash). Stale cache files
+    /// produce proving-time failures, not load-time errors.
+    pub fn with_setup_cache_path(mut self, path: impl AsRef<Path>) -> Self {
+        self.setup_cache_path = Some(path.as_ref().to_path_buf());
+        self
+    }
+
+    pub fn maybe_setup_cache_path(self, path: Option<PathBuf>) -> Self {
+        match path {
+            Some(p) => self.with_setup_cache_path(p),
+            None => self,
+        }
+    }
+
     pub fn build(self) -> Result<GpuProver> {
         GpuProver::new(
             &self.app_bin_path,
             self.worker_threads,
             self.security,
             self.level,
+            self.setup_cache_path,
         )
     }
 }
@@ -94,6 +117,7 @@ impl GpuProver {
         worker_threads: Option<usize>,
         security: SecurityLevel,
         level: ProverLevel,
+        setup_cache_path: Option<PathBuf>,
     ) -> Result<Self> {
         if matches!(worker_threads, Some(0)) {
             return Err(HostError::Prover(
@@ -102,8 +126,13 @@ impl GpuProver {
         }
 
         let app_bin_path = resolve_app_bin_path(app_bin_path)?;
-        let (command_tx, worker_handle) =
-            spawn_worker(app_bin_path, worker_threads, security, level)?;
+        let (command_tx, worker_handle) = spawn_worker(
+            app_bin_path,
+            worker_threads,
+            security,
+            level,
+            setup_cache_path,
+        )?;
 
         Ok(Self {
             command_tx,
@@ -189,6 +218,7 @@ fn spawn_worker(
     worker_threads: Option<usize>,
     security: SecurityLevel,
     level: ProverLevel,
+    setup_cache_path: Option<PathBuf>,
 ) -> Result<(mpsc::Sender<WorkerCommand>, JoinHandle<()>)> {
     let (command_tx, command_rx) = mpsc::channel();
     let (init_tx, init_rx) = mpsc::channel();
@@ -203,6 +233,7 @@ fn spawn_worker(
                 worker_threads,
                 security,
                 level,
+                setup_cache_path,
             )
         })
         .map_err(|err| {
@@ -235,6 +266,7 @@ fn gpu_worker_loop(
     worker_threads: Option<usize>,
     security: SecurityLevel,
     level: ProverLevel,
+    setup_cache_path: Option<PathBuf>,
 ) {
     // Keep all prover state inside this dedicated thread so a panic does not unwind
     // through host-call boundaries or require `AssertUnwindSafe`.
@@ -243,6 +275,7 @@ fn gpu_worker_loop(
         worker_threads,
         security,
         level.as_unrolled_level(),
+        setup_cache_path.as_deref(),
     ) {
         Ok(prover) => prover,
         Err(err) => {
@@ -294,17 +327,65 @@ fn create_unrolled_prover(
     worker_threads: Option<usize>,
     security: SecurityLevel,
     level: execution_utils::unrolled_gpu::UnrolledProverLevel,
+    setup_cache_path: Option<&Path>,
 ) -> Result<UnrolledProver> {
     let base_path = base_path(app_bin_path)?;
     let mut configuration = ExecutionProverConfiguration::default();
     if let Some(threads) = worker_threads {
         configuration.max_thread_pool_threads = Some(threads);
         configuration.replay_worker_threads_count = threads;
     }
-    Ok(UnrolledProver::new(
+
+    let Some(cache_path) = setup_cache_path else {
+        return Ok(UnrolledProver::new(
+            security.into(),
+            &base_path,
+            configuration,
+            level,
+        ));
+    };
+
+    let cache = load_cache(cache_path)?;
+    UnrolledProver::new_with_cache(
         security.into(),
         &base_path,
         configuration,
         level,
-    ))
+        &cache,
+    )
+    .map_err(|err| {
+        HostError::Prover(format!(
+            "setup cache {} is incompatible: {err}",
+            cache_path.display()
+        ))
+    })
+}
+
+fn load_cache(path: &Path) -> Result<UnrolledProverCache> {
+    if !path.exists() {
+        return Err(HostError::SetupCacheNotFound(path.display().to_string()));
+    }
+
+    let bytes = std::fs::read(path).map_err(|err| {
+        HostError::Prover(format!(
+            "failed to read setup cache {}: {err}",
+            path.display()
+        ))
+    })?;
+
+    let (cache, decoded_len): (UnrolledProverCache, usize) =
+        bincode::serde::decode_from_slice(&bytes, bincode::config::standard()).map_err(|err| {
+            HostError::Prover(format!(
+                "failed to decode setup cache {}: {err}",
+                path.display()
+            ))
+        })?;
+    if decoded_len != bytes.len() {
+        return Err(HostError::Prover(format!(
+            "setup cache {} has trailing bytes",
+            path.display()
+        )));
+    }
+
+    Ok(cache)
 }

crates/airbender-host/src/security.rs

@@ -33,6 +33,15 @@ impl From<SecurityLevel> for verifier_common::SecurityModel {
     }
 }
 
+impl From<verifier_common::SecurityModel> for SecurityLevel {
+    fn from(model: verifier_common::SecurityModel) -> Self {
+        match model {
+            verifier_common::SecurityModel::Security80 => SecurityLevel::Bits80,
+            verifier_common::SecurityModel::Security100 => SecurityLevel::Bits100,
+        }
+    }
+}
+
 impl std::fmt::Display for SecurityLevel {
     fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(formatter, "{}", self.bits())