gkr_iop config attempt (wip)

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

@@ -18,6 +18,7 @@ pub mod jump;
 pub mod logic;
 pub mod logic_imm;
 pub mod mul;
+pub mod precompiles;
 pub mod shift;
 pub mod shift_imm;
 pub mod slt;

ceno_zkvm/src/instructions/riscv.rs

@@ -0,0 +1,86 @@
+use std::{cell::RefCell, marker::PhantomData, sync::RwLock};
+
+use ceno_emul::{Change, InsnKind, KECCAK_WORDS, StepRecord, WORD_SIZE};
+use ff_ext::ExtensionField;
+use itertools::Itertools;
+
+use super::super::insn_base::WriteMEM;
+use crate::{
+    Value,
+    circuit_builder::CircuitBuilder,
+    error::ZKVMError,
+    expression::{ToExpr, WitIn},
+    instructions::{
+        Instruction,
+        riscv::{constants::UInt, insn_base::WriteRD},
+    },
+    set_val,
+    witness::LkMultiplicity,
+};
+
+use gkr_iop::{
+    ProtocolBuilder, ProtocolWitnessGenerator,
+    chip::Chip,
+    gkr::GKRCircuit,
+    precompiles::{KeccakLayout, KeccakParams},
+};
+
+pub struct KeccakF<E>(PhantomData<E>);
+pub struct KeccakFConfig<E> {
+    chip: Chip,
+    layout: KeccakLayout<E>,
+    gkr_circuit: GKRCircuit,
+    phase1_wits: 
+}
+
+impl<E: ExtensionField> Instruction<E> for KeccakF<E> {
+    type InstructionConfig = KeccakFConfig<E>;
+
+    fn name() -> String {
+        format!("keccakf")
+    }
+
+    fn construct_circuit(cb: &mut CircuitBuilder<E>) -> Result<Self::InstructionConfig, ZKVMError> {
+        let params = KeccakParams {};
+        let (layout, chip) = KeccakLayout::build(params);
+        let gkr_circuit = chip.gkr_circuit();
+
+        Ok(KeccakFConfig {
+            chip,
+            layout,
+            gkr_circuit,
+            phase1_wits:
+        })
+    }
+
+    fn assign_instance(
+        config: &Self::InstructionConfig,
+        instance: &mut [E::BaseField],
+        lk_multiplicity: &mut LkMultiplicity,
+        step: &StepRecord,
+    ) -> Result<(), ZKVMError> {
+        // let ops = &step.syscall().expect("syscall step");
+
+        // // Assign instruction.
+        // config
+        //     .dummy_insn
+        //     .assign_instance(instance, lk_multiplicity, step)?;
+
+        // set_val!(instance, config.start_addr, u64::from(ops.mem_ops[0].addr));
+
+        // // Assign registers.
+        // for ((value, writer), op) in config.reg_writes.iter().zip_eq(&ops.reg_ops) {
+        //     value.assign_value(instance, Value::new_unchecked(op.value.after));
+        //     writer.assign_op(instance, lk_multiplicity, step.cycle(), op)?;
+        // }
+
+        // // Assign memory.
+        // for ((value, writer), op) in config.mem_writes.iter().zip_eq(&ops.mem_ops) {
+        //     set_val!(instance, value.before, op.value.before as u64);
+        //     set_val!(instance, value.after, op.value.after as u64);
+        //     writer.assign_op(instance, lk_multiplicity, step.cycle(), op)?;
+        // }
+
+        Ok(())
+    }
+}

ceno_zkvm/src/instructions/riscv/precompiles/keccakf_circuit.rs

@@ -0,0 +1 @@
+pub mod keccakf_circuit;

ceno_zkvm/src/instructions/riscv/precompiles/mod.rs

@@ -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

@@ -31,10 +31,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,

gkr_iop/src/precompiles/faster_keccak.rs

@@ -1,5 +1,7 @@
 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,
+};