OMC: init heap/hint table across different shard

.github/workflows/integration.yml

@@ -98,7 +98,8 @@ jobs:
           RUST_LOG: debug
           RUSTFLAGS: "-C opt-level=3"
           MOCK_PROVING: 1
-        run: cargo run --package ceno_zkvm --features sanity-check --bin e2e -- --platform=ceno --max-cycle-per-shard=300 examples/target/riscv32im-ceno-zkvm-elf/debug/examples/ceno_rt_alloc
+        # set --max-cycle-per-shard=420 to test heap read from previous shard, probing from guest program
+        run: cargo run --package ceno_zkvm --features sanity-check --bin e2e -- --platform=ceno --max-cycle-per-shard=420 examples/target/riscv32im-ceno-zkvm-elf/debug/examples/ceno_rt_alloc
 
       - name: Run multi-shards Guest Heap Alloc (release)
         env:

Cargo.toml

@@ -25,16 +25,16 @@ version = "0.1.0"
 [workspace.dependencies]
 ceno_crypto_primitives = { git = "https://github.com/scroll-tech/ceno-patch.git", package = "ceno_crypto_primitives", branch = "main" }
 ceno_syscall = { git = "https://github.com/scroll-tech/ceno-patch.git", package = "ceno_syscall", branch = "main" }
-ff_ext = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "ff_ext", tag = "v1.0.0-alpha.18" }
-mpcs = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "mpcs", tag = "v1.0.0-alpha.18" }
-multilinear_extensions = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "multilinear_extensions", tag = "v1.0.0-alpha.18" }
-p3 = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "p3", tag = "v1.0.0-alpha.18" }
-poseidon = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "poseidon", tag = "v1.0.0-alpha.18" }
-sp1-curves = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "sp1-curves", tag = "v1.0.0-alpha.18" }
-sumcheck = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "sumcheck", tag = "v1.0.0-alpha.18" }
-transcript = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "transcript", tag = "v1.0.0-alpha.18" }
-whir = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "whir", tag = "v1.0.0-alpha.18" }
-witness = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "witness", tag = "v1.0.0-alpha.18" }
+ff_ext = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "ff_ext", rev = "58bc22f64427d19ef881c1d3c00fadf0ee9d190d" }
+mpcs = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "mpcs", rev = "58bc22f64427d19ef881c1d3c00fadf0ee9d190d" }
+multilinear_extensions = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "multilinear_extensions", rev = "58bc22f64427d19ef881c1d3c00fadf0ee9d190d" }
+p3 = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "p3", rev = "58bc22f64427d19ef881c1d3c00fadf0ee9d190d" }
+poseidon = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "poseidon", rev = "58bc22f64427d19ef881c1d3c00fadf0ee9d190d" }
+sp1-curves = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "sp1-curves", rev = "58bc22f64427d19ef881c1d3c00fadf0ee9d190d" }
+sumcheck = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "sumcheck", rev = "58bc22f64427d19ef881c1d3c00fadf0ee9d190d" }
+transcript = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "transcript", rev = "58bc22f64427d19ef881c1d3c00fadf0ee9d190d" }
+whir = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "whir", rev = "58bc22f64427d19ef881c1d3c00fadf0ee9d190d" }
+witness = { git = "https://github.com/scroll-tech/gkr-backend.git", package = "witness", rev = "58bc22f64427d19ef881c1d3c00fadf0ee9d190d" }
 
 anyhow = { version = "1.0", default-features = false }
 bincode = "1"
@@ -47,6 +47,7 @@ itertools = "0.13"
 ndarray = "*"
 num-derive = "0.4"
 num-traits = "0.2"
+once_cell = "1.21.3"
 
 metrics = "*"
 num = "*"

ceno_emul/Cargo.toml

@@ -20,6 +20,7 @@ multilinear_extensions.workspace = true
 num.workspace = true
 num-derive.workspace = true
 num-traits.workspace = true
+once_cell.workspace = true
 rayon.workspace = true
 rrs_lib = { package = "rrs-succinct", version = "0.1.0" }
 rustc-hash.workspace = true

ceno_emul/src/platform.rs

@@ -1,5 +1,6 @@
 use core::fmt::{self, Formatter};
-use std::{collections::BTreeSet, fmt::Display, ops::Range};
+use once_cell::sync::Lazy;
+use std::{collections::BTreeSet, fmt::Display, ops::Range, sync::Arc};
 
 use crate::addr::{Addr, RegIdx};
 
@@ -11,7 +12,7 @@ use crate::addr::{Addr, RegIdx};
 #[derive(Clone, Debug)]
 pub struct Platform {
     pub rom: Range<Addr>,
-    pub prog_data: BTreeSet<Addr>,
+    pub prog_data: Arc<BTreeSet<Addr>>,
     pub public_io: Range<Addr>,
 
     pub stack: Range<Addr>,
@@ -91,7 +92,7 @@ impl Display for Platform {
 // │   0x0800_0000 .. 0x1000_0000
 // │
 // └───────────────────────────── 0x8000_0000 (rom base)
-pub const CENO_PLATFORM: Platform = Platform {
+pub static CENO_PLATFORM: Lazy<Platform> = Lazy::new(|| Platform {
     rom: 0x0800_0000..0x1000_0000,       // 128 MB
     public_io: 0x1000_0000..0x1800_0000, // 128 MB
     stack: 0x1800_0000..0x2000_4000, // stack grows downward 128MB, 0x4000 reserved for debug io.
@@ -103,9 +104,9 @@ pub const CENO_PLATFORM: Platform = Platform {
     // and the real heap start from 0x3800_0000
     heap: 0x3000_0000..0x4000_0000,
     unsafe_ecall_nop: false,
-    prog_data: BTreeSet::new(),
+    prog_data: Arc::new(BTreeSet::new()),
     is_debug: false,
-};
+});
 
 impl Platform {
     // Virtual memory layout.
@@ -210,7 +211,7 @@ mod tests {
 
     #[test]
     fn test_no_overlap() {
-        let p = CENO_PLATFORM;
+        let p = CENO_PLATFORM.clone();
         // ROM and RAM do not overlap.
         assert!(!p.is_rom(p.heap.start));
         assert!(!p.is_rom(p.heap.end - WORD_SIZE as Addr));

ceno_emul/src/test_utils.rs

@@ -23,7 +23,7 @@ pub fn keccak_step() -> (StepRecord, Vec<Instruction>) {
         instructions.clone(),
         Default::default(),
     );
-    let mut vm = VMState::new(CENO_PLATFORM, program.into());
+    let mut vm = VMState::new(CENO_PLATFORM.clone(), program.into());
     let steps = vm.iter_until_halt().collect::<Result<Vec<_>>>().unwrap();
 
     (steps[2].clone(), instructions)

ceno_emul/src/tracer.rs

@@ -25,6 +25,8 @@ use std::{collections::BTreeMap, fmt, mem};
 pub struct StepRecord {
     cycle: Cycle,
     pc: Change<ByteAddr>,
+    pub heap_maxtouch_addr: Change<ByteAddr>,
+    pub hint_maxtouch_addr: Change<ByteAddr>,
     pub insn: Instruction,
 
     rs1: Option<ReadOp>,
@@ -94,6 +96,10 @@ pub trait Tracer {
 
     fn fetch(&mut self, pc: WordAddr, value: Instruction);
 
+    fn track_mmu_maxtouch_before(&mut self);
+
+    fn track_mmu_maxtouch_after(&mut self);
+
     fn load_register(&mut self, idx: RegIdx, value: Word);
 
     fn store_register(&mut self, idx: RegIdx, value: Change<Word>);
@@ -273,6 +279,8 @@ impl StepRecord {
             Some(rd),
             None,
             prev_cycle,
+            Change::default(),
+            Change::default(),
         )
     }
 
@@ -293,6 +301,8 @@ impl StepRecord {
             None,
             None,
             prev_cycle,
+            Change::default(),
+            Change::default(),
         )
     }
 
@@ -313,6 +323,8 @@ impl StepRecord {
             Some(rd),
             None,
             prev_cycle,
+            Change::default(),
+            Change::default(),
         )
     }
 
@@ -342,6 +354,8 @@ impl StepRecord {
                 previous_cycle: mem_op.previous_cycle,
             }),
             prev_cycle,
+            Change::default(),
+            Change::default(),
         )
     }
 
@@ -353,7 +367,18 @@ impl StepRecord {
         prev_cycle: Cycle,
     ) -> StepRecord {
         let pc = Change::new(pc, pc + PC_STEP_SIZE);
-        StepRecord::new_insn(cycle, pc, insn_code, None, None, Some(rd), None, prev_cycle)
+        StepRecord::new_insn(
+            cycle,
+            pc,
+            insn_code,
+            None,
+            None,
+            Some(rd),
+            None,
+            prev_cycle,
+            Change::default(),
+            Change::default(),
+        )
     }
 
     pub fn new_j_instruction(
@@ -363,7 +388,18 @@ impl StepRecord {
         rd: Change<Word>,
         prev_cycle: Cycle,
     ) -> StepRecord {
-        StepRecord::new_insn(cycle, pc, insn_code, None, None, Some(rd), None, prev_cycle)
+        StepRecord::new_insn(
+            cycle,
+            pc,
+            insn_code,
+            None,
+            None,
+            Some(rd),
+            None,
+            prev_cycle,
+            Change::default(),
+            Change::default(),
+        )
     }
 
     pub fn new_s_instruction(
@@ -385,6 +421,8 @@ impl StepRecord {
             None,
             Some(memory_op),
             prev_cycle,
+            Change::default(),
+            Change::default(),
         )
     }
 
@@ -407,6 +445,8 @@ impl StepRecord {
                 previous_cycle: 0,
             }),
             0,
+            Change::default(),
+            Change::default(),
         )
     }
 
@@ -420,6 +460,8 @@ impl StepRecord {
         rd: Option<Change<Word>>,
         memory_op: Option<WriteOp>,
         previous_cycle: Cycle,
+        heap_maxtouch_addr: Change<ByteAddr>,
+        hint_maxtouch_addr: Change<ByteAddr>,
     ) -> StepRecord {
         StepRecord {
             cycle,
@@ -442,6 +484,8 @@ impl StepRecord {
             insn,
             memory_op,
             syscall: None,
+            heap_maxtouch_addr,
+            hint_maxtouch_addr,
         }
     }
 
@@ -491,6 +535,9 @@ pub struct FullTracer {
     // 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)>>,
+    max_heap_addr_access: ByteAddr,
+    max_hint_addr_access: ByteAddr,
+    platform: Platform,
 
     // keep track of each address that the cycle when they were last accessed.
     latest_accesses: LatestAccesses,
@@ -516,8 +563,11 @@ impl FullTracer {
                 cycle: Self::SUBCYCLES_PER_INSN,
                 ..StepRecord::default()
             },
+            platform: platform.clone(),
             latest_accesses: LatestAccesses::new(platform),
             next_accesses: NextCycleAccess::new(ACCESSED_CHUNK_SIZE),
+            max_heap_addr_access: ByteAddr::from(platform.heap.start),
+            max_hint_addr_access: ByteAddr::from(platform.hints.start),
         }
     }
 
@@ -545,6 +595,18 @@ impl FullTracer {
         self.record.insn = value;
     }
 
+    #[inline(always)]
+    pub fn track_mmu_maxtouch_before(&mut self) {
+        self.record.heap_maxtouch_addr.before = self.max_heap_addr_access;
+        self.record.hint_maxtouch_addr.before = self.max_hint_addr_access;
+    }
+
+    #[inline(always)]
+    pub fn track_mmu_maxtouch_after(&mut self) {
+        self.record.heap_maxtouch_addr.after = self.max_heap_addr_access;
+        self.record.hint_maxtouch_addr.after = self.max_hint_addr_access;
+    }
+
     #[inline(always)]
     pub fn load_register(&mut self, idx: RegIdx, value: Word) {
         let addr = Platform::register_vma(idx).into();
@@ -593,7 +655,7 @@ impl FullTracer {
             unimplemented!("Only one memory access is supported");
         }
         // update min/max mmio access
-        if let Some((_, (_, end_addr, min_addr, max_addr))) = self
+        if let Some((start_addr, (_, end_addr, min_addr, max_addr))) = self
             .mmio_min_max_access
             .as_mut()
             // find the MMIO region whose start address is less than or equal to the target address
@@ -610,6 +672,19 @@ impl FullTracer {
                 if addr < *min_addr {
                     *min_addr = addr; // start is inclusive
                 }
+                if start_addr.baddr().0 == self.platform.heap.start {
+                    let access_end = addr + WordAddr::from(WORD_SIZE as u32);
+                    let access_end_byte = access_end.baddr();
+                    if access_end_byte > self.max_heap_addr_access {
+                        self.max_heap_addr_access = access_end_byte;
+                    }
+                } else if start_addr.baddr().0 == self.platform.hints.start {
+                    let access_end = addr + WordAddr::from(WORD_SIZE as u32);
+                    let access_end_byte = access_end.baddr();
+                    if access_end_byte > self.max_hint_addr_access {
+                        self.max_hint_addr_access = access_end_byte;
+                    }
+                }
             }
         }
 
@@ -763,6 +838,12 @@ impl Tracer for PreflightTracer {
     #[inline(always)]
     fn fetch(&mut self, _pc: WordAddr, _value: Instruction) {}
 
+    #[inline(always)]
+    fn track_mmu_maxtouch_before(&mut self) {}
+
+    #[inline(always)]
+    fn track_mmu_maxtouch_after(&mut self) {}
+
     #[inline(always)]
     fn load_register(&mut self, idx: RegIdx, _value: Word) {
         let addr = Platform::register_vma(idx).into();
@@ -877,6 +958,14 @@ impl Tracer for FullTracer {
         FullTracer::fetch(self, pc, value)
     }
 
+    fn track_mmu_maxtouch_before(&mut self) {
+        FullTracer::track_mmu_maxtouch_before(self)
+    }
+
+    fn track_mmu_maxtouch_after(&mut self) {
+        FullTracer::track_mmu_maxtouch_after(self)
+    }
+
     #[inline(always)]
     fn load_register(&mut self, idx: RegIdx, value: Word) {
         FullTracer::load_register(self, idx, value)

ceno_emul/src/vm_state.rs

@@ -72,7 +72,7 @@ impl<T: Tracer> VMState<T> {
     pub fn new_from_elf_with_tracer(platform: Platform, elf: &[u8]) -> Result<Self> {
         let program = Arc::new(Program::load_elf(elf, u32::MAX)?);
         let platform = Platform {
-            prog_data: program.image.keys().copied().collect(),
+            prog_data: Arc::new(program.image.keys().copied().collect()),
             ..platform
         };
         Ok(Self::new_with_tracer(platform, program))
@@ -212,6 +212,7 @@ impl<T: Tracer> EmuContext for VMState<T> {
 
     fn on_normal_end(&mut self, _decoded: &Instruction) {
         self.tracer.store_pc(ByteAddr(self.pc));
+        self.tracer.track_mmu_maxtouch_after();
     }
 
     fn get_pc(&self) -> ByteAddr {
@@ -273,6 +274,7 @@ impl<T: Tracer> EmuContext for VMState<T> {
         let idx = (relative_pc / WORD_SIZE as u32) as usize;
         let word = self.program.instructions.get(idx).copied()?;
         self.tracer.fetch(pc, word);
+        self.tracer.track_mmu_maxtouch_before();
         Some(word)
     }
 

ceno_emul/tests/test_vm_trace.rs

@@ -17,7 +17,7 @@ fn test_vm_trace() -> Result<()> {
         program_fibonacci_20(),
         Default::default(),
     );
-    let mut ctx = VMState::new(CENO_PLATFORM, Arc::new(program));
+    let mut ctx = VMState::new(CENO_PLATFORM.clone(), Arc::new(program));
 
     let steps = run(&mut ctx)?;
 
@@ -53,7 +53,7 @@ fn test_empty_program() -> Result<()> {
         vec![],
         BTreeMap::new(),
     );
-    let mut ctx = VMState::new(CENO_PLATFORM, Arc::new(empty_program));
+    let mut ctx = VMState::new(CENO_PLATFORM.clone(), Arc::new(empty_program));
     let res = run(&mut ctx);
     assert!(matches!(res, Err(e) if e.to_string().contains("InstructionAccessFault")),);
     Ok(())

ceno_host/src/lib.rs

@@ -144,7 +144,7 @@ pub fn run(
 ) -> Vec<Vec<u8>> {
     let program = Program::load_elf(elf, u32::MAX).unwrap();
     let platform = Platform {
-        prog_data: program.image.keys().copied().collect(),
+        prog_data: Arc::new(program.image.keys().copied().collect()),
         ..platform
     };