fix(perf-metrics): perf instrumentation slow (#1936)

The per-instruction call to calculate memory access adapter heights is
very slow because it iterates linearly through the entire adapter
record. I changed it now to only do this at the end. The downside is
that we won't be able to see the contribution to access adapters on a
per-instruction basis, but this could be extracted more efficiently in a
different way.

Proof it works:
https://github.com/axiom-crypto/openvm-reth-benchmark/actions/runs/16767319820

Author: jonathanpwang

ci/scripts/metric_unify/flamegraph.py

@@ -60,6 +60,9 @@ def get_stack_lines(metrics_dict, group_by_kvs, stack_keys, metric_name, sum_met
                     function_symbols = [get_function_symbol(string_table, offset) for offset in symbol_offsets]
                     stack_values.extend(function_symbols)
             else:
+                # don't make a stack frame for empty label
+                if labels[key] == '':
+                    continue
                 stack_values.append(labels[key])
         if filter:
             continue

crates/vm/src/arch/interpreter_preflight.rs

@@ -53,6 +53,7 @@ where
             self.execute_instruction(state)?;
             state.instret += 1;
         }
+
         Ok(())
     }
 
@@ -106,7 +107,7 @@ where
 
         #[cfg(feature = "metrics")]
         {
-            crate::metrics::update_instruction_metrics(state, executor, pc_entry);
+            crate::metrics::update_instruction_metrics(state, executor, pc, pc_entry);
         }
 
         Ok(())

crates/vm/src/arch/vm.rs

@@ -446,9 +446,14 @@ where
         let mut exec_state = VmExecState::new(vm_state, ctx);
         execute_spanned!("execute_e3", instance, &mut exec_state)?;
         let filtered_exec_frequencies = instance.handler.filtered_execution_frequencies();
-        let mut memory = exec_state.vm_state.memory;
-        let touched_memory = memory.finalize::<Val<E::SC>>(system_config.continuation_enabled);
+        let touched_memory = exec_state
+            .vm_state
+            .memory
+            .finalize::<Val<E::SC>>(system_config.continuation_enabled);
+        #[cfg(feature = "perf-metrics")]
+        crate::metrics::end_segment_metrics(&mut exec_state);
 
+        let memory = exec_state.vm_state.memory;
         let to_state = ExecutionState::new(exec_state.vm_state.pc, memory.timestamp());
         let public_values = system_config
             .has_public_values_chip()
@@ -507,7 +512,12 @@ where
             state.metrics.debug_infos = exe.program.debug_infos();
         }
         #[cfg(feature = "perf-metrics")]
-        state.metrics.set_pk_info(&self.pk);
+        {
+            state.metrics.set_pk_info(&self.pk);
+            state.metrics.num_sys_airs = self.config().as_ref().num_airs();
+            state.metrics.access_adapter_offset =
+                self.config().as_ref().access_adapter_air_id_offset();
+        }
         state
     }
 
@@ -1195,32 +1205,31 @@ mod vm_metrics {
     use metrics::counter;
 
     use super::*;
-    use crate::metrics::get_dyn_trace_heights_from_arenas;
+    use crate::arch::Arena;
 
     impl<E, VB> VirtualMachine<E, VB>
     where
         E: StarkEngine,
         VB: VmBuilder<E>,
     {
-        /// See [`metrics::get_trace_heights_from_arenas`](crate::metrics::get_trace_heights_from_arenas).
-        /// In addition, this function includes the constant trace heights and the used height of
-        /// the program trace.
+        /// Assumed that `record_arenas` has length equal to number of AIRs.
+        ///
+        /// Best effort calculation of the used trace heights per chip without padding to powers of
+        /// two. This is best effort because some periphery chips may not have record arenas to
+        /// instrument. This function includes the constant trace heights, and the used height of
+        /// the program trace. It does not include the memory access adapter trace heights,
+        /// which is included in `SystemChipComplex::finalize_trace_heights`.
         pub(crate) fn get_trace_heights_from_arenas(
             &self,
             system_records: &SystemRecords<Val<E::SC>>,
             record_arenas: &[VB::RecordArena],
         ) -> Vec<usize> {
             let num_airs = self.num_airs();
             assert_eq!(num_airs, record_arenas.len());
-            let sys_config = self.config().as_ref();
-            let num_sys_airs = sys_config.num_airs();
-            let access_adapter_offset = sys_config.access_adapter_air_id_offset();
-            let mut heights = get_dyn_trace_heights_from_arenas::<Val<E::SC>, _>(
-                num_sys_airs,
-                access_adapter_offset,
-                &system_records.access_adapter_records,
-                record_arenas,
-            );
+            let mut heights: Vec<usize> = record_arenas
+                .iter()
+                .map(|arena| arena.current_trace_height())
+                .collect();
             // If there are any constant trace heights, set them
             for (pk, height) in zip(&self.pk.per_air, &mut heights) {
                 if let Some(constant_height) =
@@ -1250,7 +1259,8 @@ mod vm_metrics {
             counter!("main_cells_used").absolute(main_cells_used as u64);
             counter!("total_cells_used").absolute(total_cells_used as u64);
 
-            if self.config().as_ref().profiling {
+            #[cfg(feature = "perf-metrics")]
+            {
                 for (name, value) in zip(self.air_names(), heights) {
                     let labels = [("air_name", name.to_string())];
                     counter!("rows_used", &labels).absolute(*value as u64);

crates/vm/src/metrics/mod.rs

@@ -11,14 +11,8 @@ use openvm_instructions::{
 use openvm_stark_backend::prover::{hal::ProverBackend, types::DeviceMultiStarkProvingKey};
 
 use crate::{
-    arch::{
-        execution_mode::tracegen::TracegenCtx, Arena, DenseRecordArena, PreflightExecutor,
-        VmExecState,
-    },
-    system::{
-        memory::{adapter::AccessAdapterInventory, online::TracingMemory},
-        program::PcEntry,
-    },
+    arch::{execution_mode::tracegen::TracegenCtx, Arena, PreflightExecutor, VmExecState},
+    system::{memory::online::TracingMemory, program::PcEntry},
 };
 
 pub mod cycle_tracker;
@@ -60,6 +54,7 @@ pub struct VmMetrics {
 pub fn update_instruction_metrics<F, RA, Executor>(
     state: &mut VmExecState<F, TracingMemory, TracegenCtx<RA>>,
     executor: &mut Executor,
+    prev_pc: u32, // the pc of the instruction executed, state.pc is next pc
     pc_entry: &PcEntry<F>,
 ) where
     F: Clone + Send + Sync,
@@ -80,59 +75,68 @@ pub fn update_instruction_metrics<F, RA, Executor>(
         let opcode = pc_entry.insn.opcode;
         let opcode_name = executor.get_opcode_name(opcode.as_usize());
 
-        let num_sys_airs = state.metrics.num_sys_airs;
-        let access_adapter_offset = state.metrics.access_adapter_offset;
-        let debug_info = state.metrics.debug_infos.get(pc);
+        let debug_info = state.metrics.debug_infos.get(prev_pc);
         let dsl_instr = debug_info.as_ref().map(|info| info.dsl_instruction.clone());
 
-        let now_trace_heights = get_dyn_trace_heights_from_arenas::<F, _>(
-            num_sys_airs,
-            access_adapter_offset,
-            &state.memory.access_adapter_records,
-            &state.ctx.arenas,
-        );
-        let mut now_trace_cells = now_trace_heights;
-
-        let metrics = &mut state.metrics;
-        for (main_width, cell_count) in zip(&metrics.main_widths, &mut now_trace_cells) {
-            *cell_count *= main_width;
-        }
-        metrics.update_trace_cells(now_trace_cells, opcode_name, dsl_instr);
-
-        metrics.update_current_fn(pc);
+        let now_trace_heights: Vec<usize> = state
+            .ctx
+            .arenas
+            .iter()
+            .map(|arena| arena.current_trace_height())
+            .collect();
+        let now_trace_cells = zip(&state.metrics.main_widths, &now_trace_heights)
+            .map(|(main_width, h)| main_width * h)
+            .collect_vec();
+        state
+            .metrics
+            .update_trace_cells(now_trace_cells, opcode_name, dsl_instr);
+
+        state.metrics.update_current_fn(pc);
     }
 }
 
-/// Assumed that `record_arenas` has length equal to number of AIRs.
-///
-/// Best effort calculation of the used trace heights per chip without padding to powers of
-/// two. This is best effort because some periphery chips may not have record arenas
-/// to instrument.
-///
-/// Does not include constant trace heights or the used program trace height.
-pub(crate) fn get_dyn_trace_heights_from_arenas<F, RA>(
-    num_sys_airs: usize,
-    access_adapter_offset: usize,
-    access_adapter_records: &DenseRecordArena,
-    record_arenas: &[RA],
-) -> Vec<usize>
+// Memory access adapter height calculation is slow, so only do it if this is the end of
+// execution.
+// We also clear the current trace cell counts so there aren't negative diffs at the start of the
+// next segment.
+#[cfg(feature = "perf-metrics")]
+pub fn end_segment_metrics<F, RA>(state: &mut VmExecState<F, TracingMemory, TracegenCtx<RA>>)
 where
     F: Clone + Send + Sync,
     RA: Arena,
 {
-    // First, get used heights from record arenas
-    let mut heights: Vec<usize> = record_arenas
-        .iter()
-        .map(|arena| arena.current_trace_height())
-        .collect();
-    // Memory is special case, so extract the memory AIR's trace heights from the special
-    // arena
+    use std::iter::zip;
+
+    use crate::system::memory::adapter::AccessAdapterInventory;
+
+    let access_adapter_offset = state.metrics.access_adapter_offset;
+    let num_sys_airs = state.metrics.num_sys_airs;
+    let mut now_heights = vec![0; num_sys_airs - access_adapter_offset];
     AccessAdapterInventory::<F>::compute_heights_from_arena(
-        access_adapter_records,
-        &mut heights[access_adapter_offset..num_sys_airs],
+        &state.memory.access_adapter_records,
+        &mut now_heights,
     );
-
-    heights
+    let now_trace_cells = zip(
+        &state.metrics.main_widths[access_adapter_offset..],
+        &now_heights,
+    )
+    .map(|(main_width, h)| main_width * h)
+    .collect_vec();
+    for (air_name, &now_value) in itertools::izip!(
+        &state.metrics.air_names[access_adapter_offset..],
+        &now_trace_cells,
+    ) {
+        if now_value != 0 {
+            let labels = [
+                ("air_name", air_name.clone()),
+                ("opcode", String::default()),
+                ("dsl_ir", String::default()),
+                ("cycle_tracker_span", "memory_access_adapters".to_owned()),
+            ];
+            counter!("cells_used", &labels).increment(now_value as u64);
+        }
+    }
+    state.metrics.current_trace_cells.fill(0);
 }
 
 impl VmMetrics {
@@ -151,6 +155,7 @@ impl VmMetrics {
         self.air_names = air_names;
         self.main_widths = main_widths;
         self.total_widths = total_widths;
+        self.current_trace_cells = vec![0; self.air_names.len()];
     }
 
     pub fn update_trace_cells(

crates/vm/src/system/memory/adapter/mod.rs

@@ -50,6 +50,8 @@ pub struct AccessAdapterInventory<F> {
     chips: Vec<GenericAccessAdapterChip<F>>,
     #[getset(set = "pub")]
     arena: DenseRecordArena,
+    #[cfg(feature = "metrics")]
+    pub(crate) trace_heights: Vec<usize>,
 }
 
 impl<F: Clone + Send + Sync> AccessAdapterInventory<F> {
@@ -77,6 +79,8 @@ impl<F: Clone + Send + Sync> AccessAdapterInventory<F> {
             memory_config,
             chips,
             arena: DenseRecordArena::with_byte_capacity(0),
+            #[cfg(feature = "metrics")]
+            trace_heights: Vec::new(),
         }
     }
 
@@ -177,6 +181,10 @@ impl<F: Clone + Send + Sync> AccessAdapterInventory<F> {
             .zip(heights.iter())
             .map(|(&width, &height)| RowMajorMatrix::new(vec![F::ZERO; width * height], width))
             .collect::<Vec<_>>();
+        #[cfg(feature = "metrics")]
+        {
+            self.trace_heights = heights;
+        }
 
         let mut trace_ptrs = vec![0; num_adapters];
 

crates/vm/src/system/memory/controller/mod.rs

@@ -78,7 +78,7 @@ pub struct MemoryController<F: Field> {
     pub range_checker: SharedVariableRangeCheckerChip,
     // Store separately to avoid smart pointer reference each time
     range_checker_bus: VariableRangeCheckerBus,
-    access_adapter_inventory: AccessAdapterInventory<F>,
+    pub(crate) access_adapter_inventory: AccessAdapterInventory<F>,
     pub(crate) hasher_chip: Option<Arc<Poseidon2PeripheryChip<F>>>,
 }
 

crates/vm/src/system/mod.rs

@@ -477,6 +477,7 @@ where
         use crate::system::memory::interface::MemoryInterface;
 
         let boundary_idx = PUBLIC_VALUES_AIR_ID + usize::from(self.public_values_chip.is_some());
+        let mut access_adapter_offset = boundary_idx + 1;
         match &self.memory_controller.interface_chip {
             MemoryInterface::Volatile { boundary_chip } => {
                 let boundary_height = boundary_chip
@@ -494,6 +495,7 @@ where
                 let boundary_height = 2 * boundary_chip.touched_labels.len();
                 heights[boundary_idx] = boundary_height;
                 heights[boundary_idx + 1] = merkle_chip.current_height;
+                access_adapter_offset += 1;
 
                 // Poseidon2Periphery height also varies based on memory, so set it now even though
                 // it's not a system chip:
@@ -505,6 +507,12 @@ where
                 heights[poseidon_idx] = poseidon_height;
             }
         }
+        let access_heights = &self
+            .memory_controller
+            .access_adapter_inventory
+            .trace_heights;
+        heights[access_adapter_offset..access_adapter_offset + access_heights.len()]
+            .copy_from_slice(access_heights);
     }
 }