diff --git a/crates/vm/src/arch/config.rs b/crates/vm/src/arch/config.rs index 3ffbfb74e0..13f23a45ae 100644 --- a/crates/vm/src/arch/config.rs +++ b/crates/vm/src/arch/config.rs @@ -26,9 +26,7 @@ use crate::{ Arena, ChipInventoryError, ExecutorInventory, ExecutorInventoryError, }, system::{ - memory::{ - merkle::public_values::PUBLIC_VALUES_AS, num_memory_airs, CHUNK, POINTER_MAX_BITS, - }, + memory::{merkle::public_values::PUBLIC_VALUES_AS, num_memory_airs, POINTER_MAX_BITS}, SystemChipComplex, }, }; @@ -123,6 +121,11 @@ pub const OPENVM_DEFAULT_INIT_FILE_NAME: &str = "openvm_init.rs"; const DEFAULT_U8_BLOCK_SIZE: usize = 4; const DEFAULT_NATIVE_BLOCK_SIZE: usize = 1; +/// The constant block size used for memory accesses when access adapters are disabled. +/// All memory accesses for address spaces 1-3 must use this block size. +/// This is also the block size used by the Boundary AIR for memory bus interactions. +pub const CONST_BLOCK_SIZE: usize = DEFAULT_U8_BLOCK_SIZE; + /// Trait for generating a init.rs file that contains a call to moduli_init!, /// complex_init!, sw_init! with the supported moduli and curves. /// Should be implemented by all VM config structs. @@ -183,6 +186,10 @@ pub struct MemoryConfig { pub decomp: usize, /// Maximum N AccessAdapter AIR to support. pub max_access_adapter_n: usize, + /// Whether access adapters are enabled. When disabled, all memory accesses must be of the + /// standard block size (ie, 4 for address spaces 1-3). + #[new(value = "true")] + pub access_adapters_enabled: bool, } impl Default for MemoryConfig { @@ -194,7 +201,15 @@ impl Default for MemoryConfig { addr_spaces[RV32_MEMORY_AS as usize].num_cells = MAX_CELLS; addr_spaces[PUBLIC_VALUES_AS as usize].num_cells = DEFAULT_MAX_NUM_PUBLIC_VALUES; addr_spaces[NATIVE_AS as usize].num_cells = MAX_CELLS; - Self::new(3, addr_spaces, POINTER_MAX_BITS, 29, 17, 32) + Self { + addr_space_height: 3, + addr_spaces, + pointer_max_bits: POINTER_MAX_BITS, + timestamp_max_bits: 29, + decomp: 17, + max_access_adapter_n: 32, + access_adapters_enabled: true, + } } } @@ -245,6 +260,36 @@ impl MemoryConfig { .map(|addr_sp| log2_strict_usize(addr_sp.min_block_size) as u8) .collect() } + + /// Returns true if the Native address space (AS 4) is used + /// Native AS is considered "used" if it has any allocated cells + pub fn is_native_as_used(&self) -> bool { + self.addr_spaces + .get(NATIVE_AS as usize) + .is_some_and(|config| config.num_cells > 0) + } + + /// Disables access adapters. When disabled, all memory accesses for address spaces 1-3 + /// must use the constant block size (4). Access adapters will only be used for + /// address space 4 (Native) if it is enabled. + pub fn without_access_adapters(mut self) -> Self { + self.access_adapters_enabled = false; + self + } + + /// Enables access adapters. This is the default behavior + pub fn with_access_adapters(mut self) -> Self { + self.access_adapters_enabled = true; + self + } + + /// Automatically sets `access_adapters_enabled` based on whether Native AS is used. + /// If Native AS is not used, access adapters are disabled since all other address spaces + /// use a fixed block size of 4 + pub fn with_auto_access_adapters(mut self) -> Self { + self.access_adapters_enabled = self.is_native_as_used(); + self + } } /// System-level configuration for the virtual machine. Contains all configuration parameters that @@ -375,15 +420,42 @@ impl SystemConfig { + num_memory_airs( self.continuation_enabled, self.memory_config.max_access_adapter_n, + self.memory_config.access_adapters_enabled, ) } pub fn initial_block_size(&self) -> usize { match self.continuation_enabled { - true => CHUNK, + true => CONST_BLOCK_SIZE, false => 1, } } + + /// Disables access adapters. When disabled, all memory accesses for address spaces 1-3 + /// must use the constant block size (4) + pub fn without_access_adapters(mut self) -> Self { + self.memory_config.access_adapters_enabled = false; + self + } + + /// Enables access adapters. This is the default behavior. + pub fn with_access_adapters(mut self) -> Self { + self.memory_config.access_adapters_enabled = true; + self + } + + /// Automatically sets `access_adapters_enabled` based on whether Native AS is used. + /// If Native AS is not used, access adapters are disabled since all other address spaces + /// use a fixed block size of 4. + pub fn with_auto_access_adapters(mut self) -> Self { + self.memory_config = self.memory_config.with_auto_access_adapters(); + self + } + + /// Returns true if access adapters are enabled. + pub fn access_adapters_enabled(&self) -> bool { + self.memory_config.access_adapters_enabled + } } impl Default for SystemConfig { diff --git a/crates/vm/src/arch/execution_mode/metered/ctx.rs b/crates/vm/src/arch/execution_mode/metered/ctx.rs index 8428438ca7..0b67a3b92d 100644 --- a/crates/vm/src/arch/execution_mode/metered/ctx.rs +++ b/crates/vm/src/arch/execution_mode/metered/ctx.rs @@ -64,11 +64,13 @@ impl MeteredCtx { air_names[merkle_tree_index] ); } - debug_assert!( - air_names[memory_ctx.adapter_offset].contains("AccessAdapterAir<2>"), - "air_name={}", - air_names[memory_ctx.adapter_offset] - ); + if memory_ctx.access_adapters_enabled { + debug_assert!( + air_names[memory_ctx.adapter_offset].contains("AccessAdapterAir<2>"), + "air_name={}", + air_names[memory_ctx.adapter_offset] + ); + } let segmentation_ctx = SegmentationCtx::new(air_names, widths, interactions, config.segmentation_limits); diff --git a/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs b/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs index 3429177d11..2397b1b12e 100644 --- a/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs +++ b/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs @@ -105,6 +105,7 @@ pub struct MemoryCtx { pub boundary_idx: usize, pub merkle_tree_index: Option, pub adapter_offset: usize, + pub access_adapters_enabled: bool, continuations_enabled: bool, chunk: u32, chunk_bits: u32, @@ -128,6 +129,7 @@ impl MemoryCtx { boundary_idx: config.memory_boundary_air_id(), merkle_tree_index: config.memory_merkle_air_id(), adapter_offset: config.access_adapter_air_id_offset(), + access_adapters_enabled: config.memory_config.access_adapters_enabled, chunk, chunk_bits, memory_dimensions, @@ -210,6 +212,10 @@ impl MemoryCtx { size_bits: u32, num: u32, ) { + if !self.access_adapters_enabled { + return; + } + debug_assert!((address_space as usize) < self.min_block_size_bits.len()); // SAFETY: address_space passed is usually a hardcoded constant or derived from an diff --git a/crates/vm/src/arch/execution_mode/metered_cost.rs b/crates/vm/src/arch/execution_mode/metered_cost.rs index 925bd25af2..69bfd6fe69 100644 --- a/crates/vm/src/arch/execution_mode/metered_cost.rs +++ b/crates/vm/src/arch/execution_mode/metered_cost.rs @@ -18,6 +18,7 @@ pub const DEFAULT_MAX_COST: u64 = DEFAULT_MAX_SEGMENTS * DEFAULT_SEGMENT_MAX_CEL pub struct AccessAdapterCtx { min_block_size_bits: Vec, idx_offset: usize, + enabled: bool, } impl AccessAdapterCtx { @@ -25,6 +26,7 @@ impl AccessAdapterCtx { Self { min_block_size_bits: config.memory_config.min_block_size_bits(), idx_offset: config.access_adapter_air_id_offset(), + enabled: config.memory_config.access_adapters_enabled, } } @@ -36,6 +38,10 @@ impl AccessAdapterCtx { size_bits: u32, widths: &[usize], ) { + if !self.enabled { + return; + } + debug_assert!((address_space as usize) < self.min_block_size_bits.len()); // SAFETY: address_space passed is usually a hardcoded constant or derived from an diff --git a/crates/vm/src/arch/testing/cpu.rs b/crates/vm/src/arch/testing/cpu.rs index 105962bc11..cff813801d 100644 --- a/crates/vm/src/arch/testing/cpu.rs +++ b/crates/vm/src/arch/testing/cpu.rs @@ -1,6 +1,7 @@ use std::sync::Arc; use itertools::zip_eq; +use openvm_circuit::arch::CONST_BLOCK_SIZE; use openvm_circuit_primitives::var_range::{ SharedVariableRangeCheckerChip, VariableRangeCheckerBus, VariableRangeCheckerChip, }; @@ -48,7 +49,7 @@ use crate::{ adapter::records::arena_size_bound, offline_checker::{MemoryBridge, MemoryBus}, online::TracingMemory, - MemoryAirInventory, MemoryController, SharedMemoryHelper, CHUNK, + MemoryAirInventory, MemoryController, SharedMemoryHelper, }, poseidon2::Poseidon2PeripheryChip, program::ProgramBus, @@ -323,6 +324,7 @@ impl VmChipTestBuilder { let mut mem_config = MemoryConfig::default(); mem_config.addr_spaces[RV32_REGISTER_AS as usize].num_cells = 1 << 29; mem_config.addr_spaces[NATIVE_AS as usize].num_cells = 0; + // TODO: Check if need to revert to volatile memory, after access adapters are removed Self::persistent(mem_config) } @@ -347,7 +349,7 @@ impl VmChipTestBuilder { pub fn persistent(mem_config: MemoryConfig) -> Self { setup_tracing_with_log_level(Level::INFO); - let (range_checker, memory) = Self::range_checker_and_memory(&mem_config, CHUNK); + let (range_checker, memory) = Self::range_checker_and_memory(&mem_config, CONST_BLOCK_SIZE); let hasher_chip = Arc::new(Poseidon2PeripheryChip::new( vm_poseidon2_config(), POSEIDON2_DIRECT_BUS, @@ -403,7 +405,7 @@ impl Default for VmChipTestBuilder { // removed when tests are updated. mem_config.addr_spaces[RV32_REGISTER_AS as usize].num_cells = 1 << 29; mem_config.addr_spaces[NATIVE_AS as usize].num_cells = 0; - Self::volatile(mem_config) + Self::persistent(mem_config) } } diff --git a/crates/vm/src/arch/vm.rs b/crates/vm/src/arch/vm.rs index 68555050fe..23db960d6d 100644 --- a/crates/vm/src/arch/vm.rs +++ b/crates/vm/src/arch/vm.rs @@ -622,7 +622,13 @@ where let system_config: &SystemConfig = self.config().as_ref(); let adapter_offset = system_config.access_adapter_air_id_offset(); // ATTENTION: this must agree with `num_memory_airs` - let num_adapters = log2_strict_usize(system_config.memory_config.max_access_adapter_n); + + let num_adapters = if system_config.memory_config.access_adapters_enabled { + log2_strict_usize(system_config.memory_config.max_access_adapter_n) + } else { + 0 + }; + assert_eq!(adapter_offset + num_adapters, system_config.num_airs()); let access_adapter_arena_size_bound = records::arena_size_bound( &trace_heights[adapter_offset..adapter_offset + num_adapters], diff --git a/crates/vm/src/system/memory/adapter/mod.rs b/crates/vm/src/system/memory/adapter/mod.rs index 8b0797dcf6..a9c89fc2ea 100644 --- a/crates/vm/src/system/memory/adapter/mod.rs +++ b/crates/vm/src/system/memory/adapter/mod.rs @@ -58,21 +58,26 @@ impl AccessAdapterInventory { memory_bus: MemoryBus, memory_config: MemoryConfig, ) -> Self { - let rc = range_checker; - let mb = memory_bus; - let tmb = memory_config.timestamp_max_bits; - let maan = memory_config.max_access_adapter_n; - assert!(matches!(maan, 2 | 4 | 8 | 16 | 32)); - let chips: Vec<_> = [ - Self::create_access_adapter_chip::<2>(rc.clone(), mb, tmb, maan), - Self::create_access_adapter_chip::<4>(rc.clone(), mb, tmb, maan), - Self::create_access_adapter_chip::<8>(rc.clone(), mb, tmb, maan), - Self::create_access_adapter_chip::<16>(rc.clone(), mb, tmb, maan), - Self::create_access_adapter_chip::<32>(rc.clone(), mb, tmb, maan), - ] - .into_iter() - .flatten() - .collect(); + // Only create adapter chips if access adapters are enabled + let chips: Vec<_> = if memory_config.access_adapters_enabled { + let rc = range_checker; + let mb = memory_bus; + let tmb = memory_config.timestamp_max_bits; + let maan = memory_config.max_access_adapter_n; + assert!(matches!(maan, 2 | 4 | 8 | 16 | 32)); + [ + Self::create_access_adapter_chip::<2>(rc.clone(), mb, tmb, maan), + Self::create_access_adapter_chip::<4>(rc.clone(), mb, tmb, maan), + Self::create_access_adapter_chip::<8>(rc.clone(), mb, tmb, maan), + Self::create_access_adapter_chip::<16>(rc.clone(), mb, tmb, maan), + Self::create_access_adapter_chip::<32>(rc.clone(), mb, tmb, maan), + ] + .into_iter() + .flatten() + .collect() + } else { + Vec::new() + }; Self { memory_config, chips, diff --git a/crates/vm/src/system/memory/controller/mod.rs b/crates/vm/src/system/memory/controller/mod.rs index aabe4df08d..e4733ccdf4 100644 --- a/crates/vm/src/system/memory/controller/mod.rs +++ b/crates/vm/src/system/memory/controller/mod.rs @@ -14,7 +14,6 @@ use openvm_stark_backend::{ interaction::PermutationCheckBus, p3_commit::PolynomialSpace, p3_field::{Field, PrimeField32}, - p3_maybe_rayon::prelude::{IntoParallelIterator, ParallelIterator}, p3_util::{log2_ceil_usize, log2_strict_usize}, prover::{cpu::CpuBackend, types::AirProvingContext}, Chip, @@ -24,7 +23,7 @@ use serde::{Deserialize, Serialize}; use self::interface::MemoryInterface; use super::{volatile::VolatileBoundaryChip, AddressMap}; use crate::{ - arch::{DenseRecordArena, MemoryConfig, ADDR_SPACE_OFFSET}, + arch::{DenseRecordArena, MemoryConfig, ADDR_SPACE_OFFSET, CONST_BLOCK_SIZE}, system::{ memory::{ adapter::AccessAdapterInventory, @@ -291,10 +290,30 @@ impl MemoryController { ) => { let hasher = self.hasher_chip.as_ref().unwrap(); boundary_chip.finalize(initial_memory, &final_memory, hasher.as_ref()); - let final_memory_values = final_memory - .into_par_iter() - .map(|(key, value)| (key, value.values)) - .collect(); + + // Rechunk CONST_BLOCK_SIZE blocks into CHUNK-sized blocks for merkle_chip + // Note: Equipartition key is (addr_space, ptr) where ptr is the starting pointer + let final_memory_values: Equipartition = { + use std::collections::BTreeMap; + let mut chunk_map: BTreeMap<(u32, u32), [F; CHUNK]> = BTreeMap::new(); + for ((addr_space, ptr), ts_values) in final_memory.into_iter() { + // Align to CHUNK boundary to get the chunk's starting pointer + let chunk_ptr = (ptr / CHUNK as u32) * CHUNK as u32; + let block_idx_in_chunk = + ((ptr % CHUNK as u32) / CONST_BLOCK_SIZE as u32) as usize; + let entry = chunk_map.entry((addr_space, chunk_ptr)).or_insert_with(|| { + // Initialize with values from initial memory + std::array::from_fn(|i| unsafe { + initial_memory.get_f::(addr_space, chunk_ptr + i as u32) + }) + }); + // Copy values for this block + for (i, val) in ts_values.values.into_iter().enumerate() { + entry[block_idx_in_chunk * CONST_BLOCK_SIZE + i] = val; + } + } + chunk_map + }; merkle_chip.finalize(initial_memory, &final_memory_values, hasher.as_ref()); } _ => panic!("TouchedMemory incorrect type"), diff --git a/crates/vm/src/system/memory/mod.rs b/crates/vm/src/system/memory/mod.rs index 411e7a5473..8c3f48c7f0 100644 --- a/crates/vm/src/system/memory/mod.rs +++ b/crates/vm/src/system/memory/mod.rs @@ -118,20 +118,24 @@ impl MemoryAirInventory { ); MemoryInterfaceAirs::Volatile { boundary } }; - // Memory access adapters - 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> = [ - Arc::new(AccessAdapterAir::<2> { memory_bus, lt_air }) as AirRef, - Arc::new(AccessAdapterAir::<4> { memory_bus, lt_air }) as AirRef, - Arc::new(AccessAdapterAir::<8> { memory_bus, lt_air }) as AirRef, - Arc::new(AccessAdapterAir::<16> { memory_bus, lt_air }) as AirRef, - Arc::new(AccessAdapterAir::<32> { memory_bus, lt_air }) as AirRef, - ] - .into_iter() - .take(log2_strict_usize(maan)) - .collect(); + // Memory access adapters - only create if enabled + let access_adapters: Vec> = if mem_config.access_adapters_enabled { + 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)); + [ + Arc::new(AccessAdapterAir::<2> { memory_bus, lt_air }) as AirRef, + Arc::new(AccessAdapterAir::<4> { memory_bus, lt_air }) as AirRef, + Arc::new(AccessAdapterAir::<8> { memory_bus, lt_air }) as AirRef, + Arc::new(AccessAdapterAir::<16> { memory_bus, lt_air }) as AirRef, + Arc::new(AccessAdapterAir::<32> { memory_bus, lt_air }) as AirRef, + ] + .into_iter() + .take(log2_strict_usize(maan)) + .collect() + } else { + Vec::new() + }; Self { bridge, @@ -159,7 +163,16 @@ impl MemoryAirInventory { /// This is O(1) and returns the length of /// [`MemoryAirInventory::into_airs`]. -pub fn num_memory_airs(is_persistent: bool, max_access_adapter_n: usize) -> usize { - // boundary + { merkle if is_persistent } + access_adapters - 1 + usize::from(is_persistent) + log2_strict_usize(max_access_adapter_n) +pub fn num_memory_airs( + is_persistent: bool, + max_access_adapter_n: usize, + access_adapters_enabled: bool, +) -> usize { + // boundary + { merkle if is_persistent } + access_adapters (if enabled) + let num_adapters = if access_adapters_enabled { + log2_strict_usize(max_access_adapter_n) + } else { + 0 + }; + 1 + usize::from(is_persistent) + num_adapters } diff --git a/crates/vm/src/system/memory/online.rs b/crates/vm/src/system/memory/online.rs index 6a16e0d12b..fb66845c0b 100644 --- a/crates/vm/src/system/memory/online.rs +++ b/crates/vm/src/system/memory/online.rs @@ -13,12 +13,12 @@ use tracing::instrument; use crate::{ arch::{ AddressSpaceHostConfig, AddressSpaceHostLayout, DenseRecordArena, MemoryConfig, - RecordArena, MAX_CELL_BYTE_SIZE, + RecordArena, CONST_BLOCK_SIZE, MAX_CELL_BYTE_SIZE, }, system::{ memory::{ adapter::records::{AccessLayout, AccessRecordHeader, MERGE_AND_NOT_SPLIT_FLAG}, - MemoryAddress, TimestampedEquipartition, TimestampedValues, CHUNK, + MemoryAddress, TimestampedEquipartition, TimestampedValues, }, TouchedMemory, }, @@ -944,7 +944,7 @@ impl TracingMemory { self.touched_blocks_to_equipartition::(touched_blocks), ), true => TouchedMemory::Persistent( - self.touched_blocks_to_equipartition::(touched_blocks), + self.touched_blocks_to_equipartition::(touched_blocks), ), } } diff --git a/crates/vm/src/system/memory/persistent.rs b/crates/vm/src/system/memory/persistent.rs index eeb22cbfd6..cb2be60c6b 100644 --- a/crates/vm/src/system/memory/persistent.rs +++ b/crates/vm/src/system/memory/persistent.rs @@ -20,15 +20,20 @@ use openvm_stark_backend::{ use rustc_hash::FxHashSet; use tracing::instrument; -use super::{merkle::SerialReceiver, online::INITIAL_TIMESTAMP, TimestampedValues}; +use super::{merkle::SerialReceiver, online::INITIAL_TIMESTAMP}; use crate::{ - arch::{hasher::Hasher, ADDR_SPACE_OFFSET}, + arch::{hasher::Hasher, ADDR_SPACE_OFFSET, CONST_BLOCK_SIZE}, system::memory::{ - dimensions::MemoryDimensions, offline_checker::MemoryBus, MemoryAddress, MemoryImage, - TimestampedEquipartition, + controller::CHUNK, dimensions::MemoryDimensions, offline_checker::MemoryBus, MemoryAddress, + MemoryImage, TimestampedEquipartition, }, }; +/// Number of CONST_BLOCK_SIZE blocks per CHUNK (e.g., 2 for 8/4). +/// Blocks are on the same row only for Merkle tree hashing (8 bytes at a time). +/// Memory bus interactions use per-block timestamps. +pub const BLOCKS_PER_CHUNK: usize = CHUNK / CONST_BLOCK_SIZE; + /// The values describe aligned chunk of memory of size `CHUNK`---the data together with the last /// accessed timestamp---in either the initial or final memory state. #[repr(C)] @@ -42,7 +47,10 @@ pub struct PersistentBoundaryCols { pub leaf_label: T, pub values: [T; CHUNK], pub hash: [T; CHUNK], - pub timestamp: T, + /// Per-block timestamps. Each CONST_BLOCK_SIZE block within the chunk has its own timestamp. + /// For untouched blocks, timestamp stays at 0 (balances: boundary sends at t=0 init, receives + /// at t=0 final). + pub timestamps: [T; BLOCKS_PER_CHUNK], } /// Imposes the following constraints: @@ -81,12 +89,14 @@ impl Air for PersistentBoundaryA local.expand_direction * local.expand_direction * local.expand_direction, ); - // Constrain that an "initial" row has timestamp zero. + // Constrain that an "initial" row has all timestamp zero. // Since `direction` is constrained to be in {-1, 0, 1}, we can select `direction == 1` // with the constraint below. - builder - .when(local.expand_direction * (local.expand_direction + AB::F::ONE)) - .assert_zero(local.timestamp); + let mut when_initial = + builder.when(local.expand_direction * (local.expand_direction + AB::F::ONE)); + for i in 0..BLOCKS_PER_CHUNK { + when_initial.assert_zero(local.timestamps[i]); + } let mut expand_fields = vec![ // direction = 1 => is_final = 0 @@ -109,16 +119,23 @@ impl Air for PersistentBoundaryA local.expand_direction * local.expand_direction, ); - self.memory_bus - .send( - MemoryAddress::new( - local.address_space, - local.leaf_label * AB::F::from_canonical_usize(CHUNK), - ), - local.values.to_vec(), - local.timestamp, - ) - .eval(builder, local.expand_direction); + let chunk_size_f = AB::F::from_canonical_usize(CHUNK); + for block_idx in 0..BLOCKS_PER_CHUNK { + let offset = AB::F::from_canonical_usize(block_idx * CONST_BLOCK_SIZE); + // Split the 1xCHUNK leaf into CONST_BLOCK_SIZE-sized bus messages. + // Each block uses its own timestamp - untouched blocks stay at t=0. + self.memory_bus + .send( + MemoryAddress::new( + local.address_space, + local.leaf_label * chunk_size_f + offset, + ), + local.values[block_idx * CONST_BLOCK_SIZE..(block_idx + 1) * CONST_BLOCK_SIZE] + .to_vec(), + local.timestamps[block_idx], + ) + .eval(builder, local.expand_direction); + } } } @@ -142,7 +159,8 @@ pub struct FinalTouchedLabel { final_values: [F; CHUNK], init_hash: [F; CHUNK], final_hash: [F; CHUNK], - final_timestamp: u32, + /// Per-block timestamps. Each CONST_BLOCK_SIZE block has its own timestamp. + final_timestamps: [u32; BLOCKS_PER_CHUNK], } impl Default for TouchedLabels { @@ -207,34 +225,74 @@ impl PersistentBoundaryChip { } } + /// Finalize the boundary chip with per-block timestamped memory. + /// + /// `final_memory` is at CONST_BLOCK_SIZE granularity (4 bytes per entry, single timestamp + /// each). This function rechunks into CHUNK-sized (8 bytes) groups with per-block + /// timestamps. Untouched blocks within a touched chunk get values from initial_memory and + /// timestamp 0. #[instrument(name = "boundary_finalize", level = "debug", skip_all)] pub(crate) fn finalize( &mut self, initial_memory: &MemoryImage, - // Only touched stuff - final_memory: &TimestampedEquipartition, + // Touched stuff at CONST_BLOCK_SIZE granularity + final_memory: &TimestampedEquipartition, hasher: &H, ) where H: Hasher + Sync + for<'a> SerialReceiver<&'a [F]>, { - let final_touched_labels: Vec<_> = final_memory + type BlockInfo = (usize, u32, [F; CONST_BLOCK_SIZE]); // (block_idx, timestamp, values) + type EnrichedEntry = ((u32, u32), BlockInfo); // (chunk_key, block_info) + + let enriched: Vec> = final_memory .par_iter() .map(|&((addr_space, ptr), ts_values)| { - let TimestampedValues { timestamp, values } = ts_values; + let chunk_label = ptr / CHUNK as u32; + let block_idx = ((ptr % CHUNK as u32) / CONST_BLOCK_SIZE as u32) as usize; + let key = (addr_space, chunk_label); + let block_info = (block_idx, ts_values.timestamp, ts_values.values); + (key, block_info) + }) + .collect(); + + let chunk_groups: Vec<_> = enriched + .chunk_by(|a, b| a.0 == b.0) + .map(|group| { + let key = group[0].0; + let blocks: Vec> = group.iter().map(|&(_, info)| info).collect(); + (key, blocks) + }) + .collect(); + + let final_touched_labels: Vec<_> = chunk_groups + .into_par_iter() + .map(|((addr_space, chunk_label), blocks)| { + let chunk_ptr = chunk_label * CHUNK as u32; // SAFETY: addr_space from `final_memory` are all in bounds - let init_values = array::from_fn(|i| unsafe { - initial_memory.get_f::(addr_space, ptr + i as u32) + let init_values: [F; CHUNK] = array::from_fn(|i| unsafe { + initial_memory.get_f::(addr_space, chunk_ptr + i as u32) }); + + let mut final_values = init_values; + let mut timestamps = [0u32; BLOCKS_PER_CHUNK]; + + for (block_idx, ts, values) in blocks { + timestamps[block_idx] = ts; + for (i, &val) in values.iter().enumerate() { + final_values[block_idx * CONST_BLOCK_SIZE + i] = val; + } + } + let initial_hash = hasher.hash(&init_values); - let final_hash = hasher.hash(&values); + let final_hash = hasher.hash(&final_values); FinalTouchedLabel { address_space: addr_space, - label: ptr / CHUNK as u32, + label: chunk_label, init_values, - final_values: values, + final_values, init_hash: initial_hash, final_hash, - final_timestamp: timestamp, + final_timestamps: timestamps, } }) .collect(); @@ -281,7 +339,9 @@ where leaf_label: Val::::from_canonical_u32(touched_label.label), values: touched_label.init_values, hash: touched_label.init_hash, - timestamp: Val::::from_canonical_u32(INITIAL_TIMESTAMP), + // Initial timestamps are all 0 (INITIAL_TIMESTAMP) + timestamps: [Val::::from_canonical_u32(INITIAL_TIMESTAMP); + BLOCKS_PER_CHUNK], }; *final_row.borrow_mut() = PersistentBoundaryCols { @@ -290,7 +350,10 @@ where leaf_label: Val::::from_canonical_u32(touched_label.label), values: touched_label.final_values, hash: touched_label.final_hash, - timestamp: Val::::from_canonical_u32(touched_label.final_timestamp), + // Per-block timestamps - untouched blocks stay at 0 + timestamps: touched_label + .final_timestamps + .map(Val::::from_canonical_u32), }; }); Arc::new(RowMajorMatrix::new(rows, width)) diff --git a/crates/vm/src/system/mod.rs b/crates/vm/src/system/mod.rs index d1ecb2daf1..195fa6a701 100644 --- a/crates/vm/src/system/mod.rs +++ b/crates/vm/src/system/mod.rs @@ -32,7 +32,8 @@ use crate::{ 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, + VmCircuitConfig, VmExecutionConfig, CONNECTOR_AIR_ID, CONST_BLOCK_SIZE, PROGRAM_AIR_ID, + PUBLIC_VALUES_AIR_ID, }, system::{ connector::VmConnectorChip, @@ -145,7 +146,7 @@ pub struct SystemRecords { } pub enum TouchedMemory { - Persistent(TimestampedEquipartition), + Persistent(TimestampedEquipartition), Volatile(TimestampedEquipartition), } diff --git a/extensions/rv32im/circuit/src/base_alu/tests.rs b/extensions/rv32im/circuit/src/base_alu/tests.rs index 8f38dea1f5..7e42b38989 100644 --- a/extensions/rv32im/circuit/src/base_alu/tests.rs +++ b/extensions/rv32im/circuit/src/base_alu/tests.rs @@ -163,8 +163,11 @@ fn rand_rv32_alu_test(opcode: BaseAluOpcode, num_ops: usize) { // TODO(AG): make a more meaningful test for memory accesses tester.write(2, 1024, [F::ONE; 4]); tester.write(2, 1028, [F::ONE; 4]); - let sm = tester.read(2, 1024); - assert_eq!(sm, [F::ONE; 8]); + // Avoid wider-than-min-block accesses when access adapters are disabled + let sm1 = tester.read(2, 1024); + let sm2 = tester.read(2, 1028); + assert_eq!(sm1, [F::ONE; 4]); + assert_eq!(sm2, [F::ONE; 4]); for _ in 0..num_ops { set_and_execute( @@ -201,8 +204,11 @@ fn rand_rv32_alu_test_persistent(opcode: BaseAluOpcode, num_ops: usize) { // TODO(AG): make a more meaningful test for memory accesses tester.write(2, 1024, [F::ONE; 4]); tester.write(2, 1028, [F::ONE; 4]); - let sm = tester.read(2, 1024); - assert_eq!(sm, [F::ONE; 8]); + // Avoid wider-than-min-block accesses when access adapters are disabled + let sm1 = tester.read(2, 1024); + let sm2 = tester.read(2, 1028); + assert_eq!(sm1, [F::ONE; 4]); + assert_eq!(sm2, [F::ONE; 4]); for _ in 0..num_ops { set_and_execute( diff --git a/extensions/rv32im/circuit/src/common/mod.rs b/extensions/rv32im/circuit/src/common/mod.rs index 0a58b7310b..20855af15d 100644 --- a/extensions/rv32im/circuit/src/common/mod.rs +++ b/extensions/rv32im/circuit/src/common/mod.rs @@ -9,7 +9,7 @@ mod aot { use openvm_circuit::{ arch::{ execution_mode::{metered::memory_ctx::MemoryCtx, MeteredCtx}, - AotError, SystemConfig, VmExecState, ADDR_SPACE_OFFSET, + AotError, SystemConfig, VmExecState, ADDR_SPACE_OFFSET, CONST_BLOCK_SIZE, }, system::memory::{merkle::public_values::PUBLIC_VALUES_AS, online::GuestMemory, CHUNK}, }; @@ -244,12 +244,12 @@ mod aot { // Therefore the loop only iterates once for `page_id = start_page_id`. let initial_block_size: usize = config.initial_block_size(); - if initial_block_size != CHUNK { + if initial_block_size != CONST_BLOCK_SIZE { return Err(AotError::Other(format!( - "initial_block_size must be {CHUNK}, got {initial_block_size}" + "initial_block_size must be {CONST_BLOCK_SIZE}, got {initial_block_size}" ))); } - let chunk_bits = CHUNK.ilog2(); + let chunk_bits = CONST_BLOCK_SIZE.ilog2(); let as_offset = ((address_space - ADDR_SPACE_OFFSET) as u64) << (config.memory_config.memory_dimensions().address_height); diff --git a/extensions/rv32im/circuit/src/loadstore/tests.rs b/extensions/rv32im/circuit/src/loadstore/tests.rs index 240da983d0..1157012212 100644 --- a/extensions/rv32im/circuit/src/loadstore/tests.rs +++ b/extensions/rv32im/circuit/src/loadstore/tests.rs @@ -10,7 +10,9 @@ use openvm_circuit::{ }, }; use openvm_circuit_primitives::var_range::VariableRangeCheckerChip; -use openvm_instructions::{instruction::Instruction, riscv::RV32_REGISTER_AS, LocalOpcode}; +use openvm_instructions::{ + instruction::Instruction, riscv::RV32_REGISTER_AS, LocalOpcode, NATIVE_AS, +}; use openvm_rv32im_transpiler::Rv32LoadStoreOpcode::{self, *}; use openvm_stark_backend::{ p3_air::BaseAir, @@ -131,7 +133,9 @@ fn set_and_execute>( let mem_as = mem_as.unwrap_or(if is_load { 2 } else { - *[2, 3, 4].choose(rng).unwrap() + // Avoid Native AS while access adapters are disabled. + // TODO: Revert this to [2, 3, 4] when access adapters are removed + *[2, 3].choose(rng).unwrap() }); let shift_amount = ptr_val % 4; @@ -215,10 +219,13 @@ fn rand_loadstore_test(opcode: Rv32LoadStoreOpcode, num_ops: usize) { let mut rng = create_seeded_rng(); let mut mem_config = MemoryConfig::default(); mem_config.addr_spaces[RV32_REGISTER_AS as usize].num_cells = 1 << 29; + mem_config.addr_spaces[NATIVE_AS as usize].num_cells = 0; if [STOREW, STOREB, STOREH].contains(&opcode) { mem_config.addr_spaces[PUBLIC_VALUES_AS as usize].num_cells = 1 << 29; } - let mut tester = VmChipTestBuilder::volatile(mem_config); + // Use persistent memory so initial block size matches the 4-byte alignment and + // avoids access-adapter split/merge paths when adapters are disabled. + let mut tester = VmChipTestBuilder::persistent(mem_config); let mut harness = create_harness(&mut tester); for _ in 0..num_ops { @@ -268,10 +275,12 @@ fn run_negative_loadstore_test( let mut rng = create_seeded_rng(); let mut mem_config = MemoryConfig::default(); mem_config.addr_spaces[RV32_REGISTER_AS as usize].num_cells = 1 << 29; + mem_config.addr_spaces[NATIVE_AS as usize].num_cells = 0; if [STOREW, STOREB, STOREH].contains(&opcode) { mem_config.addr_spaces[PUBLIC_VALUES_AS as usize].num_cells = 1 << 29; } - let mut tester = VmChipTestBuilder::volatile(mem_config); + // Use persistent memory so the min block size matches alignment without needing adapters. + let mut tester = VmChipTestBuilder::persistent(mem_config); let mut harness = create_harness(&mut tester); set_and_execute(