From edb901b39ddb44b9a9230eec268c6b5c738b18a8 Mon Sep 17 00:00:00 2001
From: Jonathan Wang <31040440+jonathanpwang@users.noreply.github.com>
Date: Mon, 30 Dec 2024 22:24:57 -0500
Subject: [PATCH] fix(bench): verify_fibair use volatile memory

---
 Cargo.lock                          |  1 +
 benchmarks/src/bin/fib_e2e.rs       |  3 +-
 benchmarks/src/bin/verify_fibair.rs | 47 +++++++++----------
 crates/cli/src/commands/keygen.rs   |  2 +-
 crates/sdk/Cargo.toml               |  1 +
 crates/sdk/examples/sdk.rs          |  2 +-
 crates/sdk/src/keygen/mod.rs        |  3 +-
 crates/sdk/src/lib.rs               | 13 +++++-
 crates/sdk/src/prover/agg.rs        | 43 ++----------------
 crates/sdk/src/prover/app.rs        | 70 ++++++++++++++++-------------
 crates/sdk/src/prover/mod.rs        | 12 -----
 crates/sdk/src/prover/stark.rs      |  5 ---
 crates/sdk/src/prover/vm/local.rs   |  9 ++--
 13 files changed, 86 insertions(+), 125 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index f0d5f42543..2dff52eabe 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -4268,6 +4268,7 @@ dependencies = [
  "derivative",
  "derive_more 1.0.0",
  "eyre",
+ "getset",
  "itertools 0.13.0",
  "metrics",
  "openvm",
diff --git a/benchmarks/src/bin/fib_e2e.rs b/benchmarks/src/bin/fib_e2e.rs
index e1a5b328a3..50dd94414a 100644
--- a/benchmarks/src/bin/fib_e2e.rs
+++ b/benchmarks/src/bin/fib_e2e.rs
@@ -85,8 +85,7 @@ async fn main() -> Result<()> {
     stdin.write(&n);
     run_with_metric_collection("OUTPUT_PATH", || {
         let mut e2e_prover =
-            ContinuationProver::new(&halo2_params_reader, app_pk, app_committed_exe, full_agg_pk)
-                .with_profiling();
+            ContinuationProver::new(&halo2_params_reader, app_pk, app_committed_exe, full_agg_pk);
         e2e_prover.set_program_name("fib_e2e");
         let _proof = e2e_prover.generate_proof_for_evm(stdin);
     });
diff --git a/benchmarks/src/bin/verify_fibair.rs b/benchmarks/src/bin/verify_fibair.rs
index 7d0c92f901..323bff0c1e 100644
--- a/benchmarks/src/bin/verify_fibair.rs
+++ b/benchmarks/src/bin/verify_fibair.rs
@@ -1,15 +1,11 @@
 use clap::Parser;
 use eyre::Result;
-use openvm_benchmarks::utils::{bench_from_exe, BenchmarkCli};
+use openvm_benchmarks::utils::BenchmarkCli;
 use openvm_circuit::arch::instructions::program::DEFAULT_MAX_NUM_PUBLIC_VALUES;
 use openvm_native_circuit::NativeConfig;
 use openvm_native_compiler::conversion::CompilerOptions;
 use openvm_native_recursion::testing_utils::inner::build_verification_program;
-use openvm_sdk::config::AppConfig;
-/// Benchmark of aggregation VM performance.
-/// Proofs:
-/// 1. Prove Fibonacci AIR.
-/// 2. Verify the proof of 1. by execution VM program in STARK VM.
+use openvm_sdk::{config::AppConfig, prover::AppProver, Sdk};
 use openvm_stark_sdk::{
     bench::run_with_metric_collection,
     config::{baby_bear_poseidon2::BabyBearPoseidon2Engine, FriParameters},
@@ -17,8 +13,11 @@ use openvm_stark_sdk::{
     engine::StarkFriEngine,
     openvm_stark_backend::Chip,
 };
-use tracing::info_span;
 
+/// Benchmark of aggregation VM performance.
+/// Proofs:
+/// 1. Prove Fibonacci AIR.
+/// 2. Verify the proof of 1. by execution VM program in STARK VM.
 fn main() -> Result<()> {
     let cli_args = BenchmarkCli::parse();
     let app_log_blowup = cli_args.app_log_blowup.unwrap_or(2);
@@ -30,36 +29,34 @@ fn main() -> Result<()> {
         FriParameters::standard_with_100_bits_conjectured_security(app_log_blowup),
     );
 
-    run_with_metric_collection("OUTPUT_PATH", || {
+    run_with_metric_collection("OUTPUT_PATH", || -> Result<()> {
         // run_test tries to setup tracing, but it will be ignored since run_with_metric_collection already sets it.
         let vdata = engine
             .run_test(vec![fib_chip.generate_air_proof_input()])
             .unwrap();
-        let leaf_fri_params =
+        // Unlike other apps, this "app" does not have continuations enabled.
+        let app_fri_params =
             FriParameters::standard_with_100_bits_conjectured_security(agg_log_blowup);
-        // FIXME: this should be benchmarked as a single segment VM.
         let app_vm_config = NativeConfig::aggregation(
             DEFAULT_MAX_NUM_PUBLIC_VALUES,
-            leaf_fri_params.max_constraint_degree().min(7),
-        )
-        .with_continuations();
+            app_fri_params.max_constraint_degree().min(7),
+        );
         let compiler_options = CompilerOptions::default();
         let app_config = AppConfig {
-            app_fri_params: leaf_fri_params.into(),
+            app_fri_params: app_fri_params.into(),
             app_vm_config,
-            leaf_fri_params: leaf_fri_params.into(),
+            leaf_fri_params: app_fri_params.into(),
             compiler_options,
         };
-        info_span!("Verify Fibonacci AIR").in_scope(|| {
-            let (program, input_stream) = build_verification_program(vdata, compiler_options);
-            bench_from_exe(
-                "verify_fibair",
-                app_config,
-                program,
-                input_stream.into(),
-                false,
-            )
-        })
+        let (program, input_stream) = build_verification_program(vdata, compiler_options);
+        let sdk = Sdk;
+        let app_pk = sdk.app_keygen(app_config)?;
+        let app_vk = app_pk.get_app_vk();
+        let committed_exe = sdk.commit_app_exe(app_fri_params, program.into())?;
+        let prover = AppProver::new(app_pk.app_vm_pk, committed_exe);
+        let proof = prover.generate_app_proof_without_continuations(input_stream.into());
+        sdk.verify_app_proof_without_continuations(&app_vk, &proof)?;
+        Ok(())
     })?;
     Ok(())
 }
diff --git a/crates/cli/src/commands/keygen.rs b/crates/cli/src/commands/keygen.rs
index 678ed8d5ae..f21145f0bb 100644
--- a/crates/cli/src/commands/keygen.rs
+++ b/crates/cli/src/commands/keygen.rs
@@ -39,7 +39,7 @@ impl KeygenCmd {
     pub fn run(&self) -> Result<()> {
         let app_config = read_config_toml_or_default(&self.config)?;
         let app_pk = Sdk.app_keygen(app_config)?;
-        write_app_vk_to_file(app_pk.get_vk(), &self.vk_output)?;
+        write_app_vk_to_file(app_pk.get_app_vk(), &self.vk_output)?;
         write_app_pk_to_file(app_pk, &self.output)?;
         Ok(())
     }
diff --git a/crates/sdk/Cargo.toml b/crates/sdk/Cargo.toml
index 140a78df3c..ff6e307fdc 100644
--- a/crates/sdk/Cargo.toml
+++ b/crates/sdk/Cargo.toml
@@ -41,6 +41,7 @@ async-trait.workspace = true
 metrics.workspace = true
 tracing.workspace = true
 itertools.workspace = true
+getset.workspace = true
 
 [features]
 default = ["parallel"]
diff --git a/crates/sdk/examples/sdk.rs b/crates/sdk/examples/sdk.rs
index 588e8c8368..5c60053764 100644
--- a/crates/sdk/examples/sdk.rs
+++ b/crates/sdk/examples/sdk.rs
@@ -98,7 +98,7 @@ fn main() -> Result<(), Box<dyn std::error::Error>> {
 
     // ANCHOR: verification
     // 10. Verify your program
-    let app_vk = app_pk.get_vk();
+    let app_vk = app_pk.get_app_vk();
     sdk.verify_app_proof(&app_vk, &proof)?;
     // ANCHOR_END: verification
 
diff --git a/crates/sdk/src/keygen/mod.rs b/crates/sdk/src/keygen/mod.rs
index 3bdc8567d7..8f30cdcc57 100644
--- a/crates/sdk/src/keygen/mod.rs
+++ b/crates/sdk/src/keygen/mod.rs
@@ -91,7 +91,6 @@ where
                 vm_pk.max_constraint_degree
                     <= config.app_fri_params.fri_params.max_constraint_degree()
             );
-            assert!(config.app_vm_config.system().continuation_enabled);
             VmProvingKey {
                 fri_params: config.app_fri_params.fri_params,
                 vm_config: config.app_vm_config.clone(),
@@ -122,7 +121,7 @@ where
         self.app_vm_pk.vm_config.system().num_public_values
     }
 
-    pub fn get_vk(&self) -> AppVerifyingKey {
+    pub fn get_app_vk(&self) -> AppVerifyingKey {
         AppVerifyingKey {
             fri_params: self.app_vm_pk.fri_params,
             app_vm_vk: self.app_vm_pk.vm_pk.get_vk(),
diff --git a/crates/sdk/src/lib.rs b/crates/sdk/src/lib.rs
index 433d0acac7..a4038e194a 100644
--- a/crates/sdk/src/lib.rs
+++ b/crates/sdk/src/lib.rs
@@ -21,7 +21,7 @@ use openvm_native_recursion::{
     },
     types::InnerConfig,
 };
-use openvm_stark_backend::engine::StarkEngine;
+use openvm_stark_backend::{engine::StarkEngine, prover::types::Proof};
 use openvm_stark_sdk::{
     config::{
         baby_bear_poseidon2::{BabyBearPoseidon2Config, BabyBearPoseidon2Engine},
@@ -161,10 +161,19 @@ impl Sdk {
         for seg_proof in &proof.per_segment {
             e.verify(&app_vk.app_vm_vk, seg_proof)?
         }
-        // TODO: verify continuation.
+        // TODO: verify continuation: requires App VC.
         Ok(())
     }
 
+    pub fn verify_app_proof_without_continuations(
+        &self,
+        app_vk: &AppVerifyingKey,
+        proof: &Proof<SC>,
+    ) -> Result<(), VerificationError> {
+        let e = BabyBearPoseidon2Engine::new(app_vk.fri_params);
+        e.verify(&app_vk.app_vm_vk, proof)
+    }
+
     pub fn agg_keygen(
         &self,
         config: AggConfig,
diff --git a/crates/sdk/src/prover/agg.rs b/crates/sdk/src/prover/agg.rs
index 5fe5921084..faebede4d4 100644
--- a/crates/sdk/src/prover/agg.rs
+++ b/crates/sdk/src/prover/agg.rs
@@ -2,11 +2,10 @@ use std::sync::Arc;
 
 #[cfg(feature = "bench-metrics")]
 use openvm_circuit::arch::SingleSegmentVmExecutor;
-use openvm_circuit::arch::Streams;
 use openvm_native_circuit::NativeConfig;
 use openvm_native_recursion::hints::Hintable;
 use openvm_stark_sdk::{
-    config::baby_bear_poseidon2::BabyBearPoseidon2Engine, engine::StarkFriEngine,
+    config::baby_bear_poseidon2::BabyBearPoseidon2Engine,
     openvm_stark_backend::prover::types::Proof,
 };
 use tracing::info_span;
@@ -37,13 +36,10 @@ pub struct AggStarkProver {
 
     pub num_children_internal: usize,
     pub max_internal_wrapper_layers: usize,
-
-    pub profile: bool,
 }
 pub struct LeafProver {
     prover: VmLocalProver<SC, NativeConfig, BabyBearPoseidon2Engine>,
     pub num_children_leaf: usize,
-    pub profile: bool,
 }
 
 impl AggStarkProver {
@@ -63,7 +59,6 @@ impl AggStarkProver {
             root_prover,
             num_children_internal: DEFAULT_NUM_CHILDREN_INTERNAL,
             max_internal_wrapper_layers: DEFAULT_MAX_INTERNAL_WRAPPER_LAYERS,
-            profile: false,
         }
     }
 
@@ -82,16 +77,6 @@ impl AggStarkProver {
         self
     }
 
-    pub fn set_profile(&mut self, profile: bool) -> &mut Self {
-        self.profile = profile;
-        self.leaf_prover.profile = profile;
-        self
-    }
-    pub fn with_profiling(mut self) -> Self {
-        self.set_profile(true);
-        self
-    }
-
     /// Generate a proof to aggregate app proofs.
     pub fn generate_agg_proof(&self, app_proofs: ContinuationVmProof<SC>) -> Proof<RootSC> {
         let leaf_proofs = self.leaf_prover.generate_proof(&app_proofs);
@@ -156,7 +141,7 @@ impl AggStarkProver {
                             hgt = internal_node_height
                         )
                         .in_scope(|| {
-                            single_segment_prove(&self.internal_prover, input.write(), self.profile)
+                            SingleSegmentVmProver::prove(&self.internal_prover, input.write())
                         })
                     })
                     .collect()
@@ -207,17 +192,12 @@ impl LeafProver {
         Self {
             prover,
             num_children_leaf: DEFAULT_NUM_CHILDREN_LEAF,
-            profile: false,
         }
     }
     pub fn with_num_children_leaf(mut self, num_children_leaf: usize) -> Self {
         self.num_children_leaf = num_children_leaf;
         self
     }
-    pub fn with_profile(mut self) -> Self {
-        self.profile = true;
-        self
-    }
     pub fn generate_proof(&self, app_proofs: &ContinuationVmProof<SC>) -> Vec<Proof<SC>> {
         info_span!("leaf verifier", group = "leaf").in_scope(|| {
             #[cfg(feature = "bench-metrics")]
@@ -232,7 +212,7 @@ impl LeafProver {
                 .enumerate()
                 .map(|(leaf_node_idx, input)| {
                     info_span!("leaf verifier proof", idx = leaf_node_idx).in_scope(|| {
-                        single_segment_prove(&self.prover, input.write_to_stream(), self.profile)
+                        SingleSegmentVmProver::prove(&self.prover, input.write_to_stream())
                     })
                 })
                 .collect::<Vec<_>>()
@@ -240,23 +220,6 @@ impl LeafProver {
     }
 }
 
-#[allow(unused)]
-fn single_segment_prove<E: StarkFriEngine<SC>>(
-    prover: &VmLocalProver<SC, NativeConfig, E>,
-    input: impl Into<Streams<F>> + Clone,
-    profile: bool,
-) -> Proof<SC> {
-    #[cfg(feature = "bench-metrics")]
-    if profile {
-        let mut vm_config = prover.pk.vm_config.clone();
-        vm_config.system.profiling = true;
-        let vm = SingleSegmentVmExecutor::new(vm_config);
-        vm.execute(prover.committed_exe.exe.clone(), input.clone())
-            .unwrap();
-    }
-    SingleSegmentVmProver::prove(prover, input)
-}
-
 fn heights_le(a: &[usize], b: &[usize]) -> bool {
     assert_eq!(a.len(), b.len());
     a.iter().zip(b.iter()).all(|(a, b)| a <= b)
diff --git a/crates/sdk/src/prover/app.rs b/crates/sdk/src/prover/app.rs
index 2e9555ad24..6ad7a2ded2 100644
--- a/crates/sdk/src/prover/app.rs
+++ b/crates/sdk/src/prover/app.rs
@@ -1,12 +1,14 @@
 use std::sync::Arc;
 
+use getset::Getters;
 use openvm_circuit::arch::VmConfig;
 #[cfg(feature = "bench-metrics")]
 use openvm_circuit::arch::{instructions::exe::VmExe, VmExecutor};
-use openvm_stark_backend::Chip;
+use openvm_stark_backend::{prover::types::Proof, Chip};
 use openvm_stark_sdk::config::baby_bear_poseidon2::BabyBearPoseidon2Engine;
 use tracing::info_span;
 
+use super::vm::SingleSegmentVmProver;
 use crate::{
     prover::vm::{
         local::VmLocalProver, types::VmProvingKey, ContinuationVmProof, ContinuationVmProver,
@@ -14,11 +16,10 @@ use crate::{
     NonRootCommittedExe, StdIn, F, SC,
 };
 
+#[derive(Getters)]
 pub struct AppProver<VC> {
-    /// If true, will run execution once with full metric collection for
-    /// flamegraphs (WARNING: this degrades performance).
-    pub profile: bool,
     pub program_name: Option<String>,
+    #[getset(get = "pub")]
     app_prover: VmLocalProver<SC, VC, BabyBearPoseidon2Engine>,
 }
 
@@ -31,7 +32,6 @@ impl<VC> AppProver<VC> {
         VC: VmConfig<F>,
     {
         Self {
-            profile: false,
             program_name: None,
             app_prover: VmLocalProver::<SC, VC, BabyBearPoseidon2Engine>::new(
                 app_vm_pk,
@@ -39,14 +39,6 @@ impl<VC> AppProver<VC> {
             ),
         }
     }
-    pub fn set_profile(&mut self, profile: bool) -> &mut Self {
-        self.profile = profile;
-        self
-    }
-    pub fn with_profiling(mut self) -> Self {
-        self.set_profile(true);
-        self
-    }
     pub fn set_program_name(&mut self, program_name: impl AsRef<str>) -> &mut Self {
         self.program_name = Some(program_name.as_ref().to_string());
         self
@@ -56,12 +48,17 @@ impl<VC> AppProver<VC> {
         self
     }
 
+    /// Generates proof for every continuation segment
     pub fn generate_app_proof(&self, input: StdIn) -> ContinuationVmProof<SC>
     where
         VC: VmConfig<F>,
         VC::Executor: Chip<SC>,
         VC::Periphery: Chip<SC>,
     {
+        assert!(
+            self.vm_config().system().continuation_enabled,
+            "Use generate_app_proof_without_continuations instead."
+        );
         info_span!(
             "app proof",
             group = self
@@ -70,29 +67,40 @@ impl<VC> AppProver<VC> {
                 .unwrap_or(&"app_proof".to_string())
         )
         .in_scope(|| {
-            #[cfg(feature = "bench-metrics")]
-            if self.profile {
-                emit_app_execution_metrics(
-                    self.app_prover.pk.vm_config.clone(),
-                    self.app_prover.committed_exe.exe.clone(),
-                    input.clone(),
-                );
-            }
             #[cfg(feature = "bench-metrics")]
             metrics::counter!("fri.log_blowup")
                 .absolute(self.app_prover.pk.fri_params.log_blowup as u64);
             ContinuationVmProver::prove(&self.app_prover, input)
         })
     }
-}
 
-#[cfg(feature = "bench-metrics")]
-fn emit_app_execution_metrics<VC: VmConfig<F>>(mut vm_config: VC, exe: VmExe<F>, input: StdIn)
-where
-    VC::Executor: Chip<SC>,
-    VC::Periphery: Chip<SC>,
-{
-    vm_config.system_mut().profiling = true;
-    let vm = VmExecutor::new(vm_config);
-    vm.execute_segments(exe, input).unwrap();
+    pub fn generate_app_proof_without_continuations(&self, input: StdIn) -> Proof<SC>
+    where
+        VC: VmConfig<F>,
+        VC::Executor: Chip<SC>,
+        VC::Periphery: Chip<SC>,
+    {
+        assert!(
+            !self.vm_config().system().continuation_enabled,
+            "Use generate_app_proof instead."
+        );
+        info_span!(
+            "app proof",
+            group = self
+                .program_name
+                .as_ref()
+                .unwrap_or(&"app_proof".to_string())
+        )
+        .in_scope(|| {
+            #[cfg(feature = "bench-metrics")]
+            metrics::counter!("fri.log_blowup")
+                .absolute(self.app_prover.pk.fri_params.log_blowup as u64);
+            SingleSegmentVmProver::prove(&self.app_prover, input)
+        })
+    }
+
+    /// App VM config
+    pub fn vm_config(&self) -> &VC {
+        self.app_prover.vm_config()
+    }
 }
diff --git a/crates/sdk/src/prover/mod.rs b/crates/sdk/src/prover/mod.rs
index 72b6f1921c..08c602f032 100644
--- a/crates/sdk/src/prover/mod.rs
+++ b/crates/sdk/src/prover/mod.rs
@@ -54,18 +54,6 @@ impl<VC> ContinuationProver<VC> {
         }
     }
 
-    /// Flag for enabling/disabling profiling.
-    pub fn set_profile(&mut self, profile: bool) -> &mut Self {
-        self.stark_prover.set_profile(profile);
-        // halo2 profiling is set in CompilerConfig when creating Halo2ProvingKey
-        self
-    }
-
-    pub fn with_profiling(mut self) -> Self {
-        self.set_profile(true);
-        self
-    }
-
     pub fn set_program_name(&mut self, program_name: impl AsRef<str>) -> &mut Self {
         self.stark_prover.set_program_name(program_name);
         self
diff --git a/crates/sdk/src/prover/stark.rs b/crates/sdk/src/prover/stark.rs
index 8f07aab3b8..87f5aa2cd8 100644
--- a/crates/sdk/src/prover/stark.rs
+++ b/crates/sdk/src/prover/stark.rs
@@ -37,11 +37,6 @@ impl<VC> StarkProver<VC> {
             agg_prover: AggStarkProver::new(agg_stark_pk, app_pk.leaf_committed_exe.clone()),
         }
     }
-    pub fn set_profile(&mut self, profile: bool) -> &mut Self {
-        self.app_prover.set_profile(profile);
-        self.agg_prover.set_profile(profile);
-        self
-    }
     pub fn set_program_name(&mut self, program_name: impl AsRef<str>) -> &mut Self {
         self.app_prover.set_program_name(program_name);
         self
diff --git a/crates/sdk/src/prover/vm/local.rs b/crates/sdk/src/prover/vm/local.rs
index 9f435e63ab..61dc5a1513 100644
--- a/crates/sdk/src/prover/vm/local.rs
+++ b/crates/sdk/src/prover/vm/local.rs
@@ -28,10 +28,7 @@ pub struct VmLocalProver<SC: StarkGenericConfig, VC, E: StarkFriEngine<SC>> {
     _marker: PhantomData<E>,
 }
 
-impl<SC: StarkGenericConfig, VC: VmConfig<Val<SC>>, E: StarkFriEngine<SC>> VmLocalProver<SC, VC, E>
-where
-    Val<SC>: PrimeField32,
-{
+impl<SC: StarkGenericConfig, VC, E: StarkFriEngine<SC>> VmLocalProver<SC, VC, E> {
     pub fn new(pk: Arc<VmProvingKey<SC, VC>>, committed_exe: Arc<VmCommittedExe<SC>>) -> Self {
         Self {
             pk,
@@ -57,6 +54,10 @@ where
     pub fn set_override_trace_heights(&mut self, overridden_heights: VmComplexTraceHeights) {
         self.overridden_heights = Some(overridden_heights);
     }
+
+    pub fn vm_config(&self) -> &VC {
+        &self.pk.vm_config
+    }
 }
 
 impl<SC: StarkGenericConfig, VC: VmConfig<Val<SC>>, E: StarkFriEngine<SC>> ContinuationVmProver<SC>