perf: pararellize proofs verification

Author: MarcosNicolau

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

@@ -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,
 };
 
@@ -14,6 +14,11 @@ const USER_PROOFS_PROGRAM_ELF: &[u8] =
     include_bytes!("../../aggregation_programs/sp1/elf/sp1_user_proofs_aggregator_program");
 
 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,
@@ -182,7 +187,7 @@ pub enum AlignedSP1VerificationError {
 pub(crate) fn verify(
     sp1_proof_with_pub_values_and_elf: &SP1ProofWithPubValuesAndElf,
 ) -> Result<(), AlignedSP1VerificationError> {
-    let client = &*SP1_PROVER_CLIENT;
+    let client = &*SP1_PROVER_CLIENT_CPU;
 
     let (_pk, vk) = client.setup(&sp1_proof_with_pub_values_and_elf.elf);
 
@@ -202,7 +207,7 @@ pub(crate) fn verify(
 }
 
 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,7 +110,7 @@ impl ProofsFetcher {
             info!("Data downloaded from S3, number of proofs {}", data.len());
 
             // Filter compatible proofs to be aggregated and push to queue
-            let proofs_to_add: Vec<AlignedProof> = match engine {
+            let mut proofs_to_add: Vec<AlignedProof> = match engine {
                 ZKVMEngine::SP1 => data
                     .into_iter()
                     .filter_map(|p| match p.proving_system {
@@ -153,15 +154,21 @@ impl ProofsFetcher {
                 proofs_to_add.len()
             );
 
-            // 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);
-            }
+            // Try to add them to the queue
+            // We do this in parallel, as SP1 can take quite some time in verifying
+            // because of the overhead of setting up the prover
+            proofs_to_add = proofs_to_add
+                .into_par_iter()
+                .filter(|proof| match proof.verify() {
+                    Ok(_) => true,
+                    Err(err) => {
+                        error!("Could not add proof, verification failed: {:?}", err);
+                        return false;
+                    }
+                })
+                .collect();
+
+            proofs.extend(proofs_to_add);
         }
 
         Ok(proofs)