feat/private-input-2: add private IO to MMU

scroll-tech · Nov 21, 2024 · acccad7 · acccad7
1 parent 0c32eab
commit acccad7
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 9 deletions.
diff --git a/ceno_zkvm/examples/riscv_opcodes.rs b/ceno_zkvm/examples/riscv_opcodes.rs
@@ -275,6 +275,7 @@ fn main() {
                 &reg_final,
                 &mem_final,
                 &public_io_final,
+                &[],
             )
             .unwrap();
 

diff --git a/ceno_zkvm/src/bin/e2e.rs b/ceno_zkvm/src/bin/e2e.rs
@@ -260,6 +260,7 @@ fn main() {
             &reg_final,
             &mem_final,
             &io_final,
+            &[],
         )
         .unwrap();
     // assign program circuit

diff --git a/ceno_zkvm/src/instructions/riscv/rv32im/mmu.rs b/ceno_zkvm/src/instructions/riscv/rv32im/mmu.rs
@@ -8,8 +8,8 @@ use crate::{
     error::ZKVMError,
     structs::{ProgramParams, ZKVMConstraintSystem, ZKVMFixedTraces, ZKVMWitnesses},
     tables::{
-        MemFinalRecord, MemInitRecord, NonVolatileTable, PubIOCircuit, PubIOTable, RegTable,
-        RegTableCircuit, StaticMemCircuit, StaticMemTable, TableCircuit,
+        MemFinalRecord, MemInitRecord, NonVolatileTable, PrivateIOCircuit, PubIOCircuit,
+        PubIOTable, RegTable, RegTableCircuit, StaticMemCircuit, StaticMemTable, TableCircuit,
     },
 };
 
@@ -20,6 +20,8 @@ pub struct MmuConfig<E: ExtensionField> {
     pub static_mem_config: <StaticMemCircuit<E> as TableCircuit<E>>::TableConfig,
     /// Initialization of public IO.
     pub public_io_config: <PubIOCircuit<E> as TableCircuit<E>>::TableConfig,
+    /// Initialization of private IO.
+    pub private_io_config: <PrivateIOCircuit<E> as TableCircuit<E>>::TableConfig,
     pub params: ProgramParams,
 }
 
@@ -30,11 +32,13 @@ impl<E: ExtensionField> MmuConfig<E> {
         let static_mem_config = cs.register_table_circuit::<StaticMemCircuit<E>>();
 
         let public_io_config = cs.register_table_circuit::<PubIOCircuit<E>>();
+        let private_io_config = cs.register_table_circuit::<PrivateIOCircuit<E>>();
 
         Self {
             reg_config,
             static_mem_config,
             public_io_config,
+            private_io_config,
             params: cs.params.clone(),
         }
     }
@@ -48,9 +52,11 @@ impl<E: ExtensionField> MmuConfig<E> {
         io_addrs: &[Addr],
     ) {
         assert!(
-            chain(
+            chain!(
                 static_mem_init.iter().map(|record| record.addr),
                 io_addrs.iter().copied(),
+                // TODO: optimize with min_max and Range.
+                self.params.platform.private_io.clone(),
             )
             .all_unique(),
             "memory addresses must be unique"
@@ -65,6 +71,7 @@ impl<E: ExtensionField> MmuConfig<E> {
         );
 
         fixed.register_table_circuit::<PubIOCircuit<E>>(cs, &self.public_io_config, io_addrs);
+        fixed.register_table_circuit::<PrivateIOCircuit<E>>(cs, &self.private_io_config, &());
     }
 
     pub fn assign_table_circuit(
@@ -74,6 +81,7 @@ impl<E: ExtensionField> MmuConfig<E> {
         reg_final: &[MemFinalRecord],
         static_mem_final: &[MemFinalRecord],
         io_cycles: &[Cycle],
+        private_io_final: &[MemFinalRecord],
     ) -> Result<(), ZKVMError> {
         witness.assign_table_circuit::<RegTableCircuit<E>>(cs, &self.reg_config, reg_final)?;
 
@@ -85,6 +93,12 @@ impl<E: ExtensionField> MmuConfig<E> {
 
         witness.assign_table_circuit::<PubIOCircuit<E>>(cs, &self.public_io_config, io_cycles)?;
 
+        witness.assign_table_circuit::<PrivateIOCircuit<E>>(
+            cs,
+            &self.private_io_config,
+            private_io_final,
+        )?;
+
         Ok(())
     }
 

diff --git a/ceno_zkvm/src/tables/ram.rs b/ceno_zkvm/src/tables/ram.rs
@@ -34,25 +34,25 @@ impl DynVolatileRamTable for DynMemTable {
 pub type DynMemCircuit<E> = DynVolatileRamCircuit<E, DynMemTable>;
 
 #[derive(Clone)]
-pub struct PrivateMemTable;
-impl DynVolatileRamTable for PrivateMemTable {
+pub struct PrivateIOTable;
+impl DynVolatileRamTable for PrivateIOTable {
     const RAM_TYPE: RAMType = RAMType::Memory;
     const V_LIMBS: usize = 1; // See `MemoryExpr`.
     const ZERO_INIT: bool = false;
 
     fn offset_addr(params: &ProgramParams) -> Addr {
-        params.platform.ram.start
+        params.platform.private_io.start
     }
 
     fn end_addr(params: &ProgramParams) -> Addr {
-        params.platform.ram.end
+        params.platform.private_io.end
     }
 
     fn name() -> &'static str {
-        "PrivateMemTable"
+        "PrivateIOTable"
     }
 }
-pub type PrivateMemCircuit<E> = DynVolatileRamCircuit<E, PrivateMemTable>;
+pub type PrivateIOCircuit<E> = DynVolatileRamCircuit<E, PrivateIOTable>;
 
 /// RegTable, fix size without offset
 #[derive(Clone)]