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/empirical_constraints.rs b/openvm/src/empirical_constraints.rs index ae2ded80e2..901eb8454a 100644 --- a/openvm/src/empirical_constraints.rs +++ b/openvm/src/empirical_constraints.rs @@ -1,17 +1,22 @@ use itertools::Itertools; -use openvm_circuit::arch::VmCircuitConfig; use openvm_sdk::StdIn; +use openvm_stark_backend::p3_field::FieldAlgebra; use openvm_stark_backend::p3_maybe_rayon::prelude::IntoParallelIterator; use openvm_stark_backend::p3_maybe_rayon::prelude::ParallelIterator; use openvm_stark_sdk::openvm_stark_backend::p3_field::PrimeField32; +use openvm_stark_sdk::p3_baby_bear::BabyBear; +use powdr_autoprecompiles::bus_map::BusType; use powdr_autoprecompiles::empirical_constraints::{ intersect_partitions, DebugInfo, EmpiricalConstraints, }; +use powdr_autoprecompiles::expression::AlgebraicEvaluator; +use powdr_autoprecompiles::expression::RowEvaluator; use powdr_autoprecompiles::DegreeBound; use std::collections::btree_map::Entry; use std::collections::HashMap; use std::collections::{BTreeMap, BTreeSet}; +use crate::bus_map::default_openvm_bus_map; use crate::trace_generation::do_with_trace; use crate::{CompiledProgram, OriginalCompiledProgram}; @@ -69,7 +74,7 @@ pub fn detect_empirical_constraints( // If this becomes a RAM issue, we can also pass individual segments to process_trace. // The advantage of the current approach is that the percentiles can be computed more accurately. tracing::info!(" Collecting trace..."); - let (trace, new_debug_info) = collect_trace(&program, input); + let (trace, new_debug_info) = collect_trace(&program, input, degree_bound.identities); tracing::info!(" Detecting constraints..."); constraint_detector.process_trace(trace, new_debug_info); } @@ -78,63 +83,73 @@ pub fn detect_empirical_constraints( constraint_detector.finalize() } -fn collect_trace(program: &CompiledProgram, inputs: StdIn) -> (Trace, DebugInfo) { +fn collect_trace( + program: &CompiledProgram, + inputs: StdIn, + degree_bound: usize, +) -> (Trace, DebugInfo) { let mut trace = Trace::default(); let mut debug_info = DebugInfo::default(); let mut seg_idx = 0; - do_with_trace(program, inputs, |vm, _pk, ctx| { - let global_airs = vm - .config() - .create_airs() - .unwrap() - .into_airs() - .enumerate() - .collect::>(); + do_with_trace(program, inputs, |_vm, _pk, ctx| { + let airs = program.vm_config.sdk.airs(degree_bound).unwrap(); for (air_id, proving_context) in &ctx.per_air { - let air = &global_airs[air_id]; - let Some(column_names) = air.columns() else { - // Instruction chips always have column names. - continue; - }; - if !proving_context.cached_mains.is_empty() { // Instruction chips always have a cached main. continue; } let main = proving_context.common_main.as_ref().unwrap(); - assert_eq!(main.width, column_names.len()); - - // Instruction chips have a PC and time stamp - let find_col = |name: &str| -> Option { - column_names.iter().position(|col_name| { - col_name == name || col_name == &format!("inner__{}", name) + let machine = &airs.machine_by_air_id.get(air_id).unwrap().machine; + + // Find the execution bus interation + // This assumes there is exactly one, which is the case for instruction chips + let execution_bus_interaction = machine + .bus_interactions + .iter() + .find(|interaction| { + interaction.id + == default_openvm_bus_map() + .get_bus_id(&BusType::ExecutionBridge) + .unwrap() }) - }; - let Some(pc_index) = find_col("from_state__pc") else { - continue; - }; - let ts_index = find_col("from_state__timestamp").unwrap(); + .unwrap(); for row in main.row_slices() { - let row = row.iter().map(|v| v.as_canonical_u32()).collect::>(); - let pc_value = row[pc_index]; - let ts_value = row[ts_index]; + // Create an evaluator over this row + let evaluator = RowEvaluator::new(row); + + // Evaluate the execution bus interaction + let execution = evaluator.eval_bus_interaction(execution_bus_interaction); - if pc_value == 0 { - // Padding row! + // `is_valid` is the multiplicity + let is_valid = execution.mult; + if is_valid == BabyBear::ZERO { + // If `is_valid` is zero, this is a padding row continue; } + // Recover the values of the pc and timestamp + let [pc, timestamp] = execution + .args + .map(|v| v.as_canonical_u32()) + .collect_vec() + .try_into() + .unwrap(); + + // Convert the row to u32s + // TODO: is this necessary? + let row = row.iter().map(|v| v.as_canonical_u32()).collect(); + let row = Row { cells: row, - pc: pc_value, - timestamp: (seg_idx, ts_value), + pc, + timestamp: (seg_idx, timestamp), }; trace.rows.push(row); - match debug_info.air_id_by_pc.entry(pc_value) { + match debug_info.air_id_by_pc.entry(pc) { Entry::Vacant(entry) => { entry.insert(*air_id); } @@ -143,9 +158,10 @@ fn collect_trace(program: &CompiledProgram, inputs: StdIn) -> (Trace, DebugInfo) } } if !debug_info.column_names_by_air_id.contains_key(air_id) { - debug_info - .column_names_by_air_id - .insert(*air_id, column_names.clone()); + debug_info.column_names_by_air_id.insert( + *air_id, + machine.main_columns().map(|r| (*r.name).clone()).collect(), + ); } } } diff --git a/openvm/src/extraction_utils.rs b/openvm/src/extraction_utils.rs index 07a4d3ba6c..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}; @@ -50,32 +51,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 +99,62 @@ 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 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, 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 {