Skip to content

Commit

Permalink
perf: Use usize for address_space internally everywhere
Browse files Browse the repository at this point in the history
  • Loading branch information
zlangley committed Dec 16, 2024
1 parent 31c5b18 commit b4e5a78
Showing 13 changed files with 73 additions and 96 deletions.
2 changes: 1 addition & 1 deletion crates/toolchain/instructions/src/exe.rs
Original file line number Diff line number Diff line change
@@ -6,7 +6,7 @@ use serde::{Deserialize, Serialize};
use crate::program::Program;

/// Memory image is a map from (address space, address) to word.
pub type MemoryImage<F> = BTreeMap<(F, F), F>;
pub type MemoryImage<F> = BTreeMap<(usize, usize), F>;
/// Stores the starting address, end address, and name of a set of function.
pub type FnBounds = BTreeMap<u32, FnBound>;

5 changes: 1 addition & 4 deletions crates/toolchain/transpiler/src/util.rs
Original file line number Diff line number Diff line change
@@ -166,10 +166,7 @@ pub fn elf_memory_image_to_openvm_memory_image<F: PrimeField32>(
let mut result = MemoryImage::new();
for (addr, word) in memory_image {
for (i, byte) in word.to_le_bytes().into_iter().enumerate() {
result.insert(
(F::TWO, F::from_canonical_u32(addr + i as u32)),
F::from_canonical_u8(byte),
);
result.insert((2, addr as usize + i), F::from_canonical_u8(byte));
}
}
result
7 changes: 3 additions & 4 deletions crates/vm/src/system/memory/manager/dimensions.rs
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
use derive_new::new;
use openvm_stark_backend::{p3_field::PrimeField32, p3_util::log2_strict_usize};
use openvm_stark_backend::p3_util::log2_strict_usize;

use crate::{arch::MemoryConfig, system::memory::CHUNK};

@@ -20,10 +20,9 @@ impl MemoryDimensions {
self.as_height + self.address_height
}
/// Convert an address label (address space, block id) to its index in the memory merkle tree.
pub fn label_to_index<F: PrimeField32>(&self, label: (F, usize)) -> usize {
pub fn label_to_index(&self, label: (usize, usize)) -> usize {
let (addr_space, block_id) = label;
((addr_space.as_canonical_u32() as usize - self.as_offset) << self.address_height)
+ block_id
((addr_space - self.as_offset) << self.address_height) + block_id
}
}

2 changes: 1 addition & 1 deletion crates/vm/src/system/memory/manager/interface.rs
Original file line number Diff line number Diff line change
@@ -21,7 +21,7 @@ pub enum MemoryInterface<F> {
}

impl<F: PrimeField32> MemoryInterface<F> {
pub fn touch_address(&mut self, addr_space: F, pointer: F) {
pub fn touch_address(&mut self, addr_space: usize, pointer: usize) {
match self {
MemoryInterface::Volatile { boundary_chip } => {
boundary_chip.touch_address(addr_space, pointer);
19 changes: 9 additions & 10 deletions crates/vm/src/system/memory/manager/memory.rs
Original file line number Diff line number Diff line change
@@ -76,16 +76,15 @@ impl<F: PrimeField32> Memory<F> {
let mut block_data = FxHashMap::default();
let mut data = FxHashMap::default();
for (&(address_space, block_idx), values) in initial_memory {
let address_space_usize = address_space.as_canonical_u32() as usize;
let pointer = block_idx * N;
let block = BlockData {
pointer,
size: N,
timestamp: INITIAL_TIMESTAMP,
};
for (i, value) in values.iter().enumerate() {
data.insert((address_space_usize, pointer + i), *value);
block_data.insert((address_space_usize, pointer + i), block);
data.insert((address_space, pointer + i), *value);
block_data.insert((address_space, pointer + i), block);
}
}
Self {
@@ -197,7 +196,7 @@ impl<F: PrimeField32> Memory<F> {
debug_assert_eq!(block.size, N);

equipartition.insert(
(F::from_canonical_usize(address_space), pointer / N),
(address_space, pointer / N),
TimestampedValues {
timestamp: block.timestamp,
values: self.range_array::<N>(address_space, pointer),
@@ -811,31 +810,31 @@ mod tests {
let (final_memory, records) = memory.finalize::<8>();
assert_eq!(final_memory.len(), 4);
assert_eq!(
final_memory.get(&(bb!(1), 0)),
final_memory.get(&(1, 0)),
Some(&TimestampedValues {
values: bba![1, 2, 3, 4, 0, 0, 0, 0],
timestamp: 1,
})
);
// start_index = 16 corresponds to label = 2
assert_eq!(
final_memory.get(&(bb!(1), 2)),
final_memory.get(&(1, 2)),
Some(&TimestampedValues {
values: bba![1, 1, 1, 1, 1, 1, 1, 1],
timestamp: 2,
})
);
// start_index = 24 corresponds to label = 3
assert_eq!(
final_memory.get(&(bb!(1), 3)),
final_memory.get(&(1, 3)),
Some(&TimestampedValues {
values: bba![1, 1, 1, 1, 1, 1, 1, 1],
timestamp: 2,
})
);
// start_index = 64 corresponds to label = 8
assert_eq!(
final_memory.get(&(bb!(2), 8)),
final_memory.get(&(2, 8)),
Some(&TimestampedValues {
values: bba![8, 7, 6, 5, 4, 3, 2, 1],
timestamp: 3,
@@ -852,8 +851,8 @@ mod tests {

// Initialize initial memory with blocks at indices 0 and 2
let mut initial_memory = Equipartition::<F, 8>::new();
initial_memory.insert((F::ONE, 0), bba![1, 2, 3, 4, 5, 6, 7, 8]); // Block 0, pointers 0–8
initial_memory.insert((F::ONE, 2), bba![1, 2, 3, 4, 5, 6, 7, 8]); // Block 2, pointers 16–24
initial_memory.insert((1, 0), bba![1, 2, 3, 4, 5, 6, 7, 8]); // Block 0, pointers 0–8
initial_memory.insert((1, 2), bba![1, 2, 3, 4, 5, 6, 7, 8]); // Block 2, pointers 16–24

let mut memory = Memory::new(&initial_memory);

36 changes: 17 additions & 19 deletions crates/vm/src/system/memory/manager/mod.rs
Original file line number Diff line number Diff line change
@@ -74,15 +74,15 @@ pub type MemoryControllerRef<F> = Rc<RefCell<MemoryController<F>>>;
///
/// If a key is not present in the map, then the block is uninitialized (and therefore zero).
pub type TimestampedEquipartition<F, const N: usize> =
BTreeMap<(F, usize), TimestampedValues<F, N>>;
BTreeMap<(usize, usize), TimestampedValues<F, N>>;

/// A equipartition of memory values.
///
/// The key is a pair `(address_space, label)`, where `label` is the index of the block in the
/// partition. I.e., the starting address of the block is `(address_space, label * N)`.
///
/// If a key is not present in the map, then the block is uninitialized (and therefore zero).
pub type Equipartition<F, const N: usize> = BTreeMap<(F, usize), [F; N]>;
pub type Equipartition<F, const N: usize> = BTreeMap<(usize, usize), [F; N]>;

#[derive(Debug, Getters)]
pub struct MemoryController<F> {
@@ -346,6 +346,7 @@ impl<F: PrimeField32> MemoryController<F> {
}

pub fn read<const N: usize>(&mut self, address_space: F, pointer: F) -> MemoryReadRecord<F, N> {
let address_space_usize = address_space.as_canonical_u32() as usize;
let ptr_u32 = pointer.as_canonical_u32();
assert!(
address_space == F::ZERO || ptr_u32 < (1 << self.mem_config.pointer_max_bits),
@@ -367,16 +368,15 @@ impl<F: PrimeField32> MemoryController<F> {
};
}

let (record, adapter_records) = self
.memory
.read::<N>(address_space.as_canonical_u32() as usize, ptr_u32 as usize);
let (record, adapter_records) =
self.memory.read::<N>(address_space_usize, ptr_u32 as usize);
for record in adapter_records {
self.access_adapters.add_record(record);
}

for i in 0..N as u32 {
let ptr = F::from_canonical_u32(ptr_u32 + i);
self.interface_chip.touch_address(address_space, ptr);
for i in 0..N {
self.interface_chip
.touch_address(address_space_usize, ptr_u32 as usize + i);
}

record
@@ -409,24 +409,23 @@ impl<F: PrimeField32> MemoryController<F> {
data: [F; N],
) -> MemoryWriteRecord<F, N> {
assert_ne!(address_space, F::ZERO);
let address_space_usize = address_space.as_canonical_u32() as usize;
let ptr_u32 = pointer.as_canonical_u32();
assert!(
ptr_u32 < (1 << self.mem_config.pointer_max_bits),
"memory out of bounds: {ptr_u32:?}",
);

let (record, adapter_records) = self.memory.write(
address_space.as_canonical_u32() as usize,
ptr_u32 as usize,
data,
);
let (record, adapter_records) =
self.memory
.write(address_space_usize, ptr_u32 as usize, data);
for record in adapter_records {
self.access_adapters.add_record(record);
}

for i in 0..N as u32 {
let ptr = F::from_canonical_u32(ptr_u32 + i);
self.interface_chip.touch_address(address_space, ptr);
for i in 0..N {
self.interface_chip
.touch_address(address_space_usize, ptr_u32 as usize + i);
}

record
@@ -718,9 +717,8 @@ pub fn memory_image_to_equipartition<F: PrimeField32, const N: usize>(
) -> Equipartition<F, { N }> {
let mut result = Equipartition::new();
for ((addr_space, addr), word) in memory_image {
let addr_usize = addr.as_canonical_u32() as usize;
let shift = addr_usize % N;
let key = (addr_space, addr_usize / N);
let shift = addr % N;
let key = (addr_space, addr / N);
result.entry(key).or_insert([F::ZERO; N])[shift] = word;
}
result
6 changes: 3 additions & 3 deletions crates/vm/src/system/memory/merkle/mod.rs
Original file line number Diff line number Diff line change
@@ -70,11 +70,11 @@ impl<const CHUNK: usize, F: PrimeField32> MemoryMerkleChip<CHUNK, F> {
}
}

pub fn touch_address(&mut self, address_space: F, address: F) {
pub fn touch_address(&mut self, address_space: usize, address: usize) {
self.touch_node(
0,
(address_space.as_canonical_u32() as usize) - self.air.memory_dimensions.as_offset,
(address.as_canonical_u32() as usize) / CHUNK,
address_space - self.air.memory_dimensions.as_offset,
address / CHUNK,
);
}
}
20 changes: 7 additions & 13 deletions crates/vm/src/system/memory/merkle/tests/mod.rs
Original file line number Diff line number Diff line change
@@ -6,11 +6,8 @@ use std::{
};

use openvm_stark_backend::{
interaction::InteractionType,
p3_field::{AbstractField, PrimeField32},
p3_matrix::dense::RowMajorMatrix,
prover::types::AirProofInput,
Chip, ChipUsageGetter,
interaction::InteractionType, p3_field::AbstractField, p3_matrix::dense::RowMajorMatrix,
prover::types::AirProofInput, Chip, ChipUsageGetter,
};
use openvm_stark_sdk::{
config::baby_bear_poseidon2::BabyBearPoseidon2Engine,
@@ -40,7 +37,7 @@ const COMPRESSION_BUS: DirectCompressionBus = DirectCompressionBus(POSEIDON2_DIR
fn test<const CHUNK: usize>(
memory_dimensions: MemoryDimensions,
initial_memory: &Equipartition<BabyBear, CHUNK>,
touched_labels: BTreeSet<(BabyBear, usize)>,
touched_labels: BTreeSet<(usize, usize)>,
final_memory: &Equipartition<BabyBear, CHUNK>,
) {
let MemoryDimensions {
@@ -52,7 +49,7 @@ fn test<const CHUNK: usize>(

// checking validity of test data
for (&(address_space, label), value) in final_memory {
assert!((address_space.as_canonical_u32() as usize) - as_offset < (1 << as_height));
assert!(address_space - as_offset < (1 << as_height));
assert!(label < (1 << address_height));
if initial_memory.get(&(address_space, label)) != Some(value) {
assert!(touched_labels.contains(&(address_space, label)));
@@ -76,10 +73,7 @@ fn test<const CHUNK: usize>(
MemoryMerkleChip::<CHUNK, _>::new(memory_dimensions, merkle_bus, COMPRESSION_BUS);
for &(address_space, label) in touched_labels.iter() {
for i in 0..CHUNK {
chip.touch_address(
address_space,
BabyBear::from_canonical_usize(label * CHUNK + i),
);
chip.touch_address(address_space, label * CHUNK + i);
}
}

@@ -121,7 +115,7 @@ fn test<const CHUNK: usize>(
let initial_values = *initial_memory
.get(&(address_space, address_label))
.unwrap_or(&[BabyBear::ZERO; CHUNK]);
let as_label = address_space.as_canonical_u32() as usize - as_offset;
let as_label = address_space - as_offset;
interaction(
InteractionType::Send,
false,
@@ -176,7 +170,7 @@ fn random_test<const CHUNK: usize>(
let mut touched_labels = BTreeSet::new();

while num_initial_addresses != 0 || num_touched_addresses != 0 {
let address_space = BabyBear::from_canonical_usize((next_usize() & 1) + 1);
let address_space = (next_usize() & 1) + 1;
let label = next_usize() % (1 << height);

if seen_labels.insert(label) {
3 changes: 1 addition & 2 deletions crates/vm/src/system/memory/merkle/trace.rs
Original file line number Diff line number Diff line change
@@ -133,8 +133,7 @@ impl<const CHUNK: usize, F: PrimeField32> TreeHelper<'_, CHUNK, F> {
hasher: &mut impl HasherChip<CHUNK, F>,
) -> MemoryNode<CHUNK, F> {
if height == 0 {
let address_space =
F::from_canonical_usize(as_label + self.memory_dimensions.as_offset);
let address_space = as_label + self.memory_dimensions.as_offset;
let leaf_values = *self
.final_memory
.get(&(address_space, address_label))
14 changes: 7 additions & 7 deletions crates/vm/src/system/memory/persistent.rs
Original file line number Diff line number Diff line change
@@ -128,13 +128,13 @@ pub struct PersistentBoundaryChip<F, const CHUNK: usize> {

#[derive(Debug)]
enum TouchedLabels<F, const CHUNK: usize> {
Running(FxHashSet<(F, usize)>),
Running(FxHashSet<(usize, usize)>),
Final(Vec<FinalTouchedLabel<F, CHUNK>>),
}

#[derive(Debug)]
struct FinalTouchedLabel<F, const CHUNK: usize> {
address_space: F,
address_space: usize,
label: usize,
init_values: [F; CHUNK],
final_values: [F; CHUNK],
@@ -151,7 +151,7 @@ impl<F: PrimeField32, const CHUNK: usize> Default for TouchedLabels<F, CHUNK> {
}

impl<F: PrimeField32, const CHUNK: usize> TouchedLabels<F, CHUNK> {
fn touch(&mut self, address_space: F, label: usize) {
fn touch(&mut self, address_space: usize, label: usize) {
match self {
TouchedLabels::Running(touched_labels) => {
touched_labels.insert((address_space, label));
@@ -190,8 +190,8 @@ impl<const CHUNK: usize, F: PrimeField32> PersistentBoundaryChip<F, CHUNK> {
self.overridden_height = Some(overridden_height);
}

pub fn touch_address(&mut self, address_space: F, pointer: F) {
let label = pointer.as_canonical_u32() as usize / CHUNK;
pub fn touch_address(&mut self, address_space: usize, pointer: usize) {
let label = pointer / CHUNK;
self.touched_labels.touch(address_space, label);
}

@@ -272,7 +272,7 @@ where
let (initial_row, final_row) = row.split_at_mut(width);
*initial_row.borrow_mut() = PersistentBoundaryCols {
expand_direction: Val::<SC>::ONE,
address_space: touched_label.address_space,
address_space: Val::<SC>::from_canonical_usize(touched_label.address_space),
leaf_label: Val::<SC>::from_canonical_usize(touched_label.label),
values: touched_label.init_values,
hash: touched_label.init_hash,
@@ -285,7 +285,7 @@ where

*final_row.borrow_mut() = PersistentBoundaryCols {
expand_direction: Val::<SC>::NEG_ONE,
address_space: touched_label.address_space,
address_space: Val::<SC>::from_canonical_usize(touched_label.address_space),
leaf_label: Val::<SC>::from_canonical_usize(touched_label.label),
values: touched_label.final_values,
hash: touched_label.final_hash,
15 changes: 4 additions & 11 deletions crates/vm/src/system/memory/tree/public_values.rs
Original file line number Diff line number Diff line change
@@ -111,11 +111,8 @@ pub fn extract_public_values<const CHUNK: usize, F: PrimeField32>(
final_memory: &Equipartition<F, CHUNK>,
) -> Vec<F> {
// All (addr, value) pairs in the public value address space.
let f_as_start =
F::from_canonical_usize(PUBLIC_VALUES_ADDRESS_SPACE_OFFSET + memory_dimensions.as_offset);
let f_as_end = F::from_canonical_usize(
PUBLIC_VALUES_ADDRESS_SPACE_OFFSET + memory_dimensions.as_offset + 1,
);
let f_as_start = PUBLIC_VALUES_ADDRESS_SPACE_OFFSET + memory_dimensions.as_offset;
let f_as_end = PUBLIC_VALUES_ADDRESS_SPACE_OFFSET + memory_dimensions.as_offset + 1;

let used_pvs: Vec<_> = final_memory
.range((f_as_start, 0)..(f_as_end, 0))
@@ -161,13 +158,9 @@ mod tests {
vm_config.memory_config.as_height = 4;
vm_config.memory_config.pointer_max_bits = 5;
let memory_dimensions = vm_config.memory_config.memory_dimensions();
let pv_as = F::from_canonical_usize(
PUBLIC_VALUES_ADDRESS_SPACE_OFFSET + memory_dimensions.as_offset,
);
let pv_as = PUBLIC_VALUES_ADDRESS_SPACE_OFFSET + memory_dimensions.as_offset;
let num_public_values = 16;
let memory: MemoryImage<F> = [((pv_as, F::from_canonical_u32(15)), F::ONE)]
.into_iter()
.collect();
let memory: MemoryImage<F> = [((pv_as, 15), F::ONE)].into_iter().collect();
let mut expected_pvs = F::zero_vec(num_public_values);
expected_pvs[15] = F::ONE;

Loading

0 comments on commit b4e5a78

Please sign in to comment.