openvm_circuit/arch/testing/memory/
mod.rs

1use std::collections::HashMap;
2
3use air::{MemoryDummyAir, MemoryDummyChip};
4use openvm_stark_backend::p3_field::{Field, PrimeField32};
5use rand::Rng;
6
7use crate::system::memory::{online::TracingMemory, MemoryController};
8
9pub mod air;
10
11#[cfg(feature = "cuda")]
12mod cuda;
13#[cfg(feature = "cuda")]
14pub use cuda::*;
15
16/// A dummy testing chip that will add unconstrained messages into the [MemoryBus].
17/// Stores a log of raw messages to send/receive to the [MemoryBus].
18///
19/// It will create a [air::MemoryDummyAir] to add messages to MemoryBus.
20pub struct MemoryTester<F: Field> {
21    /// Map from `block_size` to [MemoryDummyChip] of that block size
22    pub chip_for_block: HashMap<usize, MemoryDummyChip<F>>,
23    pub memory: TracingMemory,
24    pub(super) controller: MemoryController<F>,
25}
26
27impl<F: PrimeField32> MemoryTester<F> {
28    pub fn new(controller: MemoryController<F>, memory: TracingMemory) -> Self {
29        let bus = controller.memory_bus;
30        let mut chip_for_block = HashMap::new();
31        for log_block_size in 0..6 {
32            let block_size = 1 << log_block_size;
33            let chip = MemoryDummyChip::new(MemoryDummyAir::new(bus, block_size));
34            chip_for_block.insert(block_size, chip);
35        }
36        Self {
37            chip_for_block,
38            memory,
39            controller,
40        }
41    }
42
43    pub fn read<const N: usize>(&mut self, addr_space: usize, ptr: usize) -> [F; N] {
44        let memory = &mut self.memory;
45        let t = memory.timestamp();
46        // TODO: this could be improved if we added a TracingMemory::get_f function
47        let (t_prev, data) = if addr_space <= 3 {
48            let (t_prev, data) = unsafe { memory.read::<u8, N, 4>(addr_space as u32, ptr as u32) };
49            (t_prev, data.map(F::from_canonical_u8))
50        } else {
51            unsafe { memory.read::<F, N, 1>(addr_space as u32, ptr as u32) }
52        };
53        self.chip_for_block.get_mut(&N).unwrap().receive(
54            addr_space as u32,
55            ptr as u32,
56            &data,
57            t_prev,
58        );
59        self.chip_for_block
60            .get_mut(&N)
61            .unwrap()
62            .send(addr_space as u32, ptr as u32, &data, t);
63
64        data
65    }
66
67    pub fn write<const N: usize>(&mut self, addr_space: usize, ptr: usize, data: [F; N]) {
68        let memory = &mut self.memory;
69        let t = memory.timestamp();
70        // TODO: this could be improved if we added a TracingMemory::write_f function
71        let (t_prev, data_prev) = if addr_space <= 3 {
72            let (t_prev, data_prev) = unsafe {
73                memory.write::<u8, N, 4>(
74                    addr_space as u32,
75                    ptr as u32,
76                    data.map(|x| x.as_canonical_u32() as u8),
77                )
78            };
79            (t_prev, data_prev.map(F::from_canonical_u8))
80        } else {
81            unsafe { memory.write::<F, N, 1>(addr_space as u32, ptr as u32, data) }
82        };
83        self.chip_for_block.get_mut(&N).unwrap().receive(
84            addr_space as u32,
85            ptr as u32,
86            &data_prev,
87            t_prev,
88        );
89        self.chip_for_block
90            .get_mut(&N)
91            .unwrap()
92            .send(addr_space as u32, ptr as u32, &data, t);
93    }
94}
95
96pub fn gen_pointer<R>(rng: &mut R, len: usize) -> usize
97where
98    R: Rng + ?Sized,
99{
100    const MAX_MEMORY: usize = 1 << 29;
101    rng.gen_range(0..MAX_MEMORY - len) / len * len
102}
openvm_circuit/arch/testing/memory/mod.rs

openvm_circuit/arch/testing/memory/
mod.rs
