aligned step cycle and prev_cycle to local version

Author: hero78119

ceno_emul/src/shards.rs

@@ -1,14 +1,14 @@
 pub struct Shards {
     pub shard_id: usize,
-    pub num_shards: usize,
+    pub max_num_shards: usize,
 }
 
 impl Shards {
-    pub fn new(shard_id: usize, num_shards: usize) -> Self {
-        assert!(shard_id < num_shards);
+    pub fn new(shard_id: usize, max_num_shards: usize) -> Self {
+        assert!(shard_id < max_num_shards);
         Self {
             shard_id,
-            num_shards,
+            max_num_shards,
         }
     }
 
@@ -17,6 +17,6 @@ impl Shards {
     }
 
     pub fn is_last_shard(&self) -> bool {
-        self.shard_id == self.num_shards - 1
+        self.shard_id == self.max_num_shards - 1
     }
 }

ceno_emul/src/tracer.rs

@@ -25,7 +25,7 @@ use crate::{
 /// - Any pair of `rs1 / rs2 / rd` **may be the same**. Then, one op will point to the other op in the same instruction but a different subcycle. The circuits may follow the operations **without special handling** of repeated registers.
 #[derive(Clone, Debug, Default, PartialEq, Eq)]
 pub struct StepRecord {
-    cycle: Cycle,
+    pub cycle: Cycle,
     pc: Change<ByteAddr>,
     pub insn: Instruction,
 

ceno_zkvm/src/bin/e2e.rs

@@ -115,7 +115,7 @@ struct Args {
 
     // number of total shards
     #[arg(long, default_value = "1")]
-    num_shards: u32,
+    max_num_shards: u32,
 }
 
 fn main() {
@@ -248,7 +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);
+    let shards = Shards::new(args.shard_id as usize, args.max_num_shards as usize);
 
     match (args.pcs, args.field) {
         (PcsKind::Basefold, FieldType::Goldilocks) => {

ceno_zkvm/src/e2e.rs

@@ -112,7 +112,7 @@ pub struct RAMRecord {
 
 pub struct ShardContext<'a> {
     shard_id: usize,
-    num_shards: usize,
+    max_num_shards: usize,
     max_cycle: Cycle,
     addr_future_accesses: Cow<'a, HashMap<(WordAddr, Cycle), Cycle>>,
     read_thread_based_record_storage: Either<
@@ -131,7 +131,7 @@ impl<'a> Default for ShardContext<'a> {
         let max_threads = max_usable_threads();
         Self {
             shard_id: 0,
-            num_shards: 1,
+            max_num_shards: 1,
             max_cycle: Cycle::default(),
             addr_future_accesses: Cow::Owned(HashMap::new()),
             read_thread_based_record_storage: Either::Left(
@@ -154,20 +154,27 @@ impl<'a> Default for ShardContext<'a> {
 impl<'a> ShardContext<'a> {
     pub fn new(
         shard_id: usize,
-        num_shards: usize,
+        max_num_shards: usize,
         executed_instructions: usize,
         addr_future_accesses: HashMap<(WordAddr, Cycle), Cycle>,
     ) -> Self {
+        // current strategy: at least each shard deal with one instruction
+        let max_num_shards = max_num_shards.min(executed_instructions);
+        assert!(
+            shard_id < max_num_shards,
+            "implement mechanism to skip current shard proof"
+        );
+
         let max_threads = max_usable_threads();
         // let max_record_per_thread = max_insts.div_ceil(max_threads as u64);
-        let expected_inst_per_shard = executed_instructions.div_ceil(num_shards) as usize;
+        let expected_inst_per_shard = executed_instructions.div_ceil(max_num_shards) as usize;
         let max_cycle = (executed_instructions + 1) * 4; // cycle start from 4
-        let cur_shard_cycle_range = (shard_id * expected_inst_per_shard * 4).max(4)
-            ..((shard_id + 1) * expected_inst_per_shard * 4).min(max_cycle);
+        let cur_shard_cycle_range = (shard_id * expected_inst_per_shard * 4 + 4)
+            ..((shard_id + 1) * expected_inst_per_shard * 4 + 4).min(max_cycle);
 
         ShardContext {
             shard_id,
-            num_shards,
+            max_num_shards,
             max_cycle: max_cycle as Cycle,
             addr_future_accesses: Cow::Owned(addr_future_accesses),
             // TODO with_capacity optimisation
@@ -201,7 +208,7 @@ impl<'a> ShardContext<'a> {
                 .zip(write_thread_based_record_storage.iter_mut())
                 .map(|(read, write)| ShardContext {
                     shard_id: self.shard_id,
-                    num_shards: self.num_shards,
+                    max_num_shards: self.max_num_shards,
                     max_cycle: self.max_cycle,
                     addr_future_accesses: Cow::Borrowed(self.addr_future_accesses.as_ref()),
                     read_thread_based_record_storage: Either::Right(read),
@@ -220,14 +227,28 @@ impl<'a> ShardContext<'a> {
 
     #[inline(always)]
     pub fn is_last_shard(&self) -> bool {
-        self.shard_id == self.num_shards - 1
+        self.shard_id == self.max_num_shards - 1
     }
 
     #[inline(always)]
     pub fn is_current_shard_cycle(&self, cycle: Cycle) -> bool {
         self.cur_shard_cycle_range.contains(&(cycle as usize))
     }
 
+    #[inline(always)]
+    pub fn aligned_prev_ts(&self, prev_cycle: Cycle) -> Cycle {
+        let mut ts = prev_cycle.saturating_sub(self.cur_shard_cycle_range.start as Cycle);
+        if ts < 4 {
+            ts = 0
+        }
+        ts
+    }
+
+    pub fn current_shard_offset_cycle(&self) -> Cycle {
+        // `-4` as cycle of each local shard start from 4
+        (self.cur_shard_cycle_range.start as Cycle) - 4
+    }
+
     #[inline(always)]
     pub fn send(
         &mut self,
@@ -475,7 +496,7 @@ pub fn emulate_program<'a>(
 
     let shard_ctx = ShardContext::new(
         shards.shard_id,
-        shards.num_shards,
+        shards.max_num_shards,
         insts,
         vm.take_tracer().next_accesses(),
     );

ceno_zkvm/src/instructions/riscv/ecall_base.rs

@@ -69,7 +69,8 @@ impl<E: ExtensionField, const REG_ID: usize, const RW: bool> OpFixedRS<E, REG_ID
         cycle: Cycle,
         op: &WriteOp,
     ) -> Result<(), ZKVMError> {
-        set_val!(instance, self.prev_ts, op.previous_cycle);
+        let shard_prev_ts = shard_ctx.aligned_prev_ts(op.previous_cycle);
+        set_val!(instance, self.prev_ts, shard_prev_ts);
 
         // Register state
         if let Some(prev_value) = self.prev_value.as_ref() {
@@ -86,7 +87,7 @@ impl<E: ExtensionField, const REG_ID: usize, const RW: bool> OpFixedRS<E, REG_ID
         };
         // Register write
         self.lt_cfg
-            .assign_instance(instance, lk_multiplicity, op.previous_cycle, cycle)?;
+            .assign_instance(instance, lk_multiplicity, shard_prev_ts, cycle)?;
 
         shard_ctx.send(
             RAMType::Register,

ceno_zkvm/src/instructions/riscv/insn_base.rs

@@ -113,14 +113,15 @@ impl<E: ExtensionField> ReadRS1<E> {
         step: &StepRecord,
     ) -> Result<(), ZKVMError> {
         let op = step.rs1().expect("rs1 op");
+        let shard_prev_ts = shard_ctx.aligned_prev_ts(op.previous_cycle);
         set_val!(instance, self.id, op.register_index() as u64);
-        set_val!(instance, self.prev_ts, op.previous_cycle);
+        set_val!(instance, self.prev_ts, shard_prev_ts);
 
         // Register read
         self.lt_cfg.assign_instance(
             instance,
             lk_multiplicity,
-            op.previous_cycle,
+            shard_prev_ts,
             step.cycle() + Tracer::SUBCYCLE_RS1,
         )?;
         shard_ctx.send(
@@ -177,14 +178,15 @@ impl<E: ExtensionField> ReadRS2<E> {
         step: &StepRecord,
     ) -> Result<(), ZKVMError> {
         let op = step.rs2().expect("rs2 op");
+        let shard_prev_ts = shard_ctx.aligned_prev_ts(op.previous_cycle);
         set_val!(instance, self.id, op.register_index() as u64);
-        set_val!(instance, self.prev_ts, op.previous_cycle);
+        set_val!(instance, self.prev_ts, shard_prev_ts);
 
         // Register read
         self.lt_cfg.assign_instance(
             instance,
             lk_multiplicity,
-            op.previous_cycle,
+            shard_prev_ts,
             step.cycle() + Tracer::SUBCYCLE_RS2,
         )?;
 
@@ -255,8 +257,9 @@ impl<E: ExtensionField> WriteRD<E> {
         cycle: Cycle,
         op: &WriteOp,
     ) -> Result<(), ZKVMError> {
+        let shard_prev_ts = shard_ctx.aligned_prev_ts(op.previous_cycle);
         set_val!(instance, self.id, op.register_index() as u64);
-        set_val!(instance, self.prev_ts, op.previous_cycle);
+        set_val!(instance, self.prev_ts, shard_prev_ts);
 
         // Register state
         self.prev_value.assign_limbs(
@@ -268,7 +271,7 @@ impl<E: ExtensionField> WriteRD<E> {
         self.lt_cfg.assign_instance(
             instance,
             lk_multiplicity,
-            op.previous_cycle,
+            shard_prev_ts,
             cycle + Tracer::SUBCYCLE_RD,
         )?;
         shard_ctx.send(
@@ -323,14 +326,15 @@ impl<E: ExtensionField> ReadMEM<E> {
         step: &StepRecord,
     ) -> Result<(), ZKVMError> {
         let op = step.memory_op().unwrap();
+        let shard_prev_ts = shard_ctx.aligned_prev_ts(op.previous_cycle);
         // Memory state
-        set_val!(instance, self.prev_ts, op.previous_cycle);
+        set_val!(instance, self.prev_ts, shard_prev_ts);
 
         // Memory read
         self.lt_cfg.assign_instance(
             instance,
             lk_multiplicity,
-            op.previous_cycle,
+            shard_prev_ts,
             step.cycle() + Tracer::SUBCYCLE_MEM,
         )?;
 
@@ -395,12 +399,13 @@ impl WriteMEM {
         cycle: Cycle,
         op: &WriteOp,
     ) -> Result<(), ZKVMError> {
-        set_val!(instance, self.prev_ts, op.previous_cycle);
+        let shard_prev_ts = shard_ctx.aligned_prev_ts(op.previous_cycle);
+        set_val!(instance, self.prev_ts, shard_prev_ts);
 
         self.lt_cfg.assign_instance(
             instance,
             lk_multiplicity,
-            op.previous_cycle,
+            shard_prev_ts,
             cycle + Tracer::SUBCYCLE_MEM,
         )?;
 

ceno_zkvm/src/instructions/riscv/rv32im.rs

@@ -414,35 +414,46 @@ impl<E: ExtensionField> Rv32imConfig<E> {
         let mut bn254_double_records = Vec::new();
         let mut secp256k1_add_records = Vec::new();
         let mut secp256k1_double_records = Vec::new();
-        steps.into_iter().for_each(|record| {
-            let insn_kind = record.insn.kind;
-            match insn_kind {
-                // ecall / halt
-                InsnKind::ECALL if record.rs1().unwrap().value == Platform::ecall_halt() => {
-                    halt_records.push(record);
+        let current_shard_offset_cycle = shard_ctx.current_shard_offset_cycle();
+        steps
+            .into_iter()
+            .filter_map(|mut step| {
+                if shard_ctx.is_current_shard_cycle(step.cycle()) {
+                    step.cycle = step.cycle() - current_shard_offset_cycle;
+                    Some(step)
+                } else {
+                    None
                 }
-                InsnKind::ECALL if record.rs1().unwrap().value == KeccakSpec::CODE => {
-                    keccak_records.push(record);
+            })
+            .for_each(|record| {
+                let insn_kind = record.insn.kind;
+                match insn_kind {
+                    // ecall / halt
+                    InsnKind::ECALL if record.rs1().unwrap().value == Platform::ecall_halt() => {
+                        halt_records.push(record);
+                    }
+                    InsnKind::ECALL if record.rs1().unwrap().value == KeccakSpec::CODE => {
+                        keccak_records.push(record);
+                    }
+                    InsnKind::ECALL if record.rs1().unwrap().value == Bn254AddSpec::CODE => {
+                        bn254_add_records.push(record);
+                    }
+                    InsnKind::ECALL if record.rs1().unwrap().value == Bn254DoubleSpec::CODE => {
+                        bn254_double_records.push(record);
+                    }
+                    InsnKind::ECALL if record.rs1().unwrap().value == Secp256k1AddSpec::CODE => {
+                        secp256k1_add_records.push(record);
+                    }
+                    InsnKind::ECALL if record.rs1().unwrap().value == Secp256k1DoubleSpec::CODE => {
+                        secp256k1_double_records.push(record);
+                    }
+                    // other type of ecalls are handled by dummy ecall instruction
+                    _ => {
+                        // it's safe to unwrap as all_records are initialized with Vec::new()
+                        all_records.get_mut(&insn_kind).unwrap().push(record);
+                    }
                 }
-                InsnKind::ECALL if record.rs1().unwrap().value == Bn254AddSpec::CODE => {
-                    bn254_add_records.push(record);
-                }
-                InsnKind::ECALL if record.rs1().unwrap().value == Bn254DoubleSpec::CODE => {
-                    bn254_double_records.push(record);
-                }
-                InsnKind::ECALL if record.rs1().unwrap().value == Secp256k1AddSpec::CODE => {
-                    secp256k1_add_records.push(record);
-                }
-                InsnKind::ECALL if record.rs1().unwrap().value == Secp256k1DoubleSpec::CODE => {
-                    secp256k1_double_records.push(record);
-                }
-                // other type of ecalls are handled by dummy ecall instruction
-                _ => {
-                    // it's safe to unwrap as all_records are initialized with Vec::new()
-                    all_records.get_mut(&insn_kind).unwrap().push(record);
-                }
-            }
-        });
+            });
 
         for (insn_kind, (_, records)) in
             izip!(InsnKind::iter(), &all_records).sorted_by_key(|(_, (_, a))| Reverse(a.len()))