add shard info to public io

Author: hero78119

ceno_emul/src/shards.rs

@@ -11,4 +11,12 @@ impl Shards {
             num_shards,
         }
     }
+
+    pub fn is_first_shard(&self) -> bool {
+        self.shard_id == 0
+    }
+
+    pub fn is_last_shard(&self) -> bool {
+        self.shard_id == self.num_shards - 1
+    }
 }

ceno_zkvm/src/chip_handler/general.rs

@@ -4,7 +4,7 @@ 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,
+        END_CYCLE_IDX, END_PC_IDX, END_SHARD_ID_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,7 +22,7 @@ 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_shard_id(&mut self) -> Result<Instance, CircuitBuilderError>;
     fn query_mem_bus_with_read(&mut self) -> Result<Instance, CircuitBuilderError>;
     fn query_mem_bus_with_write(&mut self) -> Result<Instance, CircuitBuilderError>;
 }
@@ -63,6 +63,10 @@ impl<'a, E: ExtensionField> PublicIOQuery for CircuitBuilder<'a, E> {
         self.cs.query_instance(|| "end_cycle", END_CYCLE_IDX)
     }
 
+    fn query_shard_id(&mut self) -> Result<Instance, CircuitBuilderError> {
+        self.cs.query_instance(|| "shard_id", END_SHARD_ID_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)

ceno_zkvm/src/e2e.rs

@@ -334,6 +334,9 @@ pub fn emulate_program<'a>(
         Tracer::SUBCYCLES_PER_INSN,
         vm.get_pc().into(),
         end_cycle,
+        shards.shard_id as u32,
+        !shards.is_first_shard(), // first shard disable global read
+        !shards.is_last_shard(),  // last shard disable global write
         io_init.iter().map(|rec| rec.value).collect_vec(),
     );
 

ceno_zkvm/src/instructions/riscv/constants.rs

@@ -9,9 +9,10 @@ pub const INIT_PC_IDX: usize = 2;
 pub const INIT_CYCLE_IDX: usize = 3;
 pub const END_PC_IDX: usize = 4;
 pub const END_CYCLE_IDX: usize = 5;
-pub const PUBLIC_IO_IDX: usize = 6;
+pub const END_SHARD_ID_IDX: usize = 6;
 pub const MEM_BUS_WITH_READ_IDX: usize = 7;
 pub const MEM_BUS_WITH_WRITE_IDX: usize = 8;
+pub const PUBLIC_IO_IDX: usize = 9;
 
 pub const LIMB_BITS: usize = 16;
 pub const LIMB_MASK: u32 = 0xFFFF;

ceno_zkvm/src/scheme.rs

@@ -72,6 +72,9 @@ pub struct PublicValues {
     init_cycle: u64,
     end_pc: u32,
     end_cycle: u64,
+    shard_id: u32,
+    mem_bus_with_read: bool,
+    mem_bus_with_write: bool,
     public_io: Vec<u32>,
 }
 
@@ -82,6 +85,9 @@ impl PublicValues {
         init_cycle: u64,
         end_pc: u32,
         end_cycle: u64,
+        shard_id: u32,
+        mem_bus_with_read: bool,
+        mem_bus_with_write: bool,
         public_io: Vec<u32>,
     ) -> Self {
         Self {
@@ -90,6 +96,9 @@ impl PublicValues {
             init_cycle,
             end_pc,
             end_cycle,
+            shard_id,
+            mem_bus_with_read,
+            mem_bus_with_write,
             public_io,
         }
     }
@@ -103,6 +112,9 @@ impl PublicValues {
             vec![E::BaseField::from_canonical_u64(self.init_cycle)],
             vec![E::BaseField::from_canonical_u32(self.end_pc)],
             vec![E::BaseField::from_canonical_u64(self.end_cycle)],
+            vec![E::BaseField::from_canonical_u32(self.shard_id)],
+            vec![E::BaseField::from_bool(self.mem_bus_with_read)],
+            vec![E::BaseField::from_bool(self.mem_bus_with_write)],
         ]
         .into_iter()
         .chain(