continuation of multi-shard + ram bus

.github/workflows/lints.yml

@@ -66,3 +66,5 @@ jobs:
         run: taplo --version || cargo install taplo-cli
       - name: Run taplo
         run: taplo fmt --check --diff
+      - name: Ensure Cargo.lock not modified by build
+        run: git diff --exit-code Cargo.lock

Cargo.toml

@@ -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" }

ceno_cli/src/commands/common_args/ceno.rs

@@ -78,6 +78,14 @@ pub struct CenoOptions {
     #[arg(long)]
     pub out_vk: Option<PathBuf>,
 
+    /// shard id
+    #[arg(long, default_value = "0")]
+    shard_id: u32,
+
+    /// number of total shards.
+    #[arg(long, default_value = "1")]
+    max_num_shards: u32,
+
     /// Profiling granularity.
     /// Setting any value restricts logs to profiling information
     #[arg(long)]
@@ -337,6 +345,7 @@ fn run_elf_inner<
         std::fs::read(elf_path).context(format!("failed to read {}", elf_path.display()))?;
     let program = Program::load_elf(&elf_bytes, u32::MAX).context("failed to load elf")?;
     print_cargo_message("Loaded", format_args!("{}", elf_path.display()));
+    let shards = Shards::new(options.shard_id as usize, options.max_num_shards as usize);
 
     let public_io = options
         .read_public_io()
@@ -385,6 +394,7 @@ fn run_elf_inner<
         create_prover(backend.clone()),
         program,
         platform,
+        shards,
         &hints,
         &public_io,
         options.max_steps,

ceno_emul/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/src/chunked_vec.rs

@@ -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/lib.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/tracer.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,7 +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>,
+    latest_accesses: FxHashMap<WordAddr, Cycle>,
+    next_accesses: ChunkedVec<NextAccessPair>,
 }
 
 impl Default for Tracer {
@@ -362,7 +365,8 @@ impl Tracer {
                 cycle: Self::SUBCYCLES_PER_INSN,
                 ..StepRecord::default()
             },
-            latest_accesses: HashMap::new(),
+            latest_accesses: FxHashMap::default(),
+            next_accesses: ChunkedVec::new(1 << 20),
         }
     }
 
@@ -471,16 +475,24 @@ impl Tracer {
     /// - Record the current instruction as the origin of the latest access.
     /// - Accesses within the same instruction are distinguished by `subcycle ∈ [0, 3]`.
     pub fn track_access(&mut self, addr: WordAddr, subcycle: Cycle) -> Cycle {
-        self.latest_accesses
-            .insert(addr, self.record.cycle + subcycle)
-            .unwrap_or(0)
+        let cur_cycle = self.record.cycle + subcycle;
+        let prev_cycle = self.latest_accesses.insert(addr, cur_cycle).unwrap_or(0);
+        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) -> ChunkedVec<NextAccessPair> {
+        self.next_accesses
+    }
+
     /// Return the cycle of the pending instruction (after the last completed step).
     pub fn cycle(&self) -> Cycle {
         self.record.cycle

ceno_emul/src/vm_state.rs

@@ -68,6 +68,10 @@ impl VMState {
         &self.tracer
     }
 
+    pub fn take_tracer(self) -> Tracer {
+        self.tracer
+    }
+
     pub fn platform(&self) -> &Platform {
         &self.platform
     }

ceno_emul/tests/test_vm_trace.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_zkvm/Cargo.toml

@@ -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