feat/perf(aggregation-mode): verification & sp1 aggregation

aggregation_mode/Cargo.toml

@@ -17,6 +17,7 @@ sha3 = "0.10.8"
 reqwest = { version = "0.12" }
 ciborium = "=0.2.2"
 lambdaworks-crypto = { git = "https://github.com/lambdaclass/lambdaworks.git", rev = "5f8f2cfcc8a1a22f77e8dff2d581f1166eefb80b", features = ["serde"]}
+rayon = "1.10.0"
 # Necessary for the VerificationData type
 aligned-sdk = { path = "../batcher/aligned-sdk/" }
 # zkvms

aggregation_mode/src/aggregators/mod.rs

@@ -4,13 +4,9 @@ pub mod sp1_aggregator;
 use std::fmt::Display;
 
 use lambdaworks_crypto::merkle_tree::traits::IsMerkleTreeBackend;
-use risc0_aggregator::{
-    AlignedRisc0VerificationError, Risc0AggregationError, Risc0ProofReceiptAndImageId,
-};
+use risc0_aggregator::{Risc0AggregationError, Risc0ProofReceiptAndImageId};
 use sha3::{Digest, Keccak256};
-use sp1_aggregator::{
-    AlignedSP1VerificationError, SP1AggregationError, SP1ProofWithPubValuesAndElf,
-};
+use sp1_aggregator::{SP1AggregationError, SP1ProofWithPubValuesAndElf};
 use tracing::info;
 
 #[derive(Clone, Debug)]
@@ -219,24 +215,3 @@ impl IsMerkleTreeBackend for AlignedProof {
         hasher.finalize().into()
     }
 }
-
-#[derive(Debug)]
-pub enum AlignedVerificationError {
-    Sp1(AlignedSP1VerificationError),
-    Risc0(AlignedRisc0VerificationError),
-}
-
-impl AlignedProof {
-    pub fn verify(&self) -> Result<(), AlignedVerificationError> {
-        match self {
-            AlignedProof::SP1(proof) => sp1_aggregator::verify(proof).map_err(
-                |arg0: sp1_aggregator::AlignedSP1VerificationError| {
-                    AlignedVerificationError::Sp1(arg0)
-                },
-            ),
-            AlignedProof::Risc0(proof) => {
-                risc0_aggregator::verify(proof).map_err(AlignedVerificationError::Risc0)
-            }
-        }
-    }
-}

aggregation_mode/src/aggregators/risc0_aggregator.rs

@@ -8,7 +8,35 @@ pub struct Risc0ProofReceiptAndImageId {
     pub receipt: Receipt,
 }
 
+#[derive(Debug)]
+pub enum AlignedRisc0VerificationError {
+    Verification(String),
+    UnsupportedProof,
+}
+
 impl Risc0ProofReceiptAndImageId {
+    pub fn new(image_id: [u8; 32], receipt: Receipt) -> Self {
+        Self { image_id, receipt }
+    }
+
+    pub fn new_with_verification(
+        image_id: [u8; 32],
+        receipt: Receipt,
+    ) -> Result<Self, AlignedRisc0VerificationError> {
+        let is_supported_proof =
+            receipt.inner.composite().is_ok() || receipt.inner.succinct().is_ok();
+
+        if is_supported_proof {
+            receipt
+                .verify(image_id)
+                .map_err(|e| AlignedRisc0VerificationError::Verification(e.to_string()))?;
+        } else {
+            return Err(AlignedRisc0VerificationError::UnsupportedProof);
+        }
+
+        Ok(Self { image_id, receipt })
+    }
+
     pub fn public_inputs(&self) -> &Vec<u8> {
         &self.receipt.journal.bytes
     }
@@ -22,12 +50,6 @@ pub enum Risc0AggregationError {
     Verification(String),
 }
 
-#[derive(Debug)]
-pub enum AlignedRisc0VerificationError {
-    Verification(String),
-    UnsupportedProof,
-}
-
 /// Byte representation of the user proofs aggregator image_id, converted from `[u32; 8]` to `[u8; 32]`.
 pub const RISC0_USER_PROOFS_AGGREGATOR_PROGRAM_ID_BYTES: [u8; 32] = {
     let mut res = [0u8; 32];
@@ -169,17 +191,3 @@ pub(crate) fn run_chunk_aggregator(
 
     Ok(proof)
 }
-
-pub(crate) fn verify(
-    proof: &Risc0ProofReceiptAndImageId,
-) -> Result<(), AlignedRisc0VerificationError> {
-    // only stark proofs are supported for recursion
-    if proof.receipt.inner.composite().is_ok() || proof.receipt.inner.succinct().is_ok() {
-        proof
-            .receipt
-            .verify(proof.image_id)
-            .map_err(|e| AlignedRisc0VerificationError::Verification(e.to_string()))
-    } else {
-        Err(AlignedRisc0VerificationError::UnsupportedProof)
-    }
-}

aggregation_mode/src/aggregators/sp1_aggregator.rs

@@ -3,7 +3,7 @@ use std::sync::LazyLock;
 use alloy::primitives::Keccak256;
 use sp1_aggregation_program::SP1VkAndPubInputs;
 use sp1_sdk::{
-    EnvProver, HashableKey, Prover, ProverClient, SP1ProofWithPublicValues, SP1Stdin,
+    CpuProver, EnvProver, HashableKey, Prover, ProverClient, SP1ProofWithPublicValues, SP1Stdin,
     SP1VerifyingKey,
 };
 
@@ -13,25 +13,68 @@ const CHUNK_PROGRAM_ELF: &[u8] =
 const USER_PROOFS_PROGRAM_ELF: &[u8] =
     include_bytes!("../../aggregation_programs/sp1/elf/sp1_user_proofs_aggregator_program");
 
+#[allow(dead_code)]
 static SP1_PROVER_CLIENT: LazyLock<EnvProver> = LazyLock::new(ProverClient::from_env);
 
+/// Separate prover instance configured to always use the CPU.
+/// This is used for verification, which is performed in parallel and
+/// cannot be done on the GPU.
+static SP1_PROVER_CLIENT_CPU: LazyLock<CpuProver> =
+    LazyLock::new(|| ProverClient::builder().cpu().build());
+
 pub struct SP1ProofWithPubValuesAndElf {
     pub proof_with_pub_values: SP1ProofWithPublicValues,
     pub elf: Vec<u8>,
+    pub vk: SP1VerifyingKey,
+}
+
+#[derive(Debug)]
+pub enum AlignedSP1VerificationError {
+    Verification(sp1_sdk::SP1VerificationError),
+    UnsupportedProof,
 }
 
 impl SP1ProofWithPubValuesAndElf {
+    pub fn new(proof_with_pub_values: SP1ProofWithPublicValues, elf: Vec<u8>) -> Self {
+        let vk = vk_from_elf(&elf);
+
+        Self {
+            proof_with_pub_values,
+            elf,
+            vk,
+        }
+    }
+
+    pub fn new_with_verification(
+        proof_with_pub_values: SP1ProofWithPublicValues,
+        elf: Vec<u8>,
+    ) -> Result<Self, AlignedSP1VerificationError> {
+        let client = &*SP1_PROVER_CLIENT_CPU;
+
+        let (_pk, vk) = client.setup(&elf);
+
+        // only sp1 compressed proofs are supported for aggregation now
+        match proof_with_pub_values.proof {
+            sp1_sdk::SP1Proof::Compressed(_) => client
+                .verify(&proof_with_pub_values, &vk)
+                .map_err(AlignedSP1VerificationError::Verification),
+            _ => Err(AlignedSP1VerificationError::UnsupportedProof),
+        }?;
+
+        Ok(Self {
+            proof_with_pub_values,
+            elf,
+            vk,
+        })
+    }
+
     pub fn hash_vk_and_pub_inputs(&self) -> [u8; 32] {
         let mut hasher = Keccak256::new();
-        let vk_bytes = &self.vk().hash_bytes();
+        let vk_bytes = &self.vk.hash_bytes();
         hasher.update(vk_bytes);
         hasher.update(self.proof_with_pub_values.public_values.as_slice());
         hasher.finalize().into()
     }
-
-    pub fn vk(&self) -> SP1VerifyingKey {
-        vk_from_elf(&self.elf)
-    }
 }
 
 #[derive(Debug)]
@@ -56,15 +99,15 @@ pub(crate) fn run_user_proofs_aggregator(
             .proofs_vk_and_pub_inputs
             .push(SP1VkAndPubInputs {
                 public_inputs: proof.proof_with_pub_values.public_values.to_vec(),
-                vk: proof.vk().hash_u32(),
+                vk: proof.vk.hash_u32(),
             });
     }
 
     stdin.write(&program_input);
 
     // write proofs
     for input_proof in proofs.iter() {
-        let vk = input_proof.vk().vk;
+        let vk = input_proof.vk.vk.clone();
         // we only support sp1 Compressed proofs for now
         let sp1_sdk::SP1Proof::Compressed(proof) = input_proof.proof_with_pub_values.proof.clone()
         else {
@@ -96,6 +139,7 @@ pub(crate) fn run_user_proofs_aggregator(
     let proof_and_elf = SP1ProofWithPubValuesAndElf {
         proof_with_pub_values: proof,
         elf: USER_PROOFS_PROGRAM_ELF.to_vec(),
+        vk,
     };
 
     Ok(proof_and_elf)
@@ -115,7 +159,7 @@ pub(crate) fn run_chunk_aggregator(
         program_input.proofs_and_leaves_commitment.push((
             SP1VkAndPubInputs {
                 public_inputs: proof.proof_with_pub_values.public_values.to_vec(),
-                vk: proof.vk().hash_u32(),
+                vk: proof.vk.hash_u32(),
             },
             leaves_commitment.clone(),
         ));
@@ -125,7 +169,7 @@ pub(crate) fn run_chunk_aggregator(
 
     // write proofs
     for (input_proof, _) in proofs.iter() {
-        let vk = input_proof.vk().vk;
+        let vk = input_proof.vk.vk.clone();
         // we only support sp1 Compressed proofs for now
         let sp1_sdk::SP1Proof::Compressed(proof) = input_proof.proof_with_pub_values.proof.clone()
         else {
@@ -168,41 +212,14 @@ pub(crate) fn run_chunk_aggregator(
     let proof_and_elf = SP1ProofWithPubValuesAndElf {
         proof_with_pub_values: proof,
         elf: CHUNK_PROGRAM_ELF.to_vec(),
+        vk,
     };
 
     Ok(proof_and_elf)
 }
 
-#[derive(Debug)]
-pub enum AlignedSP1VerificationError {
-    Verification(sp1_sdk::SP1VerificationError),
-    UnsupportedProof,
-}
-
-pub(crate) fn verify(
-    sp1_proof_with_pub_values_and_elf: &SP1ProofWithPubValuesAndElf,
-) -> Result<(), AlignedSP1VerificationError> {
-    let client = &*SP1_PROVER_CLIENT;
-
-    let (_pk, vk) = client.setup(&sp1_proof_with_pub_values_and_elf.elf);
-
-    // only sp1 compressed proofs are supported for aggregation now
-    match sp1_proof_with_pub_values_and_elf
-        .proof_with_pub_values
-        .proof
-    {
-        sp1_sdk::SP1Proof::Compressed(_) => client
-            .verify(
-                &sp1_proof_with_pub_values_and_elf.proof_with_pub_values,
-                &vk,
-            )
-            .map_err(AlignedSP1VerificationError::Verification),
-        _ => Err(AlignedSP1VerificationError::UnsupportedProof),
-    }
-}
-
 pub fn vk_from_elf(elf: &[u8]) -> SP1VerifyingKey {
-    let prover = &*SP1_PROVER_CLIENT;
+    let prover = &*SP1_PROVER_CLIENT_CPU;
     let (_, vk) = prover.setup(elf);
     vk
 }

aggregation_mode/src/backend/config.rs

@@ -22,6 +22,7 @@ pub struct Config {
     pub last_aggregated_block_filepath: String,
     pub ecdsa: ECDSAConfig,
     pub proofs_per_chunk: u16,
+    pub pre_verification_enabled: bool,
 }
 
 impl Config {