From cd79e9e7ab9511f3dde51c9282622ca6d56fe9eb Mon Sep 17 00:00:00 2001 From: Zach Langley Date: Wed, 8 Jan 2025 09:59:34 -0500 Subject: [PATCH 01/15] feat: PagedVec struct --- crates/vm/src/system/memory/controller/mod.rs | 3 +- crates/vm/src/system/memory/mod.rs | 1 + crates/vm/src/system/memory/paged_vec.rs | 260 ++++++++++++++++++ 3 files changed, 263 insertions(+), 1 deletion(-) create mode 100644 crates/vm/src/system/memory/paged_vec.rs diff --git a/crates/vm/src/system/memory/controller/mod.rs b/crates/vm/src/system/memory/controller/mod.rs index 88218bedf8..984170ba2a 100644 --- a/crates/vm/src/system/memory/controller/mod.rs +++ b/crates/vm/src/system/memory/controller/mod.rs @@ -25,7 +25,7 @@ use openvm_stark_backend::{ rap::AnyRap, Chip, ChipUsageGetter, }; -use rustc_hash::FxHashMap; +use rustc_hash::{FxHashMap, FxHashSet}; use serde::{Deserialize, Serialize}; use self::interface::MemoryInterface; @@ -312,6 +312,7 @@ impl MemoryController { } pub fn memory_image(&self) -> &MemoryImage { + // TODO: just return in current form rather than converting &self.memory.data } diff --git a/crates/vm/src/system/memory/mod.rs b/crates/vm/src/system/memory/mod.rs index f8c7b6451b..b1f282e905 100644 --- a/crates/vm/src/system/memory/mod.rs +++ b/crates/vm/src/system/memory/mod.rs @@ -6,6 +6,7 @@ pub mod merkle; mod offline; pub mod offline_checker; pub mod online; +mod paged_vec; mod persistent; #[cfg(test)] mod tests; diff --git a/crates/vm/src/system/memory/paged_vec.rs b/crates/vm/src/system/memory/paged_vec.rs new file mode 100644 index 0000000000..6353e371dd --- /dev/null +++ b/crates/vm/src/system/memory/paged_vec.rs @@ -0,0 +1,260 @@ +use std::array; + +#[derive(Debug, Clone)] +pub(crate) struct PagedVec { + page_size: u64, + pages: Vec>>, +} + +impl PagedVec { + pub fn new(page_size: usize, num_pages: usize) -> Self { + Self { + page_size: page_size as u64, + pages: vec![None; num_pages], + } + } + + pub fn get(&self, index: u64) -> T { + let page_idx = (index / self.page_size) as usize; + if let Some(page) = &self.pages[page_idx] { + page[(index % self.page_size) as usize] + } else { + T::default() + } + } + + pub fn set(&mut self, index: u64, value: T) { + let page_idx = (index / self.page_size) as usize; + if let Some(page) = self.pages[page_idx].as_mut() { + page[(index % self.page_size) as usize] = value; + } else { + let page = self.pages[page_idx] + .get_or_insert_with(|| vec![T::default(); self.page_size as usize]); + page[(index % self.page_size) as usize] = value; + } + } + + pub fn get_range(&self, index: u64) -> [T; N] { + let start_page_idx = (index / self.page_size) as usize; + let end_page_idx = ((index + N as u64) / self.page_size) as usize; + + if start_page_idx == end_page_idx { + if let Some(start_page) = &self.pages[start_page_idx] { + let i = (index % self.page_size) as usize; + start_page[i..i + N] + .try_into() + .expect("Slice length matches const generic N") + } else { + [T::default(); N] + } + } else { + // TODO: This can be more efficient by copying from two slices (but most queries should + // not be cross-page). + array::from_fn(|i| self.get(index + i as u64)) + } + } + + pub fn set_range(&mut self, index: u64, values: [T; N]) { + let start_page_idx = (index / self.page_size) as usize; + let end_page_idx = ((index + N as u64) / self.page_size) as usize; + + if start_page_idx == end_page_idx { + let page = self.pages[start_page_idx] + .get_or_insert_with(|| vec![T::default(); self.page_size as usize]); + let i = (index % self.page_size) as usize; + page[i..i + N].copy_from_slice(&values); + } else { + // TODO: This can be more efficient by copying into two slices (but most queries should + // not be cross-page). + for i in 0..N { + self.set(index + i as u64, values[i]); + } + } + } + + pub fn iter(&self) -> PagedVecIter<'_, T> { + PagedVecIter { + vec: self, + current_page: 0, + current_index_in_page: 0, + } + } +} + +pub struct PagedVecIter<'a, T> { + vec: &'a PagedVec, + current_page: usize, + current_index_in_page: usize, +} + +impl<'a, T: Default + Copy> Iterator for PagedVecIter<'a, T> { + type Item = (u64, T); + + fn next(&mut self) -> Option { + while self.current_page < self.vec.pages.len() + && self.vec.pages[self.current_page].is_none() + { + self.current_page += 1; + debug_assert_eq!(self.current_index_in_page, 0); + self.current_index_in_page = 0; + } + if self.current_page >= self.vec.pages.len() { + return None; + } + let global_index = + (self.current_page as u64 * self.vec.page_size) + self.current_index_in_page as u64; + + let page = self.vec.pages[self.current_page].as_ref()?; + let value = page[self.current_index_in_page]; + + self.current_index_in_page += 1; + if self.current_index_in_page == self.vec.page_size as usize { + self.current_page += 1; + self.current_index_in_page = 0; + } + Some((global_index, value)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_basic_get_set() { + let mut v = PagedVec::new(4, 3); + assert_eq!(v.get(0), 0); + v.set(0, 42); + assert_eq!(v.get(0), 42); + } + + #[test] + fn test_cross_page_operations() { + let mut v = PagedVec::new(4, 3); + v.set(3, 10); // Last element of first page + v.set(4, 20); // First element of second page + assert_eq!(v.get(3), 10); + assert_eq!(v.get(4), 20); + } + + #[test] + fn test_page_boundaries() { + let mut v = PagedVec::new(4, 2); + // Fill first page + v.set(0, 1); + v.set(1, 2); + v.set(2, 3); + v.set(3, 4); + // Fill second page + v.set(4, 5); + v.set(5, 6); + v.set(6, 7); + v.set(7, 8); + + // Verify all values + assert_eq!(v.get_range::<8>(0), [1, 2, 3, 4, 5, 6, 7, 8]); + } + + #[test] + fn test_range_cross_page_boundary() { + let mut v = PagedVec::new(4, 2); + v.set_range::<6>(2, [10, 11, 12, 13, 14, 15]); + assert_eq!(v.get_range::<6>(2), [10, 11, 12, 13, 14, 15]); + } + + #[test] + fn test_large_indices() { + let mut v = PagedVec::new(4, 100); + let large_index = 399; + v.set(large_index, 42); + assert_eq!(v.get(large_index), 42); + } + + #[test] + fn test_range_operations_with_defaults() { + let mut v = PagedVec::new(4, 3); + v.set(2, 5); + v.set(5, 10); + + // Should include both set values and defaults + assert_eq!(v.get_range::<6>(1), [0, 5, 0, 0, 10, 0]); + } + + #[test] + fn test_non_zero_default_type() { + let mut v: PagedVec = PagedVec::new(4, 2); + assert_eq!(v.get(0), false); // bool's default + v.set(0, true); + assert_eq!(v.get(0), true); + assert_eq!(v.get(1), false); + } + + #[test] + fn test_set_range_overlapping_pages() { + let mut v = PagedVec::new(4, 3); + let test_data = [1, 2, 3, 4, 5, 6]; + v.set_range::<6>(2, test_data); + + // Verify first page + assert_eq!(v.get(2), 1); + assert_eq!(v.get(3), 2); + + // Verify second page + assert_eq!(v.get(4), 3); + assert_eq!(v.get(5), 4); + assert_eq!(v.get(6), 5); + assert_eq!(v.get(7), 6); + } + + #[test] + fn test_overlapping_set_ranges() { + let mut v = PagedVec::new(4, 3); + + // Initial set_range + v.set_range::<5>(0, [1, 2, 3, 4, 5]); + assert_eq!(v.get_range::<5>(0), [1, 2, 3, 4, 5]); + + // Overlap from beginning + v.set_range::<3>(0, [10, 20, 30]); + assert_eq!(v.get_range::<5>(0), [10, 20, 30, 4, 5]); + + // Overlap in middle + v.set_range::<2>(2, [42, 43]); + assert_eq!(v.get_range::<5>(0), [10, 20, 42, 43, 5]); + + // Overlap at end + v.set_range::<2>(4, [91, 92]); + assert_eq!(v.get_range::<6>(0), [10, 20, 42, 43, 91, 92]); + } + + #[test] + fn test_overlapping_set_ranges_cross_pages() { + let mut v = PagedVec::new(4, 3); + + // Fill across first two pages + v.set_range::<8>(0, [1, 2, 3, 4, 5, 6, 7, 8]); + + // Overlap end of first page and start of second + v.set_range::<4>(2, [21, 22, 23, 24]); + assert_eq!(v.get_range::<8>(0), [1, 2, 21, 22, 23, 24, 7, 8]); + + // Overlap multiple pages + v.set_range::<6>(1, [31, 32, 33, 34, 35, 36]); + assert_eq!(v.get_range::<8>(0), [1, 31, 32, 33, 34, 35, 36, 8]); + } + + #[test] + fn test_iterator() { + let mut v = PagedVec::new(4, 3); + + v.set_range::<6>(4, [1, 2, 3, 4, 5, 6]); + let contents: Vec<_> = v.iter().collect(); + assert_eq!(contents.len(), 8); // two pages + + for i in 0..6 { + assert_eq!(contents[i], (4 + i as u64, 1 + i)); + } + assert_eq!(contents[6], (10, 0)); + assert_eq!(contents[7], (11, 0)); + } +} From 27d20963c58eb37b7c66a2d40341ad199346668a Mon Sep 17 00:00:00 2001 From: Alexander Golovanov Date: Fri, 10 Jan 2025 00:13:29 +0100 Subject: [PATCH 02/15] Make offline memory use paged vec --- crates/vm/src/arch/config.rs | 2 +- crates/vm/src/system/memory/controller/mod.rs | 11 +- crates/vm/src/system/memory/offline.rs | 160 ++++++++++++---- crates/vm/src/system/memory/paged_vec.rs | 175 ++++++++++-------- 4 files changed, 232 insertions(+), 116 deletions(-) diff --git a/crates/vm/src/arch/config.rs b/crates/vm/src/arch/config.rs index d789e040f7..2aa24a89c4 100644 --- a/crates/vm/src/arch/config.rs +++ b/crates/vm/src/arch/config.rs @@ -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) } } diff --git a/crates/vm/src/system/memory/controller/mod.rs b/crates/vm/src/system/memory/controller/mod.rs index 984170ba2a..1a08f3206e 100644 --- a/crates/vm/src/system/memory/controller/mod.rs +++ b/crates/vm/src/system/memory/controller/mod.rs @@ -25,7 +25,7 @@ use openvm_stark_backend::{ rap::AnyRap, Chip, ChipUsageGetter, }; -use rustc_hash::{FxHashMap, FxHashSet}; +use rustc_hash::FxHashMap; use serde::{Deserialize, Serialize}; use self::interface::MemoryInterface; @@ -248,6 +248,7 @@ impl MemoryController { memory_bus, range_checker.clone(), mem_config.clk_max_bits, + mem_config, ))), access_adapters: AccessAdapterInventory::new( range_checker.clone(), @@ -298,6 +299,7 @@ impl MemoryController { memory_bus, range_checker.clone(), mem_config.clk_max_bits, + mem_config, ))), access_adapters: AccessAdapterInventory::new( range_checker.clone(), @@ -347,7 +349,7 @@ impl MemoryController { 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); @@ -453,6 +455,7 @@ impl MemoryController { } fn replay_access_log(&mut self) { + let start = std::time::Instant::now(); let log = mem::take(&mut self.memory.log); let mut offline_memory = self.offline_memory.lock().unwrap(); @@ -466,6 +469,10 @@ impl MemoryController { &mut self.access_adapters, ); } + eprintln!( + "- - - - - - - - - - - - - - - replay_access_log time: {:?}", + start.elapsed() + ); } fn replay_access( diff --git a/crates/vm/src/system/memory/offline.rs b/crates/vm/src/system/memory/offline.rs index 195fe595e2..5819e50925 100644 --- a/crates/vm/src/system/memory/offline.rs +++ b/crates/vm/src/system/memory/offline.rs @@ -4,24 +4,70 @@ use openvm_circuit_primitives::{ assert_less_than::AssertLtSubAir, var_range::SharedVariableRangeCheckerChip, }; use openvm_stark_backend::p3_field::PrimeField32; -use rustc_hash::{FxHashMap, FxHashSet}; +use rustc_hash::FxHashSet; -use crate::system::memory::{ - adapter::{AccessAdapterRecord, AccessAdapterRecordKind}, - offline_checker::{MemoryBridge, MemoryBus}, - online::Address, - MemoryAuxColsFactory, MemoryImage, RecordId, TimestampedEquipartition, TimestampedValues, +use super::paged_vec::PagedVec; +use crate::{ + arch::MemoryConfig, + system::memory::{ + adapter::{AccessAdapterRecord, AccessAdapterRecordKind}, + offline_checker::{MemoryBridge, MemoryBus}, + online::Address, + MemoryAuxColsFactory, MemoryImage, RecordId, TimestampedEquipartition, TimestampedValues, + }, }; pub const INITIAL_TIMESTAMP: u32 = 0; -#[derive(Clone, Copy, PartialEq, Eq, Debug)] +const PAGE_SIZE: usize = 1 << 13; + +#[derive(Clone, Default, PartialEq, Eq, Debug)] struct BlockData { pointer: u32, size: usize, timestamp: u32, } +struct BlockStructure { + block_data: Vec>, + as_offset: u32, +} + +impl BlockStructure { + pub fn new(as_offset: u32, as_height: usize, mem_size: usize) -> Self { + Self { + block_data: vec![PagedVec::new(PAGE_SIZE, mem_size / PAGE_SIZE); 1 << as_height], + as_offset, + } + } + + pub fn items(&self) -> impl Iterator + '_ { + self.block_data + .iter() + .enumerate() + .flat_map(move |(as_idx, page)| { + page.iter().map(move |(ptr_idx, block)| { + ( + (as_idx as u32 + self.as_offset, ptr_idx as u32), + block.clone(), + ) + }) + }) + } + + pub fn get(&self, address: &Address) -> Option<&BlockData> { + self.block_data[(address.0 - self.as_offset) as usize].get(address.1 as usize) + } + + pub fn get_mut(&mut self, address: &Address) -> Option<&mut BlockData> { + self.block_data[(address.0 - self.as_offset) as usize].get_mut(address.1 as usize) + } + + pub fn insert(&mut self, address: Address, data: BlockData) { + self.block_data[(address.0 - self.as_offset) as usize].set(address.1 as usize, data); + } +} + #[derive(Debug, Clone, PartialEq)] pub struct MemoryRecord { pub address_space: T, @@ -34,8 +80,9 @@ pub struct MemoryRecord { } pub struct OfflineMemory { - block_data: FxHashMap, - data: FxHashMap, + block_data: BlockStructure, + data: Vec>, + as_offset: u32, initial_block_size: usize, timestamp: u32, timestamp_max_bits: usize, @@ -56,12 +103,18 @@ impl OfflineMemory { memory_bus: MemoryBus, range_checker: SharedVariableRangeCheckerChip, timestamp_max_bits: usize, + config: MemoryConfig, ) -> Self { assert!(initial_block_size.is_power_of_two()); Self { - block_data: FxHashMap::default(), - data: initial_memory, + block_data: BlockStructure::new( + config.as_offset, + config.as_height, + 1 << config.pointer_max_bits, + ), + data: Self::memory_image_to_paged_vec(initial_memory, config), + as_offset: config.as_offset, initial_block_size, timestamp: INITIAL_TIMESTAMP + 1, timestamp_max_bits, @@ -71,9 +124,24 @@ impl OfflineMemory { } } - pub fn set_initial_memory(&mut self, initial_memory: MemoryImage) { + pub fn set_initial_memory(&mut self, initial_memory: MemoryImage, config: MemoryConfig) { assert_eq!(self.timestamp, INITIAL_TIMESTAMP + 1); - self.data = initial_memory; + self.data = Self::memory_image_to_paged_vec(initial_memory, config); + } + + fn memory_image_to_paged_vec( + memory_image: MemoryImage, + config: MemoryConfig, + ) -> Vec> { + let mut paged_vec = + vec![ + PagedVec::new(PAGE_SIZE, (1 << config.pointer_max_bits) / PAGE_SIZE); + 1 << config.as_height + ]; + for ((addr_space, pointer), value) in memory_image { + paged_vec[(addr_space - config.as_offset) as usize].set(pointer as usize, value); + } + paged_vec } pub(super) fn set_log_capacity(&mut self, access_capacity: usize) { @@ -121,17 +189,13 @@ impl OfflineMemory { debug_assert!(prev_timestamp < self.timestamp); - let prev_data = (0..len) - .map(|i| { - self.data - .insert((address_space, pointer + i as u32), values[i]) - .unwrap_or(F::ZERO) - }) - .collect(); + let pointer = pointer as usize; + let prev_data = self.data[(address_space - self.as_offset) as usize] + .set_range(pointer..pointer + len, &values); let record = MemoryRecord { address_space: F::from_canonical_u32(address_space), - pointer: F::from_canonical_u32(pointer), + pointer: F::from_canonical_usize(pointer), timestamp: self.timestamp, prev_timestamp, data: values, @@ -150,7 +214,6 @@ impl OfflineMemory { len: usize, ) -> Vec> { assert!(len.is_power_of_two()); - if address_space == 0 { let pointer = F::from_canonical_u32(pointer); self.log.push(Some(MemoryRecord { @@ -198,8 +261,8 @@ impl OfflineMemory { // Grab all aligned pointers that need to be re-accessed. let to_access: FxHashSet<_> = self .block_data - .keys() - .map(|&(address_space, pointer)| (address_space, (pointer / N as u32) * N as u32)) + .items() + .map(|((address_space, pointer), _)| (address_space, (pointer / N as u32) * N as u32)) .collect(); for &(address_space, pointer) in to_access.iter() { @@ -270,7 +333,8 @@ impl OfflineMemory { timestamp, }; for i in 0..half_size_u32 { - self.block_data.insert((address_space, mid_ptr + i), block); + self.block_data + .insert((address_space, mid_ptr + i), block.clone()); } } if query >= cur_ptr + half_size_u32 { @@ -281,7 +345,8 @@ impl OfflineMemory { timestamp, }; for i in 0..half_size_u32 { - self.block_data.insert((address_space, cur_ptr + i), block); + self.block_data + .insert((address_space, cur_ptr + i), block.clone()); } } if mid_ptr <= query { @@ -306,10 +371,16 @@ impl OfflineMemory { let mut prev_timestamp = None; for i in 0..size as u32 { + if self.block_data.get(&(address_space, pointer + i)).is_none() { + self.block_data.insert( + (address_space, pointer + i), + Self::initial_block_data(pointer + i, self.initial_block_size), + ); + } let block = self .block_data - .entry((address_space, pointer + i)) - .or_insert_with(|| Self::initial_block_data(pointer + i, self.initial_block_size)); + .get_mut(&(address_space, pointer + i)) + .unwrap(); debug_assert!(i == 0 || prev_timestamp == Some(block.timestamp)); prev_timestamp = Some(block.timestamp); block.timestamp = self.timestamp; @@ -327,11 +398,14 @@ impl OfflineMemory { self.split_to_make_boundary(address_space, pointer, records); self.split_to_make_boundary(address_space, pointer + size as u32, records); - let block_data = self + let mut block_data = self .block_data .get(&(address_space, pointer)) - .copied() + .cloned() .unwrap_or_else(|| Self::initial_block_data(pointer, self.initial_block_size)); + if block_data.size == 0 { + block_data = Self::initial_block_data(pointer, self.initial_block_size); + } if block_data.pointer == pointer && block_data.size == size { return; @@ -399,10 +473,11 @@ impl OfflineMemory { fn block_containing(&mut self, address_space: u32, pointer: u32) -> BlockData { if let Some(block_data) = self.block_data.get(&(address_space, pointer)) { - *block_data - } else { - Self::initial_block_data(pointer, self.initial_block_size) + if block_data.size > 0 { + return block_data.clone(); + } } + Self::initial_block_data(pointer, self.initial_block_size) } fn initial_block_data(pointer: u32, initial_block_size: usize) -> BlockData { @@ -415,7 +490,10 @@ impl OfflineMemory { } pub fn get(&self, address_space: u32, pointer: u32) -> F { - *self.data.get(&(address_space, pointer)).unwrap_or(&F::ZERO) + self.data[(address_space - self.as_offset) as usize] + .get(pointer as usize) + .cloned() + .unwrap_or_default() } fn range_array(&self, address_space: u32, pointer: u32) -> [F; N] { @@ -423,9 +501,8 @@ impl OfflineMemory { } fn range_vec(&self, address_space: u32, pointer: u32, len: usize) -> Vec { - (0..len) - .map(|i| self.get(address_space, pointer + i as u32)) - .collect() + let pointer = pointer as usize; + self.data[(address_space - self.as_offset) as usize].get_range(pointer..pointer + len) } pub fn aux_cols_factory(&self) -> MemoryAuxColsFactory { @@ -488,6 +565,7 @@ mod tests { MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), 29, + Default::default(), ); assert_eq!( partition.block_containing(0, 13), @@ -535,6 +613,7 @@ mod tests { MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), 29, + Default::default(), ); let address_space = 1; @@ -561,6 +640,7 @@ mod tests { MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), 29, + Default::default(), ); let adapter_records = memory.write(1, 0, bbvec![1, 2, 3, 4]); @@ -706,6 +786,7 @@ mod tests { MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), 29, + Default::default(), ); let adapter_records = memory.write(1, 0, bbvec![1, 2, 3, 4]); @@ -821,6 +902,7 @@ mod tests { MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), 29, + Default::default(), ); memory.write(1, 0, bbvec![4, 3, 2, 1]); @@ -843,6 +925,7 @@ mod tests { MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), 29, + Default::default(), ); memory.write(1, 0, bbvec![4, 3, 2, 1]); @@ -865,6 +948,7 @@ mod tests { MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), 29, + Default::default(), ); let (memory, records) = memory.finalize::<4>(); @@ -881,6 +965,7 @@ mod tests { MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), 29, + Default::default(), ); // Make block 0:4 in address space 1 active. memory.write(1, 0, bbvec![1, 2, 3, 4]); @@ -951,6 +1036,7 @@ mod tests { MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), 29, + Default::default(), ); // Verify initial state of block 0 (pointers 0–8) diff --git a/crates/vm/src/system/memory/paged_vec.rs b/crates/vm/src/system/memory/paged_vec.rs index 6353e371dd..0ce87fa101 100644 --- a/crates/vm/src/system/memory/paged_vec.rs +++ b/crates/vm/src/system/memory/paged_vec.rs @@ -1,77 +1,97 @@ -use std::array; +use std::ops::Range; #[derive(Debug, Clone)] pub(crate) struct PagedVec { - page_size: u64, + page_size: usize, pages: Vec>>, } -impl PagedVec { +impl PagedVec { pub fn new(page_size: usize, num_pages: usize) -> Self { Self { - page_size: page_size as u64, + page_size, pages: vec![None; num_pages], } } - pub fn get(&self, index: u64) -> T { - let page_idx = (index / self.page_size) as usize; - if let Some(page) = &self.pages[page_idx] { - page[(index % self.page_size) as usize] - } else { - T::default() - } + pub fn get(&self, index: usize) -> Option<&T> { + let page_idx = index / self.page_size; + self.pages[page_idx] + .as_ref() + .map(|page| &page[index % self.page_size]) + } + + pub fn get_mut(&mut self, index: usize) -> Option<&mut T> { + let page_idx = index / self.page_size; + self.pages[page_idx] + .as_mut() + .map(|page| &mut page[index % self.page_size]) } - pub fn set(&mut self, index: u64, value: T) { - let page_idx = (index / self.page_size) as usize; + pub fn set(&mut self, index: usize, value: T) { + let page_idx = index / self.page_size; if let Some(page) = self.pages[page_idx].as_mut() { - page[(index % self.page_size) as usize] = value; + page[index % self.page_size] = value; } else { - let page = self.pages[page_idx] - .get_or_insert_with(|| vec![T::default(); self.page_size as usize]); - page[(index % self.page_size) as usize] = value; + let page = + self.pages[page_idx].get_or_insert_with(|| vec![T::default(); self.page_size]); + page[index % self.page_size] = value; } } - pub fn get_range(&self, index: u64) -> [T; N] { - let start_page_idx = (index / self.page_size) as usize; - let end_page_idx = ((index + N as u64) / self.page_size) as usize; + pub fn get_range(&self, range: Range) -> Vec { + let start_page_idx = range.start / self.page_size; + let end_page_idx = range.end / self.page_size; if start_page_idx == end_page_idx { if let Some(start_page) = &self.pages[start_page_idx] { - let i = (index % self.page_size) as usize; - start_page[i..i + N] - .try_into() - .expect("Slice length matches const generic N") + let i = range.start % self.page_size; + start_page[i..i + range.len()].to_vec() } else { - [T::default(); N] + vec![T::default(); range.len()] } } else { // TODO: This can be more efficient by copying from two slices (but most queries should // not be cross-page). - array::from_fn(|i| self.get(index + i as u64)) + range + .map(|i| self.get(i).cloned().unwrap_or_default()) + .collect() } } - pub fn set_range(&mut self, index: u64, values: [T; N]) { - let start_page_idx = (index / self.page_size) as usize; - let end_page_idx = ((index + N as u64) / self.page_size) as usize; + pub fn set_range<'a>( + &mut self, + range: Range, + values: impl IntoIterator, + ) -> Vec + where + T: 'a, + { + let start_page_idx = range.start / self.page_size; + let end_page_idx = range.end / self.page_size; if start_page_idx == end_page_idx { let page = self.pages[start_page_idx] - .get_or_insert_with(|| vec![T::default(); self.page_size as usize]); - let i = (index % self.page_size) as usize; - page[i..i + N].copy_from_slice(&values); + .get_or_insert_with(|| vec![T::default(); self.page_size]); + let page_start = range.start - range.start % self.page_size; + let result = page[range.start - page_start..range.end - page_start].to_vec(); + for (j, value) in range.zip(values.into_iter()) { + page[j - page_start] = value.clone(); + } + result } else { // TODO: This can be more efficient by copying into two slices (but most queries should // not be cross-page). - for i in 0..N { - self.set(index + i as u64, values[i]); + let result = self.get_range(range.clone()); + for (i, value) in range.zip(values.into_iter()) { + self.set(i, value.clone()); } + result } } +} +impl PagedVec { pub fn iter(&self) -> PagedVecIter<'_, T> { PagedVecIter { vec: self, @@ -87,8 +107,8 @@ pub struct PagedVecIter<'a, T> { current_index_in_page: usize, } -impl<'a, T: Default + Copy> Iterator for PagedVecIter<'a, T> { - type Item = (u64, T); +impl Iterator for PagedVecIter<'_, T> { + type Item = (usize, T); fn next(&mut self) -> Option { while self.current_page < self.vec.pages.len() @@ -101,14 +121,13 @@ impl<'a, T: Default + Copy> Iterator for PagedVecIter<'a, T> { if self.current_page >= self.vec.pages.len() { return None; } - let global_index = - (self.current_page as u64 * self.vec.page_size) + self.current_index_in_page as u64; + let global_index = self.current_page * self.vec.page_size + self.current_index_in_page; let page = self.vec.pages[self.current_page].as_ref()?; - let value = page[self.current_index_in_page]; + let value = page[self.current_index_in_page].clone(); self.current_index_in_page += 1; - if self.current_index_in_page == self.vec.page_size as usize { + if self.current_index_in_page == self.vec.page_size { self.current_page += 1; self.current_index_in_page = 0; } @@ -123,9 +142,9 @@ mod tests { #[test] fn test_basic_get_set() { let mut v = PagedVec::new(4, 3); - assert_eq!(v.get(0), 0); + assert_eq!(v.get(0), None); v.set(0, 42); - assert_eq!(v.get(0), 42); + assert_eq!(v.get(0), Some(&42)); } #[test] @@ -133,8 +152,8 @@ mod tests { let mut v = PagedVec::new(4, 3); v.set(3, 10); // Last element of first page v.set(4, 20); // First element of second page - assert_eq!(v.get(3), 10); - assert_eq!(v.get(4), 20); + assert_eq!(v.get(3), Some(&10)); + assert_eq!(v.get(4), Some(&20)); } #[test] @@ -152,14 +171,14 @@ mod tests { v.set(7, 8); // Verify all values - assert_eq!(v.get_range::<8>(0), [1, 2, 3, 4, 5, 6, 7, 8]); + assert_eq!(v.get_range(0..8), [1, 2, 3, 4, 5, 6, 7, 8]); } #[test] fn test_range_cross_page_boundary() { let mut v = PagedVec::new(4, 2); - v.set_range::<6>(2, [10, 11, 12, 13, 14, 15]); - assert_eq!(v.get_range::<6>(2), [10, 11, 12, 13, 14, 15]); + v.set_range(2..8, &[10, 11, 12, 13, 14, 15]); + assert_eq!(v.get_range(2..8), [10, 11, 12, 13, 14, 15]); } #[test] @@ -167,7 +186,7 @@ mod tests { let mut v = PagedVec::new(4, 100); let large_index = 399; v.set(large_index, 42); - assert_eq!(v.get(large_index), 42); + assert_eq!(v.get(large_index), Some(&42)); } #[test] @@ -177,33 +196,33 @@ mod tests { v.set(5, 10); // Should include both set values and defaults - assert_eq!(v.get_range::<6>(1), [0, 5, 0, 0, 10, 0]); + assert_eq!(v.get_range(1..7), [0, 5, 0, 0, 10, 0]); } #[test] fn test_non_zero_default_type() { let mut v: PagedVec = PagedVec::new(4, 2); - assert_eq!(v.get(0), false); // bool's default + assert_eq!(v.get(0), None); // bool's default v.set(0, true); - assert_eq!(v.get(0), true); - assert_eq!(v.get(1), false); + assert_eq!(v.get(0), Some(&true)); + assert_eq!(v.get(1), None); } #[test] fn test_set_range_overlapping_pages() { let mut v = PagedVec::new(4, 3); let test_data = [1, 2, 3, 4, 5, 6]; - v.set_range::<6>(2, test_data); + v.set_range(2..8, &test_data); // Verify first page - assert_eq!(v.get(2), 1); - assert_eq!(v.get(3), 2); + assert_eq!(v.get(2), Some(&1)); + assert_eq!(v.get(3), Some(&2)); // Verify second page - assert_eq!(v.get(4), 3); - assert_eq!(v.get(5), 4); - assert_eq!(v.get(6), 5); - assert_eq!(v.get(7), 6); + assert_eq!(v.get(4), Some(&3)); + assert_eq!(v.get(5), Some(&4)); + assert_eq!(v.get(6), Some(&5)); + assert_eq!(v.get(7), Some(&6)); } #[test] @@ -211,20 +230,20 @@ mod tests { let mut v = PagedVec::new(4, 3); // Initial set_range - v.set_range::<5>(0, [1, 2, 3, 4, 5]); - assert_eq!(v.get_range::<5>(0), [1, 2, 3, 4, 5]); + v.set_range(0..5, &[1, 2, 3, 4, 5]); + assert_eq!(v.get_range(0..5), [1, 2, 3, 4, 5]); // Overlap from beginning - v.set_range::<3>(0, [10, 20, 30]); - assert_eq!(v.get_range::<5>(0), [10, 20, 30, 4, 5]); + v.set_range(0..3, &[10, 20, 30]); + assert_eq!(v.get_range(0..5), [10, 20, 30, 4, 5]); // Overlap in middle - v.set_range::<2>(2, [42, 43]); - assert_eq!(v.get_range::<5>(0), [10, 20, 42, 43, 5]); + v.set_range(2..4, &[42, 43]); + assert_eq!(v.get_range(0..5), [10, 20, 42, 43, 5]); // Overlap at end - v.set_range::<2>(4, [91, 92]); - assert_eq!(v.get_range::<6>(0), [10, 20, 42, 43, 91, 92]); + v.set_range(4..6, &[91, 92]); + assert_eq!(v.get_range(0..6), [10, 20, 42, 43, 91, 92]); } #[test] @@ -232,28 +251,32 @@ mod tests { let mut v = PagedVec::new(4, 3); // Fill across first two pages - v.set_range::<8>(0, [1, 2, 3, 4, 5, 6, 7, 8]); + v.set_range(0..8, &[1, 2, 3, 4, 5, 6, 7, 8]); // Overlap end of first page and start of second - v.set_range::<4>(2, [21, 22, 23, 24]); - assert_eq!(v.get_range::<8>(0), [1, 2, 21, 22, 23, 24, 7, 8]); + v.set_range(2..6, &[21, 22, 23, 24]); + assert_eq!(v.get_range(0..8), [1, 2, 21, 22, 23, 24, 7, 8]); // Overlap multiple pages - v.set_range::<6>(1, [31, 32, 33, 34, 35, 36]); - assert_eq!(v.get_range::<8>(0), [1, 31, 32, 33, 34, 35, 36, 8]); + v.set_range(1..7, &[31, 32, 33, 34, 35, 36]); + assert_eq!(v.get_range(0..8), [1, 31, 32, 33, 34, 35, 36, 8]); } #[test] fn test_iterator() { let mut v = PagedVec::new(4, 3); - v.set_range::<6>(4, [1, 2, 3, 4, 5, 6]); + v.set_range(4..10, &[1, 2, 3, 4, 5, 6]); let contents: Vec<_> = v.iter().collect(); assert_eq!(contents.len(), 8); // two pages - for i in 0..6 { - assert_eq!(contents[i], (4 + i as u64, 1 + i)); - } + contents + .iter() + .take(6) + .enumerate() + .for_each(|(i, &(idx, val))| { + assert_eq!((idx, val), (4 + i, 1 + i)); + }); assert_eq!(contents[6], (10, 0)); assert_eq!(contents[7], (11, 0)); } From ba97321326226af2df321abb5cb8fb64e06a4d7a Mon Sep 17 00:00:00 2001 From: Alexander Golovanov Date: Thu, 16 Jan 2025 20:49:47 +0100 Subject: [PATCH 03/15] Make `PAGE_SIZE` const generic --- crates/vm/src/system/memory/offline.rs | 12 ++-- crates/vm/src/system/memory/paged_vec.rs | 77 ++++++++++++------------ 2 files changed, 43 insertions(+), 46 deletions(-) diff --git a/crates/vm/src/system/memory/offline.rs b/crates/vm/src/system/memory/offline.rs index 5819e50925..f7e24bfaeb 100644 --- a/crates/vm/src/system/memory/offline.rs +++ b/crates/vm/src/system/memory/offline.rs @@ -19,7 +19,7 @@ use crate::{ pub const INITIAL_TIMESTAMP: u32 = 0; -const PAGE_SIZE: usize = 1 << 13; +const PAGE_SIZE: usize = 1 << 12; #[derive(Clone, Default, PartialEq, Eq, Debug)] struct BlockData { @@ -29,14 +29,14 @@ struct BlockData { } struct BlockStructure { - block_data: Vec>, + block_data: Vec>, as_offset: u32, } impl BlockStructure { pub fn new(as_offset: u32, as_height: usize, mem_size: usize) -> Self { Self { - block_data: vec![PagedVec::new(PAGE_SIZE, mem_size / PAGE_SIZE); 1 << as_height], + block_data: vec![PagedVec::new(mem_size.div_ceil(PAGE_SIZE)); 1 << as_height], as_offset, } } @@ -81,7 +81,7 @@ pub struct MemoryRecord { pub struct OfflineMemory { block_data: BlockStructure, - data: Vec>, + data: Vec>, as_offset: u32, initial_block_size: usize, timestamp: u32, @@ -132,10 +132,10 @@ impl OfflineMemory { fn memory_image_to_paged_vec( memory_image: MemoryImage, config: MemoryConfig, - ) -> Vec> { + ) -> Vec> { let mut paged_vec = vec![ - PagedVec::new(PAGE_SIZE, (1 << config.pointer_max_bits) / PAGE_SIZE); + PagedVec::new((1usize << config.pointer_max_bits).div_ceil(PAGE_SIZE)); 1 << config.as_height ]; for ((addr_space, pointer), value) in memory_image { diff --git a/crates/vm/src/system/memory/paged_vec.rs b/crates/vm/src/system/memory/paged_vec.rs index 0ce87fa101..2f6c182b16 100644 --- a/crates/vm/src/system/memory/paged_vec.rs +++ b/crates/vm/src/system/memory/paged_vec.rs @@ -1,51 +1,48 @@ use std::ops::Range; #[derive(Debug, Clone)] -pub(crate) struct PagedVec { - page_size: usize, +pub(crate) struct PagedVec { pages: Vec>>, } -impl PagedVec { - pub fn new(page_size: usize, num_pages: usize) -> Self { +impl PagedVec { + pub fn new(num_pages: usize) -> Self { Self { - page_size, pages: vec![None; num_pages], } } pub fn get(&self, index: usize) -> Option<&T> { - let page_idx = index / self.page_size; + let page_idx = index / PAGE_SIZE; self.pages[page_idx] .as_ref() - .map(|page| &page[index % self.page_size]) + .map(|page| &page[index % PAGE_SIZE]) } pub fn get_mut(&mut self, index: usize) -> Option<&mut T> { - let page_idx = index / self.page_size; + let page_idx = index / PAGE_SIZE; self.pages[page_idx] .as_mut() - .map(|page| &mut page[index % self.page_size]) + .map(|page| &mut page[index % PAGE_SIZE]) } pub fn set(&mut self, index: usize, value: T) { - let page_idx = index / self.page_size; + let page_idx = index / PAGE_SIZE; if let Some(page) = self.pages[page_idx].as_mut() { - page[index % self.page_size] = value; + page[index % PAGE_SIZE] = value; } else { - let page = - self.pages[page_idx].get_or_insert_with(|| vec![T::default(); self.page_size]); - page[index % self.page_size] = value; + let page = self.pages[page_idx].get_or_insert_with(|| vec![T::default(); PAGE_SIZE]); + page[index % PAGE_SIZE] = value; } } pub fn get_range(&self, range: Range) -> Vec { - let start_page_idx = range.start / self.page_size; - let end_page_idx = range.end / self.page_size; + let start_page_idx = range.start / PAGE_SIZE; + let end_page_idx = range.end / PAGE_SIZE; if start_page_idx == end_page_idx { if let Some(start_page) = &self.pages[start_page_idx] { - let i = range.start % self.page_size; + let i = range.start % PAGE_SIZE; start_page[i..i + range.len()].to_vec() } else { vec![T::default(); range.len()] @@ -67,13 +64,13 @@ impl PagedVec { where T: 'a, { - let start_page_idx = range.start / self.page_size; - let end_page_idx = range.end / self.page_size; + let start_page_idx = range.start / PAGE_SIZE; + let end_page_idx = range.end / PAGE_SIZE; if start_page_idx == end_page_idx { - let page = self.pages[start_page_idx] - .get_or_insert_with(|| vec![T::default(); self.page_size]); - let page_start = range.start - range.start % self.page_size; + let page = + self.pages[start_page_idx].get_or_insert_with(|| vec![T::default(); PAGE_SIZE]); + let page_start = range.start - range.start % PAGE_SIZE; let result = page[range.start - page_start..range.end - page_start].to_vec(); for (j, value) in range.zip(values.into_iter()) { page[j - page_start] = value.clone(); @@ -91,8 +88,8 @@ impl PagedVec { } } -impl PagedVec { - pub fn iter(&self) -> PagedVecIter<'_, T> { +impl PagedVec { + pub fn iter(&self) -> PagedVecIter<'_, T, PAGE_SIZE> { PagedVecIter { vec: self, current_page: 0, @@ -101,13 +98,13 @@ impl PagedVec { } } -pub struct PagedVecIter<'a, T> { - vec: &'a PagedVec, +pub struct PagedVecIter<'a, T, const PAGE_SIZE: usize> { + vec: &'a PagedVec, current_page: usize, current_index_in_page: usize, } -impl Iterator for PagedVecIter<'_, T> { +impl Iterator for PagedVecIter<'_, T, PAGE_SIZE> { type Item = (usize, T); fn next(&mut self) -> Option { @@ -121,13 +118,13 @@ impl Iterator for PagedVecIter<'_, T> { if self.current_page >= self.vec.pages.len() { return None; } - let global_index = self.current_page * self.vec.page_size + self.current_index_in_page; + let global_index = self.current_page * PAGE_SIZE + self.current_index_in_page; let page = self.vec.pages[self.current_page].as_ref()?; let value = page[self.current_index_in_page].clone(); self.current_index_in_page += 1; - if self.current_index_in_page == self.vec.page_size { + if self.current_index_in_page == PAGE_SIZE { self.current_page += 1; self.current_index_in_page = 0; } @@ -141,7 +138,7 @@ mod tests { #[test] fn test_basic_get_set() { - let mut v = PagedVec::new(4, 3); + let mut v = PagedVec::<_, 4>::new(3); assert_eq!(v.get(0), None); v.set(0, 42); assert_eq!(v.get(0), Some(&42)); @@ -149,7 +146,7 @@ mod tests { #[test] fn test_cross_page_operations() { - let mut v = PagedVec::new(4, 3); + let mut v = PagedVec::<_, 4>::new(3); v.set(3, 10); // Last element of first page v.set(4, 20); // First element of second page assert_eq!(v.get(3), Some(&10)); @@ -158,7 +155,7 @@ mod tests { #[test] fn test_page_boundaries() { - let mut v = PagedVec::new(4, 2); + let mut v = PagedVec::<_, 4>::new(2); // Fill first page v.set(0, 1); v.set(1, 2); @@ -176,14 +173,14 @@ mod tests { #[test] fn test_range_cross_page_boundary() { - let mut v = PagedVec::new(4, 2); + let mut v = PagedVec::<_, 4>::new(2); v.set_range(2..8, &[10, 11, 12, 13, 14, 15]); assert_eq!(v.get_range(2..8), [10, 11, 12, 13, 14, 15]); } #[test] fn test_large_indices() { - let mut v = PagedVec::new(4, 100); + let mut v = PagedVec::<_, 4>::new(100); let large_index = 399; v.set(large_index, 42); assert_eq!(v.get(large_index), Some(&42)); @@ -191,7 +188,7 @@ mod tests { #[test] fn test_range_operations_with_defaults() { - let mut v = PagedVec::new(4, 3); + let mut v = PagedVec::<_, 4>::new(3); v.set(2, 5); v.set(5, 10); @@ -201,7 +198,7 @@ mod tests { #[test] fn test_non_zero_default_type() { - let mut v: PagedVec = PagedVec::new(4, 2); + let mut v: PagedVec = PagedVec::new(2); assert_eq!(v.get(0), None); // bool's default v.set(0, true); assert_eq!(v.get(0), Some(&true)); @@ -210,7 +207,7 @@ mod tests { #[test] fn test_set_range_overlapping_pages() { - let mut v = PagedVec::new(4, 3); + let mut v = PagedVec::<_, 4>::new(3); let test_data = [1, 2, 3, 4, 5, 6]; v.set_range(2..8, &test_data); @@ -227,7 +224,7 @@ mod tests { #[test] fn test_overlapping_set_ranges() { - let mut v = PagedVec::new(4, 3); + let mut v = PagedVec::<_, 4>::new(3); // Initial set_range v.set_range(0..5, &[1, 2, 3, 4, 5]); @@ -248,7 +245,7 @@ mod tests { #[test] fn test_overlapping_set_ranges_cross_pages() { - let mut v = PagedVec::new(4, 3); + let mut v = PagedVec::<_, 4>::new(3); // Fill across first two pages v.set_range(0..8, &[1, 2, 3, 4, 5, 6, 7, 8]); @@ -264,7 +261,7 @@ mod tests { #[test] fn test_iterator() { - let mut v = PagedVec::new(4, 3); + let mut v = PagedVec::<_, 4>::new(3); v.set_range(4..10, &[1, 2, 3, 4, 5, 6]); let contents: Vec<_> = v.iter().collect(); From 928ccf4cc6c24771d5aca1949a6f6f13a1dabc6b Mon Sep 17 00:00:00 2001 From: Alexander Golovanov Date: Thu, 16 Jan 2025 21:46:14 +0100 Subject: [PATCH 04/15] Fix tests --- crates/vm/src/system/memory/offline.rs | 34 ++++++++++++++++++------ crates/vm/src/system/memory/paged_vec.rs | 2 +- 2 files changed, 27 insertions(+), 9 deletions(-) diff --git a/crates/vm/src/system/memory/offline.rs b/crates/vm/src/system/memory/offline.rs index f7e24bfaeb..babd8f34d6 100644 --- a/crates/vm/src/system/memory/offline.rs +++ b/crates/vm/src/system/memory/offline.rs @@ -275,6 +275,9 @@ impl OfflineMemory { let mut equipartition = TimestampedEquipartition::::new(); for (address_space, pointer) in to_access { let block = self.block_data.get(&(address_space, pointer)).unwrap(); + if block.size == 0 { + continue; + } debug_assert_eq!(block.pointer % N as u32, 0); debug_assert_eq!(block.size, N); @@ -438,9 +441,12 @@ impl OfflineMemory { let left_block = self.block_data.get(&(address_space, pointer)); let left_timestamp = left_block.map(|b| b.timestamp).unwrap_or(INITIAL_TIMESTAMP); - let size = left_block + let mut size = left_block .map(|b| b.size) .unwrap_or(self.initial_block_size); + if size == 0 { + size = self.initial_block_size; + } let right_timestamp = self .block_data @@ -530,10 +536,13 @@ mod tests { use openvm_stark_sdk::p3_baby_bear::BabyBear; use super::{BlockData, MemoryRecord, OfflineMemory}; - use crate::system::memory::{ - adapter::{AccessAdapterRecord, AccessAdapterRecordKind}, - offline_checker::MemoryBus, - MemoryImage, TimestampedValues, + use crate::{ + arch::MemoryConfig, + system::memory::{ + adapter::{AccessAdapterRecord, AccessAdapterRecordKind}, + offline_checker::MemoryBus, + MemoryImage, TimestampedValues, + }, }; macro_rules! bb { @@ -565,7 +574,10 @@ mod tests { MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), 29, - Default::default(), + MemoryConfig { + as_offset: 0, + ..Default::default() + }, ); assert_eq!( partition.block_containing(0, 13), @@ -902,7 +914,10 @@ mod tests { MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), 29, - Default::default(), + MemoryConfig { + as_offset: 0, + ..Default::default() + }, ); memory.write(1, 0, bbvec![4, 3, 2, 1]); @@ -925,7 +940,10 @@ mod tests { MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), 29, - Default::default(), + MemoryConfig { + as_offset: 0, + ..Default::default() + }, ); memory.write(1, 0, bbvec![4, 3, 2, 1]); diff --git a/crates/vm/src/system/memory/paged_vec.rs b/crates/vm/src/system/memory/paged_vec.rs index 2f6c182b16..d6a47d2490 100644 --- a/crates/vm/src/system/memory/paged_vec.rs +++ b/crates/vm/src/system/memory/paged_vec.rs @@ -202,7 +202,7 @@ mod tests { assert_eq!(v.get(0), None); // bool's default v.set(0, true); assert_eq!(v.get(0), Some(&true)); - assert_eq!(v.get(1), None); + assert_eq!(v.get(1), Some(&false)); // because we created the page } #[test] From 5070149b333b1a75654d9320d34060c9f6ad24f5 Mon Sep 17 00:00:00 2001 From: Alexander Golovanov Date: Thu, 16 Jan 2025 23:07:13 +0100 Subject: [PATCH 05/15] Fix something --- crates/vm/src/system/memory/offline.rs | 5 +++++ crates/vm/src/system/memory/paged_vec.rs | 4 ++++ 2 files changed, 9 insertions(+) diff --git a/crates/vm/src/system/memory/offline.rs b/crates/vm/src/system/memory/offline.rs index babd8f34d6..3245b7f339 100644 --- a/crates/vm/src/system/memory/offline.rs +++ b/crates/vm/src/system/memory/offline.rs @@ -301,6 +301,11 @@ impl OfflineMemory { query: u32, records: &mut Vec>, ) { + let lim = (self.data[(address_space - self.as_offset) as usize].memory_size()) as u32; + if query == lim { + return; + } + assert!(query < lim); let original_block = self.block_containing(address_space, query); if original_block.pointer == query { return; diff --git a/crates/vm/src/system/memory/paged_vec.rs b/crates/vm/src/system/memory/paged_vec.rs index d6a47d2490..c0e8f39b81 100644 --- a/crates/vm/src/system/memory/paged_vec.rs +++ b/crates/vm/src/system/memory/paged_vec.rs @@ -86,6 +86,10 @@ impl PagedVec { result } } + + pub fn memory_size(&self) -> usize { + self.pages.len() * PAGE_SIZE + } } impl PagedVec { From 8bd38daf043c87d91207b621f8efb945f19faded Mon Sep 17 00:00:00 2001 From: Alexander Golovanov Date: Fri, 17 Jan 2025 00:19:28 +0100 Subject: [PATCH 06/15] Fix --- crates/vm/src/system/memory/controller/mod.rs | 6 ------ crates/vm/src/system/memory/offline.rs | 6 ++++-- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/crates/vm/src/system/memory/controller/mod.rs b/crates/vm/src/system/memory/controller/mod.rs index 1a08f3206e..56b5d951c3 100644 --- a/crates/vm/src/system/memory/controller/mod.rs +++ b/crates/vm/src/system/memory/controller/mod.rs @@ -314,7 +314,6 @@ impl MemoryController { } pub fn memory_image(&self) -> &MemoryImage { - // TODO: just return in current form rather than converting &self.memory.data } @@ -455,7 +454,6 @@ impl MemoryController { } fn replay_access_log(&mut self) { - let start = std::time::Instant::now(); let log = mem::take(&mut self.memory.log); let mut offline_memory = self.offline_memory.lock().unwrap(); @@ -469,10 +467,6 @@ impl MemoryController { &mut self.access_adapters, ); } - eprintln!( - "- - - - - - - - - - - - - - - replay_access_log time: {:?}", - start.elapsed() - ); } fn replay_access( diff --git a/crates/vm/src/system/memory/offline.rs b/crates/vm/src/system/memory/offline.rs index 3245b7f339..c5ea3a369e 100644 --- a/crates/vm/src/system/memory/offline.rs +++ b/crates/vm/src/system/memory/offline.rs @@ -126,6 +126,7 @@ impl OfflineMemory { pub fn set_initial_memory(&mut self, initial_memory: MemoryImage, config: MemoryConfig) { assert_eq!(self.timestamp, INITIAL_TIMESTAMP + 1); + self.as_offset = config.as_offset; self.data = Self::memory_image_to_paged_vec(initial_memory, config); } @@ -267,7 +268,7 @@ impl OfflineMemory { for &(address_space, pointer) in to_access.iter() { let block = self.block_data.get(&(address_space, pointer)).unwrap(); - if block.pointer != pointer || block.size != N { + if block.size > 0 && (block.pointer != pointer || block.size != N) { self.access(address_space, pointer, N, &mut adapter_records); } } @@ -379,7 +380,8 @@ impl OfflineMemory { let mut prev_timestamp = None; for i in 0..size as u32 { - if self.block_data.get(&(address_space, pointer + i)).is_none() { + let block = self.block_data.get(&(address_space, pointer + i)); + if block.is_none() || block.unwrap().size == 0 { self.block_data.insert( (address_space, pointer + i), Self::initial_block_data(pointer + i, self.initial_block_size), From 29245d43a7f85e2593d632ee0b945e12336cb533 Mon Sep 17 00:00:00 2001 From: Alexander Golovanov Date: Fri, 17 Jan 2025 16:24:31 +0100 Subject: [PATCH 07/15] Update online memory so that it now uses pagevecs --- crates/sdk/src/commit.rs | 13 ++- crates/vm/src/arch/vm.rs | 10 ++- crates/vm/src/system/memory/controller/mod.rs | 23 ++--- .../vm/src/system/memory/merkle/tests/mod.rs | 42 +++++---- crates/vm/src/system/memory/mod.rs | 2 +- crates/vm/src/system/memory/offline.rs | 70 +++------------ crates/vm/src/system/memory/online.rs | 31 +++---- crates/vm/src/system/memory/paged_vec.rs | 90 ++++++++++++++++++- crates/vm/src/system/memory/tree/mod.rs | 4 +- .../src/system/memory/tree/public_values.rs | 19 ++-- 10 files changed, 188 insertions(+), 116 deletions(-) diff --git a/crates/sdk/src/commit.rs b/crates/sdk/src/commit.rs index b1c56d28e1..db69b9a81f 100644 --- a/crates/sdk/src/commit.rs +++ b/crates/sdk/src/commit.rs @@ -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 { .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(); diff --git a/crates/vm/src/arch/vm.rs b/crates/vm/src/arch/vm.rs index da4c8fc093..9369d7b4bf 100644 --- a/crates/vm/src/arch/vm.rs +++ b/crates/vm/src/arch/vm.rs @@ -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() { diff --git a/crates/vm/src/system/memory/controller/mod.rs b/crates/vm/src/system/memory/controller/mod.rs index 56b5d951c3..7442535b74 100644 --- a/crates/vm/src/system/memory/controller/mod.rs +++ b/crates/vm/src/system/memory/controller/mod.rs @@ -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 = FxHashMap; +pub type MemoryImage = AddressMap; #[derive(Clone, Copy, Debug, PartialEq, Eq)] pub struct TimestampedValues { @@ -229,7 +232,7 @@ impl MemoryController { 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,7 +244,7 @@ impl MemoryController { 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, @@ -286,15 +289,15 @@ impl MemoryController { 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(), @@ -355,7 +358,7 @@ impl MemoryController { match &mut self.interface_chip { MemoryInterface::Volatile { .. } => { assert!( - memory.is_empty(), + memory.items().collect::>().is_empty(), "Cannot set initial memory for volatile memory" ); } diff --git a/crates/vm/src/system/memory/merkle/tests/mod.rs b/crates/vm/src/system/memory/merkle/tests/mod.rs index 2eb0f0608b..b6f99d8b98 100644 --- a/crates/vm/src/system/memory/merkle/tests/mod.rs +++ b/crates/vm/src/system/memory/merkle/tests/mod.rs @@ -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( 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); + 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( memory: &MemoryImage, ) -> Equipartition { 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( 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( 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::::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::::tree_from_memory( memory_dimensions, &memory, diff --git a/crates/vm/src/system/memory/mod.rs b/crates/vm/src/system/memory/mod.rs index b1f282e905..792847f1bb 100644 --- a/crates/vm/src/system/memory/mod.rs +++ b/crates/vm/src/system/memory/mod.rs @@ -6,7 +6,7 @@ pub mod merkle; mod offline; pub mod offline_checker; pub mod online; -mod paged_vec; +pub mod paged_vec; mod persistent; #[cfg(test)] mod tests; diff --git a/crates/vm/src/system/memory/offline.rs b/crates/vm/src/system/memory/offline.rs index c5ea3a369e..3b1639a439 100644 --- a/crates/vm/src/system/memory/offline.rs +++ b/crates/vm/src/system/memory/offline.rs @@ -6,21 +6,18 @@ use openvm_circuit_primitives::{ use openvm_stark_backend::p3_field::PrimeField32; use rustc_hash::FxHashSet; -use super::paged_vec::PagedVec; +use super::paged_vec::{AddressMap, PagedVec, PAGE_SIZE}; use crate::{ arch::MemoryConfig, system::memory::{ adapter::{AccessAdapterRecord, AccessAdapterRecordKind}, offline_checker::{MemoryBridge, MemoryBus}, - online::Address, MemoryAuxColsFactory, MemoryImage, RecordId, TimestampedEquipartition, TimestampedValues, }, }; pub const INITIAL_TIMESTAMP: u32 = 0; -const PAGE_SIZE: usize = 1 << 12; - #[derive(Clone, Default, PartialEq, Eq, Debug)] struct BlockData { pointer: u32, @@ -28,46 +25,6 @@ struct BlockData { timestamp: u32, } -struct BlockStructure { - block_data: Vec>, - as_offset: u32, -} - -impl BlockStructure { - pub fn new(as_offset: u32, as_height: usize, mem_size: usize) -> Self { - Self { - block_data: vec![PagedVec::new(mem_size.div_ceil(PAGE_SIZE)); 1 << as_height], - as_offset, - } - } - - pub fn items(&self) -> impl Iterator + '_ { - self.block_data - .iter() - .enumerate() - .flat_map(move |(as_idx, page)| { - page.iter().map(move |(ptr_idx, block)| { - ( - (as_idx as u32 + self.as_offset, ptr_idx as u32), - block.clone(), - ) - }) - }) - } - - pub fn get(&self, address: &Address) -> Option<&BlockData> { - self.block_data[(address.0 - self.as_offset) as usize].get(address.1 as usize) - } - - pub fn get_mut(&mut self, address: &Address) -> Option<&mut BlockData> { - self.block_data[(address.0 - self.as_offset) as usize].get_mut(address.1 as usize) - } - - pub fn insert(&mut self, address: Address, data: BlockData) { - self.block_data[(address.0 - self.as_offset) as usize].set(address.1 as usize, data); - } -} - #[derive(Debug, Clone, PartialEq)] pub struct MemoryRecord { pub address_space: T, @@ -80,7 +37,7 @@ pub struct MemoryRecord { } pub struct OfflineMemory { - block_data: BlockStructure, + block_data: AddressMap, data: Vec>, as_offset: u32, initial_block_size: usize, @@ -108,11 +65,7 @@ impl OfflineMemory { assert!(initial_block_size.is_power_of_two()); Self { - block_data: BlockStructure::new( - config.as_offset, - config.as_height, - 1 << config.pointer_max_bits, - ), + block_data: AddressMap::from_mem_config(&config), data: Self::memory_image_to_paged_vec(initial_memory, config), as_offset: config.as_offset, initial_block_size, @@ -139,7 +92,7 @@ impl OfflineMemory { PagedVec::new((1usize << config.pointer_max_bits).div_ceil(PAGE_SIZE)); 1 << config.as_height ]; - for ((addr_space, pointer), value) in memory_image { + for ((addr_space, pointer), value) in memory_image.items() { paged_vec[(addr_space - config.as_offset) as usize].set(pointer as usize, value); } paged_vec @@ -343,7 +296,7 @@ impl OfflineMemory { }; for i in 0..half_size_u32 { self.block_data - .insert((address_space, mid_ptr + i), block.clone()); + .insert(&(address_space, mid_ptr + i), block.clone()); } } if query >= cur_ptr + half_size_u32 { @@ -355,7 +308,7 @@ impl OfflineMemory { }; for i in 0..half_size_u32 { self.block_data - .insert((address_space, cur_ptr + i), block.clone()); + .insert(&(address_space, cur_ptr + i), block.clone()); } } if mid_ptr <= query { @@ -383,7 +336,7 @@ impl OfflineMemory { let block = self.block_data.get(&(address_space, pointer + i)); if block.is_none() || block.unwrap().size == 0 { self.block_data.insert( - (address_space, pointer + i), + &(address_space, pointer + i), Self::initial_block_data(pointer + i, self.initial_block_size), ); } @@ -464,7 +417,7 @@ impl OfflineMemory { let timestamp = max(left_timestamp, right_timestamp); for i in 0..2 * size as u32 { self.block_data.insert( - (address_space, pointer + i), + &(address_space, pointer + i), BlockData { pointer, size: 2 * size, @@ -548,6 +501,7 @@ mod tests { system::memory::{ adapter::{AccessAdapterRecord, AccessAdapterRecordKind}, offline_checker::MemoryBus, + paged_vec::AddressMap, MemoryImage, TimestampedValues, }, }; @@ -574,7 +528,7 @@ mod tests { fn test_partition() { type F = BabyBear; - let initial_memory = MemoryImage::default(); + let initial_memory = AddressMap::new(0, 1, 16); let mut partition = OfflineMemory::::new( initial_memory, 8, @@ -1051,8 +1005,8 @@ mod tests { // Initialize initial memory with blocks at indices 0 and 2 let mut initial_memory = MemoryImage::default(); for i in 0..8 { - initial_memory.insert((1, i), F::from_canonical_u32(i + 1)); - initial_memory.insert((1, 16 + i), F::from_canonical_u32(i + 1)); + initial_memory.insert(&(1, i), F::from_canonical_u32(i + 1)); + initial_memory.insert(&(1, 16 + i), F::from_canonical_u32(i + 1)); } let mut memory = OfflineMemory::::new( diff --git a/crates/vm/src/system/memory/online.rs b/crates/vm/src/system/memory/online.rs index 6c30176a51..e4e91113d8 100644 --- a/crates/vm/src/system/memory/online.rs +++ b/crates/vm/src/system/memory/online.rs @@ -1,10 +1,13 @@ -use std::{array, fmt::Debug}; +use std::fmt::Debug; use openvm_stark_backend::p3_field::PrimeField32; -use rustc_hash::FxHashMap; use serde::{Deserialize, Serialize}; -use crate::system::memory::{offline::INITIAL_TIMESTAMP, MemoryImage, RecordId}; +use super::paged_vec::{AddressMap, PAGE_SIZE}; +use crate::{ + arch::MemoryConfig, + system::memory::{offline::INITIAL_TIMESTAMP, MemoryImage, RecordId}, +}; #[derive(Debug, Clone, Serialize, Deserialize)] pub enum MemoryLogEntry { @@ -21,25 +24,22 @@ pub enum MemoryLogEntry { IncrementTimestampBy(u32), } -/// (address_space, pointer) -pub(crate) type Address = (u32, u32); - /// A simple data structure to read to/write from memory. /// /// Stores a log of memory accesses to reconstruct aspects of memory state for trace generation. #[derive(Debug)] pub struct Memory { - pub(super) data: FxHashMap, + pub(super) data: AddressMap, pub(super) log: Vec>, timestamp: u32, } impl Memory { - pub fn new(access_capacity: usize) -> Self { + pub fn new(mem_config: &MemoryConfig) -> Self { Self { - data: MemoryImage::default(), + data: AddressMap::from_mem_config(mem_config), timestamp: INITIAL_TIMESTAMP + 1, - log: Vec::with_capacity(access_capacity), + log: Vec::with_capacity(mem_config.access_capacity), } } @@ -67,11 +67,7 @@ impl Memory { ) -> (RecordId, [F; N]) { assert!(N.is_power_of_two()); - let prev_data = array::from_fn(|i| { - self.data - .insert((address_space, pointer + i as u32), values[i]) - .unwrap_or(F::ZERO) - }); + let prev_data = self.data.set_range(&(address_space, pointer), &values); self.log.push(MemoryLogEntry::Write { address_space, @@ -117,7 +113,7 @@ impl Memory { } fn range_array(&self, address_space: u32, pointer: u32) -> [F; N] { - array::from_fn(|i| self.get(address_space, pointer + i as u32)) + self.data.get_range(&(address_space, pointer)) } } @@ -127,6 +123,7 @@ mod tests { use openvm_stark_sdk::p3_baby_bear::BabyBear; use super::Memory; + use crate::arch::MemoryConfig; macro_rules! bba { [$($x:expr),*] => { @@ -136,7 +133,7 @@ mod tests { #[test] fn test_write_read() { - let mut memory = Memory::new(0); + let mut memory = Memory::new(&MemoryConfig::default()); let address_space = 1; memory.write(address_space, 0, bba![1, 2, 3, 4]); diff --git a/crates/vm/src/system/memory/paged_vec.rs b/crates/vm/src/system/memory/paged_vec.rs index c0e8f39b81..7020a2ef8d 100644 --- a/crates/vm/src/system/memory/paged_vec.rs +++ b/crates/vm/src/system/memory/paged_vec.rs @@ -1,6 +1,14 @@ use std::ops::Range; -#[derive(Debug, Clone)] +use serde::{Deserialize, Serialize}; + +use crate::arch::MemoryConfig; + +/// (address_space, pointer) +pub(crate) type Address = (u32, u32); +pub(crate) const PAGE_SIZE: usize = 1 << 12; + +#[derive(Debug, Clone, Serialize, Deserialize)] pub(crate) struct PagedVec { pages: Vec>>, } @@ -26,13 +34,14 @@ impl PagedVec { .map(|page| &mut page[index % PAGE_SIZE]) } - pub fn set(&mut self, index: usize, value: T) { + pub fn set(&mut self, index: usize, value: T) -> Option { let page_idx = index / PAGE_SIZE; if let Some(page) = self.pages[page_idx].as_mut() { - page[index % PAGE_SIZE] = value; + Some(std::mem::replace(&mut page[index % PAGE_SIZE], value)) } else { let page = self.pages[page_idx].get_or_insert_with(|| vec![T::default(); PAGE_SIZE]); page[index % PAGE_SIZE] = value; + None } } @@ -136,6 +145,81 @@ impl Iterator for PagedVecIter<'_, T, PAGE_SIZ } } +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct AddressMap { + paged_vecs: Vec>, + as_offset: u32, +} + +impl Default for AddressMap { + fn default() -> Self { + Self::from_mem_config(&MemoryConfig::default()) + } +} + +impl AddressMap { + pub fn new(as_offset: u32, as_cnt: usize, mem_size: usize) -> Self { + Self { + paged_vecs: vec![PagedVec::new(mem_size.div_ceil(PAGE_SIZE)); as_cnt], + as_offset, + } + } + pub fn from_mem_config(mem_config: &MemoryConfig) -> Self { + Self::new( + mem_config.as_offset, + 1 << mem_config.as_height, + 1 << mem_config.pointer_max_bits, + ) + } + pub fn items(&self) -> impl Iterator + '_ { + self.paged_vecs + .iter() + .enumerate() + .flat_map(move |(as_idx, page)| { + page.iter() + .map(move |(ptr_idx, x)| ((as_idx as u32 + self.as_offset, ptr_idx as u32), x)) + }) + } + pub fn get(&self, address: &Address) -> Option<&T> { + self.paged_vecs[(address.0 - self.as_offset) as usize].get(address.1 as usize) + } + pub fn get_mut(&mut self, address: &Address) -> Option<&mut T> { + self.paged_vecs[(address.0 - self.as_offset) as usize].get_mut(address.1 as usize) + } + pub fn insert(&mut self, address: &Address, data: T) -> Option { + self.paged_vecs[(address.0 - self.as_offset) as usize].set(address.1 as usize, data) + } + pub fn get_range(&self, address: &Address) -> [T; N] { + unsafe { + self.paged_vecs[(address.0 - self.as_offset) as usize] + .get_range((address.1 as usize)..(address.1 as usize + N)) + .try_into() + .unwrap_unchecked() + } + } + pub fn set_range(&mut self, address: &Address, values: &[T; N]) -> [T; N] { + unsafe { + self.paged_vecs[(address.0 - self.as_offset) as usize] + .set_range((address.1 as usize)..(address.1 as usize + N), values) + .try_into() + .unwrap_unchecked() + } + } + + pub fn from_iter( + as_offset: u32, + as_cnt: usize, + mem_size: usize, + iter: impl IntoIterator, + ) -> Self { + let mut vec = Self::new(as_offset, as_cnt, mem_size); + for (address, data) in iter { + vec.insert(&address, data); + } + vec + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/vm/src/system/memory/tree/mod.rs b/crates/vm/src/system/memory/tree/mod.rs index 11890185ff..f855212ed8 100644 --- a/crates/vm/src/system/memory/tree/mod.rs +++ b/crates/vm/src/system/memory/tree/mod.rs @@ -98,13 +98,13 @@ impl MemoryNode { // Construct a BTreeMap that includes the address space in the label calculation, // representing the entire memory tree. let mut memory_partition = BTreeMap::new(); - for (&(address_space, pointer), value) in memory { + for ((address_space, pointer), value) in memory.items() { let label = (address_space, pointer / CHUNK as u32); let index = memory_dimensions.label_to_index(label); let chunk = memory_partition .entry(index) .or_insert_with(|| [F::ZERO; CHUNK]); - chunk[(pointer % CHUNK as u32) as usize] = *value; + chunk[(pointer % CHUNK as u32) as usize] = value; } Self::from_memory( &memory_partition, diff --git a/crates/vm/src/system/memory/tree/public_values.rs b/crates/vm/src/system/memory/tree/public_values.rs index 98325aef63..6476339338 100644 --- a/crates/vm/src/system/memory/tree/public_values.rs +++ b/crates/vm/src/system/memory/tree/public_values.rs @@ -6,7 +6,7 @@ use serde::{Deserialize, Serialize}; use crate::{ arch::hasher::Hasher, system::memory::{ - dimensions::MemoryDimensions, online::Address, tree::MemoryNode, MemoryImage, + dimensions::MemoryDimensions, paged_vec::Address, tree::MemoryNode, MemoryImage, }, }; @@ -118,7 +118,7 @@ pub fn extract_public_values( // TODO: This clones the entire memory. Ideally this should run in time proportional to // the size of the PV address space, not entire memory. - let final_memory: BTreeMap = final_memory.clone().into_iter().collect(); + let final_memory: BTreeMap = final_memory.items().collect(); let used_pvs: Vec<_> = final_memory .range((f_as_start, 0)..(f_as_end, 0)) @@ -126,13 +126,15 @@ pub fn extract_public_values( .collect(); if let Some(&last_pv) = used_pvs.last() { assert!( - last_pv.0 < num_public_values, + last_pv.0 < num_public_values || last_pv.1 == F::ZERO, "Last public value is out of bounds" ); } let mut public_values = F::zero_vec(num_public_values); for (i, pv) in used_pvs { - public_values[i] = pv; + if i < num_public_values { + public_values[i] = pv; + } } public_values } @@ -148,7 +150,7 @@ mod tests { hasher::{poseidon2::vm_poseidon2_hasher, Hasher}, SystemConfig, }, - system::memory::{tree::MemoryNode, MemoryImage, CHUNK}, + system::memory::{paged_vec::AddressMap, tree::MemoryNode, CHUNK}, }; type F = BabyBear; @@ -160,7 +162,12 @@ mod tests { let memory_dimensions = vm_config.memory_config.memory_dimensions(); let pv_as = PUBLIC_VALUES_ADDRESS_SPACE_OFFSET + memory_dimensions.as_offset; let num_public_values = 16; - let memory: MemoryImage = [((pv_as, 15), F::ONE)].into_iter().collect(); + let memory = AddressMap::from_iter( + memory_dimensions.as_offset, + 1 << memory_dimensions.as_height, + 1 << memory_dimensions.address_height, + [((pv_as, 15), F::ONE)], + ); let mut expected_pvs = F::zero_vec(num_public_values); expected_pvs[15] = F::ONE; From 95b16be7c77469561aac79d28b66fa2f18659276 Mon Sep 17 00:00:00 2001 From: Alexander Golovanov Date: Fri, 17 Jan 2025 16:45:16 +0100 Subject: [PATCH 08/15] Refactor --- crates/vm/src/system/memory/offline.rs | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/crates/vm/src/system/memory/offline.rs b/crates/vm/src/system/memory/offline.rs index 3b1639a439..fd45bc3509 100644 --- a/crates/vm/src/system/memory/offline.rs +++ b/crates/vm/src/system/memory/offline.rs @@ -361,14 +361,7 @@ impl OfflineMemory { self.split_to_make_boundary(address_space, pointer, records); self.split_to_make_boundary(address_space, pointer + size as u32, records); - let mut block_data = self - .block_data - .get(&(address_space, pointer)) - .cloned() - .unwrap_or_else(|| Self::initial_block_data(pointer, self.initial_block_size)); - if block_data.size == 0 { - block_data = Self::initial_block_data(pointer, self.initial_block_size); - } + let block_data = self.block_containing(address_space, pointer); if block_data.pointer == pointer && block_data.size == size { return; From dd0230372884cf2721655040d702fbdd58a4969d Mon Sep 17 00:00:00 2001 From: Golovanov399 Date: Fri, 17 Jan 2025 17:53:13 +0100 Subject: [PATCH 09/15] Update crates/vm/src/system/memory/merkle/tests/mod.rs Co-authored-by: Zach Langley --- crates/vm/src/system/memory/merkle/tests/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/vm/src/system/memory/merkle/tests/mod.rs b/crates/vm/src/system/memory/merkle/tests/mod.rs index b6f99d8b98..3ba7852e4a 100644 --- a/crates/vm/src/system/memory/merkle/tests/mod.rs +++ b/crates/vm/src/system/memory/merkle/tests/mod.rs @@ -53,7 +53,7 @@ fn test( let label = pointer / CHUNK as u32; assert!(address_space - as_offset < (1 << as_height)); assert!(pointer < ((CHUNK << address_height).div_ceil(PAGE_SIZE) * PAGE_SIZE) as u32); - if initial_memory.get(&(address_space, pointer)).cloned() != Some(value) { + if initial_memory.get(&(address_space, pointer)) != Some(&value) { assert!(touched_labels.contains(&(address_space, label))); } } From 877e00d8a43e9122a2a0ec8a978a4b841c92995b Mon Sep 17 00:00:00 2001 From: Alexander Golovanov Date: Fri, 17 Jan 2025 19:08:43 +0100 Subject: [PATCH 10/15] Some comments --- crates/vm/src/system/memory/controller/mod.rs | 4 +- crates/vm/src/system/memory/mod.rs | 1 + crates/vm/src/system/memory/offline.rs | 16 +---- crates/vm/src/system/memory/paged_vec.rs | 70 ++++++++----------- 4 files changed, 33 insertions(+), 58 deletions(-) diff --git a/crates/vm/src/system/memory/controller/mod.rs b/crates/vm/src/system/memory/controller/mod.rs index 7442535b74..049b1b49cf 100644 --- a/crates/vm/src/system/memory/controller/mod.rs +++ b/crates/vm/src/system/memory/controller/mod.rs @@ -250,7 +250,6 @@ impl MemoryController { 1, memory_bus, range_checker.clone(), - mem_config.clk_max_bits, mem_config, ))), access_adapters: AccessAdapterInventory::new( @@ -301,7 +300,6 @@ impl MemoryController { CHUNK, memory_bus, range_checker.clone(), - mem_config.clk_max_bits, mem_config, ))), access_adapters: AccessAdapterInventory::new( @@ -358,7 +356,7 @@ impl MemoryController { match &mut self.interface_chip { MemoryInterface::Volatile { .. } => { assert!( - memory.items().collect::>().is_empty(), + memory.is_empty(), "Cannot set initial memory for volatile memory" ); } diff --git a/crates/vm/src/system/memory/mod.rs b/crates/vm/src/system/memory/mod.rs index 792847f1bb..ac6a7d85cf 100644 --- a/crates/vm/src/system/memory/mod.rs +++ b/crates/vm/src/system/memory/mod.rs @@ -15,6 +15,7 @@ mod volatile; pub use controller::*; pub use offline::*; +pub use paged_vec::*; #[derive(PartialEq, Copy, Clone, Debug, Eq)] pub enum OpType { diff --git a/crates/vm/src/system/memory/offline.rs b/crates/vm/src/system/memory/offline.rs index fd45bc3509..deb1d0107e 100644 --- a/crates/vm/src/system/memory/offline.rs +++ b/crates/vm/src/system/memory/offline.rs @@ -6,7 +6,7 @@ use openvm_circuit_primitives::{ use openvm_stark_backend::p3_field::PrimeField32; use rustc_hash::FxHashSet; -use super::paged_vec::{AddressMap, PagedVec, PAGE_SIZE}; +use super::{AddressMap, PagedVec, PAGE_SIZE}; use crate::{ arch::MemoryConfig, system::memory::{ @@ -21,8 +21,8 @@ pub const INITIAL_TIMESTAMP: u32 = 0; #[derive(Clone, Default, PartialEq, Eq, Debug)] struct BlockData { pointer: u32, - size: usize, timestamp: u32, + size: usize, } #[derive(Debug, Clone, PartialEq)] @@ -59,7 +59,6 @@ impl OfflineMemory { initial_block_size: usize, memory_bus: MemoryBus, range_checker: SharedVariableRangeCheckerChip, - timestamp_max_bits: usize, config: MemoryConfig, ) -> Self { assert!(initial_block_size.is_power_of_two()); @@ -70,7 +69,7 @@ impl OfflineMemory { as_offset: config.as_offset, initial_block_size, timestamp: INITIAL_TIMESTAMP + 1, - timestamp_max_bits, + timestamp_max_bits: config.clk_max_bits, memory_bus, range_checker, log: vec![], @@ -527,7 +526,6 @@ mod tests { 8, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - 29, MemoryConfig { as_offset: 0, ..Default::default() @@ -578,7 +576,6 @@ mod tests { 1, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - 29, Default::default(), ); let address_space = 1; @@ -605,7 +602,6 @@ mod tests { 1, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - 29, Default::default(), ); @@ -751,7 +747,6 @@ mod tests { 8, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - 29, Default::default(), ); @@ -867,7 +862,6 @@ mod tests { 1, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - 29, MemoryConfig { as_offset: 0, ..Default::default() @@ -893,7 +887,6 @@ mod tests { 8, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - 29, MemoryConfig { as_offset: 0, ..Default::default() @@ -919,7 +912,6 @@ mod tests { 4, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - 29, Default::default(), ); @@ -936,7 +928,6 @@ mod tests { 8, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - 29, Default::default(), ); // Make block 0:4 in address space 1 active. @@ -1007,7 +998,6 @@ mod tests { 8, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - 29, Default::default(), ); diff --git a/crates/vm/src/system/memory/paged_vec.rs b/crates/vm/src/system/memory/paged_vec.rs index 7020a2ef8d..065355d8f1 100644 --- a/crates/vm/src/system/memory/paged_vec.rs +++ b/crates/vm/src/system/memory/paged_vec.rs @@ -46,59 +46,42 @@ impl PagedVec { } pub fn get_range(&self, range: Range) -> Vec { - let start_page_idx = range.start / PAGE_SIZE; - let end_page_idx = range.end / PAGE_SIZE; - - if start_page_idx == end_page_idx { - if let Some(start_page) = &self.pages[start_page_idx] { - let i = range.start % PAGE_SIZE; - start_page[i..i + range.len()].to_vec() + let mut result = Vec::with_capacity(range.len()); + for page_idx in (range.start / PAGE_SIZE)..range.end.div_ceil(PAGE_SIZE) { + let in_page_start = range.start.saturating_sub(page_idx * PAGE_SIZE); + let in_page_end = (range.end - page_idx * PAGE_SIZE).min(PAGE_SIZE); + if let Some(page) = self.pages[page_idx].as_ref() { + result.extend(page[in_page_start..in_page_end].iter().cloned()); } else { - vec![T::default(); range.len()] + result.extend(vec![T::default(); in_page_end - in_page_start]); } - } else { - // TODO: This can be more efficient by copying from two slices (but most queries should - // not be cross-page). - range - .map(|i| self.get(i).cloned().unwrap_or_default()) - .collect() } + result } - pub fn set_range<'a>( - &mut self, - range: Range, - values: impl IntoIterator, - ) -> Vec - where - T: 'a, - { - let start_page_idx = range.start / PAGE_SIZE; - let end_page_idx = range.end / PAGE_SIZE; - - if start_page_idx == end_page_idx { - let page = - self.pages[start_page_idx].get_or_insert_with(|| vec![T::default(); PAGE_SIZE]); - let page_start = range.start - range.start % PAGE_SIZE; - let result = page[range.start - page_start..range.end - page_start].to_vec(); - for (j, value) in range.zip(values.into_iter()) { - page[j - page_start] = value.clone(); - } - result - } else { - // TODO: This can be more efficient by copying into two slices (but most queries should - // not be cross-page). - let result = self.get_range(range.clone()); - for (i, value) in range.zip(values.into_iter()) { - self.set(i, value.clone()); - } - result + pub fn set_range(&mut self, range: Range, values: &[T]) -> Vec { + let mut result = Vec::with_capacity(range.len()); + let mut values = values.iter(); + for page_idx in (range.start / PAGE_SIZE)..range.end.div_ceil(PAGE_SIZE) { + let in_page_start = range.start.saturating_sub(page_idx * PAGE_SIZE); + let in_page_end = (range.end - page_idx * PAGE_SIZE).min(PAGE_SIZE); + let page = self.pages[page_idx].get_or_insert_with(|| vec![T::default(); PAGE_SIZE]); + result.extend( + page[in_page_start..in_page_end] + .iter_mut() + .map(|x| std::mem::replace(x, values.next().unwrap().clone())), + ); } + result } pub fn memory_size(&self) -> usize { self.pages.len() * PAGE_SIZE } + + pub fn is_empty(&self) -> bool { + self.pages.iter().all(|page| page.is_none()) + } } impl PagedVec { @@ -205,6 +188,9 @@ impl AddressMap { .unwrap_unchecked() } } + pub fn is_empty(&self) -> bool { + self.paged_vecs.iter().all(|page| page.is_empty()) + } pub fn from_iter( as_offset: u32, From 0db1666a1df80dde6d9c50dcddcf084ce4e1e6ef Mon Sep 17 00:00:00 2001 From: Alexander Golovanov Date: Fri, 17 Jan 2025 19:46:32 +0100 Subject: [PATCH 11/15] Try removing `initial_memory` from the arguments of `OfflineMemory::new` --- crates/vm/src/system/memory/controller/mod.rs | 3 - crates/vm/src/system/memory/offline.rs | 66 +++++++++++-------- 2 files changed, 37 insertions(+), 32 deletions(-) diff --git a/crates/vm/src/system/memory/controller/mod.rs b/crates/vm/src/system/memory/controller/mod.rs index 049b1b49cf..1629f39ae3 100644 --- a/crates/vm/src/system/memory/controller/mod.rs +++ b/crates/vm/src/system/memory/controller/mod.rs @@ -232,7 +232,6 @@ impl MemoryController { range_checker: SharedVariableRangeCheckerChip, ) -> Self { let range_checker_bus = range_checker.bus(); - let initial_memory = AddressMap::from_mem_config(&mem_config); Self { memory_bus, mem_config, @@ -246,7 +245,6 @@ impl MemoryController { }, memory: Memory::new(&mem_config), offline_memory: Arc::new(Mutex::new(OfflineMemory::new( - initial_memory, 1, memory_bus, range_checker.clone(), @@ -296,7 +294,6 @@ impl MemoryController { interface_chip, memory: Memory::new(&mem_config), // it is expected that the memory will be set later offline_memory: Arc::new(Mutex::new(OfflineMemory::new( - AddressMap::from_mem_config(&mem_config), CHUNK, memory_bus, range_checker.clone(), diff --git a/crates/vm/src/system/memory/offline.rs b/crates/vm/src/system/memory/offline.rs index deb1d0107e..f4d313babb 100644 --- a/crates/vm/src/system/memory/offline.rs +++ b/crates/vm/src/system/memory/offline.rs @@ -55,7 +55,6 @@ impl OfflineMemory { /// /// Panics if the initial block size is not a power of two. pub fn new( - initial_memory: MemoryImage, initial_block_size: usize, memory_bus: MemoryBus, range_checker: SharedVariableRangeCheckerChip, @@ -65,7 +64,7 @@ impl OfflineMemory { Self { block_data: AddressMap::from_mem_config(&config), - data: Self::memory_image_to_paged_vec(initial_memory, config), + data: vec![], as_offset: config.as_offset, initial_block_size, timestamp: INITIAL_TIMESTAMP + 1, @@ -521,16 +520,17 @@ mod tests { type F = BabyBear; let initial_memory = AddressMap::new(0, 1, 16); + let mem_config = MemoryConfig { + as_offset: 0, + ..Default::default() + }; let mut partition = OfflineMemory::::new( - initial_memory, 8, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - MemoryConfig { - as_offset: 0, - ..Default::default() - }, + mem_config, ); + partition.set_initial_memory(initial_memory, mem_config); assert_eq!( partition.block_containing(0, 13), BlockData { @@ -571,13 +571,14 @@ mod tests { #[test] fn test_write_read_initial_block_len_1() { let initial_memory = MemoryImage::default(); + let mem_config = Default::default(); let mut memory = OfflineMemory::::new( - initial_memory, 1, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - Default::default(), + mem_config, ); + memory.set_initial_memory(initial_memory, mem_config); let address_space = 1; memory.write(address_space, 0, bbvec![1, 2, 3, 4]); @@ -597,13 +598,14 @@ mod tests { fn test_records_initial_block_len_1() { let initial_memory = MemoryImage::default(); // TODO: Ideally we don't need to instantiate all this stuff since we are just testing the data structure. + let mem_config = Default::default(); let mut memory = OfflineMemory::::new( - initial_memory, 1, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - Default::default(), + mem_config, ); + memory.set_initial_memory(initial_memory, mem_config); let adapter_records = memory.write(1, 0, bbvec![1, 2, 3, 4]); @@ -742,13 +744,14 @@ mod tests { #[test] fn test_records_initial_block_len_8() { let initial_memory = MemoryImage::default(); + let mem_config = Default::default(); let mut memory = OfflineMemory::::new( - initial_memory, 8, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - Default::default(), + mem_config, ); + memory.set_initial_memory(initial_memory, mem_config); let adapter_records = memory.write(1, 0, bbvec![1, 2, 3, 4]); let write_record = memory.last_record(); @@ -857,16 +860,17 @@ mod tests { #[test] fn test_get_initial_block_len_1() { let initial_memory = MemoryImage::default(); + let mem_config = MemoryConfig { + as_offset: 0, + ..Default::default() + }; let mut memory = OfflineMemory::::new( - initial_memory, 1, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - MemoryConfig { - as_offset: 0, - ..Default::default() - }, + mem_config, ); + memory.set_initial_memory(initial_memory, mem_config); memory.write(1, 0, bbvec![4, 3, 2, 1]); @@ -882,16 +886,17 @@ mod tests { #[test] fn test_get_initial_block_len_8() { let initial_memory = MemoryImage::default(); + let mem_config = MemoryConfig { + as_offset: 0, + ..Default::default() + }; let mut memory = OfflineMemory::::new( - initial_memory, 8, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - MemoryConfig { - as_offset: 0, - ..Default::default() - }, + mem_config, ); + memory.set_initial_memory(initial_memory, mem_config); memory.write(1, 0, bbvec![4, 3, 2, 1]); @@ -907,13 +912,14 @@ mod tests { #[test] fn test_finalize_empty() { let initial_memory = MemoryImage::default(); + let mem_config = Default::default(); let mut memory = OfflineMemory::::new( - initial_memory, 4, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - Default::default(), + mem_config, ); + memory.set_initial_memory(initial_memory, mem_config); let (memory, records) = memory.finalize::<4>(); assert_eq!(memory.len(), 0); @@ -923,13 +929,14 @@ mod tests { #[test] fn test_finalize_block_len_8() { let initial_memory = MemoryImage::default(); + let mem_config = Default::default(); let mut memory = OfflineMemory::::new( - initial_memory, 8, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - Default::default(), + mem_config, ); + memory.set_initial_memory(initial_memory, mem_config); // Make block 0:4 in address space 1 active. memory.write(1, 0, bbvec![1, 2, 3, 4]); @@ -993,13 +1000,14 @@ mod tests { initial_memory.insert(&(1, 16 + i), F::from_canonical_u32(i + 1)); } + let mem_config = Default::default(); let mut memory = OfflineMemory::::new( - initial_memory, 8, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - Default::default(), + mem_config, ); + memory.set_initial_memory(initial_memory, mem_config); // Verify initial state of block 0 (pointers 0–8) memory.read(1, 0, 8); From 82ecab5e3aaf0baf1aaaeecf8e73011f819edd81 Mon Sep 17 00:00:00 2001 From: Alexander Golovanov Date: Fri, 17 Jan 2025 19:57:53 +0100 Subject: [PATCH 12/15] Revert "Try removing `initial_memory` from the arguments of `OfflineMemory::new`" This reverts commit 0db1666a1df80dde6d9c50dcddcf084ce4e1e6ef. --- crates/vm/src/system/memory/controller/mod.rs | 3 + crates/vm/src/system/memory/offline.rs | 66 ++++++++----------- 2 files changed, 32 insertions(+), 37 deletions(-) diff --git a/crates/vm/src/system/memory/controller/mod.rs b/crates/vm/src/system/memory/controller/mod.rs index 1629f39ae3..049b1b49cf 100644 --- a/crates/vm/src/system/memory/controller/mod.rs +++ b/crates/vm/src/system/memory/controller/mod.rs @@ -232,6 +232,7 @@ impl MemoryController { range_checker: SharedVariableRangeCheckerChip, ) -> Self { let range_checker_bus = range_checker.bus(); + let initial_memory = AddressMap::from_mem_config(&mem_config); Self { memory_bus, mem_config, @@ -245,6 +246,7 @@ impl MemoryController { }, memory: Memory::new(&mem_config), offline_memory: Arc::new(Mutex::new(OfflineMemory::new( + initial_memory, 1, memory_bus, range_checker.clone(), @@ -294,6 +296,7 @@ impl MemoryController { interface_chip, memory: Memory::new(&mem_config), // it is expected that the memory will be set later offline_memory: Arc::new(Mutex::new(OfflineMemory::new( + AddressMap::from_mem_config(&mem_config), CHUNK, memory_bus, range_checker.clone(), diff --git a/crates/vm/src/system/memory/offline.rs b/crates/vm/src/system/memory/offline.rs index f4d313babb..deb1d0107e 100644 --- a/crates/vm/src/system/memory/offline.rs +++ b/crates/vm/src/system/memory/offline.rs @@ -55,6 +55,7 @@ impl OfflineMemory { /// /// Panics if the initial block size is not a power of two. pub fn new( + initial_memory: MemoryImage, initial_block_size: usize, memory_bus: MemoryBus, range_checker: SharedVariableRangeCheckerChip, @@ -64,7 +65,7 @@ impl OfflineMemory { Self { block_data: AddressMap::from_mem_config(&config), - data: vec![], + data: Self::memory_image_to_paged_vec(initial_memory, config), as_offset: config.as_offset, initial_block_size, timestamp: INITIAL_TIMESTAMP + 1, @@ -520,17 +521,16 @@ mod tests { type F = BabyBear; let initial_memory = AddressMap::new(0, 1, 16); - let mem_config = MemoryConfig { - as_offset: 0, - ..Default::default() - }; let mut partition = OfflineMemory::::new( + initial_memory, 8, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - mem_config, + MemoryConfig { + as_offset: 0, + ..Default::default() + }, ); - partition.set_initial_memory(initial_memory, mem_config); assert_eq!( partition.block_containing(0, 13), BlockData { @@ -571,14 +571,13 @@ mod tests { #[test] fn test_write_read_initial_block_len_1() { let initial_memory = MemoryImage::default(); - let mem_config = Default::default(); let mut memory = OfflineMemory::::new( + initial_memory, 1, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - mem_config, + Default::default(), ); - memory.set_initial_memory(initial_memory, mem_config); let address_space = 1; memory.write(address_space, 0, bbvec![1, 2, 3, 4]); @@ -598,14 +597,13 @@ mod tests { fn test_records_initial_block_len_1() { let initial_memory = MemoryImage::default(); // TODO: Ideally we don't need to instantiate all this stuff since we are just testing the data structure. - let mem_config = Default::default(); let mut memory = OfflineMemory::::new( + initial_memory, 1, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - mem_config, + Default::default(), ); - memory.set_initial_memory(initial_memory, mem_config); let adapter_records = memory.write(1, 0, bbvec![1, 2, 3, 4]); @@ -744,14 +742,13 @@ mod tests { #[test] fn test_records_initial_block_len_8() { let initial_memory = MemoryImage::default(); - let mem_config = Default::default(); let mut memory = OfflineMemory::::new( + initial_memory, 8, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - mem_config, + Default::default(), ); - memory.set_initial_memory(initial_memory, mem_config); let adapter_records = memory.write(1, 0, bbvec![1, 2, 3, 4]); let write_record = memory.last_record(); @@ -860,17 +857,16 @@ mod tests { #[test] fn test_get_initial_block_len_1() { let initial_memory = MemoryImage::default(); - let mem_config = MemoryConfig { - as_offset: 0, - ..Default::default() - }; let mut memory = OfflineMemory::::new( + initial_memory, 1, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - mem_config, + MemoryConfig { + as_offset: 0, + ..Default::default() + }, ); - memory.set_initial_memory(initial_memory, mem_config); memory.write(1, 0, bbvec![4, 3, 2, 1]); @@ -886,17 +882,16 @@ mod tests { #[test] fn test_get_initial_block_len_8() { let initial_memory = MemoryImage::default(); - let mem_config = MemoryConfig { - as_offset: 0, - ..Default::default() - }; let mut memory = OfflineMemory::::new( + initial_memory, 8, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - mem_config, + MemoryConfig { + as_offset: 0, + ..Default::default() + }, ); - memory.set_initial_memory(initial_memory, mem_config); memory.write(1, 0, bbvec![4, 3, 2, 1]); @@ -912,14 +907,13 @@ mod tests { #[test] fn test_finalize_empty() { let initial_memory = MemoryImage::default(); - let mem_config = Default::default(); let mut memory = OfflineMemory::::new( + initial_memory, 4, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - mem_config, + Default::default(), ); - memory.set_initial_memory(initial_memory, mem_config); let (memory, records) = memory.finalize::<4>(); assert_eq!(memory.len(), 0); @@ -929,14 +923,13 @@ mod tests { #[test] fn test_finalize_block_len_8() { let initial_memory = MemoryImage::default(); - let mem_config = Default::default(); let mut memory = OfflineMemory::::new( + initial_memory, 8, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - mem_config, + Default::default(), ); - memory.set_initial_memory(initial_memory, mem_config); // Make block 0:4 in address space 1 active. memory.write(1, 0, bbvec![1, 2, 3, 4]); @@ -1000,14 +993,13 @@ mod tests { initial_memory.insert(&(1, 16 + i), F::from_canonical_u32(i + 1)); } - let mem_config = Default::default(); let mut memory = OfflineMemory::::new( + initial_memory, 8, MemoryBus(0), SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new(1, 29)), - mem_config, + Default::default(), ); - memory.set_initial_memory(initial_memory, mem_config); // Verify initial state of block 0 (pointers 0–8) memory.read(1, 0, 8); From 73891277eaecd41b04e73fc1c033b382b409d124 Mon Sep 17 00:00:00 2001 From: Alexander Golovanov Date: Tue, 21 Jan 2025 21:41:22 +0100 Subject: [PATCH 13/15] Change the memory config in the test builder to `default()` --- crates/vm/src/arch/testing/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/crates/vm/src/arch/testing/mod.rs b/crates/vm/src/arch/testing/mod.rs index 0636147c57..a6177a0a09 100644 --- a/crates/vm/src/arch/testing/mod.rs +++ b/crates/vm/src/arch/testing/mod.rs @@ -260,7 +260,7 @@ impl VmChipTestBuilder { impl Default for VmChipTestBuilder { fn default() -> Self { - let mem_config = MemoryConfig::new(2, 1, 29, 29, 17, 64, 1 << 22); + let mem_config = MemoryConfig::default(); let range_checker = SharedVariableRangeCheckerChip::new(VariableRangeCheckerBus::new( RANGE_CHECKER_BUS, mem_config.decomp, From 3f8a365b18fd27d377840d4f425925b34635d02a Mon Sep 17 00:00:00 2001 From: Alexander Golovanov Date: Wed, 22 Jan 2025 15:39:33 +0100 Subject: [PATCH 14/15] Hopefully speed something up using unsafe trickery --- crates/vm/src/system/memory/offline.rs | 2 +- crates/vm/src/system/memory/online.rs | 2 + crates/vm/src/system/memory/paged_vec.rs | 94 +++++++++++++++++++----- 3 files changed, 78 insertions(+), 20 deletions(-) diff --git a/crates/vm/src/system/memory/offline.rs b/crates/vm/src/system/memory/offline.rs index deb1d0107e..95df415b49 100644 --- a/crates/vm/src/system/memory/offline.rs +++ b/crates/vm/src/system/memory/offline.rs @@ -460,7 +460,7 @@ impl OfflineMemory { fn range_vec(&self, address_space: u32, pointer: u32, len: usize) -> Vec { let pointer = pointer as usize; - self.data[(address_space - self.as_offset) as usize].get_range(pointer..pointer + len) + self.data[(address_space - self.as_offset) as usize].range_vec(pointer..pointer + len) } pub fn aux_cols_factory(&self) -> MemoryAuxColsFactory { diff --git a/crates/vm/src/system/memory/online.rs b/crates/vm/src/system/memory/online.rs index e4e91113d8..a5bf663e4c 100644 --- a/crates/vm/src/system/memory/online.rs +++ b/crates/vm/src/system/memory/online.rs @@ -108,10 +108,12 @@ impl Memory { self.timestamp } + #[inline(always)] pub fn get(&self, address_space: u32, pointer: u32) -> F { *self.data.get(&(address_space, pointer)).unwrap_or(&F::ZERO) } + #[inline(always)] fn range_array(&self, address_space: u32, pointer: u32) -> [F; N] { self.data.get_range(&(address_space, pointer)) } diff --git a/crates/vm/src/system/memory/paged_vec.rs b/crates/vm/src/system/memory/paged_vec.rs index 065355d8f1..a84bc9d777 100644 --- a/crates/vm/src/system/memory/paged_vec.rs +++ b/crates/vm/src/system/memory/paged_vec.rs @@ -1,4 +1,4 @@ -use std::ops::Range; +use std::{mem::MaybeUninit, ops::Range, ptr}; use serde::{Deserialize, Serialize}; @@ -45,7 +45,8 @@ impl PagedVec { } } - pub fn get_range(&self, range: Range) -> Vec { + #[inline(always)] + pub fn range_vec(&self, range: Range) -> Vec { let mut result = Vec::with_capacity(range.len()); for page_idx in (range.start / PAGE_SIZE)..range.end.div_ceil(PAGE_SIZE) { let in_page_start = range.start.saturating_sub(page_idx * PAGE_SIZE); @@ -84,6 +85,63 @@ impl PagedVec { } } +impl PagedVec { + #[inline(always)] + pub fn range_array(&self, from: usize) -> [T; N] { + // Step 1: Create an uninitialized array of MaybeUninit + let mut result: [MaybeUninit; N] = unsafe { + // SAFETY: An uninitialized `[MaybeUninit; N]` is valid. + MaybeUninit::uninit().assume_init() + }; + + // Step 2: Get a mutable slice of T references from the MaybeUninit array. + let result_slice = unsafe { + // SAFETY: We are converting a pointer from MaybeUninit to T. + // This is safe because we will fully initialize every element before reading. + std::slice::from_raw_parts_mut(result.as_mut_ptr() as *mut T, N) + }; + + let start_page = from / PAGE_SIZE; + let end_page = (from + N - 1) / PAGE_SIZE; + + if start_page == end_page { + if let Some(page) = self.pages[start_page].as_ref() { + // Copy data from the page into result_slice. + let page_offset = from - start_page * PAGE_SIZE; + let src = &page[page_offset..page_offset + N]; + result_slice.copy_from_slice(src); + } else { + // If no page available, fill with default values. + result_slice.fill(T::default()); + } + } else { + debug_assert!(start_page + 1 == end_page); + let first_part = PAGE_SIZE - (from - start_page * PAGE_SIZE); + if let Some(page) = self.pages[start_page].as_ref() { + let page_offset = from - start_page * PAGE_SIZE; + let src = &page[page_offset..]; + result_slice[..first_part].copy_from_slice(src); + } else { + result_slice[..first_part].fill(T::default()); + } + if let Some(page) = self.pages[end_page].as_ref() { + let second_part = N - first_part; + let src = &page[0..second_part]; + result_slice[first_part..].copy_from_slice(src); + } else { + result_slice[first_part..].fill(T::default()); + } + } + + // Step 4: Convert the fully initialized array of MaybeUninit to [T; N]. + // SAFETY: We have initialized every element of `result`. + unsafe { + // Transmute the array; at this point each element is initialized. + ptr::read(&result as *const _ as *const [T; N]) + } + } +} + impl PagedVec { pub fn iter(&self) -> PagedVecIter<'_, T, PAGE_SIZE> { PagedVecIter { @@ -172,14 +230,6 @@ impl AddressMap { pub fn insert(&mut self, address: &Address, data: T) -> Option { self.paged_vecs[(address.0 - self.as_offset) as usize].set(address.1 as usize, data) } - pub fn get_range(&self, address: &Address) -> [T; N] { - unsafe { - self.paged_vecs[(address.0 - self.as_offset) as usize] - .get_range((address.1 as usize)..(address.1 as usize + N)) - .try_into() - .unwrap_unchecked() - } - } pub fn set_range(&mut self, address: &Address, values: &[T; N]) -> [T; N] { unsafe { self.paged_vecs[(address.0 - self.as_offset) as usize] @@ -206,6 +256,12 @@ impl AddressMap { } } +impl AddressMap { + pub fn get_range(&self, address: &Address) -> [T; N] { + self.paged_vecs[(address.0 - self.as_offset) as usize].range_array(address.1 as usize) + } +} + #[cfg(test)] mod tests { use super::*; @@ -242,14 +298,14 @@ mod tests { v.set(7, 8); // Verify all values - assert_eq!(v.get_range(0..8), [1, 2, 3, 4, 5, 6, 7, 8]); + assert_eq!(v.range_vec(0..8), [1, 2, 3, 4, 5, 6, 7, 8]); } #[test] fn test_range_cross_page_boundary() { let mut v = PagedVec::<_, 4>::new(2); v.set_range(2..8, &[10, 11, 12, 13, 14, 15]); - assert_eq!(v.get_range(2..8), [10, 11, 12, 13, 14, 15]); + assert_eq!(v.range_vec(2..8), [10, 11, 12, 13, 14, 15]); } #[test] @@ -267,7 +323,7 @@ mod tests { v.set(5, 10); // Should include both set values and defaults - assert_eq!(v.get_range(1..7), [0, 5, 0, 0, 10, 0]); + assert_eq!(v.range_vec(1..7), [0, 5, 0, 0, 10, 0]); } #[test] @@ -302,19 +358,19 @@ mod tests { // Initial set_range v.set_range(0..5, &[1, 2, 3, 4, 5]); - assert_eq!(v.get_range(0..5), [1, 2, 3, 4, 5]); + assert_eq!(v.range_vec(0..5), [1, 2, 3, 4, 5]); // Overlap from beginning v.set_range(0..3, &[10, 20, 30]); - assert_eq!(v.get_range(0..5), [10, 20, 30, 4, 5]); + assert_eq!(v.range_vec(0..5), [10, 20, 30, 4, 5]); // Overlap in middle v.set_range(2..4, &[42, 43]); - assert_eq!(v.get_range(0..5), [10, 20, 42, 43, 5]); + assert_eq!(v.range_vec(0..5), [10, 20, 42, 43, 5]); // Overlap at end v.set_range(4..6, &[91, 92]); - assert_eq!(v.get_range(0..6), [10, 20, 42, 43, 91, 92]); + assert_eq!(v.range_vec(0..6), [10, 20, 42, 43, 91, 92]); } #[test] @@ -326,11 +382,11 @@ mod tests { // Overlap end of first page and start of second v.set_range(2..6, &[21, 22, 23, 24]); - assert_eq!(v.get_range(0..8), [1, 2, 21, 22, 23, 24, 7, 8]); + assert_eq!(v.range_vec(0..8), [1, 2, 21, 22, 23, 24, 7, 8]); // Overlap multiple pages v.set_range(1..7, &[31, 32, 33, 34, 35, 36]); - assert_eq!(v.get_range(0..8), [1, 31, 32, 33, 34, 35, 36, 8]); + assert_eq!(v.range_vec(0..8), [1, 31, 32, 33, 34, 35, 36, 8]); } #[test] From 8ce871d5d0a29a489fe51cdec8d5f0a04b5c8890 Mon Sep 17 00:00:00 2001 From: Alexander Golovanov Date: Wed, 22 Jan 2025 15:49:16 +0100 Subject: [PATCH 15/15] Also (hopefully) speed up `set_range_array` --- crates/vm/src/system/memory/paged_vec.rs | 65 +++++++++++++++++++++--- 1 file changed, 57 insertions(+), 8 deletions(-) diff --git a/crates/vm/src/system/memory/paged_vec.rs b/crates/vm/src/system/memory/paged_vec.rs index a84bc9d777..efb4374582 100644 --- a/crates/vm/src/system/memory/paged_vec.rs +++ b/crates/vm/src/system/memory/paged_vec.rs @@ -140,6 +140,59 @@ impl PagedVec { ptr::read(&result as *const _ as *const [T; N]) } } + + #[inline(always)] + pub fn set_range_array(&mut self, from: usize, values: &[T; N]) -> [T; N] { + // Step 1: Create an uninitialized array for old values. + let mut result: [MaybeUninit; N] = unsafe { MaybeUninit::uninit().assume_init() }; + let result_slice = unsafe { + // SAFETY: We will fully initialize `result_slice` before reading. + std::slice::from_raw_parts_mut(result.as_mut_ptr() as *mut T, N) + }; + + let start_page = from / PAGE_SIZE; + let end_page = (from + N - 1) / PAGE_SIZE; + + if start_page == end_page { + // Ensure the page exists; if not, allocate and fill with defaults. + if self.pages[start_page].is_none() { + self.pages[start_page] = Some(vec![T::default(); PAGE_SIZE]); + } + let page = self.pages[start_page].as_mut().unwrap(); + let page_offset = from - start_page * PAGE_SIZE; + + // Copy old values from the page into result_slice. + result_slice.copy_from_slice(&page[page_offset..page_offset + N]); + // Write the new values from input into the page. + page[page_offset..page_offset + N].copy_from_slice(&values[..]); + } else { + debug_assert!(start_page + 1 == end_page); + let first_part = PAGE_SIZE - (from - start_page * PAGE_SIZE); + + // Handle the first page. + if self.pages[start_page].is_none() { + self.pages[start_page] = Some(vec![T::default(); PAGE_SIZE]); + } + let page0 = self.pages[start_page].as_mut().unwrap(); + let page_offset = from - start_page * PAGE_SIZE; + + result_slice[..first_part].copy_from_slice(&page0[page_offset..]); + page0[page_offset..].copy_from_slice(&values[..first_part]); + + // Handle the second page. + let second_part = N - first_part; + if self.pages[end_page].is_none() { + self.pages[end_page] = Some(vec![T::default(); PAGE_SIZE]); + } + let page1 = self.pages[end_page].as_mut().unwrap(); + + result_slice[first_part..].copy_from_slice(&page1[0..second_part]); + page1[0..second_part].copy_from_slice(&values[first_part..]); + } + + // Step 4: Convert the fully initialized result array to [T; N]. + unsafe { ptr::read(&result as *const _ as *const [T; N]) } + } } impl PagedVec { @@ -230,14 +283,6 @@ impl AddressMap { pub fn insert(&mut self, address: &Address, data: T) -> Option { self.paged_vecs[(address.0 - self.as_offset) as usize].set(address.1 as usize, data) } - pub fn set_range(&mut self, address: &Address, values: &[T; N]) -> [T; N] { - unsafe { - self.paged_vecs[(address.0 - self.as_offset) as usize] - .set_range((address.1 as usize)..(address.1 as usize + N), values) - .try_into() - .unwrap_unchecked() - } - } pub fn is_empty(&self) -> bool { self.paged_vecs.iter().all(|page| page.is_empty()) } @@ -260,6 +305,10 @@ impl AddressMap { pub fn get_range(&self, address: &Address) -> [T; N] { self.paged_vecs[(address.0 - self.as_offset) as usize].range_array(address.1 as usize) } + pub fn set_range(&mut self, address: &Address, values: &[T; N]) -> [T; N] { + self.paged_vecs[(address.0 - self.as_offset) as usize] + .set_range_array(address.1 as usize, values) + } } #[cfg(test)]