Skip to content
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.

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

[feat] Make online and offline memory use vectors of pages instead of hashmaps #1224

Merged
merged 17 commits into from
Jan 22, 2025
Merged
Changes from 8 commits
Commits
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
13 changes: 11 additions & 2 deletions crates/sdk/src/commit.rs
Original file line number Diff line number Diff line change
@@ -6,7 +6,10 @@ use openvm_circuit::{
instructions::exe::VmExe,
VmConfig,
},
system::{memory::tree::MemoryNode, program::trace::VmCommittedExe},
system::{
memory::{paged_vec::AddressMap, tree::MemoryNode},
program::trace::VmCommittedExe,
},
};
use openvm_native_compiler::{conversion::CompilerOptions, ir::DIGEST_SIZE};
use openvm_stark_backend::{config::StarkGenericConfig, p3_field::PrimeField32};
@@ -60,9 +63,15 @@ impl AppExecutionCommit<F> {
.commit
.into();

let mem_config = app_vm_config.system().memory_config;
let init_memory_commit = MemoryNode::tree_from_memory(
memory_dimensions,
&app_exe.exe.init_memory.clone().into_iter().collect(),
&AddressMap::from_iter(
mem_config.as_offset,
1 << mem_config.as_height,
1 << mem_config.pointer_max_bits,
app_exe.exe.init_memory.clone(),
),
&hasher,
)
.hash();
2 changes: 1 addition & 1 deletion crates/vm/src/arch/config.rs
Original file line number Diff line number Diff line change
@@ -60,7 +60,7 @@ pub struct MemoryConfig {

impl Default for MemoryConfig {
fn default() -> Self {
Self::new(29, 1, 29, 29, 17, 64, 1 << 24)
Self::new(3, 1, 29, 29, 17, 64, 1 << 24)
}
}

10 changes: 8 additions & 2 deletions crates/vm/src/arch/vm.rs
Original file line number Diff line number Diff line change
@@ -20,7 +20,7 @@ use crate::{
arch::segment::ExecutionSegment,
system::{
connector::{VmConnectorPvs, DEFAULT_SUSPEND_EXIT_CODE},
memory::{merkle::MemoryMerklePvs, MemoryImage, CHUNK},
memory::{merkle::MemoryMerklePvs, paged_vec::AddressMap, MemoryImage, CHUNK},
program::trace::VmCommittedExe,
},
};
@@ -113,11 +113,17 @@ where
let exe = exe.into();
let streams = input.into();
let mut segments = vec![];
let mem_config = self.config.system().memory_config;
let mut segment = ExecutionSegment::new(
&self.config,
exe.program.clone(),
streams,
Some(exe.init_memory.into_iter().collect()),
Some(AddressMap::from_iter(
mem_config.as_offset,
1 << mem_config.as_height,
1 << mem_config.pointer_max_bits,
exe.init_memory.clone(),
)),
exe.fn_bounds.clone(),
);
if let Some(overridden_heights) = self.overridden_heights.as_ref() {
27 changes: 16 additions & 11 deletions crates/vm/src/system/memory/controller/mod.rs
Original file line number Diff line number Diff line change
@@ -25,11 +25,14 @@ use openvm_stark_backend::{
rap::AnyRap,
Chip, ChipUsageGetter,
};
use rustc_hash::FxHashMap;
use serde::{Deserialize, Serialize};

use self::interface::MemoryInterface;
use super::{merkle::DirectCompressionBus, volatile::VolatileBoundaryChip};
use super::{
merkle::DirectCompressionBus,
paged_vec::{AddressMap, PAGE_SIZE},
volatile::VolatileBoundaryChip,
};
use crate::{
arch::{hasher::HasherChip, MemoryConfig},
system::memory::{
@@ -41,7 +44,7 @@ use crate::{
MemoryBridge, MemoryBus, MemoryReadAuxCols, MemoryReadOrImmediateAuxCols,
MemoryWriteAuxCols, AUX_LEN,
},
online::{Address, Memory, MemoryLogEntry},
online::{Memory, MemoryLogEntry},
persistent::PersistentBoundaryChip,
tree::MemoryNode,
},
@@ -59,7 +62,7 @@ pub const BOUNDARY_AIR_OFFSET: usize = 0;
#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
pub struct RecordId(pub usize);

pub type MemoryImage<F> = FxHashMap<Address, F>;
pub type MemoryImage<F> = AddressMap<F, PAGE_SIZE>;

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct TimestampedValues<T, const N: usize> {
@@ -229,7 +232,7 @@ impl<F: PrimeField32> MemoryController<F> {
range_checker: SharedVariableRangeCheckerChip,
) -> Self {
let range_checker_bus = range_checker.bus();
let initial_memory = MemoryImage::default();
let initial_memory = AddressMap::from_mem_config(&mem_config);
Self {
memory_bus,
mem_config,
@@ -241,13 +244,14 @@ impl<F: PrimeField32> MemoryController<F> {
range_checker.clone(),
),
},
memory: Memory::new(mem_config.access_capacity),
memory: Memory::new(&mem_config),
offline_memory: Arc::new(Mutex::new(OfflineMemory::new(
initial_memory,
1,
memory_bus,
range_checker.clone(),
mem_config.clk_max_bits,
mem_config,
))),
access_adapters: AccessAdapterInventory::new(
range_checker.clone(),
@@ -285,19 +289,20 @@ impl<F: PrimeField32> MemoryController<F> {
compression_bus,
),
merkle_chip: MemoryMerkleChip::new(memory_dims, merkle_bus, compression_bus),
initial_memory: MemoryImage::default(),
initial_memory: AddressMap::from_mem_config(&mem_config),
};
Self {
memory_bus,
mem_config,
interface_chip,
memory: Memory::new(0), // it is expected that the memory will be set later
memory: Memory::new(&mem_config), // it is expected that the memory will be set later
offline_memory: Arc::new(Mutex::new(OfflineMemory::new(
MemoryImage::default(),
AddressMap::from_mem_config(&mem_config),
CHUNK,
memory_bus,
range_checker.clone(),
mem_config.clk_max_bits,
mem_config,
))),
access_adapters: AccessAdapterInventory::new(
range_checker.clone(),
@@ -346,14 +351,14 @@ impl<F: PrimeField32> MemoryController<F> {
panic!("Cannot set initial memory after first timestamp");
}
let mut offline_memory = self.offline_memory.lock().unwrap();
offline_memory.set_initial_memory(memory.clone());
offline_memory.set_initial_memory(memory.clone(), self.mem_config);

self.memory = Memory::from_image(memory.clone(), self.mem_config.access_capacity);

match &mut self.interface_chip {
MemoryInterface::Volatile { .. } => {
assert!(
memory.is_empty(),
memory.items().collect::<Vec<_>>().is_empty(),
"Cannot set initial memory for volatile memory"
);
}
42 changes: 27 additions & 15 deletions crates/vm/src/system/memory/merkle/tests/mod.rs
Original file line number Diff line number Diff line change
@@ -24,6 +24,7 @@ use crate::{
columns::MemoryMerkleCols, tests::util::HashTestChip, MemoryDimensions,
MemoryMerkleBus, MemoryMerkleChip,
},
paged_vec::{AddressMap, PAGE_SIZE},
tree::MemoryNode,
Equipartition, MemoryImage,
},
@@ -48,21 +49,24 @@ fn test<const CHUNK: usize>(
let merkle_bus = MemoryMerkleBus(MEMORY_MERKLE_BUS);

// checking validity of test data
for (&(address_space, pointer), value) in final_memory {
for ((address_space, pointer), value) in final_memory.items() {
let label = pointer / CHUNK as u32;
assert!(address_space - as_offset < (1 << as_height));
assert!(label < (1 << address_height));
if initial_memory.get(&(address_space, pointer)) != Some(value) {
assert!(pointer < ((CHUNK << address_height).div_ceil(PAGE_SIZE) * PAGE_SIZE) as u32);
Copy link
Contributor

Choose a reason for hiding this comment

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

Isn't the bound just pointer < (CHUNK << address_height)? The boundary chip won't support pointer outside this range.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Unless pointer is from the last page that was padded to PAGE_SIZE elements and therefore it exceeds the supposed memory

Copy link
Contributor

Choose a reason for hiding this comment

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

But that would still be an invalid pointer?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Why?

Copy link
Contributor

Choose a reason for hiding this comment

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

If address_height is 27 and CHUNK is 8, then the max addressable cell is 2^30 - 1. If I understand correctly, this assertion might allow me to address into 2^30 or higher, if the PAGE_SIZE doesn't evenly divide 2^30. From the perspective of PagedVec, that's a fine index. But from the perspective of memory, that's not a valid pointer.

Or what am I missing?

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Yes, but how does it break anything here?

Copy link
Contributor

Choose a reason for hiding this comment

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

This is a test, so I'm not really that concerned. But the point of the assertion is to assert that pointer is a valid pointer. And valid pointers are in the range 0..CHUNK << address_height and, in particular, have nothing to do with the PAGE_SIZE of the underlying PagedVec.

if initial_memory.get(&(address_space, pointer)).cloned() != Some(value) {
assert!(touched_labels.contains(&(address_space, label)));
}
}
for key in initial_memory.keys() {
assert!(final_memory.contains_key(key));
for key in initial_memory.items().map(|(key, _)| key) {
assert!(final_memory.get(&key).is_some());
}
for &(address_space, label) in touched_labels.iter() {
let mut contains_some_key = false;
for i in 0..CHUNK {
if final_memory.contains_key(&(address_space, label * CHUNK as u32 + i as u32)) {
if final_memory
.get(&(address_space, label * CHUNK as u32 + i as u32))
.is_some()
{
contains_some_key = true;
break;
}
@@ -173,12 +177,12 @@ fn memory_to_partition<F: Default + Copy, const N: usize>(
memory: &MemoryImage<F>,
) -> Equipartition<F, N> {
let mut memory_partition = Equipartition::new();
for (&(address_space, pointer), value) in memory {
for ((address_space, pointer), value) in memory.items() {
let label = (address_space, pointer / N as u32);
let chunk = memory_partition
.entry(label)
.or_insert_with(|| [F::default(); N]);
chunk[(pointer % N as u32) as usize] = *value;
chunk[(pointer % N as u32) as usize] = value;
}
memory_partition
}
@@ -192,8 +196,8 @@ fn random_test<const CHUNK: usize>(
let mut rng = create_seeded_rng();
let mut next_u32 = || rng.next_u64() as u32;

let mut initial_memory = MemoryImage::default();
let mut final_memory = MemoryImage::default();
let mut initial_memory = AddressMap::new(1, 2, CHUNK << height);
let mut final_memory = AddressMap::new(1, 2, CHUNK << height);
let mut seen = HashSet::new();
let mut touched_labels = BTreeSet::new();

@@ -210,15 +214,15 @@ fn random_test<const CHUNK: usize>(
if is_initial && num_initial_addresses != 0 {
num_initial_addresses -= 1;
let value = BabyBear::from_canonical_u32(next_u32() % max_value);
initial_memory.insert((address_space, pointer), value);
final_memory.insert((address_space, pointer), value);
initial_memory.insert(&(address_space, pointer), value);
final_memory.insert(&(address_space, pointer), value);
}
if is_touched && num_touched_addresses != 0 {
num_touched_addresses -= 1;
touched_labels.insert((address_space, label));
if value_changes || !is_initial {
let value = BabyBear::from_canonical_u32(next_u32() % max_value);
final_memory.insert((address_space, pointer), value);
final_memory.insert(&(address_space, pointer), value);
}
}
}
@@ -260,7 +264,11 @@ fn expand_test_no_accesses() {
};
let mut hash_test_chip = HashTestChip::new();

let memory = MemoryImage::default();
let memory = AddressMap::new(
memory_dimensions.as_offset,
1 << memory_dimensions.as_height,
1 << memory_dimensions.address_height,
);
let tree = MemoryNode::<DEFAULT_CHUNK, _>::tree_from_memory(
memory_dimensions,
&memory,
@@ -293,7 +301,11 @@ fn expand_test_negative() {

let mut hash_test_chip = HashTestChip::new();

let memory = MemoryImage::default();
let memory = AddressMap::new(
memory_dimensions.as_offset,
1 << memory_dimensions.as_height,
1 << memory_dimensions.address_height,
);
let tree = MemoryNode::<DEFAULT_CHUNK, _>::tree_from_memory(
memory_dimensions,
&memory,
1 change: 1 addition & 0 deletions crates/vm/src/system/memory/mod.rs
Original file line number Diff line number Diff line change
@@ -6,6 +6,7 @@ pub mod merkle;
mod offline;
pub mod offline_checker;
pub mod online;
pub mod paged_vec;
mod persistent;
#[cfg(test)]
mod tests;
Loading