From 8d39402dac4e27038857fab82a380cd55639ecac Mon Sep 17 00:00:00 2001 From: Ayush Shukla Date: Wed, 24 Dec 2025 01:52:04 +0100 Subject: [PATCH 01/13] fix: keep track of unique memory accesses since last checkpoint --- .../vm/src/arch/execution_mode/metered/ctx.rs | 20 ++- .../arch/execution_mode/metered/memory_ctx.rs | 166 +++++++++++++++--- .../execution_mode/metered/segment_ctx.rs | 39 ++-- 3 files changed, 175 insertions(+), 50 deletions(-) diff --git a/crates/vm/src/arch/execution_mode/metered/ctx.rs b/crates/vm/src/arch/execution_mode/metered/ctx.rs index 8428438ca7..49345fca23 100644 --- a/crates/vm/src/arch/execution_mode/metered/ctx.rs +++ b/crates/vm/src/arch/execution_mode/metered/ctx.rs @@ -122,12 +122,6 @@ impl MeteredCtx { self.segmentation_ctx.segments } - fn reset_segment(&mut self) { - self.memory_ctx.clear(); - // Add merkle height contributions for all registers - self.memory_ctx.add_register_merkle_heights(); - } - #[inline(always)] pub fn check_and_segment(&mut self) -> bool { // We track the segmentation check by instrets_until_check instead of instret in order to @@ -147,7 +141,19 @@ impl MeteredCtx { ); if did_segment { - self.reset_segment(); + // Reset segment context for new segment + self.segmentation_ctx + .reset_segment(&mut self.trace_heights, &self.is_trace_height_constant); + // Reset memory context for new segment + self.memory_ctx.reset_segment(&mut self.trace_heights); + // Add merkle height contributions for all registers + self.memory_ctx.add_register_merkle_heights(); + } else { + // Update checkpoint for trace heights + self.segmentation_ctx + .update_checkpoint(self.segmentation_ctx.instret, &self.trace_heights); + // Update checkpoint for page indices + self.memory_ctx.update_checkpoint(); } did_segment } 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..d1e10ada85 100644 --- a/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs +++ b/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs @@ -32,6 +32,52 @@ impl BitSet { !was_set } + #[inline(always)] + pub fn merge_from(&mut self, other: &Self) { + debug_assert_eq!(self.words.len(), other.words.len()); + for (word, other_word) in self.words.iter_mut().zip(other.words.iter()) { + *word |= *other_word; + } + } + + /// Count bits that are set in self but not in other, and call a function for each. + /// Returns the total count of such bits. + #[inline(always)] + pub fn count_diff_with(&self, other: &Self, mut f: F) -> usize + where + F: FnMut(usize), + { + debug_assert_eq!(self.words.len(), other.words.len()); + let mut count = 0; + + for (word_idx, (&word, &other_word)) in + self.words.iter().zip(other.words.iter()).enumerate() + { + if word == other_word { + continue; + } + + let diff_bits = word & !other_word; + if diff_bits == 0 { + continue; + } + + count += diff_bits.count_ones() as usize; + + // Call function for each bit in diff + let mut remaining_bits = diff_bits; + while remaining_bits != 0 { + let bit_pos = remaining_bits.trailing_zeros(); + remaining_bits &= remaining_bits - 1; // Clear lowest set bit + + let bit_index = (word_idx << 6) + bit_pos as usize; + f(bit_index); + } + } + + count + } + /// Set all bits within [start, end) to 1, return the number of flipped bits. /// Assumes start < end and end <= self.words.len() * 64. #[inline(always)] @@ -99,7 +145,6 @@ impl BitSet { #[derive(Clone, Debug)] pub struct MemoryCtx { - pub page_indices: BitSet, memory_dimensions: MemoryDimensions, min_block_size_bits: Vec, pub boundary_idx: usize, @@ -108,9 +153,10 @@ pub struct MemoryCtx { continuations_enabled: bool, chunk: u32, chunk_bits: u32, - pub page_access_count: usize, - // Note: 32 is the maximum access adapter size. - pub addr_space_access_count: RVec, + pub(crate) page_indices_checkpoint: BitSet, + pub(crate) page_indices: BitSet, + page_access_count: usize, + addr_space_access_count: RVec, } impl MemoryCtx { @@ -121,9 +167,10 @@ impl MemoryCtx { let memory_dimensions = config.memory_config.memory_dimensions(); let merkle_height = memory_dimensions.overall_height(); + let bitset_size = 1 << (merkle_height.saturating_sub(PAGE_BITS)); + let addr_space_size = (1 << memory_dimensions.addr_space_height) + 1; + Self { - // Address height already considers `chunk_bits`. - page_indices: BitSet::new(1 << (merkle_height.saturating_sub(PAGE_BITS))), min_block_size_bits: config.memory_config.min_block_size_bits(), boundary_idx: config.memory_boundary_air_id(), merkle_tree_index: config.memory_merkle_air_id(), @@ -132,16 +179,13 @@ impl MemoryCtx { chunk_bits, memory_dimensions, continuations_enabled: config.continuation_enabled, + page_indices_checkpoint: BitSet::new(bitset_size), + page_indices: BitSet::new(bitset_size), page_access_count: 0, - addr_space_access_count: vec![0; (1 << memory_dimensions.addr_space_height) + 1].into(), + addr_space_access_count: vec![0; addr_space_size].into(), } } - #[inline(always)] - pub fn clear(&mut self) { - self.page_indices.clear(); - } - #[inline(always)] pub(crate) fn add_register_merkle_heights(&mut self) { if self.continuations_enabled { @@ -235,13 +279,79 @@ impl MemoryCtx { } } - /// Resolve all lazy updates of each memory access for memory adapters/poseidon2/merkle chip. + /// Reset memory context state for a new segment #[inline(always)] - pub(crate) fn lazy_update_boundary_heights(&mut self, trace_heights: &mut [u32]) { - debug_assert!(self.boundary_idx < trace_heights.len()); + pub(crate) fn reset_segment(&mut self, trace_heights: &mut [u32]) { + // Update trace heights with all pages accessed since last checkpoint + self.apply_height_updates( + trace_heights, + self.page_access_count, + &self.addr_space_access_count, + ); + + // Replace checkpoint with current pages and clear current + std::mem::swap(&mut self.page_indices_checkpoint, &mut self.page_indices); + self.page_indices.clear(); + self.page_access_count = 0; + // SAFETY: Resetting array elements to 0 is always safe + unsafe { + std::ptr::write_bytes( + self.addr_space_access_count.as_mut_ptr(), + 0, + self.addr_space_access_count.len(), + ); + } + } + + /// Updates the checkpoint with current safe state + #[inline(always)] + pub(crate) fn update_checkpoint(&mut self) { + // Merge current pages into checkpoint + self.page_indices_checkpoint.merge_from(&self.page_indices); + // Clear current pages to track next checkpoint period + self.page_indices.clear(); + // Reset access counts + self.page_access_count = 0; + // SAFETY: Resetting array elements to 0 is always safe + unsafe { + std::ptr::write_bytes( + self.addr_space_access_count.as_mut_ptr(), + 0, + self.addr_space_access_count.len(), + ); + } + } + + /// Calculate pages per address space relative to a checkpoint + #[inline(always)] + fn calculate_new_pages_per_addr_space(&self, checkpoint: &BitSet) -> (usize, Box<[usize]>) { + let mut addr_space_counts = + vec![0usize; self.addr_space_access_count.len()].into_boxed_slice(); + + let total_pages = self.page_indices.count_diff_with(checkpoint, |page_id| { + let block_id = (page_id << PAGE_BITS) as u64; + let (addr_space, _) = self.memory_dimensions.index_to_label(block_id); + + debug_assert!((addr_space as usize) < addr_space_counts.len()); + // SAFETY: addr_space is bounds checked in debug mode + unsafe { + *addr_space_counts.get_unchecked_mut(addr_space as usize) += 1; + } + }); + (total_pages, addr_space_counts) + } + + /// Apply height updates given page counts + #[inline(always)] + fn apply_height_updates( + &self, + trace_heights: &mut [u32], + page_count: usize, + addr_space_counts: &[usize], + ) { // On page fault, assume we add all leaves in a page - let leaves = (self.page_access_count << PAGE_BITS) as u32; + let leaves = (page_count << PAGE_BITS) as u32; // SAFETY: boundary_idx is a compile time constant within bounds unsafe { *trace_heights.get_unchecked_mut(self.boundary_idx) += leaves; @@ -261,15 +371,12 @@ impl MemoryCtx { let nodes = (((1 << PAGE_BITS) - 1) + (merkle_height - PAGE_BITS)) as u32; // SAFETY: merkle_tree_idx is guaranteed to be in bounds unsafe { - *trace_heights.get_unchecked_mut(poseidon2_idx) += nodes * 2; - *trace_heights.get_unchecked_mut(merkle_tree_idx) += nodes * 2; + *trace_heights.get_unchecked_mut(poseidon2_idx) += nodes * page_count as u32 * 2; + *trace_heights.get_unchecked_mut(merkle_tree_idx) += nodes * page_count as u32 * 2; } } - self.page_access_count = 0; - for address_space in 0..self.addr_space_access_count.len() { - // SAFETY: address_space is from 0 to len(), guaranteed to be in bounds - let x = unsafe { *self.addr_space_access_count.get_unchecked(address_space) }; + for (address_space, &x) in addr_space_counts.iter().enumerate() { if x > 0 { // Initial **and** final handling of touched pages requires send (resp. receive) in // chunk-sized units for the merkle chip @@ -281,15 +388,18 @@ impl MemoryCtx { self.chunk_bits, (x << (PAGE_BITS + 1)) as u32, ); - // SAFETY: address_space is from 0 to len(), guaranteed to be in bounds - unsafe { - *self - .addr_space_access_count - .get_unchecked_mut(address_space) = 0; - } } } } + + /// Resolve all lazy updates of each memory access for memory adapters/poseidon2/merkle chip. + #[inline(always)] + pub(crate) fn lazy_update_boundary_heights(&mut self, trace_heights: &mut [u32]) { + // Calculate diff between current and checkpoint + let (page_count, addr_space_counts) = + self.calculate_new_pages_per_addr_space(&self.page_indices_checkpoint); + self.apply_height_updates(trace_heights, page_count, &addr_space_counts); + } } #[cfg(test)] diff --git a/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs b/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs index db1aecfa31..d0fb748443 100644 --- a/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs +++ b/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs @@ -176,7 +176,7 @@ impl SegmentationCtx { } #[inline(always)] - fn should_segment( + pub(crate) fn should_segment( &self, instret: u64, trace_heights: &[u32], @@ -257,20 +257,13 @@ impl SegmentationCtx { let should_seg = self.should_segment(instret, trace_heights, is_trace_height_constant); if should_seg { - self.create_segment_from_checkpoint(instret, trace_heights, is_trace_height_constant); - } else { - self.update_checkpoint(instret, trace_heights); + self.create_segment_from_checkpoint(instret, trace_heights); } should_seg } #[inline(always)] - fn create_segment_from_checkpoint( - &mut self, - instret: u64, - trace_heights: &mut [u32], - is_trace_height_constant: &[bool], - ) { + fn create_segment_from_checkpoint(&mut self, instret: u64, trace_heights: &mut [u32]) { let instret_start = self .segments .last() @@ -296,14 +289,30 @@ impl SegmentationCtx { (instret, trace_heights.to_vec()) }; - // Reset current trace heights and checkpoint - self.reset_trace_heights(trace_heights, &segment_heights, is_trace_height_constant); - self.checkpoint_instret = 0; - let num_insns = segment_instret - instret_start; self.create_segment::(instret_start, num_insns, segment_heights); } + /// Reset segment context state for a new segment + #[inline(always)] + pub(crate) fn reset_segment( + &mut self, + trace_heights: &mut [u32], + is_trace_height_constant: &[bool], + ) { + // Get the segment heights from the last created segment + if let Some(last_segment) = self.segments.last() { + // Reset trace heights by subtracting the segment's heights + self.reset_trace_heights( + trace_heights, + &last_segment.trace_heights, + is_trace_height_constant, + ); + } + // Reset checkpoint since we're starting a new segment + self.checkpoint_instret = 0; + } + /// Resets trace heights by subtracting segment heights #[inline(always)] fn reset_trace_heights( @@ -325,7 +334,7 @@ impl SegmentationCtx { /// Updates the checkpoint with current safe state #[inline(always)] - fn update_checkpoint(&mut self, instret: u64, trace_heights: &[u32]) { + pub(crate) fn update_checkpoint(&mut self, instret: u64, trace_heights: &[u32]) { self.checkpoint_trace_heights.copy_from_slice(trace_heights); self.checkpoint_instret = instret; } From b844d94495e447b2fca50f90fa2ed97b01acc355 Mon Sep 17 00:00:00 2001 From: Ayush Shukla Date: Wed, 24 Dec 2025 01:56:31 +0100 Subject: [PATCH 02/13] fix: visibility of memory ctx members --- crates/vm/src/arch/execution_mode/metered/memory_ctx.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) 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 d1e10ada85..ded79570d5 100644 --- a/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs +++ b/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs @@ -154,9 +154,9 @@ pub struct MemoryCtx { chunk: u32, chunk_bits: u32, pub(crate) page_indices_checkpoint: BitSet, - pub(crate) page_indices: BitSet, - page_access_count: usize, - addr_space_access_count: RVec, + pub page_indices: BitSet, + pub page_access_count: usize, + pub addr_space_access_count: RVec, } impl MemoryCtx { From 0f087d409d4d6ec22259b4b8ecec7b7d5cfca16f Mon Sep 17 00:00:00 2001 From: Ayush Shukla Date: Wed, 24 Dec 2025 03:00:25 +0100 Subject: [PATCH 03/13] fix: rename to initialize_segment --- .../vm/src/arch/execution_mode/metered/ctx.rs | 8 +++---- .../arch/execution_mode/metered/memory_ctx.rs | 23 ++++++++++++++++--- .../execution_mode/metered/segment_ctx.rs | 4 ++-- 3 files changed, 26 insertions(+), 9 deletions(-) diff --git a/crates/vm/src/arch/execution_mode/metered/ctx.rs b/crates/vm/src/arch/execution_mode/metered/ctx.rs index 49345fca23..8a5c58f437 100644 --- a/crates/vm/src/arch/execution_mode/metered/ctx.rs +++ b/crates/vm/src/arch/execution_mode/metered/ctx.rs @@ -141,11 +141,11 @@ impl MeteredCtx { ); if did_segment { - // Reset segment context for new segment + // Initialize segment context for new segment self.segmentation_ctx - .reset_segment(&mut self.trace_heights, &self.is_trace_height_constant); - // Reset memory context for new segment - self.memory_ctx.reset_segment(&mut self.trace_heights); + .initialize_segment(&mut self.trace_heights, &self.is_trace_height_constant); + // Initialize memory context for new segment + self.memory_ctx.initialize_segment(&mut self.trace_heights); // Add merkle height contributions for all registers self.memory_ctx.add_register_merkle_heights(); } else { 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 ded79570d5..6414dcac3f 100644 --- a/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs +++ b/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs @@ -279,10 +279,27 @@ impl MemoryCtx { } } - /// Reset memory context state for a new segment + /// Initialize state for a new segment #[inline(always)] - pub(crate) fn reset_segment(&mut self, trace_heights: &mut [u32]) { - // Update trace heights with all pages accessed since last checkpoint + pub(crate) fn initialize_segment(&mut self, trace_heights: &mut [u32]) { + // Reset trace heights for memory chips + // SAFETY: boundary_idx is a compile time constant within bounds + unsafe { + *trace_heights.get_unchecked_mut(self.boundary_idx) = 0; + } + if let Some(merkle_tree_idx) = self.merkle_tree_index { + // SAFETY: merkle_tree_idx is guaranteed to be in bounds + unsafe { + *trace_heights.get_unchecked_mut(merkle_tree_idx) = 0; + } + let poseidon2_idx = trace_heights.len() - 2; + // SAFETY: poseidon2_idx is trace_heights.len() - 2, guaranteed to be in bounds + unsafe { + *trace_heights.get_unchecked_mut(poseidon2_idx) = 0; + } + } + + // Apply height updates for all pages accessed since last checkpoint self.apply_height_updates( trace_heights, self.page_access_count, diff --git a/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs b/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs index d0fb748443..72a20a6baa 100644 --- a/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs +++ b/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs @@ -293,9 +293,9 @@ impl SegmentationCtx { self.create_segment::(instret_start, num_insns, segment_heights); } - /// Reset segment context state for a new segment + /// Initialize state for a new segment #[inline(always)] - pub(crate) fn reset_segment( + pub(crate) fn initialize_segment( &mut self, trace_heights: &mut [u32], is_trace_height_constant: &[bool], From 8952ee952280c3823458620f3a8fc47f334fc6f4 Mon Sep 17 00:00:00 2001 From: Ayush Shukla Date: Wed, 24 Dec 2025 03:59:55 +0100 Subject: [PATCH 04/13] fix: try a hack --- .../arch/execution_mode/metered/memory_ctx.rs | 89 +++++++++---------- extensions/rv32im/circuit/src/common/mod.rs | 22 ++++- 2 files changed, 59 insertions(+), 52 deletions(-) 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 6414dcac3f..db46e5da71 100644 --- a/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs +++ b/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs @@ -153,10 +153,10 @@ pub struct MemoryCtx { continuations_enabled: bool, chunk: u32, chunk_bits: u32, - pub(crate) page_indices_checkpoint: BitSet, pub page_indices: BitSet, pub page_access_count: usize, pub addr_space_access_count: RVec, + pub addr_space_access_count_pending: RVec, } impl MemoryCtx { @@ -179,10 +179,10 @@ impl MemoryCtx { chunk_bits, memory_dimensions, continuations_enabled: config.continuation_enabled, - page_indices_checkpoint: BitSet::new(bitset_size), page_indices: BitSet::new(bitset_size), page_access_count: 0, addr_space_access_count: vec![0; addr_space_size].into(), + addr_space_access_count_pending: vec![0; addr_space_size].into(), } } @@ -232,6 +232,14 @@ impl MemoryCtx { .addr_space_access_count .get_unchecked_mut(address_space as usize) += 1; } + } else { + // SAFETY: address_space passed is usually a hardcoded constant or derived from an + // Instruction where it is bounds checked before passing + unsafe { + *self + .addr_space_access_count_pending + .get_unchecked_mut(address_space as usize) += 1; + } } } } @@ -300,75 +308,46 @@ impl MemoryCtx { } // Apply height updates for all pages accessed since last checkpoint + let page_access_count = self.addr_space_access_count_pending.iter().sum(); self.apply_height_updates( trace_heights, - self.page_access_count, - &self.addr_space_access_count, + page_access_count, + &self.addr_space_access_count_pending, ); - - // Replace checkpoint with current pages and clear current - std::mem::swap(&mut self.page_indices_checkpoint, &mut self.page_indices); - self.page_indices.clear(); - self.page_access_count = 0; // SAFETY: Resetting array elements to 0 is always safe unsafe { std::ptr::write_bytes( - self.addr_space_access_count.as_mut_ptr(), + self.addr_space_access_count_pending.as_mut_ptr(), 0, - self.addr_space_access_count.len(), + self.addr_space_access_count_pending.len(), ); } + self.page_indices.clear(); } /// Updates the checkpoint with current safe state #[inline(always)] pub(crate) fn update_checkpoint(&mut self) { - // Merge current pages into checkpoint - self.page_indices_checkpoint.merge_from(&self.page_indices); - // Clear current pages to track next checkpoint period - self.page_indices.clear(); - // Reset access counts - self.page_access_count = 0; // SAFETY: Resetting array elements to 0 is always safe unsafe { std::ptr::write_bytes( - self.addr_space_access_count.as_mut_ptr(), + self.addr_space_access_count_pending.as_mut_ptr(), 0, - self.addr_space_access_count.len(), + self.addr_space_access_count_pending.len(), ); } } - /// Calculate pages per address space relative to a checkpoint - #[inline(always)] - fn calculate_new_pages_per_addr_space(&self, checkpoint: &BitSet) -> (usize, Box<[usize]>) { - let mut addr_space_counts = - vec![0usize; self.addr_space_access_count.len()].into_boxed_slice(); - - let total_pages = self.page_indices.count_diff_with(checkpoint, |page_id| { - let block_id = (page_id << PAGE_BITS) as u64; - let (addr_space, _) = self.memory_dimensions.index_to_label(block_id); - - debug_assert!((addr_space as usize) < addr_space_counts.len()); - // SAFETY: addr_space is bounds checked in debug mode - unsafe { - *addr_space_counts.get_unchecked_mut(addr_space as usize) += 1; - } - }); - - (total_pages, addr_space_counts) - } - /// Apply height updates given page counts #[inline(always)] fn apply_height_updates( &self, trace_heights: &mut [u32], - page_count: usize, - addr_space_counts: &[usize], + page_access_count: usize, + addr_space_access_count: &[usize], ) { // On page fault, assume we add all leaves in a page - let leaves = (page_count << PAGE_BITS) as u32; + let leaves = (page_access_count << PAGE_BITS) as u32; // SAFETY: boundary_idx is a compile time constant within bounds unsafe { *trace_heights.get_unchecked_mut(self.boundary_idx) += leaves; @@ -388,12 +367,14 @@ impl MemoryCtx { let nodes = (((1 << PAGE_BITS) - 1) + (merkle_height - PAGE_BITS)) as u32; // SAFETY: merkle_tree_idx is guaranteed to be in bounds unsafe { - *trace_heights.get_unchecked_mut(poseidon2_idx) += nodes * page_count as u32 * 2; - *trace_heights.get_unchecked_mut(merkle_tree_idx) += nodes * page_count as u32 * 2; + *trace_heights.get_unchecked_mut(poseidon2_idx) += + nodes * page_access_count as u32 * 2; + *trace_heights.get_unchecked_mut(merkle_tree_idx) += + nodes * page_access_count as u32 * 2; } } - for (address_space, &x) in addr_space_counts.iter().enumerate() { + for (address_space, &x) in addr_space_access_count.iter().enumerate() { if x > 0 { // Initial **and** final handling of touched pages requires send (resp. receive) in // chunk-sized units for the merkle chip @@ -412,10 +393,20 @@ impl MemoryCtx { /// Resolve all lazy updates of each memory access for memory adapters/poseidon2/merkle chip. #[inline(always)] pub(crate) fn lazy_update_boundary_heights(&mut self, trace_heights: &mut [u32]) { - // Calculate diff between current and checkpoint - let (page_count, addr_space_counts) = - self.calculate_new_pages_per_addr_space(&self.page_indices_checkpoint); - self.apply_height_updates(trace_heights, page_count, &addr_space_counts); + self.apply_height_updates( + trace_heights, + self.page_access_count, + &self.addr_space_access_count, + ); + self.page_access_count = 0; + // SAFETY: Resetting array elements to 0 is always safe + unsafe { + std::ptr::write_bytes( + self.addr_space_access_count.as_mut_ptr(), + 0, + self.addr_space_access_count.len(), + ); + } } } diff --git a/extensions/rv32im/circuit/src/common/mod.rs b/extensions/rv32im/circuit/src/common/mod.rs index 0a58b7310b..1ef1d3dc8b 100644 --- a/extensions/rv32im/circuit/src/common/mod.rs +++ b/extensions/rv32im/circuit/src/common/mod.rs @@ -270,7 +270,13 @@ mod aot { memory_ctx_offset + offset_of!(MemoryCtx, page_access_count); let addr_space_access_count_ptr_offset = memory_ctx_offset + offset_of!(MemoryCtx, addr_space_access_count); + let addr_space_access_count_pending_ptr_offset = memory_ctx_offset + + offset_of!( + MemoryCtx, + addr_space_access_count_pending + ); let inserted_label = format!(".asm_execute_pc_{pc}_inserted"); + let not_inserted_label = format!(".asm_execute_pc_{pc}_not_inserted"); // The next section is the implementation of `BitSet::insert` in ASM. // pub fn insert(&mut self, index: usize) -> bool { // let word_index = index >> 6; @@ -302,8 +308,8 @@ mod aot { // `test (*word & mask)` asm_str += &format!(" test {ptr_reg}, {reg2}\n"); - asm_str += &format!(" jnz {inserted_label}\n"); - // When (*word & mask) == 0 + asm_str += &format!(" jnz {not_inserted_label}\n"); + // When (*word & mask) == 0 (new page inserted) // `*word += mask` asm_str += &format!(" add {ptr_reg}, {reg2}\n"); asm_str += &format!(" mov [{reg1}], {ptr_reg}\n"); @@ -319,8 +325,18 @@ mod aot { asm_str += &format!(" mov {reg1}, [{reg1}]\n"); // self.addr_space_access_count[address_space] += 1; asm_str += &format!(" add dword ptr [{reg1} + {address_space} * 4], 1\n"); + asm_str += &format!(" jmp {inserted_label}\n"); + asm_str += &format!("{not_inserted_label}:\n"); + // When (*word & mask) != 0 (page already exists) + // reg1 = &addr_space_access_count_pending.as_ptr() + asm_str += &format!( + " lea {reg1}, [{REG_EXEC_STATE_PTR} + {addr_space_access_count_pending_ptr_offset}]\n" + ); + asm_str += &format!(" mov {reg1}, [{reg1}]\n"); + // self.addr_space_access_count_pending[address_space] += 1; + asm_str += &format!(" add dword ptr [{reg1} + {address_space} * 4], 1\n"); asm_str += &format!("{inserted_label}:\n"); - // Inserted, do nothing + // Done Ok(asm_str) } From f3bfac6f496059dfaac28fcded8f6d9127e208ea Mon Sep 17 00:00:00 2001 From: Ayush Shukla Date: Wed, 24 Dec 2025 04:02:48 +0100 Subject: [PATCH 05/13] chore: cleanup --- .../arch/execution_mode/metered/memory_ctx.rs | 46 ------------------- 1 file changed, 46 deletions(-) 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 db46e5da71..275f4d6591 100644 --- a/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs +++ b/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs @@ -32,52 +32,6 @@ impl BitSet { !was_set } - #[inline(always)] - pub fn merge_from(&mut self, other: &Self) { - debug_assert_eq!(self.words.len(), other.words.len()); - for (word, other_word) in self.words.iter_mut().zip(other.words.iter()) { - *word |= *other_word; - } - } - - /// Count bits that are set in self but not in other, and call a function for each. - /// Returns the total count of such bits. - #[inline(always)] - pub fn count_diff_with(&self, other: &Self, mut f: F) -> usize - where - F: FnMut(usize), - { - debug_assert_eq!(self.words.len(), other.words.len()); - let mut count = 0; - - for (word_idx, (&word, &other_word)) in - self.words.iter().zip(other.words.iter()).enumerate() - { - if word == other_word { - continue; - } - - let diff_bits = word & !other_word; - if diff_bits == 0 { - continue; - } - - count += diff_bits.count_ones() as usize; - - // Call function for each bit in diff - let mut remaining_bits = diff_bits; - while remaining_bits != 0 { - let bit_pos = remaining_bits.trailing_zeros(); - remaining_bits &= remaining_bits - 1; // Clear lowest set bit - - let bit_index = (word_idx << 6) + bit_pos as usize; - f(bit_index); - } - } - - count - } - /// Set all bits within [start, end) to 1, return the number of flipped bits. /// Assumes start < end and end <= self.words.len() * 64. #[inline(always)] From 77e370b31ab380fcff4d73618face613225679cb Mon Sep 17 00:00:00 2001 From: Ayush Shukla Date: Wed, 24 Dec 2025 17:02:04 +0000 Subject: [PATCH 06/13] fix: try with two bitsets --- .../arch/execution_mode/metered/memory_ctx.rs | 10 ++- extensions/rv32im/circuit/src/common/mod.rs | 84 ++++++++++++------- 2 files changed, 59 insertions(+), 35 deletions(-) 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 275f4d6591..262e4169b1 100644 --- a/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs +++ b/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs @@ -108,6 +108,7 @@ pub struct MemoryCtx { chunk: u32, chunk_bits: u32, pub page_indices: BitSet, + pub page_indices_pending: BitSet, pub page_access_count: usize, pub addr_space_access_count: RVec, pub addr_space_access_count_pending: RVec, @@ -134,6 +135,7 @@ impl MemoryCtx { memory_dimensions, continuations_enabled: config.continuation_enabled, page_indices: BitSet::new(bitset_size), + page_indices_pending: BitSet::new(bitset_size), page_access_count: 0, addr_space_access_count: vec![0; addr_space_size].into(), addr_space_access_count_pending: vec![0; addr_space_size].into(), @@ -178,7 +180,6 @@ impl MemoryCtx { for page_id in start_page_id..end_page_id { if self.page_indices.insert(page_id as usize) { - self.page_access_count += 1; // SAFETY: address_space passed is usually a hardcoded constant or derived from an // Instruction where it is bounds checked before passing unsafe { @@ -186,7 +187,7 @@ impl MemoryCtx { .addr_space_access_count .get_unchecked_mut(address_space as usize) += 1; } - } else { + } else if self.page_indices_pending.insert(page_id as usize) { // SAFETY: address_space passed is usually a hardcoded constant or derived from an // Instruction where it is bounds checked before passing unsafe { @@ -276,6 +277,7 @@ impl MemoryCtx { self.addr_space_access_count_pending.len(), ); } + self.page_indices_pending.clear(); self.page_indices.clear(); } @@ -290,6 +292,7 @@ impl MemoryCtx { self.addr_space_access_count_pending.len(), ); } + self.page_indices_pending.clear(); } /// Apply height updates given page counts @@ -347,9 +350,10 @@ impl MemoryCtx { /// Resolve all lazy updates of each memory access for memory adapters/poseidon2/merkle chip. #[inline(always)] pub(crate) fn lazy_update_boundary_heights(&mut self, trace_heights: &mut [u32]) { + let page_access_count = self.addr_space_access_count.iter().sum(); self.apply_height_updates( trace_heights, - self.page_access_count, + page_access_count, &self.addr_space_access_count, ); self.page_access_count = 0; diff --git a/extensions/rv32im/circuit/src/common/mod.rs b/extensions/rv32im/circuit/src/common/mod.rs index 1ef1d3dc8b..c7ce921bb6 100644 --- a/extensions/rv32im/circuit/src/common/mod.rs +++ b/extensions/rv32im/circuit/src/common/mod.rs @@ -275,67 +275,87 @@ mod aot { MemoryCtx, addr_space_access_count_pending ); - let inserted_label = format!(".asm_execute_pc_{pc}_inserted"); - let not_inserted_label = format!(".asm_execute_pc_{pc}_not_inserted"); - // The next section is the implementation of `BitSet::insert` in ASM. - // pub fn insert(&mut self, index: usize) -> bool { - // let word_index = index >> 6; - // let bit_index = index & 63; - // let mask = 1u64 << bit_index; - // let word = unsafe { self.words.get_unchecked_mut(word_index) }; - // let was_set = (*word & mask) != 0; - // *word |= mask; - // !was_set + let page_indices_pending_ptr_offset = + memory_ctx_offset + offset_of!(MemoryCtx, page_indices_pending); + let done_label = format!(".asm_execute_pc_{pc}_done"); + let pending_label = format!(".asm_execute_pc_{pc}_pending"); + // The next section implements the two-bitset insert logic in ASM. + // Optimized for the common case where both inserts return false (bit already set). + // + // if self.page_indices.insert(page_id as usize) { + // self.addr_space_access_count[address_space] += 1; + // } else if self.page_indices_pending.insert(page_id as usize) { + // self.addr_space_access_count_pending[address_space] += 1; // } - // Start with `ptr_reg = index` - // `reg1 = word_index` + // Start with `ptr_reg = page_id` + // `reg1 = word_index = page_id >> 6` asm_str += &format!(" mov {reg1}, {ptr_reg}\n"); asm_str += &format!(" shr {reg1}, 6\n"); - // `ptr_reg = bit_index = index & 63` + // Save word_index on stack for potential second attempt + asm_str += &format!(" push {reg1}\n"); + // `ptr_reg = bit_index = page_id & 63` asm_str += &format!(" and {ptr_reg}, 63\n"); // `reg2 = mask = 1u64 << bit_index` asm_str += &format!(" mov {reg2}, 1\n"); asm_str += &format!(" shlx {reg2}, {reg2}, {ptr_reg}\n"); + + // First attempt: try to insert into page_indices // `ptr_reg = self.page_indices.ptr` asm_str += &format!(" mov {ptr_reg}, [{REG_EXEC_STATE_PTR} + {page_indices_ptr_offset}]\n"); - - // `reg1 = word_ptr = &self.words.get_unchecked_mut(word_index)` + // `reg1 = word_ptr = &page_indices.words[word_index]` asm_str += &format!(" lea {reg1}, [{ptr_reg} + {reg1} * 8]\n"); // `ptr_reg = word = *word_ptr` asm_str += &format!(" mov {ptr_reg}, [{reg1}]\n"); - - // `test (*word & mask)` + // Test if bit is already set asm_str += &format!(" test {ptr_reg}, {reg2}\n"); - asm_str += &format!(" jnz {not_inserted_label}\n"); - // When (*word & mask) == 0 (new page inserted) - // `*word += mask` - asm_str += &format!(" add {ptr_reg}, {reg2}\n"); + asm_str += &format!(" jnz {pending_label}\n"); + + // Bit not set in page_indices, set it + asm_str += &format!(" or {ptr_reg}, {reg2}\n"); asm_str += &format!(" mov [{reg1}], {ptr_reg}\n"); - // reg1 = &self.page_access_count` + // Clean up stack + asm_str += &format!(" add rsp, 8\n"); + // Update page_access_count asm_str += &format!(" lea {reg1}, [{REG_EXEC_STATE_PTR} + {page_access_count_offset}]\n"); - // self.page_access_count += 1; asm_str += &format!(" add dword ptr [{reg1}], 1\n"); - // reg1 = &addr_space_access_count.as_ptr() + // Update addr_space_access_count[address_space] asm_str += &format!( " lea {reg1}, [{REG_EXEC_STATE_PTR} + {addr_space_access_count_ptr_offset}]\n" ); asm_str += &format!(" mov {reg1}, [{reg1}]\n"); - // self.addr_space_access_count[address_space] += 1; asm_str += &format!(" add dword ptr [{reg1} + {address_space} * 4], 1\n"); - asm_str += &format!(" jmp {inserted_label}\n"); - asm_str += &format!("{not_inserted_label}:\n"); - // When (*word & mask) != 0 (page already exists) - // reg1 = &addr_space_access_count_pending.as_ptr() + asm_str += &format!(" jmp {done_label}\n"); + + // Second attempt: try to insert into page_indices_pending + asm_str += &format!("{pending_label}:\n"); + // Restore word_index from stack + asm_str += &format!(" pop {reg1}\n"); + // `ptr_reg = self.page_indices_pending.ptr` + asm_str += &format!( + " mov {ptr_reg}, [{REG_EXEC_STATE_PTR} + {page_indices_pending_ptr_offset}]\n" + ); + // `reg1 = word_ptr = &page_indices_pending.words[word_index]` + asm_str += &format!(" lea {reg1}, [{ptr_reg} + {reg1} * 8]\n"); + // `ptr_reg = word = *word_ptr` + asm_str += &format!(" mov {ptr_reg}, [{reg1}]\n"); + // Test if bit is already set (reg2 still contains mask) + asm_str += &format!(" test {ptr_reg}, {reg2}\n"); + asm_str += &format!(" jnz {done_label}\n"); + + // Bit not set in page_indices_pending, set it + asm_str += &format!(" or {ptr_reg}, {reg2}\n"); + asm_str += &format!(" mov [{reg1}], {ptr_reg}\n"); + // Update addr_space_access_count_pending[address_space] asm_str += &format!( " lea {reg1}, [{REG_EXEC_STATE_PTR} + {addr_space_access_count_pending_ptr_offset}]\n" ); asm_str += &format!(" mov {reg1}, [{reg1}]\n"); - // self.addr_space_access_count_pending[address_space] += 1; asm_str += &format!(" add dword ptr [{reg1} + {address_space} * 4], 1\n"); - asm_str += &format!("{inserted_label}:\n"); + + asm_str += &format!("{done_label}:\n"); // Done Ok(asm_str) From 922e489875f9ae473c8ad80e62a9caddb6271049 Mon Sep 17 00:00:00 2001 From: Ayush Shukla Date: Thu, 25 Dec 2025 02:11:50 +0000 Subject: [PATCH 07/13] fix: append only array instead of two bitsets --- .../vm/src/arch/execution_mode/metered/ctx.rs | 20 ++-- .../arch/execution_mode/metered/memory_ctx.rs | 89 ++++++++------ extensions/rv32im/circuit/src/common/mod.rs | 113 +++++++++--------- 3 files changed, 116 insertions(+), 106 deletions(-) diff --git a/crates/vm/src/arch/execution_mode/metered/ctx.rs b/crates/vm/src/arch/execution_mode/metered/ctx.rs index 8a5c58f437..646aea1387 100644 --- a/crates/vm/src/arch/execution_mode/metered/ctx.rs +++ b/crates/vm/src/arch/execution_mode/metered/ctx.rs @@ -49,30 +49,30 @@ impl MeteredCtx { }) .unzip(); - let memory_ctx = MemoryCtx::new(config); + // Assert that the indices are correct + let segmentation_ctx = + SegmentationCtx::new(air_names, widths, interactions, config.segmentation_limits); + let memory_ctx = MemoryCtx::new(config, segmentation_ctx.segment_check_insns); // Assert that the indices are correct debug_assert!( - air_names[memory_ctx.boundary_idx].contains("Boundary"), + segmentation_ctx.air_names[memory_ctx.boundary_idx].contains("Boundary"), "air_name={}", - air_names[memory_ctx.boundary_idx] + segmentation_ctx.air_names[memory_ctx.boundary_idx] ); if let Some(merkle_tree_index) = memory_ctx.merkle_tree_index { debug_assert!( - air_names[merkle_tree_index].contains("Merkle"), + segmentation_ctx.air_names[merkle_tree_index].contains("Merkle"), "air_name={}", - air_names[merkle_tree_index] + segmentation_ctx.air_names[merkle_tree_index] ); } debug_assert!( - air_names[memory_ctx.adapter_offset].contains("AccessAdapterAir<2>"), + segmentation_ctx.air_names[memory_ctx.adapter_offset].contains("AccessAdapterAir<2>"), "air_name={}", - air_names[memory_ctx.adapter_offset] + segmentation_ctx.air_names[memory_ctx.adapter_offset] ); - let segmentation_ctx = - SegmentationCtx::new(air_names, widths, interactions, config.segmentation_limits); - let mut ctx = Self { trace_heights, is_trace_height_constant, 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 262e4169b1..a052061b15 100644 --- a/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs +++ b/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs @@ -3,6 +3,8 @@ use openvm_instructions::riscv::{RV32_NUM_REGISTERS, RV32_REGISTER_AS, RV32_REGI use crate::{arch::SystemConfig, system::memory::dimensions::MemoryDimensions}; +const MAX_MEMORY_OPS_PER_INSN: usize = 1 << 16; + #[derive(Clone, Debug)] pub struct BitSet { words: Box<[u64]>, @@ -104,18 +106,18 @@ pub struct MemoryCtx { pub boundary_idx: usize, pub merkle_tree_index: Option, pub adapter_offset: usize, - continuations_enabled: bool, + pub continuations_enabled: bool, chunk: u32, chunk_bits: u32, pub page_indices: BitSet, - pub page_indices_pending: BitSet, pub page_access_count: usize, pub addr_space_access_count: RVec, - pub addr_space_access_count_pending: RVec, + pub page_indices_since_checkpoint: Box<[u32]>, + pub page_indices_since_checkpoint_len: usize, } impl MemoryCtx { - pub fn new(config: &SystemConfig) -> Self { + pub fn new(config: &SystemConfig, segment_check_insns: u64) -> Self { let chunk = config.initial_block_size() as u32; let chunk_bits = chunk.ilog2(); @@ -124,6 +126,12 @@ impl MemoryCtx { let bitset_size = 1 << (merkle_height.saturating_sub(PAGE_BITS)); let addr_space_size = (1 << memory_dimensions.addr_space_height) + 1; + let segment_check_insns = if segment_check_insns == u64::MAX { + 0 + } else { + segment_check_insns as usize + }; + let page_indices_since_checkpoint_cap = segment_check_insns * MAX_MEMORY_OPS_PER_INSN; Self { min_block_size_bits: config.memory_config.min_block_size_bits(), @@ -135,10 +143,11 @@ impl MemoryCtx { memory_dimensions, continuations_enabled: config.continuation_enabled, page_indices: BitSet::new(bitset_size), - page_indices_pending: BitSet::new(bitset_size), page_access_count: 0, addr_space_access_count: vec![0; addr_space_size].into(), - addr_space_access_count_pending: vec![0; addr_space_size].into(), + page_indices_since_checkpoint: vec![0; page_indices_since_checkpoint_cap] + .into_boxed_slice(), + page_indices_since_checkpoint_len: 0, } } @@ -187,14 +196,17 @@ impl MemoryCtx { .addr_space_access_count .get_unchecked_mut(address_space as usize) += 1; } - } else if self.page_indices_pending.insert(page_id as usize) { - // SAFETY: address_space passed is usually a hardcoded constant or derived from an - // Instruction where it is bounds checked before passing + } + + if self.continuations_enabled { + // Append page_id to page_indices_since_checkpoint + let len = self.page_indices_since_checkpoint_len; + debug_assert!(len < self.page_indices_since_checkpoint.len()); + // SAFETY: len is within bounds, and we extend length by 1 after writing. unsafe { - *self - .addr_space_access_count_pending - .get_unchecked_mut(address_space as usize) += 1; + *self.page_indices_since_checkpoint.as_mut_ptr().add(len) = page_id; } + self.page_indices_since_checkpoint_len = len + 1; } } } @@ -262,37 +274,40 @@ impl MemoryCtx { } } - // Apply height updates for all pages accessed since last checkpoint - let page_access_count = self.addr_space_access_count_pending.iter().sum(); - self.apply_height_updates( - trace_heights, - page_access_count, - &self.addr_space_access_count_pending, - ); - // SAFETY: Resetting array elements to 0 is always safe - unsafe { - std::ptr::write_bytes( - self.addr_space_access_count_pending.as_mut_ptr(), - 0, - self.addr_space_access_count_pending.len(), - ); - } - self.page_indices_pending.clear(); + // Apply height updates for all pages accessed since last checkpoint, and + // initialize page_indices for the new segment. + let mut addr_space_access_count = vec![0; self.addr_space_access_count.len()]; + let mut page_access_count = 0; self.page_indices.clear(); + + let pages_len = self.page_indices_since_checkpoint_len; + for i in 0..pages_len { + // SAFETY: i is within 0..pages_len and pages_len is the slice length. + let page_id = unsafe { *self.page_indices_since_checkpoint.get_unchecked(i) } as usize; + if self.page_indices.insert(page_id) { + page_access_count += 1; + let (addr_space, _) = self + .memory_dimensions + .index_to_label((page_id as u64) << PAGE_BITS); + let addr_space_idx = addr_space as usize; + debug_assert!(addr_space_idx < addr_space_access_count.len()); + // SAFETY: addr_space_idx is bounds checked in debug and derived from a valid page + // id. + unsafe { + *addr_space_access_count.get_unchecked_mut(addr_space_idx) += 1; + } + } + } + + self.apply_height_updates(trace_heights, page_access_count, &addr_space_access_count); + + self.page_indices_since_checkpoint_len = 0; } /// Updates the checkpoint with current safe state #[inline(always)] pub(crate) fn update_checkpoint(&mut self) { - // SAFETY: Resetting array elements to 0 is always safe - unsafe { - std::ptr::write_bytes( - self.addr_space_access_count_pending.as_mut_ptr(), - 0, - self.addr_space_access_count_pending.len(), - ); - } - self.page_indices_pending.clear(); + self.page_indices_since_checkpoint_len = 0; } /// Apply height updates given page counts diff --git a/extensions/rv32im/circuit/src/common/mod.rs b/extensions/rv32im/circuit/src/common/mod.rs index c7ce921bb6..6c9a9bca9a 100644 --- a/extensions/rv32im/circuit/src/common/mod.rs +++ b/extensions/rv32im/circuit/src/common/mod.rs @@ -266,97 +266,92 @@ mod aot { + offset_of!(MeteredCtx, memory_ctx); let page_indices_ptr_offset = memory_ctx_offset + offset_of!(MemoryCtx, page_indices); - let page_access_count_offset = - memory_ctx_offset + offset_of!(MemoryCtx, page_access_count); let addr_space_access_count_ptr_offset = memory_ctx_offset + offset_of!(MemoryCtx, addr_space_access_count); - let addr_space_access_count_pending_ptr_offset = memory_ctx_offset + let page_indices_since_checkpoint_ptr_offset = memory_ctx_offset + + offset_of!(MemoryCtx, page_indices_since_checkpoint); + let page_indices_since_checkpoint_len_offset = memory_ctx_offset + offset_of!( MemoryCtx, - addr_space_access_count_pending + page_indices_since_checkpoint_len ); - let page_indices_pending_ptr_offset = - memory_ctx_offset + offset_of!(MemoryCtx, page_indices_pending); + let continuations_enabled_offset = + memory_ctx_offset + offset_of!(MemoryCtx, continuations_enabled); + let inserted_label = format!(".asm_execute_pc_{pc}_inserted"); let done_label = format!(".asm_execute_pc_{pc}_done"); - let pending_label = format!(".asm_execute_pc_{pc}_pending"); - // The next section implements the two-bitset insert logic in ASM. - // Optimized for the common case where both inserts return false (bit already set). - // - // if self.page_indices.insert(page_id as usize) { - // self.addr_space_access_count[address_space] += 1; - // } else if self.page_indices_pending.insert(page_id as usize) { - // self.addr_space_access_count_pending[address_space] += 1; + let skip_append_label = format!(".asm_execute_pc_{pc}_skip_append"); + // The next section is the implementation of `BitSet::insert` in ASM. + // pub fn insert(&mut self, index: usize) -> bool { + // let word_index = index >> 6; + // let bit_index = index & 63; + // let mask = 1u64 << bit_index; + // let word = unsafe { self.words.get_unchecked_mut(word_index) }; + // let was_set = (*word & mask) != 0; + // *word |= mask; + // !was_set // } + // self.page_indices_since_checkpoint.push(page_id); - // Start with `ptr_reg = page_id` - // `reg1 = word_index = page_id >> 6` + // Start with `ptr_reg = index` + asm_str += &format!(" push {ptr_reg}\n"); + // `reg1 = word_index` asm_str += &format!(" mov {reg1}, {ptr_reg}\n"); asm_str += &format!(" shr {reg1}, 6\n"); - // Save word_index on stack for potential second attempt - asm_str += &format!(" push {reg1}\n"); - // `ptr_reg = bit_index = page_id & 63` + // `ptr_reg = bit_index = index & 63` asm_str += &format!(" and {ptr_reg}, 63\n"); // `reg2 = mask = 1u64 << bit_index` asm_str += &format!(" mov {reg2}, 1\n"); asm_str += &format!(" shlx {reg2}, {reg2}, {ptr_reg}\n"); - - // First attempt: try to insert into page_indices // `ptr_reg = self.page_indices.ptr` asm_str += &format!(" mov {ptr_reg}, [{REG_EXEC_STATE_PTR} + {page_indices_ptr_offset}]\n"); - // `reg1 = word_ptr = &page_indices.words[word_index]` + + // `reg1 = word_ptr = &self.words.get_unchecked_mut(word_index)` asm_str += &format!(" lea {reg1}, [{ptr_reg} + {reg1} * 8]\n"); // `ptr_reg = word = *word_ptr` asm_str += &format!(" mov {ptr_reg}, [{reg1}]\n"); - // Test if bit is already set - asm_str += &format!(" test {ptr_reg}, {reg2}\n"); - asm_str += &format!(" jnz {pending_label}\n"); - // Bit not set in page_indices, set it - asm_str += &format!(" or {ptr_reg}, {reg2}\n"); + // `test (*word & mask)` + asm_str += &format!(" test {ptr_reg}, {reg2}\n"); + asm_str += &format!(" jnz {inserted_label}\n"); + // When (*word & mask) == 0 + // `*word += mask` + asm_str += &format!(" add {ptr_reg}, {reg2}\n"); asm_str += &format!(" mov [{reg1}], {ptr_reg}\n"); - // Clean up stack - asm_str += &format!(" add rsp, 8\n"); - // Update page_access_count - asm_str += - &format!(" lea {reg1}, [{REG_EXEC_STATE_PTR} + {page_access_count_offset}]\n"); - asm_str += &format!(" add dword ptr [{reg1}], 1\n"); - // Update addr_space_access_count[address_space] + // reg1 = &addr_space_access_count.as_ptr() asm_str += &format!( " lea {reg1}, [{REG_EXEC_STATE_PTR} + {addr_space_access_count_ptr_offset}]\n" ); asm_str += &format!(" mov {reg1}, [{reg1}]\n"); + // self.addr_space_access_count[address_space] += 1; asm_str += &format!(" add dword ptr [{reg1} + {address_space} * 4], 1\n"); - asm_str += &format!(" jmp {done_label}\n"); - - // Second attempt: try to insert into page_indices_pending - asm_str += &format!("{pending_label}:\n"); - // Restore word_index from stack - asm_str += &format!(" pop {reg1}\n"); - // `ptr_reg = self.page_indices_pending.ptr` + asm_str += &format!("{inserted_label}:\n"); + // Skip append if continuations are disabled asm_str += &format!( - " mov {ptr_reg}, [{REG_EXEC_STATE_PTR} + {page_indices_pending_ptr_offset}]\n" + " cmp byte ptr [{REG_EXEC_STATE_PTR} + {continuations_enabled_offset}], 0\n" ); - // `reg1 = word_ptr = &page_indices_pending.words[word_index]` - asm_str += &format!(" lea {reg1}, [{ptr_reg} + {reg1} * 8]\n"); - // `ptr_reg = word = *word_ptr` - asm_str += &format!(" mov {ptr_reg}, [{reg1}]\n"); - // Test if bit is already set (reg2 still contains mask) - asm_str += &format!(" test {ptr_reg}, {reg2}\n"); - asm_str += &format!(" jnz {done_label}\n"); - - // Bit not set in page_indices_pending, set it - asm_str += &format!(" or {ptr_reg}, {reg2}\n"); - asm_str += &format!(" mov [{reg1}], {ptr_reg}\n"); - // Update addr_space_access_count_pending[address_space] + asm_str += &format!(" je {skip_append_label}\n"); + // Append page_id to page_indices_since_checkpoint + asm_str += &format!(" pop {ptr_reg}\n"); asm_str += &format!( - " lea {reg1}, [{REG_EXEC_STATE_PTR} + {addr_space_access_count_pending_ptr_offset}]\n" + " mov {reg1}, [{REG_EXEC_STATE_PTR} + {page_indices_since_checkpoint_len_offset}]\n" ); - asm_str += &format!(" mov {reg1}, [{reg1}]\n"); - asm_str += &format!(" add dword ptr [{reg1} + {address_space} * 4], 1\n"); - + asm_str += &format!( + " mov {reg2}, [{REG_EXEC_STATE_PTR} + {page_indices_since_checkpoint_ptr_offset}]\n" + ); + let ptr_reg_32 = convert_x86_reg(ptr_reg, Width::W32).ok_or_else(|| { + AotError::Other(format!("unsupported ptr_reg for 32-bit store: {ptr_reg}")) + })?; + asm_str += &format!(" mov dword ptr [{reg2} + {reg1} * 4], {ptr_reg_32}\n"); + asm_str += &format!(" add {reg1}, 1\n"); + asm_str += &format!( + " mov [{REG_EXEC_STATE_PTR} + {page_indices_since_checkpoint_len_offset}], {reg1}\n" + ); + asm_str += &format!(" jmp {done_label}\n"); + asm_str += &format!("{skip_append_label}:\n"); + // Discard pushed page_id + asm_str += &format!(" pop {ptr_reg}\n"); asm_str += &format!("{done_label}:\n"); - // Done Ok(asm_str) } From 04f11f72a01a103cb1752acbfdbbe44ad4d7f6df Mon Sep 17 00:00:00 2001 From: Ayush Shukla Date: Tue, 30 Dec 2025 20:49:43 +0000 Subject: [PATCH 08/13] fix: remove if continuations enabled check --- .../vm/src/arch/execution_mode/metered/ctx.rs | 23 ++++++++----- .../arch/execution_mode/metered/memory_ctx.rs | 32 +++++++++---------- .../execution_mode/metered/segment_ctx.rs | 3 +- extensions/rv32im/circuit/src/common/mod.rs | 14 -------- 4 files changed, 32 insertions(+), 40 deletions(-) diff --git a/crates/vm/src/arch/execution_mode/metered/ctx.rs b/crates/vm/src/arch/execution_mode/metered/ctx.rs index 646aea1387..6a36ad4338 100644 --- a/crates/vm/src/arch/execution_mode/metered/ctx.rs +++ b/crates/vm/src/arch/execution_mode/metered/ctx.rs @@ -49,7 +49,6 @@ impl MeteredCtx { }) .unzip(); - // Assert that the indices are correct let segmentation_ctx = SegmentationCtx::new(air_names, widths, interactions, config.segmentation_limits); let memory_ctx = MemoryCtx::new(config, segmentation_ctx.segment_check_insns); @@ -80,11 +79,6 @@ impl MeteredCtx { segmentation_ctx, suspend_on_segment: false, }; - if !config.continuation_enabled { - // force single segment - ctx.segmentation_ctx.segment_check_insns = u64::MAX; - ctx.segmentation_ctx.instrets_until_check = u64::MAX; - } // Add merkle height contributions for all registers ctx.memory_ctx.add_register_merkle_heights(); @@ -98,9 +92,8 @@ impl MeteredCtx { self.segmentation_ctx.set_max_trace_height(max_trace_height); let max_check_freq = (max_trace_height / 2) as u64; if max_check_freq < self.segmentation_ctx.segment_check_insns { - self.segmentation_ctx.segment_check_insns = max_check_freq; + self = self.with_segment_check_insns(max_check_freq); } - self.segmentation_ctx.instrets_until_check = self.segmentation_ctx.segment_check_insns; self } @@ -114,6 +107,20 @@ impl MeteredCtx { self } + pub fn with_segment_check_insns(mut self, segment_check_insns: u64) -> Self { + self.segmentation_ctx.segment_check_insns = segment_check_insns; + self.segmentation_ctx.instrets_until_check = segment_check_insns; + + // Update memory context with new segment check instructions + let page_indices_since_checkpoint_cap = + MemoryCtx::::calculate_checkpoint_capacity(segment_check_insns); + + self.memory_ctx.page_indices_since_checkpoint = + vec![0; page_indices_since_checkpoint_cap].into_boxed_slice(); + self.memory_ctx.page_indices_since_checkpoint_len = 0; + self + } + pub fn segments(&self) -> &[Segment] { &self.segmentation_ctx.segments } 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 a052061b15..5772b0a162 100644 --- a/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs +++ b/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs @@ -3,7 +3,8 @@ use openvm_instructions::riscv::{RV32_NUM_REGISTERS, RV32_REGISTER_AS, RV32_REGI use crate::{arch::SystemConfig, system::memory::dimensions::MemoryDimensions}; -const MAX_MEMORY_OPS_PER_INSN: usize = 1 << 16; +/// Upper bound on memory operations per instruction. Used for buffer allocation. +pub const MAX_MEMORY_OPS_PER_INSN: usize = 1 << 16; #[derive(Clone, Debug)] pub struct BitSet { @@ -126,12 +127,8 @@ impl MemoryCtx { let bitset_size = 1 << (merkle_height.saturating_sub(PAGE_BITS)); let addr_space_size = (1 << memory_dimensions.addr_space_height) + 1; - let segment_check_insns = if segment_check_insns == u64::MAX { - 0 - } else { - segment_check_insns as usize - }; - let page_indices_since_checkpoint_cap = segment_check_insns * MAX_MEMORY_OPS_PER_INSN; + let page_indices_since_checkpoint_cap = + Self::calculate_checkpoint_capacity(segment_check_insns); Self { min_block_size_bits: config.memory_config.min_block_size_bits(), @@ -151,6 +148,11 @@ impl MemoryCtx { } } + #[inline(always)] + pub(super) fn calculate_checkpoint_capacity(segment_check_insns: u64) -> usize { + segment_check_insns as usize * MAX_MEMORY_OPS_PER_INSN + } + #[inline(always)] pub(crate) fn add_register_merkle_heights(&mut self) { if self.continuations_enabled { @@ -198,16 +200,14 @@ impl MemoryCtx { } } - if self.continuations_enabled { - // Append page_id to page_indices_since_checkpoint - let len = self.page_indices_since_checkpoint_len; - debug_assert!(len < self.page_indices_since_checkpoint.len()); - // SAFETY: len is within bounds, and we extend length by 1 after writing. - unsafe { - *self.page_indices_since_checkpoint.as_mut_ptr().add(len) = page_id; - } - self.page_indices_since_checkpoint_len = len + 1; + // Append page_id to page_indices_since_checkpoint + let len = self.page_indices_since_checkpoint_len; + debug_assert!(len < self.page_indices_since_checkpoint.len()); + // SAFETY: len is within bounds, and we extend length by 1 after writing. + unsafe { + *self.page_indices_since_checkpoint.as_mut_ptr().add(len) = page_id; } + self.page_indices_since_checkpoint_len = len + 1; } } diff --git a/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs b/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs index 72a20a6baa..19fc0ad223 100644 --- a/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs +++ b/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs @@ -68,8 +68,7 @@ pub struct SegmentationCtx { pub(crate) segmentation_limits: SegmentationLimits, pub instret: u64, pub instrets_until_check: u64, - #[getset(set_with = "pub")] - pub segment_check_insns: u64, + pub(super) segment_check_insns: u64, /// Checkpoint of trace heights at last known state where all thresholds satisfied pub(crate) checkpoint_trace_heights: Vec, /// Instruction count at the checkpoint diff --git a/extensions/rv32im/circuit/src/common/mod.rs b/extensions/rv32im/circuit/src/common/mod.rs index 6c9a9bca9a..d26d7e4853 100644 --- a/extensions/rv32im/circuit/src/common/mod.rs +++ b/extensions/rv32im/circuit/src/common/mod.rs @@ -275,11 +275,7 @@ mod aot { MemoryCtx, page_indices_since_checkpoint_len ); - let continuations_enabled_offset = - memory_ctx_offset + offset_of!(MemoryCtx, continuations_enabled); let inserted_label = format!(".asm_execute_pc_{pc}_inserted"); - let done_label = format!(".asm_execute_pc_{pc}_done"); - let skip_append_label = format!(".asm_execute_pc_{pc}_skip_append"); // The next section is the implementation of `BitSet::insert` in ASM. // pub fn insert(&mut self, index: usize) -> bool { // let word_index = index >> 6; @@ -326,11 +322,6 @@ mod aot { // self.addr_space_access_count[address_space] += 1; asm_str += &format!(" add dword ptr [{reg1} + {address_space} * 4], 1\n"); asm_str += &format!("{inserted_label}:\n"); - // Skip append if continuations are disabled - asm_str += &format!( - " cmp byte ptr [{REG_EXEC_STATE_PTR} + {continuations_enabled_offset}], 0\n" - ); - asm_str += &format!(" je {skip_append_label}\n"); // Append page_id to page_indices_since_checkpoint asm_str += &format!(" pop {ptr_reg}\n"); asm_str += &format!( @@ -347,11 +338,6 @@ mod aot { asm_str += &format!( " mov [{REG_EXEC_STATE_PTR} + {page_indices_since_checkpoint_len_offset}], {reg1}\n" ); - asm_str += &format!(" jmp {done_label}\n"); - asm_str += &format!("{skip_append_label}:\n"); - // Discard pushed page_id - asm_str += &format!(" pop {ptr_reg}\n"); - asm_str += &format!("{done_label}:\n"); Ok(asm_str) } From 6607435fc5b43bc38ffc62a4640a7db8ae5eaed4 Mon Sep 17 00:00:00 2001 From: Ayush Shukla Date: Tue, 30 Dec 2025 21:20:12 +0000 Subject: [PATCH 09/13] fix: remove native test expecting single segment --- .../native/circuit/tests/integration_test.rs | 36 ------------------- 1 file changed, 36 deletions(-) diff --git a/extensions/native/circuit/tests/integration_test.rs b/extensions/native/circuit/tests/integration_test.rs index d0d56f6d4c..6cb64077e8 100644 --- a/extensions/native/circuit/tests/integration_test.rs +++ b/extensions/native/circuit/tests/integration_test.rs @@ -11,7 +11,6 @@ use openvm_circuit::system::cuda::extensions::SystemGpuBuilder as SystemBuilder; use openvm_circuit::{arch::RowMajorMatrixArena, system::SystemCpuBuilder as SystemBuilder}; use openvm_circuit::{ arch::{ - execution_mode::metered::segment_ctx::{SegmentationLimits, DEFAULT_SEGMENT_CHECK_INSNS}, hasher::{poseidon2::vm_poseidon2_hasher, Hasher}, verify_segments, verify_single, AirInventory, ContinuationVmProver, PreflightExecutionOutput, SingleSegmentVmProver, VirtualMachine, VmCircuitConfig, @@ -948,41 +947,6 @@ fn test_vm_execute_native_chips() { .expect("Failed to execute"); } -// This test ensures that metered execution never segments when continuations is disabled -#[test] -fn test_single_segment_executor_no_segmentation() { - setup_tracing(); - - let mut config = test_native_config(); - config - .system - .set_segmentation_limits(SegmentationLimits::default().with_max_trace_height(1)); - - let engine = TestEngine::new(FriParameters::new_for_testing(3)); - let (vm, _) = - VirtualMachine::new_with_keygen(engine, NativeBuilder::default(), config).unwrap(); - let instructions: Vec<_> = (0..2 * DEFAULT_SEGMENT_CHECK_INSNS) - .map(|_| Instruction::large_from_isize(ADD.global_opcode(), 0, 0, 1, 4, 0, 0, 0)) - .chain(std::iter::once(Instruction::from_isize( - TERMINATE.global_opcode(), - 0, - 0, - 0, - 0, - 0, - ))) - .collect(); - - let exe = VmExe::new(Program::from_instructions(&instructions)); - let executor_idx_to_air_idx = vm.executor_idx_to_air_idx(); - let metered_ctx = vm.build_metered_ctx(&exe); - vm.executor() - .metered_instance(&exe, &executor_idx_to_air_idx) - .unwrap() - .execute_metered(vec![], metered_ctx) - .unwrap(); -} - #[test] fn test_vm_execute_metered_cost_native_chips() { type F = BabyBear; From dad85eec9554181a6a3d13677c58cce0bbbf730b Mon Sep 17 00:00:00 2001 From: Ayush Shukla Date: Tue, 30 Dec 2025 22:11:50 +0000 Subject: [PATCH 10/13] fix: avoid stack push/pop --- .../arch/execution_mode/metered/memory_ctx.rs | 18 +++++----- extensions/rv32im/circuit/src/common/mod.rs | 34 +++++++++---------- 2 files changed, 25 insertions(+), 27 deletions(-) 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 5772b0a162..32ee4cd9f3 100644 --- a/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs +++ b/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs @@ -190,6 +190,15 @@ impl MemoryCtx { let end_page_id = ((end_block_id - 1) >> PAGE_BITS) + 1; for page_id in start_page_id..end_page_id { + // Append page_id to page_indices_since_checkpoint + let len = self.page_indices_since_checkpoint_len; + debug_assert!(len < self.page_indices_since_checkpoint.len()); + // SAFETY: len is within bounds, and we extend length by 1 after writing. + unsafe { + *self.page_indices_since_checkpoint.as_mut_ptr().add(len) = page_id; + } + self.page_indices_since_checkpoint_len = len + 1; + if self.page_indices.insert(page_id as usize) { // SAFETY: address_space passed is usually a hardcoded constant or derived from an // Instruction where it is bounds checked before passing @@ -199,15 +208,6 @@ impl MemoryCtx { .get_unchecked_mut(address_space as usize) += 1; } } - - // Append page_id to page_indices_since_checkpoint - let len = self.page_indices_since_checkpoint_len; - debug_assert!(len < self.page_indices_since_checkpoint.len()); - // SAFETY: len is within bounds, and we extend length by 1 after writing. - unsafe { - *self.page_indices_since_checkpoint.as_mut_ptr().add(len) = page_id; - } - self.page_indices_since_checkpoint_len = len + 1; } } diff --git a/extensions/rv32im/circuit/src/common/mod.rs b/extensions/rv32im/circuit/src/common/mod.rs index d26d7e4853..eec334509a 100644 --- a/extensions/rv32im/circuit/src/common/mod.rs +++ b/extensions/rv32im/circuit/src/common/mod.rs @@ -288,8 +288,22 @@ mod aot { // } // self.page_indices_since_checkpoint.push(page_id); - // Start with `ptr_reg = index` - asm_str += &format!(" push {ptr_reg}\n"); + // Append page_id to page_indices_since_checkpoint + asm_str += &format!( + " mov {reg1}, [{REG_EXEC_STATE_PTR} + {page_indices_since_checkpoint_len_offset}]\n" + ); + asm_str += &format!( + " mov {reg2}, [{REG_EXEC_STATE_PTR} + {page_indices_since_checkpoint_ptr_offset}]\n" + ); + let ptr_reg_32 = convert_x86_reg(ptr_reg, Width::W32).ok_or_else(|| { + AotError::Other(format!("unsupported ptr_reg for 32-bit store: {ptr_reg}")) + })?; + asm_str += &format!(" mov dword ptr [{reg2} + {reg1} * 4], {ptr_reg_32}\n"); + asm_str += &format!(" add {reg1}, 1\n"); + asm_str += &format!( + " mov [{REG_EXEC_STATE_PTR} + {page_indices_since_checkpoint_len_offset}], {reg1}\n" + ); + // `reg1 = word_index` asm_str += &format!(" mov {reg1}, {ptr_reg}\n"); asm_str += &format!(" shr {reg1}, 6\n"); @@ -322,22 +336,6 @@ mod aot { // self.addr_space_access_count[address_space] += 1; asm_str += &format!(" add dword ptr [{reg1} + {address_space} * 4], 1\n"); asm_str += &format!("{inserted_label}:\n"); - // Append page_id to page_indices_since_checkpoint - asm_str += &format!(" pop {ptr_reg}\n"); - asm_str += &format!( - " mov {reg1}, [{REG_EXEC_STATE_PTR} + {page_indices_since_checkpoint_len_offset}]\n" - ); - asm_str += &format!( - " mov {reg2}, [{REG_EXEC_STATE_PTR} + {page_indices_since_checkpoint_ptr_offset}]\n" - ); - let ptr_reg_32 = convert_x86_reg(ptr_reg, Width::W32).ok_or_else(|| { - AotError::Other(format!("unsupported ptr_reg for 32-bit store: {ptr_reg}")) - })?; - asm_str += &format!(" mov dword ptr [{reg2} + {reg1} * 4], {ptr_reg_32}\n"); - asm_str += &format!(" add {reg1}, 1\n"); - asm_str += &format!( - " mov [{REG_EXEC_STATE_PTR} + {page_indices_since_checkpoint_len_offset}], {reg1}\n" - ); Ok(asm_str) } From cdf29cb3ec4bb0152784144c4319ac229de5a7de Mon Sep 17 00:00:00 2001 From: Ayush Shukla Date: Wed, 31 Dec 2025 15:46:03 +0000 Subject: [PATCH 11/13] chore: clean up --- .../vm/src/arch/execution_mode/metered/ctx.rs | 4 +- .../arch/execution_mode/metered/memory_ctx.rs | 42 ++++++++----------- extensions/rv32im/circuit/src/common/mod.rs | 33 +++++++++------ 3 files changed, 40 insertions(+), 39 deletions(-) diff --git a/crates/vm/src/arch/execution_mode/metered/ctx.rs b/crates/vm/src/arch/execution_mode/metered/ctx.rs index 6a36ad4338..3a87d15fef 100644 --- a/crates/vm/src/arch/execution_mode/metered/ctx.rs +++ b/crates/vm/src/arch/execution_mode/metered/ctx.rs @@ -82,6 +82,8 @@ impl MeteredCtx { // Add merkle height contributions for all registers ctx.memory_ctx.add_register_merkle_heights(); + ctx.memory_ctx + .lazy_update_boundary_heights(&mut ctx.trace_heights); ctx } @@ -153,8 +155,6 @@ impl MeteredCtx { .initialize_segment(&mut self.trace_heights, &self.is_trace_height_constant); // Initialize memory context for new segment self.memory_ctx.initialize_segment(&mut self.trace_heights); - // Add merkle height contributions for all registers - self.memory_ctx.add_register_merkle_heights(); } else { // Update checkpoint for trace heights self.segmentation_ctx 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 32ee4cd9f3..3b9f652309 100644 --- a/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs +++ b/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs @@ -107,11 +107,10 @@ pub struct MemoryCtx { pub boundary_idx: usize, pub merkle_tree_index: Option, pub adapter_offset: usize, - pub continuations_enabled: bool, + continuations_enabled: bool, chunk: u32, chunk_bits: u32, pub page_indices: BitSet, - pub page_access_count: usize, pub addr_space_access_count: RVec, pub page_indices_since_checkpoint: Box<[u32]>, pub page_indices_since_checkpoint_len: usize, @@ -140,7 +139,6 @@ impl MemoryCtx { memory_dimensions, continuations_enabled: config.continuation_enabled, page_indices: BitSet::new(bitset_size), - page_access_count: 0, addr_space_access_count: vec![0; addr_space_size].into(), page_indices_since_checkpoint: vec![0; page_indices_since_checkpoint_cap] .into_boxed_slice(), @@ -257,7 +255,10 @@ impl MemoryCtx { /// Initialize state for a new segment #[inline(always)] pub(crate) fn initialize_segment(&mut self, trace_heights: &mut [u32]) { - // Reset trace heights for memory chips + // Clear page indices for the new segment + self.page_indices.clear(); + + // Reset trace heights for memory chips as 0 // SAFETY: boundary_idx is a compile time constant within bounds unsafe { *trace_heights.get_unchecked_mut(self.boundary_idx) = 0; @@ -277,15 +278,11 @@ impl MemoryCtx { // Apply height updates for all pages accessed since last checkpoint, and // initialize page_indices for the new segment. let mut addr_space_access_count = vec![0; self.addr_space_access_count.len()]; - let mut page_access_count = 0; - self.page_indices.clear(); - let pages_len = self.page_indices_since_checkpoint_len; for i in 0..pages_len { // SAFETY: i is within 0..pages_len and pages_len is the slice length. let page_id = unsafe { *self.page_indices_since_checkpoint.get_unchecked(i) } as usize; if self.page_indices.insert(page_id) { - page_access_count += 1; let (addr_space, _) = self .memory_dimensions .index_to_label((page_id as u64) << PAGE_BITS); @@ -298,10 +295,12 @@ impl MemoryCtx { } } } - - self.apply_height_updates(trace_heights, page_access_count, &addr_space_access_count); - + self.apply_height_updates(trace_heights, &addr_space_access_count); self.page_indices_since_checkpoint_len = 0; + + // Add merkle height contributions for all registers + self.add_register_merkle_heights(); + self.lazy_update_boundary_heights(trace_heights); } /// Updates the checkpoint with current safe state @@ -312,12 +311,9 @@ impl MemoryCtx { /// Apply height updates given page counts #[inline(always)] - fn apply_height_updates( - &self, - trace_heights: &mut [u32], - page_access_count: usize, - addr_space_access_count: &[usize], - ) { + fn apply_height_updates(&self, trace_heights: &mut [u32], addr_space_access_count: &[usize]) { + let page_access_count: usize = addr_space_access_count.iter().sum(); + // On page fault, assume we add all leaves in a page let leaves = (page_access_count << PAGE_BITS) as u32; // SAFETY: boundary_idx is a compile time constant within bounds @@ -346,7 +342,9 @@ impl MemoryCtx { } } - for (address_space, &x) in addr_space_access_count.iter().enumerate() { + for address_space in 0..addr_space_access_count.len() { + // SAFETY: address_space is from 0 to len(), guaranteed to be in bounds + let x = unsafe { *addr_space_access_count.get_unchecked(address_space) }; if x > 0 { // Initial **and** final handling of touched pages requires send (resp. receive) in // chunk-sized units for the merkle chip @@ -365,13 +363,7 @@ impl MemoryCtx { /// Resolve all lazy updates of each memory access for memory adapters/poseidon2/merkle chip. #[inline(always)] pub(crate) fn lazy_update_boundary_heights(&mut self, trace_heights: &mut [u32]) { - let page_access_count = self.addr_space_access_count.iter().sum(); - self.apply_height_updates( - trace_heights, - page_access_count, - &self.addr_space_access_count, - ); - self.page_access_count = 0; + self.apply_height_updates(trace_heights, &self.addr_space_access_count); // SAFETY: Resetting array elements to 0 is always safe unsafe { std::ptr::write_bytes( diff --git a/extensions/rv32im/circuit/src/common/mod.rs b/extensions/rv32im/circuit/src/common/mod.rs index eec334509a..827eb6b76f 100644 --- a/extensions/rv32im/circuit/src/common/mod.rs +++ b/extensions/rv32im/circuit/src/common/mod.rs @@ -212,8 +212,15 @@ mod aot { // let end_page_id = ((end_block_id - 1) >> PAGE_BITS) + 1; // for page_id in start_page_id..end_page_id { + // // Append page_id to page_indices_since_checkpoint + // let len = self.page_indices_since_checkpoint_len; + // // SAFETY: len is within bounds, and we extend length by 1 after writing. + // unsafe { + // *self.page_indices_since_checkpoint.as_mut_ptr().add(len) = page_id; + // } + // self.page_indices_since_checkpoint_len = len + 1; + // // if self.page_indices.insert(page_id as usize) { - // self.page_access_count += 1; // // SAFETY: address_space passed is usually a hardcoded constant or derived // from an // Instruction where it is bounds checked before passing // unsafe { @@ -276,17 +283,6 @@ mod aot { page_indices_since_checkpoint_len ); let inserted_label = format!(".asm_execute_pc_{pc}_inserted"); - // The next section is the implementation of `BitSet::insert` in ASM. - // pub fn insert(&mut self, index: usize) -> bool { - // let word_index = index >> 6; - // let bit_index = index & 63; - // let mask = 1u64 << bit_index; - // let word = unsafe { self.words.get_unchecked_mut(word_index) }; - // let was_set = (*word & mask) != 0; - // *word |= mask; - // !was_set - // } - // self.page_indices_since_checkpoint.push(page_id); // Append page_id to page_indices_since_checkpoint asm_str += &format!( @@ -304,6 +300,18 @@ mod aot { " mov [{REG_EXEC_STATE_PTR} + {page_indices_since_checkpoint_len_offset}], {reg1}\n" ); + // The next section is the implementation of `BitSet::insert` in ASM. + // pub fn insert(&mut self, index: usize) -> bool { + // let word_index = index >> 6; + // let bit_index = index & 63; + // let mask = 1u64 << bit_index; + // let word = unsafe { self.words.get_unchecked_mut(word_index) }; + // let was_set = (*word & mask) != 0; + // *word |= mask; + // !was_set + // } + + // Start with `ptr_reg = index` // `reg1 = word_index` asm_str += &format!(" mov {reg1}, {ptr_reg}\n"); asm_str += &format!(" shr {reg1}, 6\n"); @@ -336,6 +344,7 @@ mod aot { // self.addr_space_access_count[address_space] += 1; asm_str += &format!(" add dword ptr [{reg1} + {address_space} * 4], 1\n"); asm_str += &format!("{inserted_label}:\n"); + // Inserted, do nothing Ok(asm_str) } From d55c7295a3879019985365bd13c8f3d3262177cb Mon Sep 17 00:00:00 2001 From: Ayush Shukla Date: Wed, 31 Dec 2025 18:18:11 +0000 Subject: [PATCH 12/13] fix: checkpoint at segmentation --- .../vm/src/arch/execution_mode/metered/ctx.rs | 38 +++++++++++++++---- .../arch/execution_mode/metered/memory_ctx.rs | 1 - .../execution_mode/metered/segment_ctx.rs | 18 ++++----- 3 files changed, 37 insertions(+), 20 deletions(-) diff --git a/crates/vm/src/arch/execution_mode/metered/ctx.rs b/crates/vm/src/arch/execution_mode/metered/ctx.rs index 3a87d15fef..ce1144bfbf 100644 --- a/crates/vm/src/arch/execution_mode/metered/ctx.rs +++ b/crates/vm/src/arch/execution_mode/metered/ctx.rs @@ -150,18 +150,40 @@ impl MeteredCtx { ); if did_segment { - // Initialize segment context for new segment + // Initialize contexts for new segment self.segmentation_ctx .initialize_segment(&mut self.trace_heights, &self.is_trace_height_constant); - // Initialize memory context for new segment self.memory_ctx.initialize_segment(&mut self.trace_heights); - } else { - // Update checkpoint for trace heights - self.segmentation_ctx - .update_checkpoint(self.segmentation_ctx.instret, &self.trace_heights); - // Update checkpoint for page indices - self.memory_ctx.update_checkpoint(); + + // Check if the new segment is within limits + if self.segmentation_ctx.should_segment( + self.segmentation_ctx.instret, + &self.trace_heights, + &self.is_trace_height_constant, + ) { + let trace_heights_str = self + .trace_heights + .iter() + .zip(self.segmentation_ctx.air_names.iter()) + .filter(|(&height, _)| height > 0) + .map(|(&height, name)| format!(" {name} = {height}")) + .collect::>() + .join("\n"); + tracing::warn!( + "Segment initialized with heights that exceed limits\n\ + instret={}\n\ + trace_heights=[\n{}\n]", + self.segmentation_ctx.instret, + trace_heights_str + ); + } } + + // Update checkpoints + self.segmentation_ctx + .update_checkpoint(self.segmentation_ctx.instret, &self.trace_heights); + self.memory_ctx.update_checkpoint(); + did_segment } 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 3b9f652309..9eeec8ad07 100644 --- a/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs +++ b/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs @@ -296,7 +296,6 @@ impl MemoryCtx { } } self.apply_height_updates(trace_heights, &addr_space_access_count); - self.page_indices_since_checkpoint_len = 0; // Add merkle height contributions for all registers self.add_register_merkle_heights(); diff --git a/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs b/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs index 19fc0ad223..0b34ad3f99 100644 --- a/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs +++ b/crates/vm/src/arch/execution_mode/metered/segment_ctx.rs @@ -299,17 +299,13 @@ impl SegmentationCtx { trace_heights: &mut [u32], is_trace_height_constant: &[bool], ) { - // Get the segment heights from the last created segment - if let Some(last_segment) = self.segments.last() { - // Reset trace heights by subtracting the segment's heights - self.reset_trace_heights( - trace_heights, - &last_segment.trace_heights, - is_trace_height_constant, - ); - } - // Reset checkpoint since we're starting a new segment - self.checkpoint_instret = 0; + // Reset trace heights by subtracting the last segment's heights + let last_segment = self.segments.last().unwrap(); + self.reset_trace_heights( + trace_heights, + &last_segment.trace_heights, + is_trace_height_constant, + ); } /// Resets trace heights by subtracting segment heights From f1c9fba114b59c80e8dc94100230b98b18c294e8 Mon Sep 17 00:00:00 2001 From: Jonathan Wang <31040440+jonathanpwang@users.noreply.github.com> Date: Sun, 4 Jan 2026 17:46:51 -0800 Subject: [PATCH 13/13] fix: add safety assert to ensure buffer doesn't overflow --- .../vm/src/arch/execution_mode/metered/memory_ctx.rs | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) 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 9eeec8ad07..68e2c6a665 100644 --- a/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs +++ b/crates/vm/src/arch/execution_mode/metered/memory_ctx.rs @@ -3,8 +3,8 @@ use openvm_instructions::riscv::{RV32_NUM_REGISTERS, RV32_REGISTER_AS, RV32_REGI use crate::{arch::SystemConfig, system::memory::dimensions::MemoryDimensions}; -/// Upper bound on memory operations per instruction. Used for buffer allocation. -pub const MAX_MEMORY_OPS_PER_INSN: usize = 1 << 16; +/// Upper bound on number of memory pages accessed per instruction. Used for buffer allocation. +pub const MAX_MEM_PAGE_OPS_PER_INSN: usize = 1 << 16; #[derive(Clone, Debug)] pub struct BitSet { @@ -148,7 +148,7 @@ impl MemoryCtx { #[inline(always)] pub(super) fn calculate_checkpoint_capacity(segment_check_insns: u64) -> usize { - segment_check_insns as usize * MAX_MEMORY_OPS_PER_INSN + segment_check_insns as usize * MAX_MEM_PAGE_OPS_PER_INSN } #[inline(always)] @@ -186,6 +186,11 @@ impl MemoryCtx { let end_block_id = start_block_id + num_blocks; let start_page_id = start_block_id >> PAGE_BITS; let end_page_id = ((end_block_id - 1) >> PAGE_BITS) + 1; + assert!( + self.page_indices_since_checkpoint_len + (end_page_id - start_page_id) as usize + <= self.page_indices_since_checkpoint.len(), + "more than {MAX_MEM_PAGE_OPS_PER_INSN} memory pages accessed in a single instruction" + ); for page_id in start_page_id..end_page_id { // Append page_id to page_indices_since_checkpoint