[WIP] continuation of multi-shard + ram bus

ceno_emul/src/lib.rs

@@ -45,3 +45,4 @@ pub mod utils;
 pub mod test_utils;
 
 pub mod host_utils;
+pub mod shards;

ceno_emul/src/shards.rs

@@ -0,0 +1,14 @@
+pub struct Shards {
+    pub shard_id: usize,
+    pub num_shards: usize,
+}
+
+impl Shards {
+    pub fn new(shard_id: usize, num_shards: usize) -> Self {
+        assert!(shard_id < num_shards);
+        Self {
+            shard_id,
+            num_shards,
+        }
+    }
+}

ceno_emul/src/tracer.rs

@@ -306,6 +306,7 @@ pub struct Tracer {
     // (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>,
 }
 
 impl Default for Tracer {
@@ -363,6 +364,7 @@ impl Tracer {
                 ..StepRecord::default()
             },
             latest_accesses: HashMap::new(),
+            next_accesses: HashMap::new(),
         }
     }
 
@@ -471,16 +473,22 @@ 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.insert((addr, prev_cycle), 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> {
         &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> {
+        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_zkvm/src/bin/e2e.rs

@@ -1,4 +1,4 @@
-use ceno_emul::{IterAddresses, Platform, Program, WORD_SIZE, Word};
+use ceno_emul::{IterAddresses, Platform, Program, WORD_SIZE, Word, shards::Shards};
 use ceno_host::{CenoStdin, memory_from_file};
 #[cfg(all(feature = "jemalloc", unix, not(test)))]
 use ceno_zkvm::print_allocated_bytes;
@@ -108,6 +108,14 @@ struct Args {
     /// The security level to use.
     #[arg(short, long, value_enum, default_value_t = SecurityLevel::default())]
     security_level: SecurityLevel,
+
+    // shard id
+    #[arg(long, default_value = "0")]
+    shard_id: u32,
+
+    // number of total shards
+    #[arg(long, default_value = "1")]
+    num_shards: u32,
 }
 
 fn main() {
@@ -240,6 +248,7 @@ fn main() {
         .unwrap_or_default();
 
     let max_steps = args.max_steps.unwrap_or(usize::MAX);
+    let shards = Shards::new(args.shard_id as usize, args.num_shards as usize);
 
     match (args.pcs, args.field) {
         (PcsKind::Basefold, FieldType::Goldilocks) => {
@@ -249,6 +258,7 @@ fn main() {
                 prover,
                 program,
                 platform,
+                shards,
                 &hints,
                 &public_io,
                 max_steps,
@@ -264,6 +274,7 @@ fn main() {
                 prover,
                 program,
                 platform,
+                shards,
                 &hints,
                 &public_io,
                 max_steps,
@@ -279,6 +290,7 @@ fn main() {
                 prover,
                 program,
                 platform,
+                shards,
                 &hints,
                 &public_io,
                 max_steps,
@@ -294,6 +306,7 @@ fn main() {
                 prover,
                 program,
                 platform,
+                shards,
                 &hints,
                 &public_io,
                 max_steps,
@@ -320,6 +333,7 @@ fn run_inner<
     pd: PD,
     program: Program,
     platform: Platform,
+    shards: Shards,
     hints: &[u32],
     public_io: &[u32],
     max_steps: usize,
@@ -328,7 +342,7 @@ fn run_inner<
     checkpoint: Checkpoint,
 ) {
     let result = run_e2e_with_checkpoint::<E, PCS, _, _>(
-        pd, program, platform, hints, public_io, max_steps, checkpoint,
+        pd, program, platform, shards, hints, public_io, max_steps, checkpoint,
     );
 
     let zkvm_proof = result

ceno_zkvm/src/chip_handler/general.rs

@@ -4,8 +4,8 @@ use gkr_iop::{error::CircuitBuilderError, tables::LookupTable};
 use crate::{
     circuit_builder::CircuitBuilder,
     instructions::riscv::constants::{
-        END_CYCLE_IDX, END_PC_IDX, EXIT_CODE_IDX, INIT_CYCLE_IDX, INIT_PC_IDX, PUBLIC_IO_IDX,
-        UINT_LIMBS,
+        END_CYCLE_IDX, END_PC_IDX, EXIT_CODE_IDX, INIT_CYCLE_IDX, INIT_PC_IDX,
+        MEM_BUS_WITH_READ_IDX, MEM_BUS_WITH_WRITE_IDX, PUBLIC_IO_IDX, UINT_LIMBS,
     },
     tables::InsnRecord,
 };
@@ -22,6 +22,9 @@ pub trait PublicIOQuery {
     fn query_end_pc(&mut self) -> Result<Instance, CircuitBuilderError>;
     fn query_end_cycle(&mut self) -> Result<Instance, CircuitBuilderError>;
     fn query_public_io(&mut self) -> Result<[Instance; UINT_LIMBS], CircuitBuilderError>;
+
+    fn query_mem_bus_with_read(&mut self) -> Result<Instance, CircuitBuilderError>;
+    fn query_mem_bus_with_write(&mut self) -> Result<Instance, CircuitBuilderError>;
 }
 
 impl<'a, E: ExtensionField> InstFetch<E> for CircuitBuilder<'a, E> {
@@ -60,6 +63,16 @@ impl<'a, E: ExtensionField> PublicIOQuery for CircuitBuilder<'a, E> {
         self.cs.query_instance(|| "end_cycle", END_CYCLE_IDX)
     }
 
+    fn query_mem_bus_with_read(&mut self) -> Result<Instance, CircuitBuilderError> {
+        self.cs
+            .query_instance(|| "mem_bus_with_read", MEM_BUS_WITH_READ_IDX)
+    }
+
+    fn query_mem_bus_with_write(&mut self) -> Result<Instance, CircuitBuilderError> {
+        self.cs
+            .query_instance(|| "mem_bus_with_write", MEM_BUS_WITH_WRITE_IDX)
+    }
+
     fn query_public_io(&mut self) -> Result<[Instance; UINT_LIMBS], CircuitBuilderError> {
         Ok([
             self.cs.query_instance(|| "public_io_low", PUBLIC_IO_IDX)?,