keccakf precompile + GKR-IOP integration

Cargo.toml

@@ -36,6 +36,7 @@ cfg-if = "1.0"
 criterion = { version = "0.5", features = ["html_reports"] }
 crossbeam-channel = "0.5"
 itertools = "0.13"
+ndarray = "*"
 num-bigint = { version = "0.4.6" }
 num-derive = "0.4"
 num-traits = "0.2"

ceno_zkvm/Cargo.toml

@@ -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/src/structs.rs

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

gkr_iop/Cargo.toml

@@ -12,12 +12,26 @@ 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
 p3-field.workspace = true
 p3-goldilocks.workspace = true
 rand.workspace = true
 rayon.workspace = true
 subprotocols = { path = "../subprotocols" }
 thiserror = "1"
+tiny-keccak.workspace = true
 transcript = { path = "../transcript" }
+
+[dev-dependencies]
+criterion.workspace = true
+
+[[bench]]
+harness = false
+name = "keccak_f"
+
+[[bench]]
+harness = false
+name = "faster_keccak"

gkr_iop/benches/faster_keccak.rs

@@ -0,0 +1,38 @@
+use std::time::Duration;
+
+use criterion::*;
+use gkr_iop::precompiles::run_faster_keccakf;
+
+use rand::{Rng, SeedableRng};
+criterion_group!(benches, keccak_f_fn);
+criterion_main!(benches);
+
+const NUM_SAMPLES: usize = 10;
+
+fn keccak_f_fn(c: &mut Criterion) {
+    // expand more input size once runtime is acceptable
+    let mut group = c.benchmark_group(format!("keccak_f"));
+    group.sample_size(NUM_SAMPLES);
+
+    // Benchmark the proving time
+    group.bench_function(BenchmarkId::new("keccak_f", format!("keccak_f")), |b| {
+        b.iter_custom(|iters| {
+            let mut time = Duration::new(0, 0);
+            for _ in 0..iters {
+                // Use seeded rng for debugging convenience
+                let mut rng = rand::rngs::StdRng::seed_from_u64(42);
+                let state1: [u64; 25] = std::array::from_fn(|_| rng.gen());
+                let state2: [u64; 25] = std::array::from_fn(|_| rng.gen());
+
+                let instant = std::time::Instant::now();
+                let _ = black_box(run_faster_keccakf(vec![state1, state2], false, false));
+                let elapsed = instant.elapsed();
+                time += elapsed;
+            }
+
+            time
+        });
+    });
+
+    group.finish();
+}

gkr_iop/benches/keccak_f.rs

@@ -0,0 +1,38 @@
+use std::time::Duration;
+
+use criterion::*;
+use gkr_iop::precompiles::{run_faster_keccakf, run_keccakf};
+use p3_field::extension::BinomialExtensionField;
+use p3_goldilocks::Goldilocks;
+use rand::{Rng, SeedableRng};
+criterion_group!(benches, keccak_f_fn);
+criterion_main!(benches);
+
+const NUM_SAMPLES: usize = 10;
+
+fn keccak_f_fn(c: &mut Criterion) {
+    // expand more input size once runtime is acceptable
+    let mut group = c.benchmark_group(format!("keccak_f"));
+    group.sample_size(NUM_SAMPLES);
+
+    // Benchmark the proving time
+    group.bench_function(BenchmarkId::new("keccak_f", format!("keccak_f")), |b| {
+        b.iter_custom(|iters| {
+            let mut time = Duration::new(0, 0);
+            for _ in 0..iters {
+                // Use seeded rng for debugging convenience
+                let mut rng = rand::rngs::StdRng::seed_from_u64(42);
+                let state: [u64; 25] = std::array::from_fn(|_| rng.gen());
+
+                let instant = std::time::Instant::now();
+                let _ = black_box(run_keccakf(state, false, false));
+                let elapsed = instant.elapsed();
+                time += elapsed;
+            }
+
+            time
+        });
+    });
+
+    group.finish();
+}