[precompile] mem read/write rlc within precompile circuit (#973)

hero78119 · web-flow · commit d6d2c87ca32b · 2025-07-04T09:21:52.000Z
This PR add the missing witness: current ts, mem read ts, mem start addr
in keccak precompile.
Those information are needed for consistent check.
Beside, mem read/write are rlc with those record

&gt; This PR is not fully sound, as we need to revamp current eq with
selector logic

### benchmark

The observed performance regression is expected, as the additional
workload, whether handled inside or outside the precompile circuit, is
necessary and unavoidable.

| Benchmark | Median Time (s) | Median Change (%) |

|----------------------------------|------------------|----------------------------------------|
| keccak_lookup_f_4096 | 1.0320 | +4.7694% (Performance has regressed) |
| keccak_lookup_f_8192 | 2.0793 | +9.6826% (Performance has regressed) |
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/ceno_zkvm/src/chip_handler/global_state.rs b/ceno_zkvm/src/chip_handler/global_state.rs
@@ -12,7 +12,6 @@ impl<E: ExtensionField> GlobalStateRegisterMachineChipOperations<E> for CircuitB
             pc,
             ts,
         ];
-
         self.read_record(|| "state_in", RAMType::GlobalState, record)
     }
 
diff --git a/ceno_zkvm/src/e2e.rs b/ceno_zkvm/src/e2e.rs
@@ -706,7 +706,7 @@ pub fn run_e2e_with_checkpoint<
 
     // proving
     let backend: CpuBackend<E, PCS> = CpuBackend::new();
-    let device: CpuProver<CpuBackend<E, PCS>> = CpuProver::new(backend);
+    let device = CpuProver::new(backend);
     let mut prover = ZKVMProver::new(pk, device);
 
     if is_mock_proving {
@@ -780,7 +780,7 @@ pub fn run_e2e_proof<
 
     // proving
     let backend: CpuBackend<E, PCS> = CpuBackend::new();
-    let device: CpuProver<CpuBackend<E, PCS>> = CpuProver::new(backend);
+    let device = CpuProver::new(backend);
     let mut prover = ZKVMProver::new(pk, device);
 
     if is_mock_proving {
diff --git a/ceno_zkvm/src/instructions/riscv/dummy/dummy_ecall.rs b/ceno_zkvm/src/instructions/riscv/dummy/dummy_ecall.rs
@@ -1,6 +1,6 @@
 use std::marker::PhantomData;
 
-use ceno_emul::{Change, InsnKind, KeccakSpec, StepRecord, SyscallSpec};
+use ceno_emul::{ByteAddr, Change, InsnKind, KeccakSpec, StepRecord, SyscallSpec};
 use ff_ext::{ExtensionField, SmallField};
 use itertools::{Itertools, zip_eq};
 use witness::RowMajorMatrix;
@@ -22,7 +22,9 @@ use multilinear_extensions::{ToExpr, WitIn};
 
 use gkr_iop::{
     ProtocolWitnessGenerator,
-    precompiles::{AND_LOOKUPS, KeccakLayout, KeccakTrace, RANGE_LOOKUPS, XOR_LOOKUPS},
+    precompiles::{
+        AND_LOOKUPS, KECCAK_INPUT32_SIZE, KeccakLayout, KeccakTrace, RANGE_LOOKUPS, XOR_LOOKUPS,
+    },
 };
 
 /// LargeEcallDummy can handle any instruction and produce its effects,
@@ -208,8 +210,14 @@ impl<E: ExtensionField> GKRIOPInstruction<E> for LargeEcallDummy<E, KeccakSpec>
                     .unwrap()
             })
             .collect_vec();
+        let num_instances = instances.len();
 
-        layout.phase1_witness_group(KeccakTrace { instances })
+        layout.phase1_witness_group(KeccakTrace {
+            instances,
+            ram_start_addr: vec![ByteAddr::from(0); num_instances],
+            cur_ts: vec![0; num_instances],
+            read_ts: vec![[0; KECCAK_INPUT32_SIZE]; num_instances],
+        })
     }
 
     fn assign_instance_with_gkr_iop(
diff --git a/ceno_zkvm/src/scheme/tests.rs b/ceno_zkvm/src/scheme/tests.rs
@@ -13,9 +13,7 @@ use crate::{
         prover::ZkVMCpuProver,
     },
     set_val,
-    structs::{
-        PointAndEval, RAMType::Register, ZKVMConstraintSystem, ZKVMFixedTraces, ZKVMWitnesses,
-    },
+    structs::{PointAndEval, RAMType, ZKVMConstraintSystem, ZKVMFixedTraces, ZKVMWitnesses},
     tables::{ProgramTableCircuit, U16TableCircuit},
     witness::LkMultiplicity,
 };
@@ -64,8 +62,8 @@ impl<E: ExtensionField, const L: usize, const RW: usize> Instruction<E> for Test
         let reg_id = cb.create_witin(|| "reg_id");
         (0..RW).try_for_each(|_| {
             let record = vec![1.into(), reg_id.expr()];
-            cb.read_record(|| "read", Register, record.clone())?;
-            cb.write_record(|| "write", Register, record)?;
+            cb.read_record(|| "read", RAMType::Register, record.clone())?;
+            cb.write_record(|| "write", RAMType::Register, record)?;
             Result::<(), ZKVMError>::Ok(())
         })?;
         (0..L).try_for_each(|_| {
diff --git a/ceno_zkvm/src/structs.rs b/ceno_zkvm/src/structs.rs
@@ -7,12 +7,11 @@ use crate::{
     witness::LkMultiplicity,
 };
 use ceno_emul::{CENO_PLATFORM, KeccakSpec, Platform, StepRecord, SyscallSpec};
-use either::Either;
 use ff_ext::ExtensionField;
 use gkr_iop::{LookupTable, gkr::GKRCircuit, precompiles::KeccakLayout};
 use itertools::Itertools;
 use mpcs::{Point, PolynomialCommitmentScheme};
-use multilinear_extensions::{Expression, impl_expr_from_unsigned};
+use multilinear_extensions::Expression;
 use serde::{Deserialize, Serialize, de::DeserializeOwned};
 use std::{
     collections::{BTreeMap, HashMap},
@@ -45,14 +44,7 @@ pub type ChallengeId = u16;
 
 pub type ROMType = LookupTable;
 
-#[derive(Clone, Debug, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
-pub enum RAMType {
-    GlobalState,
-    Register,
-    Memory,
-}
-
-impl_expr_from_unsigned!(RAMType);
+pub type RAMType = gkr_iop::RAMType;
 
 pub type PointAndEval<F> = multilinear_extensions::mle::PointAndEval<F>;
 
diff --git a/gkr_iop/Cargo.toml b/gkr_iop/Cargo.toml
@@ -12,6 +12,7 @@ version.workspace = true
 [dependencies]
 ark-std = { version = "0.5" }
 bincode.workspace = true
+ceno_emul = { path = "../ceno_emul" }
 clap.workspace = true
 either.workspace = true
 ff_ext = { path = "../ff_ext" }
diff --git a/gkr_iop/src/gkr/layer_constraint_system.rs b/gkr_iop/src/gkr/layer_constraint_system.rs
@@ -1,15 +1,15 @@
 use std::{cmp::Ordering, collections::BTreeMap};
 
-use ff_ext::ExtensionField;
-use itertools::{Itertools, chain, izip};
-use multilinear_extensions::{Expression, Fixed, ToExpr, WitnessId, rlc_chip_record};
-use p3_field::FieldAlgebra;
-
 use crate::{
-    LookupTable,
+    LookupTable, RAMType,
     evaluation::EvalExpression,
     gkr::layer::{Layer, LayerType, ROTATION_OPENING_COUNT},
 };
+use ceno_emul::Tracer;
+use ff_ext::ExtensionField;
+use itertools::{Itertools, chain, izip};
+use multilinear_extensions::{Expression, Fixed, ToExpr, WitnessId, rlc_chip_record};
+use p3_field::FieldAlgebra;
 
 #[derive(Clone, Debug, Default)]
 pub struct RotationParams<E: ExtensionField> {
@@ -39,6 +39,9 @@ pub struct LayerConstraintSystem<E: ExtensionField> {
     pub xor_lookups: Vec<Expression<E>>,
     pub range_lookups: Vec<Expression<E>>,
 
+    pub ram_read: Vec<Expression<E>>,
+    pub ram_write: Vec<Expression<E>>,
+
     // global challenge
     pub alpha: Expression<E>,
     pub beta: Expression<E>,
@@ -66,6 +69,8 @@ impl<E: ExtensionField> LayerConstraintSystem<E> {
             and_lookups: vec![],
             xor_lookups: vec![],
             range_lookups: vec![],
+            ram_read: vec![],
+            ram_write: vec![],
             alpha,
             beta,
         }
@@ -89,29 +94,31 @@ impl<E: ExtensionField> LayerConstraintSystem<E> {
         self.expr_names.push(name);
     }
 
-    fn rlc_table_lookup(
-        &self,
-        lookup_table: LookupTable,
-        values: impl Iterator<Item = Expression<E>>,
-    ) -> Expression<E> {
-        rlc_chip_record(
-            chain!(
-                std::iter::once(E::BaseField::from_canonical_u64(lookup_table as u64).expr()),
-                values
-            )
-            .collect_vec(),
+    pub fn lookup_and8(&mut self, a: Expression<E>, b: Expression<E>, c: Expression<E>) {
+        let rlc_record = rlc_chip_record(
+            vec![
+                E::BaseField::from_canonical_u64(LookupTable::And as u64).expr(),
+                a,
+                b,
+                c,
+            ],
             self.alpha.clone(),
             self.beta.clone(),
-        )
-    }
-
-    pub fn lookup_and8(&mut self, a: Expression<E>, b: Expression<E>, c: Expression<E>) {
-        let rlc_record = self.rlc_table_lookup(LookupTable::And, [a, b, c].iter().cloned());
+        );
         self.and_lookups.push(rlc_record);
     }
 
     pub fn lookup_xor8(&mut self, a: Expression<E>, b: Expression<E>, c: Expression<E>) {
-        let rlc_record = self.rlc_table_lookup(LookupTable::Xor, [a, b, c].iter().cloned());
+        let rlc_record = rlc_chip_record(
+            vec![
+                E::BaseField::from_canonical_u64(LookupTable::Xor as u64).expr(),
+                a,
+                b,
+                c,
+            ],
+            self.alpha.clone(),
+            self.beta.clone(),
+        );
         self.xor_lookups.push(rlc_record);
     }
 
@@ -120,17 +127,90 @@ impl<E: ExtensionField> LayerConstraintSystem<E> {
     /// `value << (16 - size)`.
     pub fn lookup_range(&mut self, value: Expression<E>, size: usize) {
         assert!(size <= 16);
-        self.range_lookups
-            .push(self.rlc_table_lookup(LookupTable::U16, vec![value.clone()].into_iter()));
+        let rlc_record = rlc_chip_record(
+            vec![
+                E::BaseField::from_canonical_u64(LookupTable::U16 as u64).expr(),
+                value.clone(),
+            ],
+            self.alpha.clone(),
+            self.beta.clone(),
+        );
+        self.range_lookups.push(rlc_record);
         if size < 16 {
-            self.range_lookups.push(
-                self.rlc_table_lookup(
-                    LookupTable::U16,
-                    [value * E::BaseField::from_canonical_u64(1 << (16 - size)).expr()]
-                        .iter()
-                        .cloned(),
-                ),
-            )
+            let rlc_record = rlc_chip_record(
+                vec![
+                    E::BaseField::from_canonical_u64(LookupTable::U16 as u64).expr(),
+                    value * E::BaseField::from_canonical_u64(1 << (16 - size)).expr(),
+                ],
+                self.alpha.clone(),
+                self.beta.clone(),
+            );
+            self.range_lookups.push(rlc_record)
+        }
+    }
+
+    /// records RAM write operations into the `ram_write` trace vector using RLC encoding.
+    ///
+    /// this function appends one RLC-encoded record per word written to memory,
+    /// starting from `mem_start_addr`. Each record includes:
+    /// - The operation type (assumed to be `RAMType::Memory`)
+    /// - The memory address of the word (adjusted by `4 * index`)
+    /// - The value being written
+    /// - The timestamp of the write
+    pub fn ram_write_record(
+        &mut self,
+        mem_start_addr: Expression<E>,
+        values: Vec<Expression<E>>,
+        cur_ts: Expression<E>,
+    ) {
+        for (idx, value) in values.into_iter().enumerate() {
+            let rlc_record = rlc_chip_record(
+                vec![
+                    E::BaseField::from_canonical_u64(RAMType::Memory as u64).expr(),
+                    mem_start_addr.clone()
+                    // `4` is num bytes per word
+                    // TODO fetch from constant
+                        + E::BaseField::from_canonical_u64((4 * idx) as u64).expr(),
+                    value,
+                    cur_ts.clone() + E::BaseField::from_canonical_u64(Tracer::SUBCYCLE_MEM).expr(),
+                ],
+                self.alpha.clone(),
+                self.beta.clone(),
+            );
+            self.ram_write.push(rlc_record);
+        }
+    }
+
+    /// records RAM read operations into the `ram_read` trace vector using RLC encoding.
+    ///
+    /// this function appends one RLC-encoded record per word written to memory,
+    /// starting from `mem_start_addr`. Each record includes:
+    /// - The operation type (assumed to be `RAMType::Memory`)
+    /// - The memory address of the word (adjusted by `4 * index`)
+    /// - The value being read
+    /// - The ts corresponding to each value
+    pub fn ram_read_record(
+        &mut self,
+        mem_start_addr: Expression<E>,
+        values: Vec<Expression<E>>,
+        ts: Vec<Expression<E>>,
+    ) {
+        assert_eq!(values.len(), ts.len());
+        for (idx, (value, ts)) in izip!(values, ts).enumerate() {
+            let rlc_record = rlc_chip_record(
+                vec![
+                    E::BaseField::from_canonical_u64(RAMType::Memory as u64).expr(),
+                    mem_start_addr.clone()
+                    // `4` is num bytes per word
+                    // TODO fetch from constant
+                        + E::BaseField::from_canonical_u64((4 * idx) as u64).expr(),
+                    value,
+                    ts,
+                ],
+                self.alpha.clone(),
+                self.beta.clone(),
+            );
+            self.ram_read.push(rlc_record);
         }
     }
 
@@ -270,8 +350,27 @@ impl<E: ExtensionField> LayerConstraintSystem<E> {
         layer_name: String,
         in_expr_evals: Vec<usize>,
         n_challenges: usize,
+        ram_write_evals: impl ExactSizeIterator<Item = (Option<Expression<E>>, usize)>,
+        ram_read_evals: impl ExactSizeIterator<Item = (Option<Expression<E>>, usize)>,
         lookup_evals: impl ExactSizeIterator<Item = (Option<Expression<E>>, usize)>,
     ) -> Layer<E> {
+        // process ram read/write record
+        assert_eq!(ram_write_evals.len(), self.ram_write.len(),);
+        assert_eq!(ram_read_evals.len(), self.ram_read.len(),);
+
+        for (idx, ram_expr, ram_eval) in izip!(
+            0..,
+            chain!(self.ram_write.clone(), self.ram_read.clone(),),
+            ram_write_evals.chain(ram_read_evals)
+        ) {
+            self.add_non_zero_constraint(
+                ram_expr - E::BaseField::ONE.expr(), // ONE is for padding
+                (ram_eval.0, EvalExpression::Single(ram_eval.1)),
+                format!("round 0th: {idx}th ram read/write operation"),
+            );
+        }
+
+        // process lookup records
         assert_eq!(
             lookup_evals.len(),
             self.and_lookups.len() + self.xor_lookups.len() + self.range_lookups.len()
@@ -336,6 +435,7 @@ impl<E: ExtensionField> LayerConstraintSystem<E> {
                     is_linear_so_far && t.is_linear()
                 });
 
+        // process evaluation group by eq expression
         let mut eq_map = BTreeMap::new();
         izip!(
             evals.into_iter(),
diff --git a/gkr_iop/src/lib.rs b/gkr_iop/src/lib.rs
@@ -2,9 +2,10 @@ use std::marker::PhantomData;
 
 use crate::hal::{ProtocolWitnessGeneratorProver, ProverDevice};
 use chip::Chip;
+use either::Either;
 use ff_ext::ExtensionField;
 use gkr::{GKRCircuit, GKRCircuitOutput, GKRCircuitWitness, layer::LayerWitness};
-use multilinear_extensions::mle::ArcMultilinearExtension;
+use multilinear_extensions::{Expression, impl_expr_from_unsigned, mle::ArcMultilinearExtension};
 use strum_macros::EnumIter;
 use transcript::Transcript;
 use utils::infer_layer_witness;
@@ -103,3 +104,13 @@ pub enum LookupTable {
     Pow,         // a ** b where a is 2 and b is 5-bit value
     Instruction, // Decoded instruction from the fixed program.
 }
+
+#[derive(Clone, Debug, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
+#[repr(usize)]
+pub enum RAMType {
+    GlobalState,
+    Register,
+    Memory,
+}
+
+impl_expr_from_unsigned!(RAMType);
diff --git a/gkr_iop/src/precompiles/bitwise_keccakf.rs b/gkr_iop/src/precompiles/bitwise_keccakf.rs
diff --git a/gkr_iop/src/precompiles/lookup_keccakf.rs b/gkr_iop/src/precompiles/lookup_keccakf.rs
diff --git a/gkr_iop/src/precompiles/mod.rs b/gkr_iop/src/precompiles/mod.rs

Original file line number	Diff line number	Diff line change
`@@ -12,7 +12,6 @@ impl<E: ExtensionField> GlobalStateRegisterMachineChipOperations<E> for CircuitB`
`12`	`12`	`pc,`
`13`	`13`	`ts,`
`14`	`14`	`];`
`15`		`-`
`16`	`15`	`self.read_record(\|\| "state_in", RAMType::GlobalState, record)`
`17`	`16`	`}`
`18`	`17`