feat: Multibinary prover

Cargo.toml

@@ -1,5 +1,6 @@
 [workspace]
 members = [ "crates/sequencer_proof_client",
+    "crates/zksync_os_prover_common",
     "crates/zksync_os_fri_prover",
     "crates/zksync_os_snark_prover",
     "crates/zksync_os_prover_service"

crates/zksync_os_fri_prover/Cargo.toml

@@ -14,6 +14,7 @@ categories.workspace = true
 zksync_airbender_cli.workspace = true
 zksync_airbender_execution_utils.workspace = true
 zksync_sequencer_proof_client = { path = "../sequencer_proof_client" }
+zksync_os_prover_common = { path = "../zksync_os_prover_common" }
 
 anyhow.workspace = true
 base64.workspace = true
@@ -29,4 +30,4 @@ vise.workspace = true
 vise-exporter.workspace = true
 
 [features]
-gpu = ["zksync_airbender_cli/gpu"]
+gpu = ["zksync_airbender_cli/gpu", "zksync_os_prover_common/gpu"]

crates/zksync_os_fri_prover/src/lib.rs

@@ -9,15 +9,14 @@ use clap::Parser;
 use tracing_subscriber::{EnvFilter, FmtSubscriber};
 use zksync_airbender_cli::prover_utils::{
     create_proofs_internal, create_recursion_proofs, load_binary_from_path, serialize_to_file,
-    GpuSharedState,
 };
 use zksync_airbender_execution_utils::{Machine, ProgramProof, RecursionStrategy};
 use zksync_sequencer_proof_client::{sequencer_proof_client::SequencerProofClient, ProofClient};
 
 use crate::metrics::FRI_PROVER_METRICS;
 
 pub mod metrics;
-
+pub use zksync_os_prover_common::MultiBinaryProver;
 /// Command-line arguments for the Zksync OS prover
 #[derive(Parser, Debug)]
 #[command(name = "Zksync OS Prover")]
@@ -58,7 +57,7 @@ pub fn create_proof(
     prover_input: Vec<u32>,
     binary: &Vec<u32>,
     circuit_limit: usize,
-    _gpu_state: &mut GpuSharedState,
+    _multi_prover: &mut MultiBinaryProver<usize>,
 ) -> ProgramProof {
     let mut timing = Some(0f64);
     let (proof_list, proof_metadata) = create_proofs_internal(
@@ -68,7 +67,7 @@ pub fn create_proof(
         circuit_limit,
         None,
         #[cfg(feature = "gpu")]
-        &mut Some(_gpu_state),
+        &mut Some(_multi_prover.execution_prover_mut()),
         #[cfg(not(feature = "gpu"))]
         &mut None,
         &mut timing, // timing info
@@ -80,7 +79,7 @@ pub fn create_proof(
         RecursionStrategy::UseReducedLog23Machine,
         &None,
         #[cfg(feature = "gpu")]
-        &mut Some(_gpu_state),
+        &mut Some(_multi_prover.execution_prover_mut()),
         #[cfg(not(feature = "gpu"))]
         &mut None,
         &mut timing, // timing info
@@ -101,14 +100,23 @@ pub async fn run(args: Args) {
         .app_bin_path
         .unwrap_or_else(|| Path::new(&manifest_path).join("../../multiblock_batch.bin"));
     let binary = load_binary_from_path(&binary_path.to_str().unwrap().to_string());
-    // For regular fri proving, we keep using reduced RiscV machine.
+
+    #[cfg(feature = "gpu")]
+    use zksync_airbender_gpu_prover::circuit_type::MainCircuitType;
+
     #[cfg(feature = "gpu")]
-    let mut gpu_state = GpuSharedState::new(
-        &binary,
-        zksync_airbender_cli::prover_utils::MainCircuitType::ReducedRiscVMachine,
+    let main_binaries = vec![(0, MainCircuitType::RiscVCycles, binary.clone())];
+
+    #[cfg(feature = "gpu")]
+    let mut multi_prover = MultiBinaryProver::new(
+        1,                                    // max_concurrent_batches
+        main_binaries,                        // your main binaries
+        MainCircuitType::ReducedRiscVMachine, // recursion circuit type
+        1,                                    // recursion key
     );
+
     #[cfg(not(feature = "gpu"))]
-    let mut gpu_state = GpuSharedState::new(&binary);
+    let mut multi_prover = MultiBinaryProver::new();
 
     tracing::info!(
         "Starting Zksync OS FRI prover for {}",
@@ -122,7 +130,7 @@ pub async fn run(args: Args) {
             &client,
             &binary,
             args.circuit_limit,
-            &mut gpu_state,
+            &mut multi_prover,
             args.path.clone(),
         )
         .await
@@ -144,8 +152,8 @@ pub async fn run_inner<P: ProofClient>(
     client: &P,
     binary: &Vec<u32>,
     circuit_limit: usize,
-    #[cfg(feature = "gpu")] gpu_state: &mut GpuSharedState,
-    #[cfg(not(feature = "gpu"))] gpu_state: &mut GpuSharedState<'_>,
+    #[cfg(feature = "gpu")] multi_prover: &mut MultiBinaryProver<usize>,
+    #[cfg(not(feature = "gpu"))] multi_prover: &mut MultiBinaryProver<usize>,
     path: Option<PathBuf>,
 ) -> anyhow::Result<bool> {
     let (block_number, prover_input) = match client.pick_fri_job().await {
@@ -172,7 +180,7 @@ pub async fn run_inner<P: ProofClient>(
 
     tracing::info!("Starting proving block number {}", block_number);
 
-    let proof = create_proof(prover_input, binary, circuit_limit, gpu_state);
+    let proof = create_proof(prover_input, binary, circuit_limit, multi_prover);
 
     tracing::info!("Finished proving block number {}", block_number);
     let proof_bytes: Vec<u8> = bincode::serde::encode_to_vec(&proof, bincode::config::standard())

crates/zksync_os_prover_common/Cargo.toml

@@ -0,0 +1,13 @@
+[package]
+name = "zksync_os_prover_common"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+# Execution utils for universal verifier binary
+zksync_airbender_execution_utils = { workspace = true, optional = true }
+
+[features]
+default = []
+gpu = ["zksync_airbender_execution_utils"]
+

crates/zksync_os_prover_common/src/lib.rs

@@ -0,0 +1,3 @@
+pub mod multi_binary_prover;
+
+pub use multi_binary_prover::MultiBinaryProver;

crates/zksync_os_prover_common/src/multi_binary_prover.rs

@@ -0,0 +1,92 @@
+use std::fmt::Debug;
+use std::hash::Hash;
+
+#[cfg(feature = "gpu")]
+use zksync_airbender_gpu_prover::circuit_type::MainCircuitType;
+
+#[cfg(feature = "gpu")]
+use zksync_airbender_execution_utils::{get_padded_binary, UNIVERSAL_CIRCUIT_VERIFIER};
+
+#[cfg(feature = "gpu")]
+pub use zksync_airbender_execution_utils::execution::prover::{ExecutableBinary, ExecutionProver};
+
+#[cfg(feature = "gpu")]
+pub struct MultiBinaryProver<K: Clone + Debug + Eq + Hash> {
+    execution_prover: ExecutionProver<K>,
+    recursion_key: K,
+}
+
+#[cfg(feature = "gpu")]
+impl<K: Clone + Debug + Eq + Hash> MultiBinaryProver<K> {
+    pub fn new(
+        max_concurrent_batches: usize,
+        main_binaries: Vec<(K, MainCircuitType, Vec<u32>)>,
+        recursion_circuit_type: MainCircuitType,
+        recursion_key: K,
+    ) -> Self {
+        // Validate recursion circuit type
+        assert!(
+            recursion_circuit_type == MainCircuitType::ReducedRiscVMachine
+                || recursion_circuit_type == MainCircuitType::ReducedRiscVLog23Machine,
+            "Recursion circuit type must be ReducedRiscVMachine or ReducedRiscVLog23Machine"
+        );
+
+        // Convert main binaries to ExecutableBinary
+        let mut all_binaries: Vec<ExecutableBinary<K, Vec<u32>>> = main_binaries
+            .into_iter()
+            .map(|(key, circuit_type, bytecode)| ExecutableBinary {
+                key,
+                circuit_type,
+                bytecode,
+            })
+            .collect();
+
+        // Add recursion binary (universal circuit verifier)
+        let recursion_binary = ExecutableBinary {
+            key: recursion_key.clone(),
+            circuit_type: recursion_circuit_type,
+            bytecode: get_padded_binary(UNIVERSAL_CIRCUIT_VERIFIER),
+        };
+        all_binaries.push(recursion_binary);
+
+        let execution_prover = ExecutionProver::new(max_concurrent_batches, all_binaries);
+
+        Self {
+            execution_prover,
+            recursion_key,
+        }
+    }
+
+    pub fn execution_prover(&self) -> &ExecutionProver<K> {
+        &self.execution_prover
+    }
+
+    pub fn execution_prover_mut(&mut self) -> &mut ExecutionProver<K> {
+        &mut self.execution_prover
+    }
+
+    pub fn recursion_key(&self) -> &K {
+        &self.recursion_key
+    }
+}
+
+/// Non-GPU version of MultiBinaryProver (placeholder for compatibility)
+#[cfg(not(feature = "gpu"))]
+pub struct MultiBinaryProver<K: Clone + Debug + Eq + Hash> {
+    _phantom: std::marker::PhantomData<K>,
+}
+
+#[cfg(not(feature = "gpu"))]
+impl<K: Clone + Debug + Eq + Hash> MultiBinaryProver<K> {
+    pub fn new() -> Self {
+        Self {
+            _phantom: std::marker::PhantomData,
+        }
+    }
+}
+
+impl<K: Clone + Debug + Eq + Hash> Default for MultiBinaryProver<K> {
+    fn default() -> Self {
+        Self::new()
+    }
+}

crates/zksync_os_prover_service/src/lib.rs

@@ -8,11 +8,10 @@ use tracing_subscriber::{EnvFilter, FmtSubscriber};
 #[cfg(feature = "gpu")]
 use zkos_wrapper::gpu::snark::gpu_create_snark_setup_data;
 use zksync_airbender_cli::prover_utils::load_binary_from_path;
-#[cfg(not(feature = "gpu"))]
-use zksync_airbender_cli::prover_utils::GpuSharedState;
 #[cfg(feature = "gpu")]
 use zksync_airbender_cli::prover_utils::GpuSharedState;
 use zksync_airbender_execution_utils::{get_padded_binary, UNIVERSAL_CIRCUIT_VERIFIER};
+use zksync_os_fri_prover::MultiBinaryProver;
 #[cfg(feature = "gpu")]
 use zksync_os_snark_prover::compute_compression_vk;
 use zksync_sequencer_proof_client::sequencer_proof_client::SequencerProofClient;
@@ -93,12 +92,12 @@ pub async fn run(args: Args) {
 
         // For regular fri proving, we keep using reduced RiscV machine.
         #[cfg(feature = "gpu")]
-        let mut gpu_state = GpuSharedState::new(
-            &binary,
-            zksync_airbender_cli::prover_utils::MainCircuitType::ReducedRiscVMachine,
-        );
+        let main_binaries = vec![(0, MainCircuitType::RiscVCycles, binary.clone())];
+        #[cfg(feature = "gpu")]
+        let mut multi_prover =
+            MultiBinaryProver::new(1, main_binaries, MainCircuitType::ReducedRiscVMachine, 1);
         #[cfg(not(feature = "gpu"))]
-        let mut gpu_state = GpuSharedState::new(&binary);
+        let mut multi_prover = MultiBinaryProver::new();
 
         // Run FRI prover until we hit one of the limits
         tracing::info!("Running FRI prover");
@@ -107,7 +106,7 @@ pub async fn run(args: Args) {
                 &client,
                 &binary,
                 args.circuit_limit,
-                &mut gpu_state,
+                &mut multi_prover,
                 args.fri_path.clone(),
             )
             .await
@@ -129,7 +128,7 @@ pub async fn run(args: Args) {
             }
         }
         #[cfg(feature = "gpu")]
-        drop(gpu_state);
+        drop(multi_prover);
 
         // Here we do exactly one SNARK proof
         tracing::info!("Running SNARK prover");

crates/zksync_os_snark_prover/Cargo.toml

@@ -13,6 +13,7 @@ categories.workspace = true
 [dependencies]
 #zksync_prover_job_processor.workspace = true
 zksync_sequencer_proof_client = { path = "../sequencer_proof_client" }
+zksync_os_prover_common = { path = "../zksync_os_prover_common" }
 
 zksync_airbender_cli.workspace = true
 zksync_airbender_execution_utils.workspace = true
@@ -36,4 +37,4 @@ vise.workspace = true
 vise-exporter.workspace = true
 
 [features]
-gpu = ["zksync_airbender_cli/gpu", "zkos_wrapper/gpu", "proof-compression/gpu"]
+gpu = ["zksync_airbender_cli/gpu", "zkos_wrapper/gpu", "proof-compression/gpu", "zksync_os_prover_common/gpu"]

crates/zksync_os_snark_prover/src/lib.rs

@@ -11,13 +11,14 @@ use zkos_wrapper::{
 };
 use zkos_wrapper::{prove, serialize_to_file, SnarkWrapperProof};
 use zksync_airbender_cli::prover_utils::{
-    create_final_proofs_from_program_proof, create_proofs_internal, GpuSharedState,
+    create_final_proofs_from_program_proof, create_proofs_internal,
 };
 use zksync_airbender_execution_utils::{
     generate_oracle_data_for_universal_verifier, generate_oracle_data_from_metadata_and_proof_list,
     get_padded_binary, Machine, ProgramProof, RecursionStrategy, VerifierCircuitsIdentifiers,
     UNIVERSAL_CIRCUIT_VERIFIER,
 };
+use zksync_os_prover_common::MultiBinaryProver;
 use zksync_sequencer_proof_client::{
     sequencer_proof_client::SequencerProofClient, ProofClient, SnarkProofInputs,
 };
@@ -61,7 +62,7 @@ pub fn generate_verification_key(
 pub fn merge_fris(
     snark_proof_input: SnarkProofInputs,
     verifier_binary: &Vec<u32>,
-    gpu_state: &mut GpuSharedState,
+    _multi_prover: &mut MultiBinaryProver<usize>,
 ) -> ProgramProof {
     SNARK_PROVER_METRICS
         .fri_proofs_merged
@@ -104,7 +105,10 @@ pub fn merge_fris(
             &zksync_airbender_execution_utils::Machine::Reduced,
             100, // Guessing - FIXME!!
             Some(first_metadata.create_prev_metadata()),
-            &mut Some(gpu_state),
+            #[cfg(feature = "gpu")]
+            &mut Some(_multi_prover.execution_prover_mut()),
+            #[cfg(not(feature = "gpu"))]
+            &mut None,
             &mut Some(0f64),
         );
         // Let's do recursion.
@@ -122,7 +126,10 @@ pub fn merge_fris(
                 &Machine::Reduced,
                 proof_metadata.total_proofs(),
                 Some(proof_metadata.create_prev_metadata()),
-                &mut Some(gpu_state),
+                #[cfg(feature = "gpu")]
+                &mut Some(_multi_prover.execution_prover_mut()),
+                #[cfg(not(feature = "gpu"))]
+                &mut None,
                 &mut Some(0f64),
             );
         }
@@ -255,18 +262,18 @@ pub async fn run_inner<P: ProofClient>(
     );
     tracing::info!("Initializing GPU state");
     #[cfg(feature = "gpu")]
-    let mut gpu_state = GpuSharedState::new(
-        verifier_binary,
-        zksync_airbender_cli::prover_utils::MainCircuitType::ReducedRiscVMachine,
-    );
+    let main_binaries = vec![(0, MainCircuitType::RiscVCycles, verifier_binary.clone())];
+    #[cfg(feature = "gpu")]
+    let mut multi_prover =
+        MultiBinaryProver::new(1, main_binaries, MainCircuitType::ReducedRiscVMachine, 1);
     #[cfg(not(feature = "gpu"))]
-    let mut gpu_state = GpuSharedState::new(verifier_binary);
+    let mut multi_prover = MultiBinaryProver::new();
     tracing::info!("Finished initializing GPU state");
-    let proof = merge_fris(snark_proof_input, verifier_binary, &mut gpu_state);
+    let proof = merge_fris(snark_proof_input, verifier_binary, &mut multi_prover);
 
-    // Drop GPU state to release the airbender GPU resources (as now Final Proof will be taking them).
+    // Drop multi_prover to release the airbender GPU resources (as now Final Proof will be taking them).
     #[cfg(feature = "gpu")]
-    drop(gpu_state);
+    drop(multi_prover);
 
     tracing::info!("Creating final proof before SNARKification");