add gkr_iop state to ConstraintSystem/ZKVMWitness (wip, Apr 2)

Author: mcalancea

Cargo.lock

@@ -19,6 +19,7 @@ serde_json.workspace = true
 base64 = "0.22"
 ceno_emul = { path = "../ceno_emul" }
 ff_ext = { path = "../ff_ext" }
+gkr_iop = { path = "../gkr_iop" }
 mpcs = { path = "../mpcs" }
 multilinear_extensions = { version = "0", path = "../multilinear_extensions" }
 p3 = { path = "../p3" }

ceno_zkvm/Cargo.toml

@@ -2,13 +2,17 @@ use crate::{
     circuit_builder::{CircuitBuilder, ConstraintSystem},
     error::ZKVMError,
     expression::Expression,
-    instructions::Instruction,
+    instructions::{Instruction, riscv::dummy::LargeEcallDummy},
     state::StateCircuit,
     tables::{RMMCollections, TableCircuit},
     witness::LkMultiplicity,
 };
-use ceno_emul::{CENO_PLATFORM, Platform, StepRecord};
+use ceno_emul::{CENO_PLATFORM, KeccakSpec, Platform, StepRecord};
 use ff_ext::ExtensionField;
+use gkr_iop::{
+    ProtocolWitnessGenerator,
+    precompiles::{KeccakLayout, KeccakTrace},
+};
 use itertools::{Itertools, chain};
 use mpcs::PolynomialCommitmentScheme;
 use multilinear_extensions::{
@@ -130,6 +134,13 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> VerifyingKey<E, PCS>
     }
 }
 
+#[derive(Clone)]
+pub struct KeccakGKRIOP<E> {
+    pub chip: gkr_iop::chip::Chip,
+    pub layout: KeccakLayout<E>,
+    pub circuit: gkr_iop::gkr::GKRCircuit,
+}
+
 #[derive(Clone, Debug)]
 pub struct ProgramParams {
     pub platform: Platform,
@@ -154,6 +165,7 @@ pub struct ZKVMConstraintSystem<E: ExtensionField> {
     pub(crate) circuit_css: BTreeMap<String, ConstraintSystem<E>>,
     pub(crate) initial_global_state_expr: Expression<E>,
     pub(crate) finalize_global_state_expr: Expression<E>,
+    pub keccak_gkr_iop: Option<KeccakGKRIOP<E>>,
     pub params: ProgramParams,
 }
 
@@ -164,6 +176,7 @@ impl<E: ExtensionField> Default for ZKVMConstraintSystem<E> {
             initial_global_state_expr: Expression::ZERO,
             finalize_global_state_expr: Expression::ZERO,
             params: ProgramParams::default(),
+            keccak_gkr_iop: None,
         }
     }
 }
@@ -175,6 +188,25 @@ impl<E: ExtensionField> ZKVMConstraintSystem<E> {
             ..Default::default()
         }
     }
+
+    pub fn register_keccakf_circuit(
+        &mut self,
+    ) -> <LargeEcallDummy<E, KeccakSpec> as Instruction<E>>::InstructionConfig {
+        // Add GKR-IOP instance
+        let params = gkr_iop::precompiles::KeccakParams {};
+        let (layout, chip) = <KeccakLayout<E> as gkr_iop::ProtocolBuilder>::build(params);
+        let circuit = chip.gkr_circuit();
+
+        assert!(self.keccak_gkr_iop.is_none());
+        self.keccak_gkr_iop = Some(KeccakGKRIOP {
+            layout,
+            chip,
+            circuit,
+        });
+
+        self.register_opcode_circuit::<LargeEcallDummy<E, KeccakSpec>>()
+    }
+
     pub fn register_opcode_circuit<OC: Instruction<E>>(&mut self) -> OC::InstructionConfig {
         let mut cs = ConstraintSystem::new(|| format!("riscv_opcode/{}", OC::name()));
         let mut circuit_builder =
@@ -220,6 +252,14 @@ pub struct ZKVMFixedTraces<E: ExtensionField> {
 }
 
 impl<E: ExtensionField> ZKVMFixedTraces<E> {
+    pub fn register_keccakf_circuit(&mut self, _cs: &ZKVMConstraintSystem<E>) {
+        assert!(
+            self.circuit_fixed_traces
+                .insert(LargeEcallDummy::<E, KeccakSpec>::name(), None)
+                .is_none()
+        );
+    }
+
     pub fn register_opcode_circuit<OC: Instruction<E>>(&mut self, _cs: &ZKVMConstraintSystem<E>) {
         assert!(self.circuit_fixed_traces.insert(OC::name(), None).is_none());
     }
@@ -244,6 +284,7 @@ impl<E: ExtensionField> ZKVMFixedTraces<E> {
 
 #[derive(Default, Clone)]
 pub struct ZKVMWitnesses<E: ExtensionField> {
+    keccak_trace: <KeccakLayout<E> as gkr_iop::ProtocolWitnessGenerator<E>>::Trace,
     witnesses_opcodes: BTreeMap<String, RowMajorMatrix<E::BaseField>>,
     witnesses_tables: BTreeMap<String, RMMCollections<E::BaseField>>,
     lk_mlts: BTreeMap<String, LkMultiplicity>,
@@ -263,6 +304,44 @@ impl<E: ExtensionField> ZKVMWitnesses<E> {
         self.lk_mlts.get(name)
     }
 
+    pub fn assign_keccakf_circuit(
+        &mut self,
+        css: &mut ZKVMConstraintSystem<E>,
+        config: &<LargeEcallDummy<E, KeccakSpec> as Instruction<E>>::InstructionConfig,
+        records: Vec<StepRecord>,
+    ) -> Result<(), ZKVMError> {
+        // Ugly copy paste from assign_opcode_circuit, but we need to use the row major matrix
+        let cs = css
+            .get_cs(&LargeEcallDummy::<E, KeccakSpec>::name())
+            .unwrap();
+        let (witness, logup_multiplicity) = LargeEcallDummy::<E, KeccakSpec>::assign_instances(
+            config,
+            cs.num_witin as usize,
+            records,
+        )?;
+
+        // GKR-IOP-specific trace from row major witness
+        self.keccak_trace = KeccakTrace::from(witness.clone());
+
+        assert!(
+            self.witnesses_opcodes
+                .insert(LargeEcallDummy::<E, KeccakSpec>::name(), witness)
+                .is_none()
+        );
+        assert!(
+            !self
+                .witnesses_tables
+                .contains_key(&LargeEcallDummy::<E, KeccakSpec>::name())
+        );
+        assert!(
+            self.lk_mlts
+                .insert(LargeEcallDummy::<E, KeccakSpec>::name(), logup_multiplicity)
+                .is_none()
+        );
+
+        Ok(())
+    }
+
     pub fn assign_opcode_circuit<OC: Instruction<E>>(
         &mut self,
         cs: &ZKVMConstraintSystem<E>,

ceno_zkvm/src/structs.rs

@@ -12,6 +12,7 @@ version.workspace = true
 [dependencies]
 ark-std.workspace = true
 ff_ext = { path = "../ff_ext" }
+witness = { path = "../witness" }
 itertools.workspace = true
 multilinear_extensions = { version = "0.1.0", path = "../multilinear_extensions" }
 ndarray.workspace = true

gkr_iop/Cargo.toml

@@ -4,13 +4,13 @@ use super::Chip;
 
 impl Chip {
     /// Extract information from Chip that required in the GKR phase.
-    pub fn gkr_circuit(&'_ self) -> GKRCircuit<'_> {
+    pub fn gkr_circuit(&self) -> GKRCircuit {
         GKRCircuit {
-            layers: &self.layers,
+            layers: self.layers.clone(),
             n_challenges: self.n_challenges,
             n_evaluations: self.n_evaluations,
-            base_openings: &self.base_openings,
-            ext_openings: &self.ext_openings,
+            base_openings: self.base_openings.clone(),
+            ext_openings: self.ext_openings.clone(),
         }
     }
 }

gkr_iop/src/chip/protocol.rs

@@ -1,5 +1,5 @@
 use ff_ext::ExtensionField;
-use itertools::{Itertools, chain, izip};
+use itertools::{chain, izip, Itertools};
 use layer::{Layer, LayerWitness};
 use subprotocols::{expression::Point, sumcheck::SumcheckProof};
 use transcript::Transcript;
@@ -13,13 +13,13 @@ pub mod layer;
 pub mod mock;
 
 #[derive(Clone, Debug)]
-pub struct GKRCircuit<'a> {
-    pub layers: &'a [Layer],
+pub struct GKRCircuit {
+    pub layers: Vec<Layer>,
 
     pub n_challenges: usize,
     pub n_evaluations: usize,
-    pub base_openings: &'a [(usize, EvalExpression)],
-    pub ext_openings: &'a [(usize, EvalExpression)],
+    pub base_openings: Vec<(usize, EvalExpression)>,
+    pub ext_openings: Vec<(usize, EvalExpression)>,
 }
 
 #[derive(Clone, Debug)]
@@ -42,7 +42,7 @@ pub struct Evaluation<E: ExtensionField> {
 
 pub struct GKRClaims<Evaluation>(pub Vec<Evaluation>);
 
-impl GKRCircuit<'_> {
+impl GKRCircuit {
     pub fn prove<E>(
         &self,
         circuit_wit: GKRCircuitWitness<E>,
@@ -56,7 +56,7 @@ impl GKRCircuit<'_> {
         let mut evaluations = out_evals.to_vec();
         evaluations.resize(self.n_evaluations, PointAndEval::default());
         let mut challenges = challenges.to_vec();
-        let sumcheck_proofs = izip!(self.layers, circuit_wit.layers)
+        let sumcheck_proofs = izip!(&self.layers, circuit_wit.layers)
             .map(|(layer, layer_wit)| {
                 layer.prove(layer_wit, &mut evaluations, &mut challenges, transcript)
             })
@@ -85,7 +85,7 @@ impl GKRCircuit<'_> {
         let mut challenges = challenges.to_vec();
         let mut evaluations = out_evals.to_vec();
         evaluations.resize(self.n_evaluations, PointAndEval::default());
-        for (layer, layer_proof) in izip!(self.layers, sumcheck_proofs) {
+        for (layer, layer_proof) in izip!(self.layers.clone(), sumcheck_proofs) {
             layer.verify(layer_proof, &mut evaluations, &mut challenges, transcript)?;
         }
 
@@ -99,7 +99,7 @@ impl GKRCircuit<'_> {
         evaluations: &[PointAndEval<E>],
         challenges: &[E],
     ) -> Vec<Evaluation<E>> {
-        chain!(self.base_openings, self.ext_openings)
+        chain!(&self.base_openings, &self.ext_openings)
             .map(|(poly, eval)| {
                 let poly = *poly;
                 let PointAndEval { point, eval: value } = eval.evaluate(evaluations, challenges);

gkr_iop/src/gkr.rs

@@ -1,7 +1,7 @@
 use std::marker::PhantomData;
 
 use ff_ext::ExtensionField;
-use itertools::{Itertools, izip};
+use itertools::{izip, Itertools};
 use rand::rngs::OsRng;
 use subprotocols::{
     expression::{Expression, VectorType},
@@ -12,7 +12,7 @@ use thiserror::Error;
 
 use crate::{evaluation::EvalExpression, utils::SliceIterator};
 
-use super::{GKRCircuit, GKRCircuitWitness, layer::LayerType};
+use super::{layer::LayerType, GKRCircuit, GKRCircuitWitness};
 
 pub struct MockProver<E: ExtensionField>(PhantomData<E>);
 
@@ -35,7 +35,7 @@ pub enum MockProverError<F: ExtensionField> {
 
 impl<E: ExtensionField> MockProver<E> {
     pub fn check(
-        circuit: GKRCircuit<'_>,
+        circuit: GKRCircuit,
         circuit_wit: &GKRCircuitWitness<E>,
         mut evaluations: Vec<VectorType<E>>,
         mut challenges: Vec<E>,

gkr_iop/src/gkr/mock.rs

@@ -18,6 +18,8 @@ use crate::{
 };
 use ndarray::{range, s, ArrayView, Ix2, Ix3};
 
+use witness::RowMajorMatrix;
+
 use super::utils::{u64s_to_felts, zero_eval, CenoLookup};
 use p3_field::{extension::BinomialExtensionField, PrimeCharacteristicRing};
 
@@ -31,10 +33,10 @@ use transcript::BasicTranscript;
 type E = BinomialExtensionField<Goldilocks, 2>;
 
 #[derive(Clone, Debug, Default)]
-struct KeccakParams {}
+pub struct KeccakParams {}
 
 #[derive(Clone, Debug, Default)]
-struct KeccakLayout<E> {
+pub struct KeccakLayout<E> {
     params: KeccakParams,
 
     committed_bits_id: usize,
@@ -727,10 +729,19 @@ impl<E: ExtensionField> ProtocolBuilder for KeccakLayout<E> {
     }
 }
 
+#[derive(Clone, Default)]
 pub struct KeccakTrace {
     pub instances: Vec<[u32; KECCAK_INPUT_SIZE]>,
 }
 
+use p3_field::Field;
+
+impl<F: Field> From<RowMajorMatrix<F>> for KeccakTrace {
+    fn from(value: RowMajorMatrix<F>) -> Self {
+        unimplemented!();
+    }
+}
+
 impl<E> ProtocolWitnessGenerator<E> for KeccakLayout<E>
 where
     E: ExtensionField,

gkr_iop/src/precompiles/faster_keccak.rs

@@ -1,5 +1,8 @@
 mod faster_keccak;
 mod keccak_f;
 mod utils;
-pub use faster_keccak::run_faster_keccakf;
 pub use keccak_f::run_keccakf;
+pub use {
+    faster_keccak::run_faster_keccakf, faster_keccak::KeccakLayout, faster_keccak::KeccakParams,
+    faster_keccak::KeccakTrace,
+};