From e75d450d29cf3eef88e621f3cdeb4d33fbaeed29 Mon Sep 17 00:00:00 2001 From: schaeff Date: Thu, 4 Dec 2025 22:56:37 +0100 Subject: [PATCH 1/3] use insertion idx as air id --- autoprecompiles/src/adapter.rs | 1 - openvm/src/customize_exe.rs | 1 - openvm/src/extraction_utils.rs | 95 +++++++++++-------- openvm/src/powdr_extension/executor/mod.rs | 65 +++++++------ .../trace_generator/cpu/mod.rs | 12 +-- 5 files changed, 90 insertions(+), 84 deletions(-) diff --git a/autoprecompiles/src/adapter.rs b/autoprecompiles/src/adapter.rs index 351cdf9fbd..c014aecb89 100644 --- a/autoprecompiles/src/adapter.rs +++ b/autoprecompiles/src/adapter.rs @@ -87,7 +87,6 @@ where >; type CustomBusTypes: Clone + Display + Sync + Eq + PartialEq; type ApcStats: Send + Sync; - type AirId: Eq + Hash + Send + Sync; fn into_field(e: Self::PowdrField) -> Self::Field; diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs index c3b88bacab..44337ee0ae 100644 --- a/openvm/src/customize_exe.rs +++ b/openvm/src/customize_exe.rs @@ -63,7 +63,6 @@ impl<'a> Adapter for BabyBearOpenVmApcAdapter<'a> { OpenVmMemoryBusInteraction; type CustomBusTypes = OpenVmBusType; type ApcStats = OvmApcStats; - type AirId = String; fn into_field(e: Self::PowdrField) -> Self::Field { openvm_stark_sdk::p3_baby_bear::BabyBear::from_canonical_u32( diff --git a/openvm/src/extraction_utils.rs b/openvm/src/extraction_utils.rs index 07a4d3ba6c..baa56e3dbc 100644 --- a/openvm/src/extraction_utils.rs +++ b/openvm/src/extraction_utils.rs @@ -50,32 +50,35 @@ use crate::utils::symbolic_to_algebraic; // TODO: Use ` as FieldExtensionAlgebra>>::D` instead after fixing p3 dependency const EXT_DEGREE: usize = 4; +#[derive(Clone, Serialize, Deserialize)] +pub struct OriginalMachine { + pub machine: SymbolicMachine, + pub name: String, + pub metrics: AirMetrics, +} + #[derive(Clone, Serialize, Deserialize, Default)] pub struct OriginalAirs { - pub(crate) opcode_to_air: HashMap, - pub(crate) air_name_to_machine: BTreeMap, AirMetrics)>, + pub(crate) air_id_by_opcode: HashMap, + pub(crate) machine_by_air_id: BTreeMap>, } impl InstructionHandler for OriginalAirs { type Field = F; type Instruction = Instr; - type AirId = String; + type AirId = usize; fn get_instruction_air_and_id( &self, instruction: &Self::Instruction, ) -> (Self::AirId, &SymbolicMachine) { - let id = self - .opcode_to_air - .get(&instruction.0.opcode) - .unwrap() - .clone(); - let air = &self.air_name_to_machine.get(&id).unwrap().0; + let id = *self.air_id_by_opcode.get(&instruction.0.opcode).unwrap(); + let air = &self.machine_by_air_id.get(&id).unwrap().machine; (id, air) } fn is_allowed(&self, instruction: &Self::Instruction) -> bool { - self.opcode_to_air.contains_key(&instruction.0.opcode) + self.air_id_by_opcode.contains_key(&instruction.0.opcode) } fn is_branching(&self, instruction: &Self::Instruction) -> bool { @@ -95,59 +98,63 @@ impl OriginalAirs { pub fn insert_opcode( &mut self, opcode: VmOpcode, - air_name: String, - machine: impl Fn() -> Result<(SymbolicMachine, AirMetrics), UnsupportedOpenVmReferenceError>, + air_id: usize, + machine: impl Fn() -> Result, UnsupportedOpenVmReferenceError>, ) -> Result<(), UnsupportedOpenVmReferenceError> { - if self.opcode_to_air.contains_key(&opcode) { + if self.air_id_by_opcode.contains_key(&opcode) { panic!("Opcode {opcode} already exists"); } // Insert the machine only if `air_name` isn't already present - if !self.air_name_to_machine.contains_key(&air_name) { + if let std::collections::btree_map::Entry::Vacant(e) = self.machine_by_air_id.entry(air_id) + { let machine_instance = machine()?; - self.air_name_to_machine - .insert(air_name.clone(), machine_instance); + e.insert(machine_instance); } - self.opcode_to_air.insert(opcode, air_name); + self.air_id_by_opcode.insert(opcode, air_id); + Ok(()) } pub fn get_instruction_metrics(&self, opcode: VmOpcode) -> Option<&AirMetrics> { - self.opcode_to_air.get(&opcode).and_then(|air_name| { - self.air_name_to_machine - .get(air_name) - .map(|(_, metrics)| metrics) + self.air_id_by_opcode.get(&opcode).and_then(|air_id| { + self.machine_by_air_id + .get(air_id) + .map(|machine| &machine.metrics) }) } pub fn allow_list(&self) -> Vec { - self.opcode_to_air.keys().cloned().collect() + self.air_id_by_opcode.keys().cloned().collect() } pub fn airs_by_name(&self) -> impl Iterator)> { - self.air_name_to_machine - .iter() - .map(|(name, (machine, _))| (name, machine)) + self.machine_by_air_id + .values() + .map(|machine| (&machine.name, &machine.machine)) } } /// For each air name, the dimension of a record arena needed to store the /// records for a single APC call. -pub fn record_arena_dimension_by_air_name_per_apc_call( +pub fn record_arena_dimension_by_air_id_per_apc_call( apc: &Apc>, air_by_opcode_id: &OriginalAirs, -) -> BTreeMap { +) -> BTreeMap { apc.instructions().iter().map(|instr| &instr.0.opcode).fold( 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(); + let air_id = air_by_opcode_id.air_id_by_opcode.get(opcode).unwrap(); // Increment the height for this air name, initializing if necessary - acc.entry(air_name.clone()) + acc.entry(*air_id) .or_insert_with(|| { - let (_, air_metrics) = - air_by_opcode_id.air_name_to_machine.get(air_name).unwrap(); + let air_metrics = &air_by_opcode_id + .machine_by_air_id + .get(air_id) + .unwrap() + .metrics; RecordArenaDimension { height: 0, @@ -300,11 +307,12 @@ impl OriginalVmConfig { }) .map(|(op, executor_id)| { let insertion_index = chip_inventory.executor_idx_to_insertion_idx[executor_id]; - let air_ref = &chip_inventory.airs().ext_airs()[insertion_index]; - (op, air_ref) + (op, insertion_index) }) // find executor for opcode - .try_fold(OriginalAirs::default(), |mut airs, (op, air_ref)| { - airs.insert_opcode(op, air_ref.name(), || { + .try_fold(OriginalAirs::default(), |mut airs, (op, insertion_idx)| { + airs.insert_opcode(op, insertion_idx, || { + let air_ref = &chip_inventory.airs().ext_airs()[insertion_idx]; + let name = air_ref.name(); let columns = get_columns(air_ref.clone()); let constraints = get_constraints(air_ref.clone()); let metrics = get_air_metrics(air_ref.clone(), max_degree); @@ -321,14 +329,17 @@ impl OriginalVmConfig { .map(|expr| openvm_bus_interaction_to_powdr(expr, &columns)) .collect::>()?; - Ok(( - SymbolicMachine { - constraints: powdr_exprs.into_iter().map(Into::into).collect(), - bus_interactions: powdr_bus_interactions, - derived_columns: vec![], - }, + let machine = SymbolicMachine { + constraints: powdr_exprs.into_iter().map(Into::into).collect(), + bus_interactions: powdr_bus_interactions, + derived_columns: vec![], + }; + + Ok(OriginalMachine { + machine, + name, metrics, - )) + }) })?; Ok(airs) diff --git a/openvm/src/powdr_extension/executor/mod.rs b/openvm/src/powdr_extension/executor/mod.rs index ef9613e407..3792e939d7 100644 --- a/openvm/src/powdr_extension/executor/mod.rs +++ b/openvm/src/powdr_extension/executor/mod.rs @@ -8,7 +8,7 @@ use std::{ use crate::{ extraction_utils::{ - record_arena_dimension_by_air_name_per_apc_call, OriginalAirs, OriginalVmConfig, + record_arena_dimension_by_air_id_per_apc_call, OriginalAirs, OriginalVmConfig, }, Instr, }; @@ -94,10 +94,10 @@ impl OriginalArenas { /// Returns the arena of the given air name. /// - Panics if the arenas are not initialized. - pub fn take_arena(&mut self, air_name: &str) -> Option { + pub fn take_arena(&mut self, air_id: usize) -> Option { match self { OriginalArenas::Uninitialized => panic!("original arenas are uninitialized"), - OriginalArenas::Initialized(initialized) => initialized.take_arena(air_name), + OriginalArenas::Initialized(initialized) => initialized.take_arena(air_id), } } @@ -125,7 +125,7 @@ impl OriginalArenas { #[derive(Default)] pub struct InitializedOriginalArenas { arenas: Vec>, - air_name_to_arena_index: HashMap, + air_id_to_arena_index: HashMap, pub number_of_calls: usize, } @@ -137,33 +137,32 @@ impl InitializedOriginalArenas { apc: &Arc>>, ) -> Self { let record_arena_dimensions = - record_arena_dimension_by_air_name_per_apc_call(apc, original_airs); - 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_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) - }, - ); + record_arena_dimension_by_air_id_per_apc_call(apc, original_airs); + let (air_id_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_id, + RecordArenaDimension { + height: num_calls, + width: air_width, + }, + ), + )| { + air_name_to_arena_index.insert(air_id, 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, + air_id_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, } @@ -177,8 +176,8 @@ impl InitializedOriginalArenas { .expect("arena missing for index") } - fn take_arena(&mut self, air_name: &str) -> Option { - let index = *self.air_name_to_arena_index.get(air_name)?; + fn take_arena(&mut self, air_id: usize) -> Option { + let index = *self.air_id_to_arena_index.get(&air_id)?; self.arenas[index].take() } } @@ -574,10 +573,10 @@ impl PowdrExecutor { 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) + record_arena_dimension_by_air_id_per_apc_call(apc.as_ref(), &air_by_opcode_id) .iter() .enumerate() - .map(|(idx, (name, _))| (name.clone(), idx)) + .map(|(idx, (name, _))| (*name, idx)) .collect::>(); let cached_instructions_meta = apc @@ -590,7 +589,7 @@ impl PowdrExecutor { .expect("missing executor for opcode") as usize; let air_name = air_by_opcode_id - .opcode_to_air + .air_id_by_opcode .get(&instruction.0.opcode) .expect("missing air for opcode"); let arena_index = *arena_index_by_name diff --git a/openvm/src/powdr_extension/trace_generator/cpu/mod.rs b/openvm/src/powdr_extension/trace_generator/cpu/mod.rs index 94d9ef6d3e..546f5b380d 100644 --- a/openvm/src/powdr_extension/trace_generator/cpu/mod.rs +++ b/openvm/src/powdr_extension/trace_generator/cpu/mod.rs @@ -119,16 +119,14 @@ impl PowdrTraceGeneratorCpu { .inventory }; - let dummy_trace_by_air_name: HashMap> = chip_inventory + let dummy_trace_by_air_id: HashMap> = chip_inventory .chips() .iter() .enumerate() .rev() - .filter_map(|(insertion_idx, chip)| { - let air_name = chip_inventory.airs().ext_airs()[insertion_idx].name(); - + .filter_map(|(air_id, chip)| { let record_arena = { - match original_arenas.take_arena(&air_name) { + match original_arenas.take_arena(air_id) { Some(ra) => ra, None => return None, // skip this iteration, because we only have record arena for chips that are used } @@ -136,7 +134,7 @@ impl PowdrTraceGeneratorCpu { let shared_trace = chip.generate_proving_ctx(record_arena).common_main.unwrap(); - Some((air_name, SharedCpuTrace::from(shared_trace))) + Some((air_id, SharedCpuTrace::from(shared_trace))) }) .collect(); @@ -146,7 +144,7 @@ impl PowdrTraceGeneratorCpu { apc_poly_id_to_index, columns_to_compute, } = generate_trace( - &dummy_trace_by_air_name, + &dummy_trace_by_air_id, &self.original_airs, num_apc_calls, &self.apc, From 136b0ce8642e595429f8939f9ac384847165dfec Mon Sep 17 00:00:00 2001 From: schaeff Date: Thu, 4 Dec 2025 23:01:20 +0100 Subject: [PATCH 2/3] import entry --- openvm/src/extraction_utils.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/openvm/src/extraction_utils.rs b/openvm/src/extraction_utils.rs index baa56e3dbc..93d361eddb 100644 --- a/openvm/src/extraction_utils.rs +++ b/openvm/src/extraction_utils.rs @@ -1,3 +1,4 @@ +use std::collections::btree_map::Entry; use std::collections::{BTreeMap, HashMap}; use std::sync::{Arc, Mutex}; @@ -105,8 +106,7 @@ impl OriginalAirs { panic!("Opcode {opcode} already exists"); } // Insert the machine only if `air_name` isn't already present - if let std::collections::btree_map::Entry::Vacant(e) = self.machine_by_air_id.entry(air_id) - { + if let Entry::Vacant(e) = self.machine_by_air_id.entry(air_id) { let machine_instance = machine()?; e.insert(machine_instance); } From 20eb88e8ca9f010af838719beb327045d454ceeb Mon Sep 17 00:00:00 2001 From: Schaeff Date: Thu, 4 Dec 2025 23:20:07 +0100 Subject: [PATCH 3/3] fix gpu --- .../trace_generator/cuda/mod.rs | 18 ++++++++---------- 1 file changed, 8 insertions(+), 10 deletions(-) diff --git a/openvm/src/powdr_extension/trace_generator/cuda/mod.rs b/openvm/src/powdr_extension/trace_generator/cuda/mod.rs index e5adb49399..e2b1054c96 100644 --- a/openvm/src/powdr_extension/trace_generator/cuda/mod.rs +++ b/openvm/src/powdr_extension/trace_generator/cuda/mod.rs @@ -221,16 +221,14 @@ impl PowdrTraceGeneratorGpu { .inventory }; - let dummy_trace_by_air_name: HashMap> = chip_inventory + let dummy_trace_by_air_name: HashMap> = chip_inventory .chips() .iter() .enumerate() .rev() - .filter_map(|(insertion_idx, chip)| { - let air_name = chip_inventory.airs().ext_airs()[insertion_idx].name(); - + .filter_map(|(air_id, chip)| { let record_arena = { - match original_arenas.take_arena(&air_name) { + match original_arenas.take_arena(air_id) { Some(ra) => ra, None => return None, // skip this iteration, because we only have record arena for chips that are used } @@ -238,7 +236,7 @@ impl PowdrTraceGeneratorGpu { let shared_trace = chip.generate_proving_ctx(record_arena).common_main.unwrap(); - Some((air_name, shared_trace)) + Some((air_id, shared_trace)) }) .collect(); @@ -264,8 +262,8 @@ impl PowdrTraceGeneratorGpu { .iter() // along with their substitutions .zip_eq(self.apc.subs()) - // map to `(air_name, substitutions)` - .map(|(instr, subs)| (&self.original_airs.opcode_to_air[&instr.0.opcode], subs)) + // map to `(air_id, substitutions)` + .map(|(instr, subs)| (&self.original_airs.air_id_by_opcode[&instr.0.opcode], subs)) // group by air name. This results in `HashMap>` where the length of the vector is the number of rows which are created in this air, per apc call .into_group_map() // go through each air and its substitutions @@ -273,7 +271,7 @@ impl PowdrTraceGeneratorGpu { .enumerate() .fold( (Vec::new(), Vec::new()), - |(mut airs, mut substitutions), (air_index, (air_name, subs_by_row))| { + |(mut airs, mut substitutions), (air_index, (air_id, subs_by_row))| { // Find the substitutions that map to an apc column let new_substitutions: Vec = subs_by_row .iter() @@ -293,7 +291,7 @@ impl PowdrTraceGeneratorGpu { .collect(); // get the device dummy trace for this air - let dummy_trace = &dummy_trace_by_air_name[*air_name]; + let dummy_trace = &dummy_trace_by_air_name[*air_id]; use openvm_stark_backend::prover::hal::MatrixDimensions; airs.push(OriginalAir {