optimise tracer performance

Author: hero78119

Cargo.lock

@@ -57,9 +57,17 @@ rand_chacha = { version = "0.3", features = ["serde1"] }
 rand_core = "0.6"
 rayon = "1.10"
 rkyv = { version = "0.8", features = ["pointer_width_32"] }
+rustc-hash = "2.0.0"
 secp = "0.4.1"
 serde = { version = "1.0", features = ["derive", "rc"] }
 serde_json = "1.0"
+smallvec = { version = "1.13.2", features = [
+  "const_generics",
+  "const_new",
+  "serde",
+  "union",
+  "write",
+] }
 strum = "0.26"
 strum_macros = "0.26"
 substrate-bn = { version = "0.6.0" }

Cargo.toml

@@ -19,9 +19,12 @@ itertools.workspace = true
 multilinear_extensions.workspace = true
 num-derive.workspace = true
 num-traits.workspace = true
+rayon.workspace = true
 rrs_lib = { package = "rrs-succinct", version = "0.1.0" }
+rustc-hash.workspace = true
 secp.workspace = true
 serde.workspace = true
+smallvec.workspace = true
 strum.workspace = true
 strum_macros.workspace = true
 substrate-bn.workspace = true

ceno_emul/Cargo.toml

@@ -0,0 +1,89 @@
+use rayon::iter::{IntoParallelIterator, ParallelIterator};
+use std::ops::{Index, IndexMut};
+
+/// a chunked vector that grows in fixed-size chunks.
+#[derive(Default, Debug, Clone)]
+pub struct ChunkedVec<T> {
+    chunks: Vec<Vec<T>>,
+    chunk_size: usize,
+    len: usize,
+}
+
+impl<T: Default + Send> ChunkedVec<T> {
+    /// create a new ChunkedVec with a given chunk size.
+    pub fn new(chunk_size: usize) -> Self {
+        assert!(chunk_size > 0, "chunk_size must be > 0");
+        Self {
+            chunks: Vec::new(),
+            chunk_size,
+            len: 0,
+        }
+    }
+
+    /// get the current number of elements.
+    pub fn len(&self) -> usize {
+        self.len
+    }
+
+    /// returns true if the vector is empty.
+    pub fn is_empty(&self) -> bool {
+        self.len == 0
+    }
+
+    /// access element by index (immutable).
+    pub fn get(&self, index: usize) -> Option<&T> {
+        if index >= self.len {
+            return None;
+        }
+        let chunk_idx = index / self.chunk_size;
+        let within_idx = index % self.chunk_size;
+        self.chunks.get(chunk_idx)?.get(within_idx)
+    }
+
+    /// access element by index (mutable).
+    /// get mutable reference to element at index, auto-creating chunks as needed
+    pub fn get_or_create(&mut self, index: usize) -> &mut T {
+        let chunk_idx = index / self.chunk_size;
+        let within_idx = index % self.chunk_size;
+
+        // Ensure enough chunks exist
+        if chunk_idx >= self.chunks.len() {
+            let to_create = chunk_idx + 1 - self.chunks.len();
+
+            // Use rayon to create all missing chunks in parallel
+            let mut new_chunks: Vec<Vec<T>> = (0..to_create)
+                .map(|_| {
+                    (0..self.chunk_size)
+                        .into_par_iter()
+                        .map(|_| Default::default())
+                        .collect::<Vec<_>>()
+                })
+                .collect();
+
+            self.chunks.append(&mut new_chunks);
+        }
+
+        let chunk = &mut self.chunks[chunk_idx];
+
+        // Update the overall length
+        if index >= self.len {
+            self.len = index + 1;
+        }
+
+        &mut chunk[within_idx]
+    }
+}
+
+impl<T: Default + Send> Index<usize> for ChunkedVec<T> {
+    type Output = T;
+
+    fn index(&self, index: usize) -> &Self::Output {
+        self.get(index).expect("index out of bounds")
+    }
+}
+
+impl<T: Default + Send> IndexMut<usize> for ChunkedVec<T> {
+    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
+        self.get_or_create(index)
+    }
+}

ceno_emul/src/chunked_vec.rs

@@ -7,7 +7,7 @@ mod platform;
 pub use platform::{CENO_PLATFORM, Platform};
 
 mod tracer;
-pub use tracer::{Change, MemOp, ReadOp, StepRecord, Tracer, WriteOp};
+pub use tracer::{Change, MemOp, NextAccessPair, ReadOp, StepRecord, Tracer, WriteOp};
 
 mod vm_state;
 pub use vm_state::VMState;
@@ -44,4 +44,6 @@ pub mod utils;
 
 pub mod test_utils;
 
+mod chunked_vec;
+pub use chunked_vec::ChunkedVec as NextCycleAccess;
 pub mod host_utils;

ceno_emul/src/lib.rs

@@ -1,13 +1,13 @@
-use std::{
-    collections::{BTreeMap, HashMap},
-    fmt, mem,
-};
+use rustc_hash::FxHashMap;
+use smallvec::SmallVec;
+use std::{collections::BTreeMap, fmt, mem};
 
 use ceno_rt::WORD_SIZE;
 
 use crate::{
     CENO_PLATFORM, InsnKind, Instruction, PC_STEP_SIZE, Platform,
     addr::{ByteAddr, Cycle, RegIdx, Word, WordAddr},
+    chunked_vec::ChunkedVec,
     encode_rv32,
     syscalls::{SyscallEffects, SyscallWitness},
 };
@@ -39,6 +39,8 @@ pub struct StepRecord {
     syscall: Option<SyscallWitness>,
 }
 
+pub type NextAccessPair = SmallVec<[(WordAddr, Cycle); 1]>;
+
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
 pub struct MemOp<T> {
     /// Virtual Memory Address.
@@ -305,8 +307,8 @@ pub struct Tracer {
     // record each section max access address
     // (start_addr -> (start_addr, end_addr, min_access_addr, max_access_addr))
     mmio_min_max_access: Option<BTreeMap<WordAddr, (WordAddr, WordAddr, WordAddr, WordAddr)>>,
-    latest_accesses: HashMap<WordAddr, Cycle>,
-    next_accesses: HashMap<(WordAddr, Cycle), Cycle>,
+    latest_accesses: FxHashMap<WordAddr, Cycle>,
+    next_accesses: ChunkedVec<NextAccessPair>,
 }
 
 impl Default for Tracer {
@@ -363,8 +365,8 @@ impl Tracer {
                 cycle: Self::SUBCYCLES_PER_INSN,
                 ..StepRecord::default()
             },
-            latest_accesses: HashMap::new(),
-            next_accesses: HashMap::new(),
+            latest_accesses: FxHashMap::default(),
+            next_accesses: ChunkedVec::new(1 << 20),
         }
     }
 
@@ -475,17 +477,19 @@ impl Tracer {
     pub fn track_access(&mut self, addr: WordAddr, subcycle: Cycle) -> Cycle {
         let cur_cycle = self.record.cycle + subcycle;
         let prev_cycle = self.latest_accesses.insert(addr, cur_cycle).unwrap_or(0);
-        self.next_accesses.insert((addr, prev_cycle), cur_cycle);
+        self.next_accesses
+            .get_or_create(prev_cycle as usize)
+            .push((addr, cur_cycle));
         prev_cycle
     }
 
     /// Return all the addresses that were accessed and the cycle when they were last accessed.
-    pub fn final_accesses(&self) -> &HashMap<WordAddr, Cycle> {
+    pub fn final_accesses(&self) -> &FxHashMap<WordAddr, Cycle> {
         &self.latest_accesses
     }
 
     /// Return all the addresses that were accessed and the cycle when they were last accessed.
-    pub fn next_accesses(self) -> HashMap<(WordAddr, Cycle), Cycle> {
+    pub fn next_accesses(self) -> ChunkedVec<NextAccessPair> {
         self.next_accesses
     }
 

ceno_emul/src/tracer.rs

@@ -1,9 +1,7 @@
 #![allow(clippy::unusual_byte_groupings)]
 use anyhow::Result;
-use std::{
-    collections::{BTreeMap, HashMap},
-    sync::Arc,
-};
+use rustc_hash::FxHashMap;
+use std::{collections::BTreeMap, sync::Arc};
 
 use ceno_emul::{
     CENO_PLATFORM, Cycle, EmuContext, InsnKind, Instruction, Platform, Program, StepRecord, Tracer,
@@ -111,8 +109,8 @@ fn expected_ops_fibonacci_20() -> Vec<InsnKind> {
 }
 
 /// Reconstruct the last access of each register.
-fn expected_final_accesses_fibonacci_20() -> HashMap<WordAddr, Cycle> {
-    let mut accesses = HashMap::new();
+fn expected_final_accesses_fibonacci_20() -> FxHashMap<WordAddr, Cycle> {
+    let mut accesses = FxHashMap::default();
     let x = |i| WordAddr::from(Platform::register_vma(i));
     const C: Cycle = Tracer::SUBCYCLES_PER_INSN;
 

ceno_emul/tests/test_vm_trace.rs

@@ -34,6 +34,7 @@ witness.workspace = true
 itertools.workspace = true
 ndarray.workspace = true
 prettytable-rs.workspace = true
+rustc-hash.workspace = true
 strum.workspace = true
 strum_macros.workspace = true
 tracing.workspace = true

ceno_zkvm/Cargo.toml

@@ -16,8 +16,9 @@ use crate::{
     tables::{MemFinalRecord, MemInitRecord, ProgramTableCircuit, ProgramTableConfig},
 };
 use ceno_emul::{
-    Addr, ByteAddr, CENO_PLATFORM, Cycle, EmuContext, InsnKind, IterAddresses, Platform, Program,
-    StepRecord, Tracer, VMState, WORD_SIZE, Word, WordAddr, host_utils::read_all_messages,
+    Addr, ByteAddr, CENO_PLATFORM, Cycle, EmuContext, InsnKind, IterAddresses, NextAccessPair,
+    NextCycleAccess, Platform, Program, StepRecord, Tracer, VMState, WORD_SIZE, Word, WordAddr,
+    host_utils::read_all_messages,
 };
 use clap::ValueEnum;
 use either::Either;
@@ -147,7 +148,7 @@ pub struct ShardContext<'a> {
     shards: Shards,
     max_cycle: Cycle,
     // TODO optimize this map as it's super huge
-    addr_future_accesses: Cow<'a, HashMap<(WordAddr, Cycle), Cycle>>,
+    addr_future_accesses: Cow<'a, NextCycleAccess<NextAccessPair>>,
     read_thread_based_record_storage:
         Either<Vec<BTreeMap<WordAddr, RAMRecord>>, &'a mut BTreeMap<WordAddr, RAMRecord>>,
     write_thread_based_record_storage:
@@ -161,7 +162,7 @@ impl<'a> Default for ShardContext<'a> {
         Self {
             shards: Shards::default(),
             max_cycle: Cycle::default(),
-            addr_future_accesses: Cow::Owned(HashMap::new()),
+            addr_future_accesses: Cow::Owned(Default::default()),
             read_thread_based_record_storage: Either::Left(
                 (0..max_threads)
                     .into_par_iter()
@@ -183,7 +184,7 @@ impl<'a> ShardContext<'a> {
     pub fn new(
         shards: Shards,
         executed_instructions: usize,
-        addr_future_accesses: HashMap<(WordAddr, Cycle), Cycle>,
+        addr_future_accesses: NextCycleAccess<NextAccessPair>,
     ) -> Self {
         // current strategy: at least each shard deal with one instruction
         let max_num_shards = shards.max_num_shards.min(executed_instructions);
@@ -329,8 +330,21 @@ impl<'a> ShardContext<'a> {
         }
 
         // check write to external mem bus
-        if let Some(future_touch_cycle) = self.addr_future_accesses.get(&(addr, cycle))
-            && *future_touch_cycle >= self.cur_shard_cycle_range.end as Cycle
+        if let Some(future_touch_cycle) =
+            self.addr_future_accesses
+                .get(cycle as usize)
+                .and_then(|res| {
+                    if res.len() == 1 {
+                        Some(res[0].1)
+                    } else if res.len() > 1 {
+                        res.iter()
+                            .find(|(m_addr, _)| *m_addr == addr)
+                            .map(|(_, cycle)| *cycle)
+                    } else {
+                        None
+                    }
+                })
+            && future_touch_cycle >= self.cur_shard_cycle_range.end as Cycle
             && self.is_current_shard_cycle(cycle)
         {
             let ram_record = self