Skip to content

fix(aot): x86 assembly for addr_space_access_count update #2340

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Merged
jonathanpwang merged 3 commits into from
Jan 5, 2026
Merged

fix(aot): x86 assembly for addr_space_access_count update #2340

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
4208c0b
fix(aot): x86 assembly for `addr_space_access_count` update
jonathanpwang Jan 5, 2026
38e51e3
Revert "fix(aot): x86 assembly for `addr_space_access_count` update"
jonathanpwang Jan 5, 2026
8398a0b
use RVec u32
jonathanpwang Jan 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 7 additions & 9 deletions crates/vm/src/arch/execution_mode/metered/memory_ctx.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,7 +111,7 @@ pub struct MemoryCtx<const PAGE_BITS: usize> {
chunk: u32,
chunk_bits: u32,
pub page_indices: BitSet,
pub addr_space_access_count: RVec<usize>,
pub addr_space_access_count: RVec<u32>,
pub page_indices_since_checkpoint: Box<[u32]>,
pub page_indices_since_checkpoint_len: usize,
}
Expand Down Expand Up @@ -315,11 +315,11 @@ impl<const PAGE_BITS: usize> MemoryCtx<PAGE_BITS> {

/// Apply height updates given page counts
#[inline(always)]
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();
fn apply_height_updates(&self, trace_heights: &mut [u32], addr_space_access_count: &[u32]) {
let page_access_count: u32 = 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;
let leaves = page_access_count << PAGE_BITS;
// SAFETY: boundary_idx is a compile time constant within bounds
unsafe {
*trace_heights.get_unchecked_mut(self.boundary_idx) += leaves;
Expand All @@ -339,10 +339,8 @@ impl<const PAGE_BITS: usize> MemoryCtx<PAGE_BITS> {
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_access_count as u32 * 2;
*trace_heights.get_unchecked_mut(merkle_tree_idx) +=
nodes * page_access_count as u32 * 2;
*trace_heights.get_unchecked_mut(poseidon2_idx) += nodes * page_access_count * 2;
*trace_heights.get_unchecked_mut(merkle_tree_idx) += nodes * page_access_count * 2;
}
}

Expand All @@ -358,7 +356,7 @@ impl<const PAGE_BITS: usize> MemoryCtx<PAGE_BITS> {
trace_heights,
address_space as u32,
self.chunk_bits,
(x << (PAGE_BITS + 1)) as u32,
x << (PAGE_BITS + 1),
);
}
}
Expand Down