fix: take program air height into account in metered execution (#2137)

- add `ProgramAir` height as a constant trace height in metered
execution

doesn't seem to affect the
[benchmark](https://github.com/axiom-crypto/openvm-reth-benchmark/actions/runs/17747443362)
which is good

Author: shuklaayush

CHANGELOG.md

@@ -5,6 +5,11 @@ All notable changes to OpenVM will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project follows a versioning principles documented in [VERSIONING.md](./VERSIONING.md).
 
+## Unreleased
+
+### Changed
+- (Executor) Modified `VirtualMachine::build_metered_ctx` to take the program (`&VmExe<Val<E::SC>>`) as an argument.
+
 ## v1.4.0 (2025-09-01)
 
 ### Added

benchmarks/execute/benches/execute.rs

@@ -16,9 +16,7 @@ use openvm_bigint_circuit::{Int256, Int256CpuProverExt, Int256Executor};
 use openvm_bigint_transpiler::Int256TranspilerExtension;
 use openvm_circuit::{
     arch::{
-        execution_mode::{MeteredCostCtx, MeteredCtx},
-        instructions::exe::VmExe,
-        interpreter::InterpretedInstance,
+        execution_mode::MeteredCostCtx, instructions::exe::VmExe, interpreter::InterpretedInstance,
         ContinuationVmProof, *,
     },
     derive::VmConfig,
@@ -57,6 +55,7 @@ use openvm_stark_sdk::{
     openvm_stark_backend::{
         self,
         config::{StarkGenericConfig, Val},
+        keygen::types::MultiStarkProvingKey,
         p3_field::PrimeField32,
         proof::Proof,
         prover::{
@@ -86,7 +85,7 @@ const APP_PROGRAMS: &[&str] = &[
 const LEAF_VERIFIER_PROGRAMS: &[&str] = &["kitchen-sink"];
 const INTERNAL_VERIFIER_PROGRAMS: &[&str] = &["fibonacci"];
 
-static METERED_CTX: OnceLock<(MeteredCtx, Vec<usize>)> = OnceLock::new();
+static VM_PROVING_KEY: OnceLock<MultiStarkProvingKey<SC>> = OnceLock::new();
 static METERED_COST_CTX: OnceLock<(MeteredCostCtx, Vec<usize>)> = OnceLock::new();
 static EXECUTOR: OnceLock<VmExecutor<BabyBear, ExecuteConfig>> = OnceLock::new();
 
@@ -231,22 +230,22 @@ fn load_program_executable(program: &str) -> Result<VmExe<BabyBear>> {
     Ok(VmExe::from_elf(elf, transpiler)?)
 }
 
-fn metering_setup() -> &'static (MeteredCtx, Vec<usize>) {
-    METERED_CTX.get_or_init(|| {
+fn vm_proving_key() -> &'static MultiStarkProvingKey<SC> {
+    VM_PROVING_KEY.get_or_init(|| {
         let config = ExecuteConfig::default();
         let engine = BabyBearPoseidon2Engine::new(FriParameters::standard_fast());
-        let (vm, _) = VirtualMachine::new_with_keygen(engine, ExecuteBuilder, config).unwrap();
-        let ctx = vm.build_metered_ctx();
-        let executor_idx_to_air_idx = vm.executor_idx_to_air_idx();
-        (ctx, executor_idx_to_air_idx)
+        let circuit = config.create_airs().expect("Failed to create AIRs");
+        circuit.keygen(&engine)
     })
 }
 
 fn metered_cost_setup() -> &'static (MeteredCostCtx, Vec<usize>) {
     METERED_COST_CTX.get_or_init(|| {
         let config = ExecuteConfig::default();
         let engine = BabyBearPoseidon2Engine::new(FriParameters::standard_fast());
-        let (vm, _) = VirtualMachine::new_with_keygen(engine, ExecuteBuilder, config).unwrap();
+        let pk = vm_proving_key();
+        let d_pk = engine.device().transport_pk_to_device(pk);
+        let vm = VirtualMachine::new(engine, ExecuteBuilder, config, d_pk).unwrap();
         let ctx = vm.build_metered_cost_ctx();
         let executor_idx_to_air_idx = vm.executor_idx_to_air_idx();
         (ctx, executor_idx_to_air_idx)
@@ -280,9 +279,16 @@ fn benchmark_execute_metered(bencher: Bencher, program: &str) {
     bencher
         .with_inputs(|| {
             let exe = load_program_executable(program).expect("Failed to load program executable");
-            let (ctx, executor_idx_to_air_idx) = metering_setup();
+            let config = ExecuteConfig::default();
+            let engine = BabyBearPoseidon2Engine::new(FriParameters::standard_fast());
+            let pk = vm_proving_key();
+            let d_pk = engine.device().transport_pk_to_device(pk);
+            let vm = VirtualMachine::new(engine, ExecuteBuilder, config, d_pk).unwrap();
+            let executor_idx_to_air_idx = vm.executor_idx_to_air_idx();
+
+            let ctx = vm.build_metered_ctx(&exe);
             let interpreter = executor()
-                .metered_instance(&exe, executor_idx_to_air_idx)
+                .metered_instance(&exe, &executor_idx_to_air_idx)
                 .unwrap();
             (interpreter, vec![], ctx.clone())
         })
@@ -403,7 +409,7 @@ fn benchmark_leaf_verifier_execute_metered(bencher: Bencher, program: &str) {
     bencher
         .with_inputs(|| {
             let (vm, leaf_exe, input_stream) = setup_leaf_verifier(program);
-            let ctx = vm.build_metered_ctx();
+            let ctx = vm.build_metered_ctx(&leaf_exe);
             let executor_idx_to_air_idx = vm.executor_idx_to_air_idx();
             let interpreter = vm
                 .executor()
@@ -459,7 +465,7 @@ fn benchmark_internal_verifier_execute_metered(bencher: Bencher, program: &str)
     bencher
         .with_inputs(|| {
             let (vm, internal_exe, input_stream) = setup_internal_verifier(program);
-            let ctx = vm.build_metered_ctx();
+            let ctx = vm.build_metered_ctx(&internal_exe);
             let executor_idx_to_air_idx = vm.executor_idx_to_air_idx();
             let interpreter = vm
                 .executor()

benchmarks/execute/src/execute-verifier.rs

@@ -332,7 +332,7 @@ fn execute_verifier(
         }
         ExecutionMode::Metered => {
             tracing::info!("Running metered execute...");
-            let ctx = vm.build_metered_ctx();
+            let ctx = vm.build_metered_ctx(exe);
             let interpreter = vm.metered_interpreter(exe)?;
             interpreter.execute_metered(input_stream, ctx)?;
         }

benchmarks/prove/src/bin/kitchen_sink.rs

@@ -31,7 +31,7 @@ fn verify_native_max_trace_heights(
     for leaf_input in leaf_inputs {
         let exe = leaf_prover.exe().clone();
         let vm = &mut leaf_prover.vm;
-        let metered_ctx = vm.build_metered_ctx();
+        let metered_ctx = vm.build_metered_ctx(&exe);
         let (segments, _) = vm
             .metered_interpreter(&exe)?
             .execute_metered(leaf_input.write_to_stream(), metered_ctx)?;

crates/sdk/src/lib.rs

@@ -375,7 +375,7 @@ where
         let vm = app_prover.vm();
         let exe = app_prover.exe();
 
-        let ctx = vm.build_metered_ctx();
+        let ctx = vm.build_metered_ctx(&exe);
         let interpreter = vm
             .metered_interpreter(&exe)
             .map_err(VirtualMachineError::from)?;

crates/vm/src/arch/execution_mode/metered/ctx.rs

@@ -1,5 +1,6 @@
 use std::num::NonZero;
 
+use itertools::Itertools;
 use openvm_instructions::riscv::{RV32_IMM_AS, RV32_REGISTER_AS};
 
 use super::{
@@ -152,17 +153,20 @@ impl<const PAGE_BITS: usize> MeteredCtx<PAGE_BITS> {
     }
 
     #[allow(dead_code)]
-    pub fn print_heights(&self) {
-        println!("{:>10} {:<30}", "Height", "Air Name");
-        println!("{}", "-".repeat(42));
-        for (i, height) in self.trace_heights.iter().enumerate() {
-            let air_name = self
-                .segmentation_ctx
-                .air_names
-                .get(i)
-                .map(|s| s.as_str())
-                .unwrap_or("Unknown");
-            println!("{:>10} {:<30}", height, air_name);
+    pub fn print_segment(&self) {
+        println!("{}", "-".repeat(80));
+        println!("Segment {}", self.segmentation_ctx.segments.len() - 1);
+        println!("{}", "-".repeat(80));
+        println!("{:>10} {:>10} {:<30}", "Width", "Height", "Air Name");
+        println!("{}", "-".repeat(80));
+        for ((&width, &height), air_name) in self
+            .segmentation_ctx
+            .widths
+            .iter()
+            .zip_eq(self.trace_heights.iter())
+            .zip_eq(self.segmentation_ctx.air_names.iter())
+        {
+            println!("{:>10} {:>10} {:<30}", width, height, air_name.as_str());
         }
     }
 }

crates/vm/src/arch/execution_mode/metered/segment_ctx.rs

@@ -40,7 +40,7 @@ impl Default for SegmentationLimits {
 pub struct SegmentationCtx {
     pub segments: Vec<Segment>,
     pub(crate) air_names: Vec<String>,
-    widths: Vec<usize>,
+    pub(crate) widths: Vec<usize>,
     interactions: Vec<usize>,
     pub(crate) segmentation_limits: SegmentationLimits,
     pub instret_last_segment_check: u64,
@@ -152,13 +152,13 @@ impl SegmentationCtx {
             return false;
         }
 
-        for (i, (height, is_constant)) in trace_heights
+        for (i, (&height, is_constant)) in trace_heights
             .iter()
             .zip(is_trace_height_constant.iter())
             .enumerate()
         {
             // Only segment if the height is not constant and exceeds the maximum height
-            if !is_constant && *height > self.segmentation_limits.max_trace_height {
+            if !is_constant && height > self.segmentation_limits.max_trace_height {
                 let air_name = &self.air_names[i];
                 tracing::info!(
                     "Segment {:2} | instret {:9} | chip {} ({}) height ({:8}) > max ({:8})",

crates/vm/src/arch/vm.rs

@@ -780,8 +780,11 @@ where
     }
 
     /// Convenience method to construct a [MeteredCtx] using data from the stored proving key.
-    pub fn build_metered_ctx(&self) -> MeteredCtx {
-        let (constant_trace_heights, air_names, widths, interactions): (
+    pub fn build_metered_ctx(&self, exe: &VmExe<Val<E::SC>>) -> MeteredCtx {
+        let config = self.config().as_ref();
+        let program_len = exe.program.num_defined_instructions();
+
+        let (mut constant_trace_heights, air_names, widths, interactions): (
             Vec<_>,
             Vec<_>,
             Vec<_>,
@@ -804,6 +807,13 @@ where
             })
             .multiunzip();
 
+        // Program trace is the same for all segments
+        constant_trace_heights[PROGRAM_AIR_ID] = Some(program_len);
+        if config.has_public_values_chip() {
+            // Public values chip is only present when there's a single segment
+            constant_trace_heights[PUBLIC_VALUES_AIR_ID] = Some(config.num_public_values);
+        }
+
         self.executor().build_metered_ctx(
             &constant_trace_heights,
             &air_names,
@@ -975,7 +985,7 @@ where
         let input = input.into();
         self.reset_state(input.clone());
         let vm = &mut self.vm;
-        let metered_ctx = vm.build_metered_ctx();
+        let metered_ctx = vm.build_metered_ctx(&self.exe);
         let metered_interpreter = vm.metered_interpreter(&self.exe)?;
         let (segments, _) = metered_interpreter.execute_metered(input, metered_ctx)?;
         let mut proofs = Vec::with_capacity(segments.len());

crates/vm/src/utils/stark_utils.rs

@@ -114,7 +114,7 @@ where
     let vk = pk.get_vk();
     let exe = exe.into();
     let input = input.into();
-    let metered_ctx = vm.build_metered_ctx();
+    let metered_ctx = vm.build_metered_ctx(&exe);
     let (segments, _) = vm
         .metered_interpreter(&exe)?
         .execute_metered(input.clone(), metered_ctx)?;

extensions/native/circuit/src/utils.rs

@@ -97,9 +97,9 @@ pub mod test_utils {
         let input = input_stream.into();
 
         let engine = E::new(FriParameters::new_for_testing(1));
-        let (vm, _) = VirtualMachine::new_with_keygen(engine, builder, config)?;
-        let ctx = vm.build_metered_ctx();
         let exe = VmExe::new(program);
+        let (vm, _) = VirtualMachine::new_with_keygen(engine, builder, config)?;
+        let ctx = vm.build_metered_ctx(&exe);
         let (mut segments, _) = vm
             .metered_interpreter(&exe)?
             .execute_metered(input.clone(), ctx)?;