openvm_circuit/system/memory/
online.rs

1use std::{array::from_fn, fmt::Debug};
2
3use getset::Getters;
4use openvm_instructions::exe::SparseMemoryImage;
5use openvm_stark_backend::{
6    p3_field::{Field, PrimeField32},
7    p3_maybe_rayon::prelude::*,
8};
9use tracing::instrument;
10
11use crate::{
12    arch::{AddressSpaceHostConfig, AddressSpaceHostLayout, MemoryConfig, DEFAULT_BLOCK_SIZE},
13    system::{memory::TimestampedValues, TouchedMemory},
14};
15
16mod basic;
17#[cfg(any(unix, windows))]
18mod memmap;
19mod paged_vec;
20
21#[cfg(not(any(unix, windows)))]
22pub use basic::*;
23#[cfg(any(unix, windows))]
24pub use memmap::*;
25pub use paged_vec::PagedVec;
26
27#[cfg(all(any(unix, windows), not(feature = "basic-memory")))]
28pub type MemoryBackend = memmap::MmapMemory;
29#[cfg(any(not(any(unix, windows)), feature = "basic-memory"))]
30pub type MemoryBackend = basic::BasicMemory;
31
32pub const INITIAL_TIMESTAMP: u32 = 0;
33/// Default mmap page size. Change this if using THB.
34pub const PAGE_SIZE: usize = 4096;
35
36/// (address_space, pointer)
37pub type Address = (u32, u32);
38
39/// API for any memory implementation that allocates a contiguous region of memory.
40pub trait LinearMemory {
41    /// Create instance of `Self` with `size` bytes.
42    fn new(size: usize) -> Self;
43    /// Allocated size of the memory in bytes.
44    fn size(&self) -> usize;
45    /// Returns the entire memory as a raw byte slice.
46    fn as_slice(&self) -> &[u8];
47    /// Returns the entire memory as a raw byte slice.
48    fn as_mut_slice(&mut self) -> &mut [u8];
49    /// Fill the memory with zeros.
50    fn fill_zero(&mut self) {
51        self.as_mut_slice().fill(0);
52    }
53    /// Read `BLOCK` from `self` at `from` address without moving it.
54    ///
55    /// Panics or segfaults if `from..from + size_of::<BLOCK>()` is out of bounds.
56    ///
57    /// # Safety
58    /// - `BLOCK` should be "plain old data" (see [`Pod`](https://docs.rs/bytemuck/latest/bytemuck/trait.Pod.html)).
59    ///   We do not add a trait bound due to Plonky3 types not implementing the trait.
60    /// - See [`core::ptr::read`] for similar considerations.
61    /// - Memory at `from` must be properly aligned for `BLOCK`. Use [`Self::read_unaligned`] if
62    ///   alignment is not guaranteed.
63    unsafe fn read<BLOCK: Copy>(&self, from: usize) -> BLOCK;
64    /// Read `BLOCK` from `self` at `from` address without moving it.
65    /// Same as [`Self::read`] except that it does not require alignment.
66    ///
67    /// Panics or segfaults if `from..from + size_of::<BLOCK>()` is out of bounds.
68    ///
69    /// # Safety
70    /// - `BLOCK` should be "plain old data" (see [`Pod`](https://docs.rs/bytemuck/latest/bytemuck/trait.Pod.html)).
71    ///   We do not add a trait bound due to Plonky3 types not implementing the trait.
72    /// - See [`core::ptr::read`] for similar considerations.
73    unsafe fn read_unaligned<BLOCK: Copy>(&self, from: usize) -> BLOCK;
74    /// Write `BLOCK` to `self` at `start` address without reading the old value. Does not drop
75    /// `values`. Semantically, `values` is moved into the location pointed to by `start`.
76    ///
77    /// Panics or segfaults if `start..start + size_of::<BLOCK>()` is out of bounds.
78    ///
79    /// # Safety
80    /// - See [`core::ptr::write`] for similar considerations.
81    /// - Memory at `start` must be properly aligned for `BLOCK`. Use [`Self::write_unaligned`] if
82    ///   alignment is not guaranteed.
83    unsafe fn write<BLOCK: Copy>(&mut self, start: usize, values: BLOCK);
84    /// Write `BLOCK` to `self` at `start` address without reading the old value. Does not drop
85    /// `values`. Semantically, `values` is moved into the location pointed to by `start`.
86    /// Same as [`Self::write`] but without alignment requirement.
87    ///
88    /// Panics or segfaults if `start..start + size_of::<BLOCK>()` is out of bounds.
89    ///
90    /// # Safety
91    /// - See [`core::ptr::write`] for similar considerations.
92    unsafe fn write_unaligned<BLOCK: Copy>(&mut self, start: usize, values: BLOCK);
93    /// Swaps `values` with memory at `start..start + size_of::<BLOCK>()`.
94    ///
95    /// Panics or segfaults if `start..start + size_of::<BLOCK>()` is out of bounds.
96    ///
97    /// # Safety
98    /// - `BLOCK` should be "plain old data" (see [`Pod`](https://docs.rs/bytemuck/latest/bytemuck/trait.Pod.html)).
99    ///   We do not add a trait bound due to Plonky3 types not implementing the trait.
100    /// - Memory at `start` must be properly aligned for `BLOCK`.
101    /// - The data in `values` should not overlap with memory in `self`.
102    unsafe fn swap<BLOCK: Copy>(&mut self, start: usize, values: &mut BLOCK);
103    /// Copies `data` into memory at `to` address.
104    ///
105    /// Panics or segfaults if `to..to + size_of_val(data)` is out of bounds.
106    ///
107    /// # Safety
108    /// - `T` should be "plain old data" (see [`Pod`](https://docs.rs/bytemuck/latest/bytemuck/trait.Pod.html)).
109    ///   We do not add a trait bound due to Plonky3 types not implementing the trait.
110    /// - The underlying memory of `data` should not overlap with `self`.
111    /// - The starting pointer of `self` should be aligned to `T`.
112    /// - The memory pointer at `to` should be aligned to `T`.
113    unsafe fn copy_nonoverlapping<T: Copy>(&mut self, to: usize, data: &[T]);
114    /// Returns a slice `&[T]` for the memory region `start..start + len`.
115    ///
116    /// Panics or segfaults if `start..start + len * size_of::<T>()` is out of bounds.
117    ///
118    /// # Safety
119    /// - `T` should be "plain old data" (see [`Pod`](https://docs.rs/bytemuck/latest/bytemuck/trait.Pod.html)).
120    ///   We do not add a trait bound due to Plonky3 types not implementing the trait.
121    /// - Memory at `start` must be properly aligned for `T`.
122    unsafe fn get_aligned_slice<T: Copy>(&self, start: usize, len: usize) -> &[T];
123}
124
125/// Map from address space to linear memory.
126/// The underlying memory is typeless, stored as raw bytes, but usage implicitly assumes that each
127/// address space has memory cells of a fixed type (e.g., `u8, F`). We do not use a typemap for
128/// performance reasons, and it is up to the user to enforce types. Needless to say, this is a very
129/// `unsafe` API.
130#[derive(Debug, Clone)]
131#[repr(C)]
132pub struct AddressMap<M: LinearMemory = MemoryBackend> {
133    /// Underlying memory data.
134    pub mem: Vec<M>,
135    /// Host configuration for each address space.
136    pub config: Vec<AddressSpaceHostConfig>,
137}
138
139impl Default for AddressMap {
140    fn default() -> Self {
141        Self::from_mem_config(&MemoryConfig::default())
142    }
143}
144
145impl<M: LinearMemory> AddressMap<M> {
146    pub fn new(config: Vec<AddressSpaceHostConfig>) -> Self {
147        assert_eq!(config[0].num_cells, 0, "Address space 0 must have 0 cells");
148        let mem = config
149            .iter()
150            .map(|config| M::new(config.num_cells.checked_mul(config.layout.size()).unwrap()))
151            .collect();
152        Self { mem, config }
153    }
154
155    pub fn from_mem_config(mem_config: &MemoryConfig) -> Self {
156        Self::new(mem_config.addr_spaces.clone())
157    }
158
159    #[inline(always)]
160    pub fn get_memory(&self) -> &Vec<M> {
161        &self.mem
162    }
163
164    #[inline(always)]
165    pub fn get_memory_mut(&mut self) -> &mut Vec<M> {
166        &mut self.mem
167    }
168
169    /// Fill each address space memory with zeros. Does not change the config.
170    pub fn fill_zero(&mut self) {
171        for mem in &mut self.mem {
172            mem.fill_zero();
173        }
174    }
175
176    /// # Safety
177    /// - Assumes `addr_space` is within the configured memory and not out of bounds
178    pub unsafe fn get_f<F: PrimeField32>(&self, addr_space: u32, ptr: u32) -> F {
179        let layout = &self.config.get_unchecked(addr_space as usize).layout;
180        let start = ptr as usize * layout.size();
181        let bytes = self.get_u8_slice(addr_space, start, layout.size());
182        layout.to_field(bytes)
183    }
184
185    /// # Safety
186    /// - `T` **must** be the correct type for a single memory cell for `addr_space`
187    /// - Assumes `addr_space` is within the configured memory and not out of bounds
188    pub unsafe fn get<T: Copy>(&self, (addr_space, ptr): Address) -> T {
189        debug_assert_eq!(
190            size_of::<T>(),
191            self.config[addr_space as usize].layout.size()
192        );
193        // SAFETY:
194        // - alignment is automatic since we multiply by `size_of::<T>()`
195        self.mem
196            .get_unchecked(addr_space as usize)
197            .read((ptr as usize) * size_of::<T>())
198    }
199
200    /// Panics or segfaults if `ptr..ptr + len` is out of bounds
201    ///
202    /// # Safety
203    /// - `T` **must** be the correct type for a single memory cell for `addr_space`
204    /// - Assumes `addr_space` is within the configured memory and not out of bounds
205    pub unsafe fn get_slice<T: Copy + Debug>(
206        &self,
207        (addr_space, ptr): Address,
208        len: usize,
209    ) -> &[T] {
210        debug_assert_eq!(
211            size_of::<T>(),
212            self.config[addr_space as usize].layout.size()
213        );
214        let start = (ptr as usize) * size_of::<T>();
215        let mem = self.mem.get_unchecked(addr_space as usize);
216        // SAFETY:
217        // - alignment is automatic since we multiply by `size_of::<T>()`
218        mem.get_aligned_slice(start, len)
219    }
220
221    /// Reads the slice at **byte** addresses `start..start + len` from address space `addr_space`
222    /// linear memory. Panics or segfaults if `start..start + len` is out of bounds
223    ///
224    /// # Safety
225    /// - Assumes `addr_space` is within the configured memory and not out of bounds
226    pub unsafe fn get_u8_slice(&self, addr_space: u32, start: usize, len: usize) -> &[u8] {
227        let mem = self.mem.get_unchecked(addr_space as usize);
228        mem.get_aligned_slice(start, len)
229    }
230
231    /// Copies `data` into the memory at `(addr_space, ptr)`.
232    ///
233    /// Panics or segfaults if `ptr + size_of_val(data)` is out of bounds.
234    ///
235    /// # Safety
236    /// - `T` **must** be the correct type for a single memory cell for `addr_space`
237    /// - The linear memory in `addr_space` is aligned to `T`.
238    pub unsafe fn copy_slice_nonoverlapping<T: Copy>(
239        &mut self,
240        (addr_space, ptr): Address,
241        data: &[T],
242    ) {
243        let start = (ptr as usize) * size_of::<T>();
244        // SAFETY:
245        // - Linear memory is aligned to `T` and `start` is multiple of `size_of::<T>()` so
246        //   alignment is satisfied.
247        // - `data` and `self.mem` are non-overlapping
248        self.mem
249            .get_unchecked_mut(addr_space as usize)
250            .copy_nonoverlapping(start, data);
251    }
252
253    /// # Safety
254    /// - `T` **must** be the correct type for a single memory cell for `addr_space`
255    /// - Assumes `addr_space` is within the configured memory and not out of bounds
256    pub fn set_from_sparse(&mut self, sparse_map: &SparseMemoryImage) {
257        for (&(addr_space, index), &data_byte) in sparse_map.iter() {
258            // SAFETY:
259            // - safety assumptions in function doc comments
260            unsafe {
261                self.mem
262                    .get_unchecked_mut(addr_space as usize)
263                    .write_unaligned(index as usize, data_byte);
264            }
265        }
266    }
267}
268
269/// API for guest memory conforming to OpenVM ISA
270// @dev Note we don't make this a trait because phantom executors currently need a concrete type for
271// guest memory
272#[derive(Debug, Clone)]
273#[repr(C)]
274pub struct GuestMemory {
275    pub memory: AddressMap,
276}
277
278impl GuestMemory {
279    pub fn new(addr: AddressMap) -> Self {
280        Self { memory: addr }
281    }
282
283    /// Returns `[pointer:BLOCK_SIZE]_{address_space}`
284    ///
285    /// # Safety
286    /// The type `T` must be stack-allocated `repr(C)` or `repr(transparent)`,
287    /// and it must be the exact type used to represent a single memory cell in
288    /// address space `address_space`. For standard usage,
289    /// `T` is either `u8` or `F` where `F` is the base field of the ZK backend.
290    #[inline(always)]
291    pub unsafe fn read<T, const BLOCK_SIZE: usize>(
292        &self,
293        addr_space: u32,
294        ptr: u32,
295    ) -> [T; BLOCK_SIZE]
296    where
297        T: Copy + Debug,
298    {
299        self.debug_assert_cell_type::<T>(addr_space);
300        // SAFETY:
301        // - `T` should be "plain old data"
302        // - alignment for `[T; BLOCK_SIZE]` is automatic since we multiply by `size_of::<T>()`
303        self.memory
304            .get_memory()
305            .get_unchecked(addr_space as usize)
306            .read((ptr as usize) * size_of::<T>())
307    }
308
309    /// Writes `values` to `[pointer:BLOCK_SIZE]_{address_space}`
310    ///
311    /// # Safety
312    /// See [`GuestMemory::read`].
313    #[inline(always)]
314    pub unsafe fn write<T, const BLOCK_SIZE: usize>(
315        &mut self,
316        addr_space: u32,
317        ptr: u32,
318        values: [T; BLOCK_SIZE],
319    ) where
320        T: Copy + Debug,
321    {
322        self.debug_assert_cell_type::<T>(addr_space);
323        // SAFETY:
324        // - alignment for `[T; BLOCK_SIZE]` is automatic since we multiply by `size_of::<T>()`
325        self.memory
326            .get_memory_mut()
327            .get_unchecked_mut(addr_space as usize)
328            .write((ptr as usize) * size_of::<T>(), values);
329    }
330
331    /// Swaps `values` with `[pointer:BLOCK_SIZE]_{address_space}`.
332    ///
333    /// # Safety
334    /// See [`GuestMemory::read`] and [`LinearMemory::swap`].
335    #[inline(always)]
336    pub unsafe fn swap<T, const BLOCK_SIZE: usize>(
337        &mut self,
338        addr_space: u32,
339        ptr: u32,
340        values: &mut [T; BLOCK_SIZE],
341    ) where
342        T: Copy + Debug,
343    {
344        self.debug_assert_cell_type::<T>(addr_space);
345        // SAFETY:
346        // - alignment for `[T; BLOCK_SIZE]` is automatic since we multiply by `size_of::<T>()`
347        self.memory
348            .get_memory_mut()
349            .get_unchecked_mut(addr_space as usize)
350            .swap((ptr as usize) * size_of::<T>(), values);
351    }
352
353    #[inline(always)]
354    #[allow(clippy::missing_safety_doc)]
355    pub unsafe fn get_slice<T: Copy + Debug>(&self, addr_space: u32, ptr: u32, len: usize) -> &[T] {
356        self.memory.get_slice((addr_space, ptr), len)
357    }
358
359    #[inline(always)]
360    fn debug_assert_cell_type<T>(&self, addr_space: u32) {
361        debug_assert_eq!(
362            size_of::<T>(),
363            self.memory.config[addr_space as usize].layout.size()
364        );
365    }
366}
367
368/// Online memory that stores additional information for trace generation purposes.
369/// In particular, keeps track of timestamp.
370#[derive(Getters)]
371pub struct TracingMemory {
372    pub timestamp: u32,
373    /// The underlying data memory, with memory cells typed by address space: see [AddressMap].
374    #[getset(get = "pub")]
375    pub data: GuestMemory,
376    /// Maps `(addr_space, ptr / DEFAULT_BLOCK_SIZE)` to the latest access timestamp.
377    /// A value of 0 means the 4-cell touched-memory slot has never been accessed.
378    pub(super) meta: Vec<PagedVec<u32, PAGE_SIZE>>,
379}
380
381impl TracingMemory {
382    pub fn new(mem_config: &MemoryConfig) -> Self {
383        let image = GuestMemory::new(AddressMap::from_mem_config(mem_config));
384        Self::from_image(image)
385    }
386
387    /// Constructor from pre-existing memory image.
388    pub fn from_image(image: GuestMemory) -> Self {
389        let meta = image
390            .memory
391            .config
392            .iter()
393            .map(|config| PagedVec::new(config.num_cells.div_ceil(DEFAULT_BLOCK_SIZE)))
394            .collect();
395        Self {
396            data: image,
397            meta,
398            timestamp: INITIAL_TIMESTAMP + 1,
399        }
400    }
401
402    #[inline(always)]
403    fn assert_valid_access(&self, block_size: usize, addr_space: u32, ptr: u32) {
404        debug_assert_ne!(addr_space, 0);
405        debug_assert!(block_size.is_power_of_two());
406        debug_assert_eq!(
407            block_size, DEFAULT_BLOCK_SIZE,
408            "TracingMemory only supports {DEFAULT_BLOCK_SIZE}-cell accesses; got {block_size}"
409        );
410        assert_eq!(
411            ptr % block_size as u32,
412            0,
413            "pointer={ptr} not aligned to block_size {block_size}"
414        );
415    }
416
417    /// Returns the previous access timestamp and updates the metadata slot.
418    /// Block size is always `DEFAULT_BLOCK_SIZE`, so this is a single-slot read/write.
419    #[inline(always)]
420    fn prev_access_time(&mut self, address_space: usize, pointer: usize) -> u32 {
421        let idx = pointer / DEFAULT_BLOCK_SIZE;
422        // SAFETY: address_space is validated during instruction decoding
423        let meta_page = unsafe { self.meta.get_unchecked_mut(address_space) };
424        let prev = meta_page.get(idx);
425        meta_page.set(idx, self.timestamp);
426        prev
427    }
428
429    /// Atomic read operation which increments the timestamp by 1.
430    /// Returns `(t_prev, [pointer:BLOCK_SIZE]_{address_space})`.
431    ///
432    /// # Safety
433    /// - `T` must be `repr(C)` or `repr(transparent)` and match the cell type for `address_space`.
434    /// - `address_space` must be valid.
435    /// - `BLOCK_SIZE` is measured in memory cells and is tracked in fixed `DEFAULT_BLOCK_SIZE`
436    ///   touched-memory slots.
437    #[inline(always)]
438    pub unsafe fn read<T, const BLOCK_SIZE: usize>(
439        &mut self,
440        address_space: u32,
441        pointer: u32,
442    ) -> (u32, [T; BLOCK_SIZE])
443    where
444        T: Copy + Debug,
445    {
446        self.assert_valid_access(BLOCK_SIZE, address_space, pointer);
447        let values = self.data.read(address_space, pointer);
448        let t_prev = self.prev_access_time(address_space as usize, pointer as usize);
449        self.timestamp += 1;
450
451        (t_prev, values)
452    }
453
454    /// Atomic write operation. Returns `(t_prev, values_prev)`.
455    ///
456    /// # Safety
457    /// - `T` must be `repr(C)` or `repr(transparent)` and match the cell type for `address_space`.
458    /// - `address_space` must be valid.
459    /// - `BLOCK_SIZE` is measured in memory cells and is tracked in fixed `DEFAULT_BLOCK_SIZE`
460    ///   touched-memory slots.
461    #[inline(always)]
462    pub unsafe fn write<T, const BLOCK_SIZE: usize>(
463        &mut self,
464        address_space: u32,
465        pointer: u32,
466        values: [T; BLOCK_SIZE],
467    ) -> (u32, [T; BLOCK_SIZE])
468    where
469        T: Copy + Debug,
470    {
471        self.assert_valid_access(BLOCK_SIZE, address_space, pointer);
472        let values_prev = self.data.read(address_space, pointer);
473        let t_prev = self.prev_access_time(address_space as usize, pointer as usize);
474        self.data.write(address_space, pointer, values);
475        self.timestamp += 1;
476
477        (t_prev, values_prev)
478    }
479
480    pub fn increment_timestamp(&mut self) {
481        self.timestamp += 1;
482    }
483
484    pub fn increment_timestamp_by(&mut self, amount: u32) {
485        self.timestamp += amount;
486    }
487
488    pub fn timestamp(&self) -> u32 {
489        self.timestamp
490    }
491
492    /// Finalize the boundary and merkle chips.
493    #[instrument(name = "memory_finalize", skip_all)]
494    pub fn finalize<F: Field>(&mut self) -> TouchedMemory<F> {
495        self.touched_blocks_to_equipartition::<F>(self.touched_blocks())
496    }
497
498    /// Returns the list of all touched blocks (address, timestamp), sorted by address.
499    fn touched_blocks(&self) -> Vec<(Address, u32)> {
500        assert_eq!(self.meta.len(), self.data.memory.config.len());
501        let mut touched_blocks: Vec<_> = self
502            .meta
503            .par_iter()
504            .enumerate()
505            .flat_map_iter(|(addr_space, meta_page)| {
506                meta_page
507                    .par_iter()
508                    .filter_map(move |(idx, timestamp)| {
509                        if timestamp > INITIAL_TIMESTAMP {
510                            let ptr = idx as u32 * DEFAULT_BLOCK_SIZE as u32;
511                            Some(((addr_space as u32, ptr), timestamp))
512                        } else {
513                            None
514                        }
515                    })
516                    .collect::<Vec<_>>()
517            })
518            .collect();
519        // This sort may not be strictly necessary, but it makes the finalize path independent of
520        // Rayon ordering.
521        touched_blocks.sort_unstable_by_key(|(addr, _)| *addr);
522        touched_blocks
523    }
524
525    /// Returns the fixed 4-byte touched memory equipartition.
526    fn touched_blocks_to_equipartition<F: Field>(
527        &self,
528        touched_blocks: Vec<((u32, u32), u32)>,
529    ) -> TouchedMemory<F> {
530        debug_assert!(touched_blocks.is_sorted_by_key(|(addr, _)| addr));
531        touched_blocks
532            .into_par_iter()
533            .map(|((addr_space, ptr), timestamp)| {
534                let addr_space_config = &self.data.memory.config[addr_space as usize];
535                let cell_size = addr_space_config.layout.size();
536                let values = from_fn(|i| unsafe {
537                    addr_space_config
538                        .layout
539                        .to_field(self.data.memory.get_u8_slice(
540                            addr_space,
541                            (ptr as usize + i) * cell_size,
542                            cell_size,
543                        ))
544                });
545                ((addr_space, ptr), TimestampedValues { timestamp, values })
546            })
547            .collect()
548    }
549}
openvm_circuit/system/memory/online.rs

openvm_circuit/system/memory/
online.rs
