extract shard id from cycle

Author: hero78119

ceno_zkvm/src/e2e.rs

@@ -106,10 +106,14 @@ pub struct RAMRecord {
     pub ram_type: RAMType,
     pub id: u64,
     pub addr: WordAddr,
+    // prev_cycle and cycle are global cycle
     pub prev_cycle: Cycle,
     pub cycle: Cycle,
     pub prev_value: Option<Word>,
     pub value: Word,
+    // for global reads, `shard_id` refers to the shard that previously produced this value.
+    // for global write, `shard_id` refers to current shard.
+    pub shard_id: usize,
 }
 
 #[derive(Clone, Debug)]
@@ -155,6 +159,7 @@ pub struct ShardContext<'a> {
     write_thread_based_record_storage:
         Either<Vec<BTreeMap<WordAddr, RAMRecord>>, &'a mut BTreeMap<WordAddr, RAMRecord>>,
     pub cur_shard_cycle_range: std::ops::Range<usize>,
+    pub expected_inst_per_shard: usize,
 }
 
 impl<'a> Default for ShardContext<'a> {
@@ -177,6 +182,7 @@ impl<'a> Default for ShardContext<'a> {
                     .collect::<Vec<_>>(),
             ),
             cur_shard_cycle_range: Tracer::SUBCYCLES_PER_INSN as usize..usize::MAX,
+            expected_inst_per_shard: usize::MAX,
         }
     }
 }
@@ -223,6 +229,7 @@ impl<'a> ShardContext<'a> {
                     .collect::<Vec<_>>(),
             ),
             cur_shard_cycle_range,
+            expected_inst_per_shard,
         }
     }
 
@@ -244,6 +251,7 @@ impl<'a> ShardContext<'a> {
                     read_thread_based_record_storage: Either::Right(read),
                     write_thread_based_record_storage: Either::Right(write),
                     cur_shard_cycle_range: self.cur_shard_cycle_range.clone(),
+                    expected_inst_per_shard: self.expected_inst_per_shard,
                 })
                 .collect_vec(),
             _ => panic!("invalid type"),
@@ -279,6 +287,14 @@ impl<'a> ShardContext<'a> {
         self.cur_shard_cycle_range.contains(&(cycle as usize))
     }
 
+    #[inline(always)]
+    pub fn extract_prev_shard_id(&self, cycle: Cycle) -> usize {
+        let subcycle_per_insn = Tracer::SUBCYCLES_PER_INSN;
+        let per_shard_cycles =
+            (self.expected_inst_per_shard as u64).saturating_mul(subcycle_per_insn);
+        ((cycle.saturating_sub(subcycle_per_insn)) / per_shard_cycles) as usize
+    }
+
     #[inline(always)]
     pub fn aligned_prev_ts(&self, prev_cycle: Cycle) -> Cycle {
         let mut ts = prev_cycle.saturating_sub(self.current_shard_offset_cycle());
@@ -311,6 +327,7 @@ impl<'a> ShardContext<'a> {
             && self.is_current_shard_cycle(cycle)
             && !self.is_first_shard()
         {
+            let prev_shard_id = self.extract_prev_shard_id(prev_cycle);
             let ram_record = self
                 .read_thread_based_record_storage
                 .as_mut()
@@ -326,6 +343,7 @@ impl<'a> ShardContext<'a> {
                     cycle,
                     prev_value,
                     value,
+                    shard_id: prev_shard_id,
                 },
             );
         }
@@ -363,6 +381,7 @@ impl<'a> ShardContext<'a> {
                     cycle,
                     prev_value,
                     value,
+                    shard_id: self.shards.shard_id,
                 },
             );
         }

ceno_zkvm/src/instructions/riscv/rv32im/mmu.rs

@@ -187,7 +187,7 @@ impl<E: ExtensionField> MmuConfig<'_, E> {
                         addr,
                         ram_type: record.ram_type,
                         value: record.value,
-                        shard: record.prev_cycle, // TODO: extract shard id properly
+                        shard: record.shard_id as u64,
                         local_clk: local_clk,
                         global_clk: global_clk,
                         is_write: is_write,