Skip to content

Proof of concept of column names without touching stark-backend #46

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Schaeff wants to merge 1 commit into
base: upstream-main
Choose a base branch
from
Open

Proof of concept of column names without touching stark-backend #46

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 15 additions & 5 deletions crates/vm/src/arch/extensions.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ use openvm_stark_backend::{
types::{AirProvingContext, ProvingContext},
},
rap::AnyRap,
AirRef, AnyChip, Chip,
AnyChip, Chip,
};
use rustc_hash::FxHashMap;
use tracing::info_span;
Expand Down Expand Up @@ -145,7 +145,7 @@ pub struct AirInventory<SC: StarkGenericConfig> {
/// Note that the system will ensure that the first AIR in the list is always the
/// [VariableRangeCheckerAir].
#[get = "pub"]
ext_airs: Vec<AirRef<SC>>,
ext_airs: Vec<AirRefWithColumnNames<SC>>,
/// `ext_start[i]` will have the starting index in `ext_airs` for extension `i`
ext_start: Vec<usize>,

Expand Down Expand Up @@ -399,6 +399,16 @@ impl<F, E> ExecutorInventoryBuilder<'_, F, E> {
}
}

pub trait ColumnNames {}
Copy link
Author

@Schaeff Schaeff Oct 31, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This trait would have the column_name method


impl<T> ColumnNames for T {}
Copy link
Author

@Schaeff Schaeff Oct 31, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Currently we implement it for all types just to prove the approach, but in practice this would be done for each air, as we currently do.


pub trait AnyRapWithColumnNames<SC: StarkGenericConfig>: AnyRap<SC> + ColumnNames {}
Copy link
Author

@Schaeff Schaeff Oct 31, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This extends the stark-backend trait with our ColumnNames trait


pub type AirRefWithColumnNames<SC> = Arc<dyn AnyRapWithColumnNames<SC>>;
Copy link
Author

@Schaeff Schaeff Oct 31, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This does the same for AirRef


impl<SC: StarkGenericConfig, R: AnyRap<SC> + ColumnNames> AnyRapWithColumnNames<SC> for R {}
Copy link
Author

@Schaeff Schaeff Oct 31, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This implementation would stay: if AnyRap and ColumnNames are implemented, then AnyRapWithColumnNames is implemented


impl<SC: StarkGenericConfig> AirInventory<SC> {
/// Outside of this crate, [AirInventory] must be constructed via [SystemConfig].
pub(crate) fn new(
Expand Down Expand Up @@ -434,11 +444,11 @@ impl<SC: StarkGenericConfig> AirInventory<SC> {
.filter_map(|air| air.as_any().downcast_ref())
}

pub fn add_air<A: AnyRap<SC> + 'static>(&mut self, air: A) {
pub fn add_air<A: AnyRapWithColumnNames<SC> + 'static>(&mut self, air: A) {
Copy link
Author

@Schaeff Schaeff Oct 31, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We replace all usages of AnyRap by our extension.

self.add_air_ref(Arc::new(air));
}

pub fn add_air_ref(&mut self, air: AirRef<SC>) {
pub fn add_air_ref(&mut self, air: AirRefWithColumnNames<SC>) {
self.ext_airs.push(air);
}

Expand All @@ -452,7 +462,7 @@ impl<SC: StarkGenericConfig> AirInventory<SC> {
/// This is the system AIRs, followed by the other AIRs in the **reverse** of the order they
/// were added in the VM extension definitions. In particular, the AIRs that have dependencies
/// appear later. The system guarantees that the last AIR is the [VariableRangeCheckerAir].
pub fn into_airs(self) -> impl Iterator<Item = AirRef<SC>> {
pub fn into_airs(self) -> impl Iterator<Item = AirRefWithColumnNames<SC>> {
self.system
.into_airs()
.into_iter()
Expand Down
4 changes: 3 additions & 1 deletion crates/vm/src/arch/vm.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1315,6 +1315,8 @@ pub fn debug_proving_ctx<E, VB>(
let (airs, pks, proof_inputs): (Vec<_>, Vec<_>, Vec<_>) =
multiunzip(ctx.per_air.iter().map(|(air_id, air_ctx)| {
// Transfer from device **back** to host so the debugger can read the data.

use openvm_stark_backend::rap::AnyRap;
let cached_mains = air_ctx
.cached_mains
.iter()
Expand All @@ -1331,7 +1333,7 @@ pub fn debug_proving_ctx<E, VB>(
public_values,
};
(
global_airs[*air_id].clone(),
global_airs[*air_id].clone() as Arc<dyn AnyRap<_>>,
Copy link
Author

@Schaeff Schaeff Oct 31, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

When we need to reach into stark-backend, we lower to the original AnyRap

Copy link
Author

@Schaeff Schaeff Oct 31, 2025 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

There are no other compile errors, so I suspect Rust is able to do this lowering automatically in other cases (for example, keygen)

pk.per_air[*air_id].clone(),
raw,
)
Expand Down
21 changes: 10 additions & 11 deletions crates/vm/src/system/memory/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@ use openvm_stark_backend::{
interaction::PermutationCheckBus,
p3_field::Field,
p3_util::{log2_ceil_usize, log2_strict_usize},
AirRef,
};

pub mod adapter;
Expand All @@ -24,7 +23,7 @@ pub use controller::*;
pub use online::{Address, AddressMap, INITIAL_TIMESTAMP};

use crate::{
arch::{MemoryConfig, ADDR_SPACE_OFFSET},
arch::{AirRefWithColumnNames, MemoryConfig, ADDR_SPACE_OFFSET},
system::memory::{
adapter::AccessAdapterAir, dimensions::MemoryDimensions, interface::MemoryInterfaceAirs,
merkle::MemoryMerkleAir, offline_checker::MemoryBridge, persistent::PersistentBoundaryAir,
Expand Down Expand Up @@ -75,7 +74,7 @@ impl<S, T> MemoryAddress<S, T> {
pub struct MemoryAirInventory<SC: StarkGenericConfig> {
pub bridge: MemoryBridge,
pub interface: MemoryInterfaceAirs,
pub access_adapters: Vec<AirRef<SC>>,
pub access_adapters: Vec<AirRefWithColumnNames<SC>>,
}

impl<SC: StarkGenericConfig> MemoryAirInventory<SC> {
Expand Down Expand Up @@ -122,12 +121,12 @@ impl<SC: StarkGenericConfig> MemoryAirInventory<SC> {
let lt_air = IsLtSubAir::new(range_bus, mem_config.timestamp_max_bits);
let maan = mem_config.max_access_adapter_n;
assert!(matches!(maan, 2 | 4 | 8 | 16 | 32));
let access_adapters: Vec<AirRef<SC>> = [
Arc::new(AccessAdapterAir::<2> { memory_bus, lt_air }) as AirRef<SC>,
Arc::new(AccessAdapterAir::<4> { memory_bus, lt_air }) as AirRef<SC>,
Arc::new(AccessAdapterAir::<8> { memory_bus, lt_air }) as AirRef<SC>,
Arc::new(AccessAdapterAir::<16> { memory_bus, lt_air }) as AirRef<SC>,
Arc::new(AccessAdapterAir::<32> { memory_bus, lt_air }) as AirRef<SC>,
let access_adapters: Vec<AirRefWithColumnNames<SC>> = [
Arc::new(AccessAdapterAir::<2> { memory_bus, lt_air }) as AirRefWithColumnNames<SC>,
Arc::new(AccessAdapterAir::<4> { memory_bus, lt_air }) as AirRefWithColumnNames<SC>,
Arc::new(AccessAdapterAir::<8> { memory_bus, lt_air }) as AirRefWithColumnNames<SC>,
Arc::new(AccessAdapterAir::<16> { memory_bus, lt_air }) as AirRefWithColumnNames<SC>,
Arc::new(AccessAdapterAir::<32> { memory_bus, lt_air }) as AirRefWithColumnNames<SC>,
]
.into_iter()
.take(log2_strict_usize(maan))
Expand All @@ -141,8 +140,8 @@ impl<SC: StarkGenericConfig> MemoryAirInventory<SC> {
}

/// The order of memory AIRs is boundary, merkle (if exists), access adapters
pub fn into_airs(self) -> Vec<AirRef<SC>> {
let mut airs: Vec<AirRef<SC>> = Vec::new();
pub fn into_airs(self) -> Vec<AirRefWithColumnNames<SC>> {
let mut airs: Vec<AirRefWithColumnNames<SC>> = Vec::new();
match self.interface {
MemoryInterfaceAirs::Volatile { boundary } => {
airs.push(Arc::new(boundary));
Expand Down
17 changes: 9 additions & 8 deletions crates/vm/src/system/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,18 +19,19 @@ use openvm_stark_backend::{
hal::{MatrixDimensions, ProverBackend},
types::{AirProvingContext, CommittedTraceData},
},
AirRef, Chip,
Chip,
};
use rustc_hash::FxHashMap;

use self::{connector::VmConnectorAir, program::ProgramAir, public_values::PublicValuesAir};
use crate::{
arch::{
vm_poseidon2_config, AirInventory, AirInventoryError, BusIndexManager, ChipInventory,
ChipInventoryError, DenseRecordArena, ExecutionBridge, ExecutionBus, ExecutionState,
ExecutorInventory, ExecutorInventoryError, MatrixRecordArena, PhantomSubExecutor,
RowMajorMatrixArena, SystemConfig, VmAirWrapper, VmBuilder, VmChipComplex, VmChipWrapper,
VmCircuitConfig, VmExecutionConfig, CONNECTOR_AIR_ID, PROGRAM_AIR_ID, PUBLIC_VALUES_AIR_ID,
vm_poseidon2_config, AirInventory, AirInventoryError, AirRefWithColumnNames,
BusIndexManager, ChipInventory, ChipInventoryError, DenseRecordArena, ExecutionBridge,
ExecutionBus, ExecutionState, ExecutorInventory, ExecutorInventoryError, MatrixRecordArena,
PhantomSubExecutor, RowMajorMatrixArena, SystemConfig, VmAirWrapper, VmBuilder,
VmChipComplex, VmChipWrapper, VmCircuitConfig, VmExecutionConfig, CONNECTOR_AIR_ID,
PROGRAM_AIR_ID, PUBLIC_VALUES_AIR_ID,
},
system::{
connector::VmConnectorChip,
Expand Down Expand Up @@ -224,8 +225,8 @@ impl<SC: StarkGenericConfig> SystemAirInventory<SC> {
}
}

pub fn into_airs(self) -> Vec<AirRef<SC>> {
let mut airs: Vec<AirRef<SC>> = Vec::new();
pub fn into_airs(self) -> Vec<AirRefWithColumnNames<SC>> {
let mut airs: Vec<AirRefWithColumnNames<SC>> = Vec::new();
airs.push(Arc::new(self.program));
airs.push(Arc::new(self.connector));
if let Some(public_values) = self.public_values {
Expand Down
9 changes: 6 additions & 3 deletions crates/vm/src/system/poseidon2/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ use openvm_stark_backend::{
config::{StarkGenericConfig, Val},
interaction::{BusIndex, LookupBus},
p3_field::{Field, PrimeField32},
AirRef, ChipUsageGetter,
ChipUsageGetter,
};

#[cfg(test)]
Expand All @@ -27,7 +27,10 @@ mod chip;
pub use chip::*;

use crate::{
arch::hasher::{Hasher, HasherChip},
arch::{
hasher::{Hasher, HasherChip},
AirRefWithColumnNames,
},
system::poseidon2::air::Poseidon2PeripheryAir,
};
pub mod columns;
Expand Down Expand Up @@ -60,7 +63,7 @@ pub fn new_poseidon2_periphery_air<SC: StarkGenericConfig>(
poseidon2_config: Poseidon2Config<Val<SC>>,
direct_bus: LookupBus,
max_constraint_degree: usize,
) -> AirRef<SC> {
) -> AirRefWithColumnNames<SC> {
if max_constraint_degree >= 7 {
Arc::new(Poseidon2PeripheryAir::<Val<SC>, 0>::new(
Arc::new(Poseidon2SubAir::new(poseidon2_config.constants.into())),
Expand Down
Loading