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

@@ -66,4 +66,8 @@ pub trait Instruction<E: ExtensionField> {
         raw_witin.padding_by_strategy();
         Ok((raw_witin, lk_multiplicity))
     }
+
+    fn add_gkr_iop() {
+        unimplemented!();
+    }
 }

ceno_zkvm/src/instructions.rs

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

@@ -16,6 +16,7 @@ pub mod e2e;
 pub mod expression;
 pub mod gadgets;
 mod keygen;
+pub mod precompiles;
 pub mod state;
 pub mod stats;
 pub mod structs;

ceno_zkvm/src/lib.rs

@@ -2,13 +2,20 @@ use crate::{
     circuit_builder::{CircuitBuilder, ConstraintSystem},
     error::ZKVMError,
     expression::Expression,
-    instructions::Instruction,
+    instructions::{
+        Instruction,
+        riscv::{dummy::LargeEcallDummy, precompiles::keccakf_circuit},
+    },
     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 +137,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 +168,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 +179,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 +191,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 +255,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 +287,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 +307,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(),
         }
     }
 }