From 5ae915d7fdbcd401f2a595962b8a200b6214f7ca Mon Sep 17 00:00:00 2001 From: Steve Wang Date: Fri, 4 Jul 2025 15:33:31 +0800 Subject: [PATCH] Calculate non-APC LogUp columns for PGO (#2960) Need to first merge https://github.com/powdr-labs/stark-backend/pull/11. And then merge https://github.com/powdr-labs/openvm/pull/32. Original comment by @Schaeff: ``` so here's my suspicion: to get the number of columns of an air, we use executor.air().width(), in fully qualified syntax openvm_stark_backend::p3_air::BaseAir::width(&executor.air()) this is the number of main trace columns, which does not include interaction columns in order to get the full number of columns, we need to do what we already do for get_constraints but additionally call InteractionPhaseAirBuilder::finalize_interactions on the symbolic builder to materialize the interactions. Sadly InteractionPhaseAirBuilder is pub(crate) so it seems we need to change stark-backend ``` This PR follows the comment above, but has additional changes: 1. Instead of using `finalize_interactions`, call lower level functions for greater efficiency. 2. Refactored `AirMetrics` related functions, so we always show the number of main constraint vs bus interaction columns. Next step: 1. We should calculate log up columns for APCs as well during PGO. --------- Co-authored-by: Thibaut Schaeffer --- Cargo.toml | 58 ++++++++++----------- openvm/src/customize_exe.rs | 14 ++--- openvm/src/extraction_utils.rs | 93 +++++++++++++++++++++++++++++----- openvm/src/lib.rs | 22 ++++---- 4 files changed, 126 insertions(+), 61 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 0531a91bf..b30e168ce 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -83,47 +83,47 @@ powdr-autoprecompiles = { path = "./autoprecompiles", version = "0.1.4" } powdr-openvm = { path = "./openvm", version = "0.1.4" } # openvm -openvm = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be" } -openvm-build = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be" } -openvm-rv32im-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be" } -openvm-rv32im-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be" } -openvm-rv32im-guest = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be", default-features = false } -openvm-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be" } -openvm-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be" } -openvm-circuit-derive = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be" } -openvm-circuit-primitives = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be" } -openvm-circuit-primitives-derive = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be" } -openvm-instructions = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be" } -openvm-instructions-derive = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be" } -openvm-sdk = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be", default-features = false, features = [ +openvm = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36" } +openvm-build = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36" } +openvm-rv32im-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36" } +openvm-rv32im-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36" } +openvm-rv32im-guest = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36", default-features = false } +openvm-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36" } +openvm-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36" } +openvm-circuit-derive = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36" } +openvm-circuit-primitives = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36" } +openvm-circuit-primitives-derive = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36" } +openvm-instructions = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36" } +openvm-instructions-derive = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36" } +openvm-sdk = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36", default-features = false, features = [ "parallel", "jemalloc", "nightly-features", "bench-metrics", ] } -openvm-ecc-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be" } -openvm-ecc-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be" } -openvm-keccak256-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be" } -openvm-keccak256-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be" } -openvm-sha256-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be" } -openvm-sha256-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be" } -openvm-algebra-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be" } -openvm-algebra-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be" } -openvm-bigint-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be" } -openvm-bigint-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be" } -openvm-pairing-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be" } -openvm-pairing-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be" } -openvm-native-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be", default-features = false } -openvm-native-recursion = { git = "https://github.com/powdr-labs/openvm.git", rev = "26bb25be", default-features = false } +openvm-ecc-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36" } +openvm-ecc-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36" } +openvm-keccak256-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36" } +openvm-keccak256-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36" } +openvm-sha256-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36" } +openvm-sha256-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36" } +openvm-algebra-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36" } +openvm-algebra-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36" } +openvm-bigint-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36" } +openvm-bigint-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36" } +openvm-pairing-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36" } +openvm-pairing-transpiler = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36" } +openvm-native-circuit = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36", default-features = false } +openvm-native-recursion = { git = "https://github.com/powdr-labs/openvm.git", rev = "d730a36", default-features = false } # stark-backend -openvm-stark-sdk = { git = "https://github.com/powdr-labs/stark-backend.git", rev = "c33fda6", default-features = false, features = [ +openvm-stark-sdk = { git = "https://github.com/powdr-labs/stark-backend.git", rev = "ee4e22b", default-features = false, features = [ "parallel", "jemalloc", "nightly-features", "bench-metrics", ] } -openvm-stark-backend = { git = "https://github.com/powdr-labs/stark-backend.git", rev = "c33fda6", default-features = false, features = [ +openvm-stark-backend = { git = "https://github.com/powdr-labs/stark-backend.git", rev = "ee4e22b", default-features = false, features = [ "parallel", "jemalloc", "bench-metrics", diff --git a/openvm/src/customize_exe.rs b/openvm/src/customize_exe.rs index bcedb4cad..0e89611a2 100644 --- a/openvm/src/customize_exe.rs +++ b/openvm/src/customize_exe.rs @@ -5,9 +5,9 @@ use std::sync::Arc; use crate::extraction_utils::{OriginalAirs, OriginalVmConfig}; use crate::opcode::{branch_opcodes_bigint_set, branch_opcodes_set}; use crate::utils::{fractional_knapsack, KnapsackItem, UnsupportedOpenVmReferenceError}; -use crate::IntoOpenVm; use crate::OpenVmField; use crate::OriginalCompiledProgram; +use crate::{AirMetrics, IntoOpenVm}; use crate::{CompiledProgram, SpecializedConfig}; use itertools::Itertools; use openvm_instructions::instruction::Instruction; @@ -519,13 +519,13 @@ fn create_apcs_with_cell_pgo( } } - let max_total_apc_columns = max_total_columns.map(|max_total_columns| { - let chip_inventory_air_widths = original_config.chip_inventory_air_widths(); - let total_non_apc_columns = chip_inventory_air_widths + let max_total_apc_columns: Option = max_total_columns.map(|max_total_columns| { + let chip_inventory_air_metrics = original_config.chip_inventory_air_metrics(); + let total_non_apc_columns = chip_inventory_air_metrics .iter() - .map(|(air_name, width)| { - tracing::debug!("Chip inventory air {} has width {}", air_name, width); - width + .map(|AirMetrics { name, widths, .. }| { + tracing::debug!("Chip inventory air {} has {}", name, widths); + widths.preprocessed + widths.main + widths.log_up }) .sum::(); max_total_columns - total_non_apc_columns diff --git a/openvm/src/extraction_utils.rs b/openvm/src/extraction_utils.rs index 536268ed3..90b608190 100644 --- a/openvm/src/extraction_utils.rs +++ b/openvm/src/extraction_utils.rs @@ -2,13 +2,17 @@ use std::collections::{BTreeMap, BTreeSet, HashMap}; use std::sync::{Arc, Mutex}; use crate::air_builder::AirKeygenBuilder; -use crate::IntoOpenVm; use crate::{opcode::instruction_allowlist, BabyBearSC, SpecializedConfig}; +use crate::{AirMetrics, IntoOpenVm}; use openvm_circuit::arch::{VmChipComplex, VmConfig, VmInventoryError}; use openvm_circuit_primitives::bitwise_op_lookup::SharedBitwiseOperationLookupChip; use openvm_circuit_primitives::range_tuple::SharedRangeTupleCheckerChip; use openvm_instructions::VmOpcode; use openvm_sdk::config::{SdkVmConfig, SdkVmConfigExecutor, SdkVmConfigPeriphery}; +use openvm_stark_backend::air_builders::symbolic::SymbolicRapBuilder; +use openvm_stark_backend::config::{PackedChallenge, Val}; +use openvm_stark_backend::interaction::fri_log_up::find_interaction_chunks; +use openvm_stark_backend::p3_field::FieldExtensionAlgebra; use openvm_stark_backend::{ air_builders::symbolic::SymbolicConstraints, config::StarkGenericConfig, rap::AnyRap, Chip, }; @@ -272,18 +276,24 @@ impl OriginalVmConfig { self.sdk_config.create_chip_complex() } - pub fn chip_inventory_air_widths(&self) -> HashMap { - self.chip_complex() - .inventory + pub fn chip_inventory_air_metrics(&self) -> Vec { + let inventory = &self.chip_complex().inventory; + + inventory .executors() .iter() - .map(|executor| { - let air: Arc> = executor.air(); - let name = air.name(); - let width = air.width(); - (name, width) + .map(|executor| executor.air()) + .chain( + inventory + .periphery() + .iter() + .map(|periphery| periphery.air()), + ) + .map(|air| { + // both executors and periphery implement the same `air()` API + get_air_metrics(air) }) - .collect::>() + .collect() } } @@ -329,12 +339,71 @@ pub fn get_name(air: Arc>) -> String { pub fn get_constraints( air: Arc>, ) -> SymbolicConstraints { + let builder = symbolic_builder_with_degree(air, None); + builder.constraints() +} + +pub fn get_air_metrics(air: Arc>) -> AirMetrics { + let app_log_blow_up = 2; + let max_degree = (1 << app_log_blow_up) + 1; + + let name = air.name(); + let main = air.width(); + + let symbolic_rap_builder = symbolic_builder_with_degree(air, Some(max_degree)); + let preprocessed = symbolic_rap_builder.width().preprocessed.unwrap_or(0); + + let SymbolicConstraints { + constraints, + interactions, + } = symbolic_rap_builder.constraints(); + + let log_up = (find_interaction_chunks(&interactions, max_degree) + .interaction_partitions() + .len() + + 1) + * as FieldExtensionAlgebra>>::D; + + AirMetrics { + name, + widths: AirWidths { + preprocessed, + main, + log_up, + }, + constraints: constraints.len(), + bus_interactions: interactions.len(), + } +} + +pub fn symbolic_builder_with_degree( + air: Arc>, + max_constraint_degree: Option, +) -> SymbolicRapBuilder { let perm = default_perm(); let security_params = SecurityParameters::standard_fast(); let config = config_from_perm(&perm, security_params); let air_keygen_builder = AirKeygenBuilder::new(config.pcs(), air); - let builder = air_keygen_builder.get_symbolic_builder(None); - builder.constraints() + air_keygen_builder.get_symbolic_builder(max_constraint_degree) +} + +pub struct AirWidths { + pub preprocessed: usize, + pub main: usize, + pub log_up: usize, +} + +impl std::fmt::Display for AirWidths { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!( + f, + "Total Width: {} (Preprocessed: {} Main: {}, Log Up: {})", + self.preprocessed + self.main + self.log_up, + self.preprocessed, + self.main, + self.log_up + ) + } } #[cfg(test)] diff --git a/openvm/src/lib.rs b/openvm/src/lib.rs index 4632b0001..4b09c227e 100644 --- a/openvm/src/lib.rs +++ b/openvm/src/lib.rs @@ -45,7 +45,7 @@ use tracing_subscriber::{ Layer, }; -use crate::extraction_utils::{export_pil, get_constraints, OriginalVmConfig}; +use crate::extraction_utils::{export_pil, get_air_metrics, AirWidths, OriginalVmConfig}; use crate::instruction_formatter::openvm_opcode_formatter; use crate::powdr_extension::PowdrPrecompile; use crate::traits::OpenVmField; @@ -508,7 +508,7 @@ pub struct OriginalCompiledProgram { pub struct AirMetrics { pub name: String, - pub width: usize, + pub widths: AirWidths, pub constraints: usize, pub bus_interactions: usize, } @@ -523,20 +523,13 @@ impl CompiledProgram { .iter() .filter_map(|executor| { let air = executor.air(); - let width = air.width(); let name = air.name(); // We actually give name "powdr_air_for_opcode_" to the AIRs, // but OpenVM uses the actual Rust type (PowdrAir) as the name in this method. // TODO this is hacky but not sure how to do it better rn. if name.starts_with("PowdrAir") || name.starts_with("PlonkAir") { - let constraints = get_constraints(air); - Some(AirMetrics { - name: name.to_string(), - width, - constraints: constraints.constraints.len(), - bus_interactions: constraints.interactions.len(), - }) + Some(get_air_metrics(air)) } else { None } @@ -1049,7 +1042,10 @@ mod tests { .powdr_airs_metrics(); assert_eq!(machines.len(), 1); let m = &machines[0]; - assert_eq!([m.width, m.constraints, m.bus_interactions], [49, 22, 31]); + assert_eq!( + [m.widths.main, m.constraints, m.bus_interactions], + [49, 22, 31] + ); } fn test_keccak_machine(pgo_config: PgoConfig) { @@ -1060,7 +1056,7 @@ mod tests { assert_eq!(machines.len(), 1); let m = &machines[0]; assert_eq!( - [m.width, m.constraints, m.bus_interactions], + [m.widths.main, m.constraints, m.bus_interactions], [2011, 166, 1783] ); } @@ -1083,7 +1079,7 @@ mod tests { .powdr_airs_metrics(); assert_eq!(machines.len(), 1); let m = &machines[0]; - assert_eq!(m.width, 26); + assert_eq!(m.widths.main, 26); assert_eq!(m.constraints, 1); assert_eq!(m.bus_interactions, 16); }