feat: Benchmark halo2 flamegraphs

benchmarks/Cargo.toml

@@ -55,7 +55,7 @@ pprof = { version = "0.13", features = [
 
 [features]
 default = ["parallel", "mimalloc", "bench-metrics"]
-bench-metrics = ["openvm-native-recursion/bench-metrics"]
+bench-metrics = ["openvm-native-recursion/bench-metrics", "openvm-native-compiler/bench-metrics"]
 profiling = ["openvm-sdk/profiling"]
 aggregation = []
 static-verifier = ["openvm-native-recursion/static-verifier"]

benchmarks/src/utils.rs

@@ -122,6 +122,7 @@ impl BenchmarkCli {
             halo2_config: Halo2Config {
                 verifier_k: self.halo2_outer_k.unwrap_or(24),
                 wrapper_k: self.halo2_wrapper_k,
+                profiling: self.profiling,
             },
         }
     }

crates/sdk/Cargo.toml

@@ -50,6 +50,7 @@ default = ["parallel"]
 bench-metrics = [
     "openvm-circuit/bench-metrics",
     "openvm-native-recursion/bench-metrics",
+    "openvm-native-compiler/bench-metrics",
 ]
 profiling = ["openvm-circuit/function-span", "openvm-transpiler/function-span"]
 parallel = ["openvm-circuit/parallel"]

crates/sdk/src/config/mod.rs

@@ -49,6 +49,8 @@ pub struct Halo2Config {
     pub verifier_k: usize,
     /// If not specified, keygen will tune wrapper_k automatically.
     pub wrapper_k: Option<usize>,
+    /// Sets the profiling mode of halo2 VM
+    pub profiling: bool,
 }
 
 impl<VC> AppConfig<VC> {
@@ -101,6 +103,7 @@ impl Default for AggConfig {
             halo2_config: Halo2Config {
                 verifier_k: 24,
                 wrapper_k: None,
+                profiling: false,
             },
         }
     }

crates/sdk/src/keygen/mod.rs

@@ -76,6 +76,8 @@ pub struct Halo2ProvingKey {
     pub verifier: Halo2VerifierProvingKey,
     /// Wrapper circuit to verify static verifier and reduce the verification costs in the final proof.
     pub wrapper: Halo2WrapperProvingKey,
+    /// Whether to collect detailed profiling metrics
+    pub profiling: bool,
 }
 
 impl<VC: VmConfig<F>> AppProvingKey<VC>
@@ -315,7 +317,11 @@ impl AggProvingKey {
         } else {
             Halo2WrapperProvingKey::keygen_auto_tune(reader, dummy_snark)
         };
-        let halo2_pk = Halo2ProvingKey { verifier, wrapper };
+        let halo2_pk = Halo2ProvingKey {
+            verifier,
+            wrapper,
+            profiling: halo2_config.profiling,
+        };
         Self {
             agg_stark_pk,
             halo2_pk,

crates/sdk/src/prover/halo2.rs

@@ -30,8 +30,11 @@ impl Halo2Prover {
     pub fn prove_for_evm(&self, root_proof: &Proof<RootSC>) -> EvmProof {
         let mut witness = Witness::default();
         root_proof.write(&mut witness);
-        let snark = info_span!("prove", group = "halo2_outer")
-            .in_scope(|| self.halo2_pk.verifier.prove(&self.verifier_srs, witness));
+        let snark = info_span!("prove", group = "halo2_outer").in_scope(|| {
+            self.halo2_pk
+                .verifier
+                .prove(&self.verifier_srs, witness, self.halo2_pk.profiling)
+        });
         info_span!("prove_for_evm", group = "halo2_wrapper").in_scope(|| {
             self.halo2_pk
                 .wrapper

crates/sdk/tests/integration_test.rs

@@ -90,6 +90,7 @@ fn agg_config_for_test() -> AggConfig {
         halo2_config: Halo2Config {
             verifier_k: 24,
             wrapper_k: None,
+            profiling: false,
         },
     }
 }

extensions/native/recursion/src/halo2/mod.rs

@@ -199,14 +199,15 @@ impl Halo2Prover {
         pk: &ProvingKey<G1Affine>,
         dsl_operations: DslOperations<C>,
         witness: Witness<C>,
+        profiling: bool,
     ) -> Snark {
         let k = config_params.k;
         #[cfg(feature = "bench-metrics")]
         let start = std::time::Instant::now();
         let builder = Self::builder(CircuitBuilderStage::Prover, k)
             .use_params(config_params)
             .use_break_points(break_points);
-        let builder = Self::populate(builder, dsl_operations, witness, false);
+        let builder = Self::populate(builder, dsl_operations, witness, profiling);
         #[cfg(feature = "bench-metrics")]
         {
             let stats = builder.statistics();

extensions/native/recursion/src/halo2/testing_utils.rs

@@ -57,7 +57,7 @@ pub fn run_static_verifier_test(
     .entered();
     let mut witness = Witness::default();
     vparams.data.proof.write(&mut witness);
-    let static_verifier_snark = stark_verifier_circuit.prove(params, witness);
+    let static_verifier_snark = stark_verifier_circuit.prove(params, witness, false);
     info_span.exit();
     (stark_verifier_circuit, static_verifier_snark)
 }

extensions/native/recursion/src/halo2/verifier.rs

@@ -39,14 +39,20 @@ pub fn generate_halo2_verifier_proving_key(
 }
 
 impl Halo2VerifierProvingKey {
-    pub fn prove(&self, params: &Halo2Params, witness: Witness<OuterConfig>) -> Snark {
+    pub fn prove(
+        &self,
+        params: &Halo2Params,
+        witness: Witness<OuterConfig>,
+        profiling: bool,
+    ) -> Snark {
         Halo2Prover::prove(
             params,
             self.pinning.metadata.config_params.clone(),
             self.pinning.metadata.break_points.clone(),
             &self.pinning.pk,
             self.dsl_ops.clone(),
             witness,
+            profiling,
         )
     }
     // TODO: Add verify method