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

Co-authored-by: Julian Arce <52429267+JuArce@users.noreply.github.com>
Co-authored-by: Mauro Toscano <12560266+MauroToscano@users.noreply.github.com>

Author: 3 people

aggregation_mode/Cargo.lock

@@ -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/Cargo.toml

@@ -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/mod.rs

@@ -8,7 +8,32 @@ pub struct Risc0ProofReceiptAndImageId {
     pub receipt: Receipt,
 }
 
+#[derive(Debug)]
+pub enum AlignedRisc0VerificationError {
+    Verification(String),
+    UnsupportedProof,
+}
+
 impl Risc0ProofReceiptAndImageId {
+    /// Constructs a new instance of the struct, verifying the provided receipt against the given image ID.
+    pub fn new(
+        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 +47,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 +188,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/risc0_aggregator.rs

@@ -2,8 +2,10 @@ use std::sync::LazyLock;
 
 use alloy::primitives::Keccak256;
 use sp1_aggregation_program::SP1VkAndPubInputs;
+#[cfg(feature = "prove")]
+use sp1_sdk::EnvProver;
 use sp1_sdk::{
-    EnvProver, HashableKey, Prover, ProverClient, SP1ProofWithPublicValues, SP1Stdin,
+    CpuProver, HashableKey, Prover, ProverClient, SP1ProofWithPublicValues, SP1Stdin,
     SP1VerifyingKey,
 };
 
@@ -13,25 +15,59 @@ const CHUNK_PROGRAM_ELF: &[u8] =
 const USER_PROOFS_PROGRAM_ELF: &[u8] =
     include_bytes!("../../aggregation_programs/sp1/elf/sp1_user_proofs_aggregator_program");
 
+#[cfg(feature = "prove")]
 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 {
+    /// Constructs a new instance of the struct by verifying a given SP1 proof with its public values.
+    pub fn new(
+        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 +92,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 +132,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 +152,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 +162,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 +205,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/aggregators/sp1_aggregator.rs

@@ -16,6 +16,7 @@ use alloy::{
     primitives::Address,
     providers::{Provider, ProviderBuilder},
 };
+use rayon::prelude::*;
 use risc0_zkvm::Receipt;
 use tracing::{error, info};
 
@@ -109,38 +110,49 @@ impl ProofsFetcher {
             // Filter compatible proofs to be aggregated and push to queue
             let proofs_to_add: Vec<AlignedProof> = match engine {
                 ZKVMEngine::SP1 => data
-                    .into_iter()
-                    .filter_map(|p| match p.proving_system {
-                        ProvingSystemId::SP1 => {
-                            let elf = p.vm_program_code?;
-                            let proof_with_pub_values = bincode::deserialize(&p.proof).ok()?;
-                            let sp1_proof = SP1ProofWithPubValuesAndElf {
-                                proof_with_pub_values,
-                                elf,
-                            };
-
-                            Some(AlignedProof::SP1(sp1_proof.into()))
+                    .into_par_iter()
+                    .filter_map(|p| {
+                        if p.proving_system != ProvingSystemId::SP1 {
+                            return None;
+                        };
+
+                        let elf = p.vm_program_code?;
+                        let proof_with_pub_values = bincode::deserialize(&p.proof).ok()?;
+                        let sp1_proof =
+                            SP1ProofWithPubValuesAndElf::new(proof_with_pub_values, elf);
+
+                        match sp1_proof {
+                            Ok(proof) => Some(AlignedProof::SP1(proof.into())),
+                            Err(err) => {
+                                error!("Could not add proof, verification failed: {:?}", err);
+                                None
+                            }
                         }
-
-                        _ => None,
                     })
                     .collect(),
                 ZKVMEngine::RISC0 => data
-                    .into_iter()
-                    .filter_map(|p| match p.proving_system {
-                        ProvingSystemId::Risc0 => {
-                            let mut image_id = [0u8; 32];
-                            image_id.copy_from_slice(p.vm_program_code?.as_slice());
-                            let public_inputs = p.pub_input?;
-                            let inner_receipt: risc0_zkvm::InnerReceipt =
-                                bincode::deserialize(&p.proof).ok()?;
-
-                            let receipt = Receipt::new(inner_receipt, public_inputs);
-                            let risc0_proof = Risc0ProofReceiptAndImageId { image_id, receipt };
-
-                            Some(AlignedProof::Risc0(risc0_proof.into()))
+                    .into_par_iter()
+                    .filter_map(|p| {
+                        if p.proving_system != ProvingSystemId::Risc0 {
+                            return None;
+                        };
+
+                        let mut image_id = [0u8; 32];
+                        image_id.copy_from_slice(p.vm_program_code?.as_slice());
+                        let public_inputs = p.pub_input?;
+                        let inner_receipt: risc0_zkvm::InnerReceipt =
+                            bincode::deserialize(&p.proof).ok()?;
+
+                        let receipt = Receipt::new(inner_receipt, public_inputs);
+                        let risc0_proof = Risc0ProofReceiptAndImageId::new(image_id, receipt);
+
+                        match risc0_proof {
+                            Ok(proof) => Some(AlignedProof::Risc0(proof.into())),
+                            Err(err) => {
+                                error!("Could not add proof, verification failed: {:?}", err);
+                                None
+                            }
                         }
-                        _ => None,
                     })
                     .collect(),
             };
@@ -163,15 +175,7 @@ impl ProofsFetcher {
                 return Ok(proofs);
             }
 
-            // try to add them to the queue
-            for proof in proofs_to_add {
-                if let Err(err) = proof.verify() {
-                    error!("Could not add proof, verification failed: {:?}", err);
-                    continue;
-                };
-
-                proofs.push(proof);
-            }
+            proofs.extend(proofs_to_add);
         }
 
         // Update last processed block after collecting logs