feat: use cached max trace heights in leaf/internal aggregation (#1839)

Author: stephenh-axiom-xyz

Cargo.lock

@@ -10,6 +10,7 @@ license.workspace = true
 [dependencies]
 openvm-benchmarks-utils.workspace = true
 openvm-circuit.workspace = true
+openvm-continuations.workspace = true
 openvm-sdk.workspace = true
 openvm-stark-backend.workspace = true
 openvm-stark-sdk.workspace = true

benchmarks/prove/Cargo.toml

@@ -5,22 +5,85 @@ use eyre::Result;
 use num_bigint::BigUint;
 use openvm_algebra_circuit::{Fp2Extension, ModularExtension};
 use openvm_benchmarks_prove::util::BenchmarkCli;
-use openvm_circuit::arch::{instructions::exe::VmExe, SystemConfig};
+use openvm_circuit::arch::{instructions::exe::VmExe, SingleSegmentVmExecutor, SystemConfig};
+use openvm_continuations::verifier::leaf::types::LeafVmVerifierInput;
 use openvm_ecc_circuit::{WeierstrassExtension, P256_CONFIG, SECP256K1_CONFIG};
+use openvm_native_circuit::{NativeConfig, NATIVE_MAX_TRACE_HEIGHTS};
 use openvm_native_recursion::halo2::utils::{CacheHalo2ParamsReader, DEFAULT_PARAMS_DIR};
 use openvm_pairing_circuit::{PairingCurve, PairingExtension};
 use openvm_pairing_guest::{
     bls12_381::BLS12_381_COMPLEX_STRUCT_NAME, bn254::BN254_COMPLEX_STRUCT_NAME,
 };
 use openvm_sdk::{
-    commit::commit_app_exe, config::SdkVmConfig, prover::EvmHalo2Prover,
-    DefaultStaticVerifierPvHandler, Sdk, StdIn,
+    commit::commit_app_exe,
+    config::SdkVmConfig,
+    keygen::AppProvingKey,
+    prover::{vm::types::VmProvingKey, EvmHalo2Prover},
+    DefaultStaticVerifierPvHandler, NonRootCommittedExe, Sdk, StdIn, SC,
 };
 use openvm_stark_sdk::{
     bench::run_with_metric_collection, config::baby_bear_poseidon2::BabyBearPoseidon2Engine,
 };
 use openvm_transpiler::FromElf;
 
+fn verify_native_max_trace_heights(
+    sdk: &Sdk,
+    app_pk: Arc<AppProvingKey<SdkVmConfig>>,
+    app_committed_exe: Arc<NonRootCommittedExe>,
+    leaf_vm_pk: Arc<VmProvingKey<SC, NativeConfig>>,
+    num_children_leaf: usize,
+) -> Result<()> {
+    let app_proof =
+        sdk.generate_app_proof(app_pk.clone(), app_committed_exe.clone(), StdIn::default())?;
+    let leaf_inputs =
+        LeafVmVerifierInput::chunk_continuation_vm_proof(&app_proof, num_children_leaf);
+    let vm_vk = leaf_vm_pk.vm_pk.get_vk();
+
+    leaf_inputs.iter().for_each(|leaf_input| {
+        let executor = {
+            let mut executor = SingleSegmentVmExecutor::new(leaf_vm_pk.vm_config.clone());
+            executor
+                .set_trace_height_constraints(leaf_vm_pk.vm_pk.trace_height_constraints.clone());
+            executor
+        };
+        let max_trace_heights = executor
+            .execute_metered(
+                app_pk.leaf_committed_exe.exe.clone(),
+                leaf_input.write_to_stream(),
+                &vm_vk.total_widths(),
+                &vm_vk.num_interactions(),
+            )
+            .expect("execute_metered failed");
+        println!("max_trace_heights: {:?}", max_trace_heights);
+
+        let actual_trace_heights = executor
+            .execute_and_generate(
+                app_pk.leaf_committed_exe.clone(),
+                leaf_input.write_to_stream(),
+                &max_trace_heights,
+            )
+            .expect("execute_and_generate failed")
+            .per_air
+            .iter()
+            .map(|(_, air)| air.raw.height())
+            .collect::<Vec<usize>>();
+        println!("actual_trace_heights: {:?}", actual_trace_heights);
+
+        actual_trace_heights
+            .iter()
+            .zip(NATIVE_MAX_TRACE_HEIGHTS)
+            .for_each(|(&actual, &expected)| {
+                assert!(
+                    actual <= (expected as usize),
+                    "Actual trace height {} exceeds expected height {}",
+                    actual,
+                    expected
+                );
+            });
+    });
+    Ok(())
+}
+
 fn main() -> Result<()> {
     let args = BenchmarkCli::parse();
 
@@ -88,6 +151,15 @@ fn main() -> Result<()> {
         &DefaultStaticVerifierPvHandler,
     )?;
 
+    // Verify that NATIVE_MAX_TRACE_HEIGHTS remains valid
+    verify_native_max_trace_heights(
+        &sdk,
+        app_pk.clone(),
+        app_committed_exe.clone(),
+        full_agg_pk.agg_stark_pk.leaf_vm_pk.clone(),
+        args.agg_tree_config.num_children_leaf,
+    )?;
+
     run_with_metric_collection("OUTPUT_PATH", || -> Result<()> {
         let mut prover = EvmHalo2Prover::<_, BabyBearPoseidon2Engine>::new(
             &halo2_params_reader,

benchmarks/prove/src/bin/kitchen_sink.rs

@@ -207,12 +207,8 @@ where
         vm_heights
     };
     // For the dummy proof, we must override the trace heights.
-    let app_prover =
-        VmLocalProver::<SC, VC, BabyBearPoseidon2Engine>::new_with_overridden_trace_heights(
-            app_vm_pk,
-            dummy_exe,
-            Some(overridden_heights),
-        );
+    let app_prover = VmLocalProver::<SC, VC, BabyBearPoseidon2Engine>::new(app_vm_pk, dummy_exe)
+        .with_overridden_continuation_trace_heights(overridden_heights);
     ContinuationVmProver::prove(&app_prover, vec![])
 }
 

crates/sdk/src/keygen/dummy.rs

@@ -6,7 +6,7 @@ use openvm_continuations::verifier::{
     leaf::types::LeafVmVerifierInput,
     root::types::RootVmVerifierInput,
 };
-use openvm_native_circuit::NativeConfig;
+use openvm_native_circuit::{NativeConfig, NATIVE_MAX_TRACE_HEIGHTS};
 use openvm_native_compiler::ir::DIGEST_SIZE;
 use openvm_native_recursion::hints::Hintable;
 use openvm_stark_sdk::{engine::StarkFriEngine, openvm_stark_backend::proof::Proof};
@@ -45,14 +45,16 @@ impl<E: StarkFriEngine<SC>> AggStarkProver<E> {
         tree_config: AggregationTreeConfig,
     ) -> Self {
         let leaf_prover =
-            VmLocalProver::<SC, NativeConfig, E>::new(agg_stark_pk.leaf_vm_pk, leaf_committed_exe);
+            VmLocalProver::<SC, NativeConfig, E>::new(agg_stark_pk.leaf_vm_pk, leaf_committed_exe)
+                .with_overridden_single_segment_trace_heights(NATIVE_MAX_TRACE_HEIGHTS.to_vec());
         let leaf_controller = LeafProvingController {
             num_children: tree_config.num_children_leaf,
         };
         let internal_prover = VmLocalProver::<SC, NativeConfig, E>::new(
             agg_stark_pk.internal_vm_pk,
             agg_stark_pk.internal_committed_exe,
-        );
+        )
+        .with_overridden_single_segment_trace_heights(NATIVE_MAX_TRACE_HEIGHTS.to_vec());
         let root_prover = RootVerifierLocalProver::new(agg_stark_pk.root_verifier_pk);
         Self {
             leaf_prover,

crates/sdk/src/prover/agg.rs

@@ -27,7 +27,8 @@ use crate::prover::vm::{
 pub struct VmLocalProver<SC: StarkGenericConfig, VC, E: StarkFriEngine<SC>> {
     pub pk: Arc<VmProvingKey<SC, VC>>,
     pub committed_exe: Arc<VmCommittedExe<SC>>,
-    overridden_heights: Option<VmComplexTraceHeights>,
+    continuation_heights: Option<VmComplexTraceHeights>,
+    single_segment_heights: Option<Vec<u32>>,
     _marker: PhantomData<E>,
 }
 
@@ -36,26 +37,26 @@ impl<SC: StarkGenericConfig, VC, E: StarkFriEngine<SC>> VmLocalProver<SC, VC, E>
         Self {
             pk,
             committed_exe,
-            overridden_heights: None,
+            continuation_heights: None,
+            single_segment_heights: None,
             _marker: PhantomData,
         }
     }
 
-    pub fn new_with_overridden_trace_heights(
-        pk: Arc<VmProvingKey<SC, VC>>,
-        committed_exe: Arc<VmCommittedExe<SC>>,
-        overridden_heights: Option<VmComplexTraceHeights>,
+    pub fn with_overridden_continuation_trace_heights(
+        mut self,
+        overridden_heights: VmComplexTraceHeights,
     ) -> Self {
-        Self {
-            pk,
-            committed_exe,
-            overridden_heights,
-            _marker: PhantomData,
-        }
+        self.continuation_heights = Some(overridden_heights);
+        self
     }
 
-    pub fn set_override_trace_heights(&mut self, overridden_heights: VmComplexTraceHeights) {
-        self.overridden_heights = Some(overridden_heights);
+    pub fn with_overridden_single_segment_trace_heights(
+        mut self,
+        overridden_heights: Vec<u32>,
+    ) -> Self {
+        self.single_segment_heights = Some(overridden_heights);
+        self
     }
 
     pub fn vm_config(&self) -> &VC {
@@ -81,7 +82,7 @@ where
         let mut vm = VirtualMachine::new_with_overridden_trace_heights(
             e,
             self.pk.vm_config.clone(),
-            self.overridden_heights.clone(),
+            self.continuation_heights.clone(),
         );
         vm.set_trace_height_constraints(trace_height_constraints.clone());
         let VmCommittedExe {
@@ -179,16 +180,20 @@ where
 
         let vm_vk = self.pk.vm_pk.get_vk();
         let input = input.into();
-        let max_trace_heights = executor
-            .execute_metered(
-                self.committed_exe.exe.clone(),
-                input.clone(),
-                &vm_vk.total_widths(),
-                &vm_vk.num_interactions(),
-            )
-            .expect("execute_metered failed");
+        let max_trace_heights = if let Some(overridden_heights) = &self.single_segment_heights {
+            overridden_heights
+        } else {
+            &executor
+                .execute_metered(
+                    self.committed_exe.exe.clone(),
+                    input.clone(),
+                    &vm_vk.total_widths(),
+                    &vm_vk.num_interactions(),
+                )
+                .expect("execute_metered failed")
+        };
         let proof_input = executor
-            .execute_and_generate(self.committed_exe.clone(), input, &max_trace_heights)
+            .execute_and_generate(self.committed_exe.clone(), input, max_trace_heights)
             .unwrap();
 
         let vm = VirtualMachine::new(e, executor.config);

crates/sdk/src/prover/vm/local.rs

@@ -175,9 +175,10 @@ where
         &self,
         state: &mut VmStateMut<F, GuestMemory, MeteredCtx>,
         instruction: &Instruction<F>,
-        _chip_index: usize,
+        chip_index: usize,
     ) -> Result<(), ExecutionError> {
         self.execute_e1(state, instruction)?;
+        state.ctx.trace_heights[chip_index] += 1;
 
         Ok(())
     }

crates/vm/src/system/phantom/mod.rs

@@ -264,10 +264,9 @@ where
         &self,
         state: &mut VmStateMut<F, GuestMemory, MeteredCtx>,
         instruction: &Instruction<F>,
-        chip_index: usize,
+        _chip_index: usize,
     ) -> Result<()> {
         self.execute_e1(state, instruction)?;
-        state.ctx.trace_heights[chip_index] += 1;
 
         Ok(())
     }

crates/vm/src/system/public_values/core.rs

@@ -471,3 +471,11 @@ impl Default for Rv32WithKernelsConfig {
 
 // Default implementation uses no init file
 impl InitFileGenerator for Rv32WithKernelsConfig {}
+
+// Pre-computed maximum trace heights for NativeConfig. Found by doubling
+// the actual trace heights of kitchen-sink leaf verification (except for
+// VariableRangeChecker, which has a fixed height).
+pub const NATIVE_MAX_TRACE_HEIGHTS: &[u32] = &[
+    4194304, 4, 128, 2097152, 8388608, 4194304, 262144, 2097152, 16777216, 2097152, 8388608,
+    262144, 2097152, 1048576, 4194304, 65536, 262144,
+];