[OVM 1.4] 70% off APC preflight execution overhead (#3437)

As seen in the test results below (first being current main, second
being this PR, and third being APC=0 base line, each run 3+ times so
should be robust), preflight execution overhead from APC should be
mostly gone (7851 --> 4699 with a baseline of 4622). However, it's
extremely perplexing that tracegen time increased by almost a full
second (7370 --> 8233 with a baseline of 7171) despite nothing much
changed there. Tried multiple ways to dissect and benchmark where the
increase comes from but with no avail.

The net effect should still be savings, as now preflight + tracegen
should have very small overhead `(4699 + 8233) / (4622 + 7171) = ~+10%.`

Disclaimer: these stats are for 100 APCs, so can be quite different for
other # of APCs. Calling this 70% off preflight because it eliminates
most preflight overhead but somehow adds to tracegen...

```
                                                 filename  num_segments  app_proof_cells  app_proof_cols  total_proof_time_ms  app_proof_time_ms  app_execute_preflight_time_ms  app_execute_metered_time_ms  app_trace_gen_time_ms  leaf_proof_time_ms  inner_recursion_proof_time_ms  normal_instruction_ratio  openvm_precompile_ratio  powdr_ratio  powdr_rows
       /home/steve/openvm-reth-benchmark/apc_100_new.json            19      13856523983          354152                31097              31097                           7851                          708                   7370                   0                              0                  0.307127                 0.540265     0.152608    14033237
 /home/steve/openvm-reth-benchmark/apc_100_preflight.json            19      13856523983          354152                31166              31166                           4699                          701                   8233                   0                              0                  0.307127                 0.540265     0.152608    14033237
               ../openvm-reth-benchmark/metrics_apc0.json            26      20019740816          216005                42660              42660                           4622                          749                   7171                   0                              0                  0.612871                 0.387129     0.000000           0

```

Author: qwang98

openvm/src/extraction_utils.rs

@@ -136,9 +136,9 @@ impl<F> OriginalAirs<F> {
 pub fn record_arena_dimension_by_air_name_per_apc_call<F>(
     apc: &Apc<F, Instr<F>>,
     air_by_opcode_id: &OriginalAirs<F>,
-) -> HashMap<String, RecordArenaDimension> {
+) -> BTreeMap<String, RecordArenaDimension> {
     apc.instructions().iter().map(|instr| &instr.0.opcode).fold(
-        HashMap::new(),
+        BTreeMap::new(),
         |mut acc, opcode| {
             // Get the air name for this opcode
             let air_name = air_by_opcode_id.opcode_to_air.get(opcode).unwrap();

openvm/src/powdr_extension/executor/mod.rs

@@ -13,6 +13,7 @@ use crate::{
     Instr,
 };
 
+use itertools::Itertools;
 use openvm_circuit::arch::{
     execution_mode::{ExecutionCtx, MeteredCtx},
     Arena, DenseRecordArena, E2PreCompute, MatrixRecordArena, PreflightExecutor,
@@ -33,7 +34,6 @@ use openvm_circuit::{
     },
     system::memory::online::{GuestMemory, TracingMemory},
 };
-use powdr_autoprecompiles::InstructionHandler;
 
 /// A struct which holds the state of the execution based on the original instructions in this block and a dummy inventory.
 /// It holds arenas for each original use for both cpu and gpu execution, so that this struct can be agnostic to the execution backend.
@@ -45,6 +45,7 @@ pub struct PowdrExecutor {
     pub original_arenas_cpu: Rc<RefCell<OriginalArenas<MatrixRecordArena<BabyBear>>>>,
     pub original_arenas_gpu: Rc<RefCell<OriginalArenas<DenseRecordArena>>>,
     pub height_change: u32,
+    cached_instructions_meta: Vec<CachedInstructionMeta>,
 }
 
 /// A shared mutable reference to the arenas used to store the traces of the original instructions, accessed during preflight execution and trace generation.
@@ -82,21 +83,21 @@ impl<A: Arena> OriginalArenas<A> {
         }
     }
 
-    /// Returns a mutable reference to the arenas.
+    /// Returns a mutable reference to the arena of the given vector index.
     /// - Panics if the arenas are not initialized.
-    pub fn arenas_mut(&mut self) -> &mut HashMap<String, A> {
+    pub fn arena_mut_by_index(&mut self, index: usize) -> &mut A {
         match self {
             OriginalArenas::Uninitialized => panic!("original arenas are uninitialized"),
-            OriginalArenas::Initialized(initialized) => &mut initialized.arenas,
+            OriginalArenas::Initialized(initialized) => initialized.arena_mut_by_index(index),
         }
     }
 
-    /// Returns a reference to the arenas.
+    /// Returns the arena of the given air name.
     /// - Panics if the arenas are not initialized.
-    pub fn arenas(&self) -> &HashMap<String, A> {
+    pub fn take_arena(&mut self, air_name: &str) -> Option<A> {
         match self {
             OriginalArenas::Uninitialized => panic!("original arenas are uninitialized"),
-            OriginalArenas::Initialized(initialized) => &initialized.arenas,
+            OriginalArenas::Initialized(initialized) => initialized.take_arena(air_name),
         }
     }
 
@@ -123,7 +124,8 @@ impl<A: Arena> OriginalArenas<A> {
 /// and how many calls to each air are made per APC call.
 #[derive(Default)]
 pub struct InitializedOriginalArenas<A> {
-    pub arenas: HashMap<String, A>,
+    arenas: Vec<Option<A>>,
+    air_name_to_arena_index: HashMap<String, usize>,
     pub number_of_calls: usize,
 }
 
@@ -136,28 +138,49 @@ impl<A: Arena> InitializedOriginalArenas<A> {
     ) -> Self {
         let record_arena_dimensions =
             record_arena_dimension_by_air_name_per_apc_call(apc, original_airs);
-        Self {
-            arenas: record_arena_dimensions
-                .iter()
-                .map(
-                    |(
+        let (air_name_to_arena_index, arenas) =
+            record_arena_dimensions.into_iter().enumerate().fold(
+                (HashMap::new(), Vec::new()),
+                |(mut air_name_to_arena_index, mut arenas),
+                 (
+                    idx,
+                    (
                         air_name,
                         RecordArenaDimension {
                             height: num_calls,
                             width: air_width,
                         },
-                    )| {
-                        (
-                            air_name.clone(),
-                            A::with_capacity(*num_calls * apc_call_count_estimate, *air_width),
-                        )
-                    },
-                )
-                .collect(),
+                    ),
+                )| {
+                    air_name_to_arena_index.insert(air_name, idx);
+                    arenas.push(Some(A::with_capacity(
+                        num_calls * apc_call_count_estimate,
+                        air_width,
+                    )));
+                    (air_name_to_arena_index, arenas)
+                },
+            );
+
+        Self {
+            arenas,
+            air_name_to_arena_index,
             // This is the actual number of calls, which we don't know yet. It will be updated during preflight execution.
             number_of_calls: 0,
         }
     }
+
+    #[inline]
+    fn arena_mut_by_index(&mut self, index: usize) -> &mut A {
+        self.arenas
+            .get_mut(index)
+            .and_then(|arena| arena.as_mut())
+            .expect("arena missing for index")
+    }
+
+    fn take_arena(&mut self, air_name: &str) -> Option<A> {
+        let index = *self.air_name_to_arena_index.get(air_name)?;
+        self.arenas[index].take()
+    }
 }
 
 /// The dimensions of a record arena for a given air name, used to initialize the arenas.
@@ -166,6 +189,12 @@ pub struct RecordArenaDimension {
     pub width: usize,
 }
 
+#[derive(Clone, Copy)]
+struct CachedInstructionMeta {
+    executor_index: usize,
+    arena_index: usize,
+}
+
 /// A struct to interpret the pre-compute data as for PowdrExecutor.
 #[derive(AlignedBytesBorrow, Clone)]
 #[repr(C)]
@@ -426,20 +455,15 @@ impl PreflightExecutor<BabyBear, MatrixRecordArena<BabyBear>> for PowdrExecutor
         let apc_call_count = || ctx.trace_buffer.len() / ctx.width;
 
         original_arenas.ensure_initialized(apc_call_count, &self.air_by_opcode_id, &self.apc);
-
-        let arenas = original_arenas.arenas_mut();
-
         // execute the original instructions one by one
-        for instruction in self.apc.instructions() {
-            let executor = self
-                .executor_inventory
-                .get_executor(instruction.0.opcode)
-                .unwrap();
-
-            let air_name = self
-                .air_by_opcode_id
-                .get_instruction_air_and_id(instruction)
-                .0;
+        for (instruction, cached_meta) in self
+            .apc
+            .instructions()
+            .iter()
+            .zip_eq(&self.cached_instructions_meta)
+        {
+            let executor = &self.executor_inventory.executors[cached_meta.executor_index];
+            let ctx_arena = original_arenas.arena_mut_by_index(cached_meta.arena_index);
 
             let state = VmStateMut {
                 pc,
@@ -448,7 +472,7 @@ impl PreflightExecutor<BabyBear, MatrixRecordArena<BabyBear>> for PowdrExecutor
                 rng,
                 custom_pvs,
                 // We execute in the context of the relevant original table
-                ctx: arenas.get_mut(&air_name).unwrap(),
+                ctx: ctx_arena,
                 // TODO: should we pass around the same metrics object, or snapshot it at the beginning of this method and apply a single update at the end?
                 #[cfg(feature = "metrics")]
                 metrics,
@@ -501,20 +525,15 @@ impl PreflightExecutor<BabyBear, DenseRecordArena> for PowdrExecutor {
         };
 
         original_arenas.ensure_initialized(apc_call_count, &self.air_by_opcode_id, &self.apc);
-
-        let arenas = original_arenas.arenas_mut();
-
         // execute the original instructions one by one
-        for instruction in self.apc.instructions() {
-            let executor = self
-                .executor_inventory
-                .get_executor(instruction.0.opcode)
-                .unwrap();
-
-            let air_name = self
-                .air_by_opcode_id
-                .get_instruction_air_and_id(instruction)
-                .0;
+        for (instruction, cached_meta) in self
+            .apc
+            .instructions()
+            .iter()
+            .zip(&self.cached_instructions_meta)
+        {
+            let executor = &self.executor_inventory.executors[cached_meta.executor_index];
+            let ctx_arena = original_arenas.arena_mut_by_index(cached_meta.arena_index);
 
             let state = VmStateMut {
                 pc,
@@ -523,7 +542,7 @@ impl PreflightExecutor<BabyBear, DenseRecordArena> for PowdrExecutor {
                 rng,
                 custom_pvs,
                 // We execute in the context of the relevant original table
-                ctx: arenas.get_mut(&air_name).unwrap(),
+                ctx: ctx_arena,
                 // TODO: should we pass around the same metrics object, or snapshot it at the beginning of this method and apply a single update at the end?
                 #[cfg(feature = "metrics")]
                 metrics,
@@ -552,13 +571,46 @@ impl PowdrExecutor {
         record_arena_by_air_name_gpu: Rc<RefCell<OriginalArenas<DenseRecordArena>>>,
         height_change: u32,
     ) -> Self {
+        let executor_inventory = base_config.sdk_config.sdk.create_executors().unwrap();
+
+        let arena_index_by_name =
+            record_arena_dimension_by_air_name_per_apc_call(apc.as_ref(), &air_by_opcode_id)
+                .iter()
+                .enumerate()
+                .map(|(idx, (name, _))| (name.clone(), idx))
+                .collect::<HashMap<_, _>>();
+
+        let cached_instructions_meta = apc
+            .instructions()
+            .iter()
+            .map(|instruction| {
+                let executor_index = *executor_inventory
+                    .instruction_lookup
+                    .get(&instruction.0.opcode)
+                    .expect("missing executor for opcode")
+                    as usize;
+                let air_name = air_by_opcode_id
+                    .opcode_to_air
+                    .get(&instruction.0.opcode)
+                    .expect("missing air for opcode");
+                let arena_index = *arena_index_by_name
+                    .get(air_name)
+                    .expect("missing arena for air");
+                CachedInstructionMeta {
+                    executor_index,
+                    arena_index,
+                }
+            })
+            .collect();
+
         Self {
             air_by_opcode_id,
-            executor_inventory: base_config.sdk_config.sdk.create_executors().unwrap(),
+            executor_inventory,
             apc,
             original_arenas_cpu: record_arena_by_air_name_cpu,
             original_arenas_gpu: record_arena_by_air_name_gpu,
             height_change,
+            cached_instructions_meta,
         }
     }
 }

openvm/src/powdr_extension/trace_generator/cpu/mod.rs

@@ -119,8 +119,6 @@ impl PowdrTraceGeneratorCpu {
             .inventory
         };
 
-        let arenas = original_arenas.arenas_mut();
-
         let dummy_trace_by_air_name: HashMap<String, SharedCpuTrace<BabyBear>> = chip_inventory
             .chips()
             .iter()
@@ -130,7 +128,7 @@ impl PowdrTraceGeneratorCpu {
                 let air_name = chip_inventory.airs().ext_airs()[insertion_idx].name();
 
                 let record_arena = {
-                    match arenas.remove(&air_name) {
+                    match original_arenas.take_arena(&air_name) {
                         Some(ra) => ra,
                         None => return None, // skip this iteration, because we only have record arena for chips that are used
                     }

openvm/src/powdr_extension/trace_generator/cuda/mod.rs

@@ -219,8 +219,6 @@ impl PowdrTraceGeneratorGpu {
             .inventory
         };
 
-        let arenas = original_arenas.arenas_mut();
-
         let dummy_trace_by_air_name: HashMap<String, DeviceMatrix<BabyBear>> = chip_inventory
             .chips()
             .iter()
@@ -230,7 +228,7 @@ impl PowdrTraceGeneratorGpu {
                 let air_name = chip_inventory.airs().ext_airs()[insertion_idx].name();
 
                 let record_arena = {
-                    match arenas.remove(&air_name) {
+                    match original_arenas.take_arena(&air_name) {
                         Some(ra) => ra,
                         None => return None, // skip this iteration, because we only have record arena for chips that are used
                     }