[precompile] part2 integrate keccak precompile into e2e flow

ceno_cli/src/commands/common_args/ceno.rs

@@ -319,7 +319,7 @@ impl CenoOptions {
 
 fn run_elf_inner<
     E: ExtensionField + LkMultiplicityKey,
-    PCS: PolynomialCommitmentScheme<E> + 'static,
+    PCS: PolynomialCommitmentScheme<E> + Serialize + 'static,
     P: AsRef<Path>,
 >(
     options: &CenoOptions,

ceno_host/tests/test_elf.rs

@@ -231,13 +231,13 @@ fn test_hashing() -> Result<()> {
 }
 
 #[test]
-fn test_ceno_rt_keccak() -> Result<()> {
-    let program_elf = ceno_examples::ceno_rt_keccak;
+fn test_keccak_syscall() -> Result<()> {
+    let program_elf = ceno_examples::keccak_syscall;
     let mut state = VMState::new_from_elf(unsafe_platform(), program_elf)?;
     let steps = run(&mut state)?;
 
     // Expect the program to have written successive states between Keccak permutations.
-    const ITERATIONS: usize = 4;
+    const ITERATIONS: usize = 100;
     let keccak_outs = sample_keccak_f(ITERATIONS);
 
     let all_messages = read_all_messages(&state);

ceno_zkvm/Cargo.toml

@@ -17,6 +17,7 @@ ceno_host = { path = "../ceno_host" }
 either.workspace = true
 ff_ext = { path = "../ff_ext" }
 gkr_iop = { path = "../gkr_iop" }
+keccakf = "0.1.2"
 mpcs = { path = "../mpcs" }
 multilinear_extensions = { version = "0", path = "../multilinear_extensions" }
 p3.workspace = true
@@ -29,6 +30,7 @@ transcript = { path = "../transcript" }
 witness = { path = "../witness" }
 
 itertools.workspace = true
+ndarray.workspace = true
 num-traits.workspace = true
 paste.workspace = true
 poseidon.workspace = true
@@ -101,3 +103,15 @@ name = "quadratic_sorting"
 [[bench]]
 harness = false
 name = "is_prime"
+
+[[bench]]
+harness = false
+name = "keccak"
+
+[[bench]]
+harness = false
+name = "bitwise_keccakf"
+
+[[bench]]
+harness = false
+name = "lookup_keccakf"

gkr_iop/benches/bitwise_keccakf.rs → ceno_zkvm/benches/bitwise_keccakf.rs

@@ -1,8 +1,8 @@
 use std::time::Duration;
 
+use ceno_zkvm::precompiles::{run_bitwise_keccakf, setup_bitwise_keccak_gkr_circuit};
 use criterion::*;
 use ff_ext::GoldilocksExt2;
-use gkr_iop::precompiles::{run_bitwise_keccakf, setup_bitwise_keccak_gkr_circuit};
 use itertools::Itertools;
 use mpcs::BasefoldDefault;
 use rand::{RngCore, SeedableRng};
@@ -36,7 +36,7 @@ fn keccak_f_fn(c: &mut Criterion) {
 
                         #[allow(clippy::unit_arg)]
                         run_bitwise_keccakf::<GoldilocksExt2, BasefoldDefault<GoldilocksExt2>>(
-                            setup_bitwise_keccak_gkr_circuit(),
+                            setup_bitwise_keccak_gkr_circuit().expect("setup circuit error"),
                             black_box(states),
                             false,
                             false,

ceno_zkvm/benches/keccak.rs

@@ -0,0 +1,107 @@
+use std::time::Duration;
+
+use ceno_emul::{Platform, Program};
+use ceno_host::CenoStdin;
+use ceno_zkvm::{
+    self,
+    e2e::{Checkpoint, Preset, run_e2e_with_checkpoint, setup_platform},
+    scheme::constants::MAX_NUM_VARIABLES,
+};
+mod alloc;
+use criterion::*;
+
+use ff_ext::GoldilocksExt2;
+use mpcs::{BasefoldDefault, SecurityLevel};
+
+criterion_group! {
+  name = keccak_prove_group;
+  config = Criterion::default().warm_up_time(Duration::from_millis(20000));
+  targets = keccak_prove,
+}
+
+criterion_main!(keccak_prove_group);
+
+const NUM_SAMPLES: usize = 10;
+
+type Pcs = BasefoldDefault<E>;
+type E = GoldilocksExt2;
+
+// Relevant init data for keccak run
+fn setup() -> (Program, Platform) {
+    let stack_size = 32768;
+    let heap_size = 2097152;
+    let pub_io_size = 16;
+    let program = Program::load_elf(ceno_examples::keccak_syscall, u32::MAX).unwrap();
+    let platform = setup_platform(Preset::Ceno, &program, stack_size, heap_size, pub_io_size);
+    (program, platform)
+}
+
+fn keccak_prove(c: &mut Criterion) {
+    let (program, platform) = setup();
+    // retrive 1 << 20th keccak element >> max_steps
+    let mut hints = CenoStdin::default();
+    let _ = hints.write(&vec![1, 2, 3]);
+    let max_steps = usize::MAX;
+    // estimate proof size data first
+    // let result = run_e2e_with_checkpoint::<E, Pcs>(
+    //     program.clone(),
+    //     platform.clone(),
+    //     &Vec::from(&hints),
+    //     &[],
+    //     max_steps,
+    //     MAX_NUM_VARIABLES,
+    //     SecurityLevel::default(),
+    //     Checkpoint::Complete,
+    // );
+    // let proof = result.proof.expect("PrepSanityCheck do not provide proof");
+    // let vk = result.vk.expect("PrepSanityCheck do not provide verifier");
+
+    // println!("e2e proof {}", proof);
+    // let transcript = BasicTranscript::new(b"riscv");
+    // let verifier = ZKVMVerifier::<E, Pcs>::new(vk);
+    // assert!(
+    //     verifier
+    //         .verify_proof_halt(proof, transcript, false)
+    //         .expect("verify proof return with error"),
+    // );
+    // println!();
+    // println!("max_steps = {}", max_steps);
+
+    // expand more input size once runtime is acceptable
+    let mut group = c.benchmark_group(format!("keccak_max_steps_{}", max_steps));
+    group.sample_size(NUM_SAMPLES);
+
+    // Benchmark the proving time
+    group.bench_function(
+        BenchmarkId::new("prove_keccak", format!("keccak_max_steps_{}", max_steps)),
+        |b| {
+            b.iter_custom(|iters| {
+                let mut time = Duration::new(0, 0);
+                for _ in 0..iters {
+                    let result = run_e2e_with_checkpoint::<E, Pcs>(
+                        program.clone(),
+                        platform.clone(),
+                        &Vec::from(&hints),
+                        &[],
+                        max_steps,
+                        MAX_NUM_VARIABLES,
+                        SecurityLevel::default(),
+                        Checkpoint::PrepE2EProving,
+                    );
+                    let instant = std::time::Instant::now();
+                    result.next_step();
+                    let elapsed = instant.elapsed();
+                    println!(
+                        "Keccak::create_proof, max_steps = {}, time = {}",
+                        max_steps,
+                        elapsed.as_secs_f64()
+                    );
+                    time += elapsed;
+                }
+                time
+            });
+        },
+    );
+
+    group.finish();
+}

gkr_iop/benches/lookup_keccakf.rs → ceno_zkvm/benches/lookup_keccakf.rs

@@ -1,8 +1,8 @@
 use std::time::Duration;
 
+use ceno_zkvm::precompiles::{run_faster_keccakf, setup_lookup_keccak_gkr_circuit};
 use criterion::*;
 use ff_ext::GoldilocksExt2;
-use gkr_iop::precompiles::{run_faster_keccakf, setup_lookup_keccak_gkr_circuit};
 
 use itertools::Itertools;
 use mpcs::BasefoldDefault;
@@ -38,7 +38,8 @@ fn keccak_f_fn(c: &mut Criterion) {
 
                         let instant = std::time::Instant::now();
 
-                        let circuit = setup_lookup_keccak_gkr_circuit();
+                        let circuit =
+                            setup_lookup_keccak_gkr_circuit().expect("setup circuit error");
                         #[allow(clippy::unit_arg)]
                         let _ =
                             run_faster_keccakf::<GoldilocksExt2, BasefoldDefault<GoldilocksExt2>>(

ceno_zkvm/benches/riscv_add.rs

@@ -41,9 +41,9 @@ const NUM_SAMPLES: usize = 10;
 fn bench_add(c: &mut Criterion) {
     type Pcs = BasefoldDefault<E>;
     let mut zkvm_cs = ZKVMConstraintSystem::default();
-    let _ = zkvm_cs.register_opcode_circuit::<AddInstruction<E>>();
+    let config = zkvm_cs.register_opcode_circuit::<AddInstruction<E>>();
     let mut zkvm_fixed_traces = ZKVMFixedTraces::default();
-    zkvm_fixed_traces.register_opcode_circuit::<AddInstruction<E>>(&zkvm_cs);
+    zkvm_fixed_traces.register_opcode_circuit::<AddInstruction<E>>(&zkvm_cs, &config);
 
     let param = Pcs::setup(1 << MAX_NUM_VARIABLES, SecurityLevel::default()).unwrap();
     let (pp, vp) = Pcs::trim(param, 1 << MAX_NUM_VARIABLES).unwrap();

gkr_iop/src/bin/bitwise_keccak.rs → ceno_zkvm/src/bin/bitwise_keccak.rs

@@ -1,6 +1,6 @@
+use ceno_zkvm::precompiles::{run_bitwise_keccakf, setup_bitwise_keccak_gkr_circuit};
 use clap::{Parser, command};
 use ff_ext::GoldilocksExt2;
-use gkr_iop::precompiles::{run_bitwise_keccakf, setup_bitwise_keccak_gkr_circuit};
 use itertools::Itertools;
 use mpcs::BasefoldDefault;
 use rand::{RngCore, SeedableRng};
@@ -70,5 +70,10 @@ fn main() {
     let states: Vec<[u64; 25]> = (0..num_instance)
         .map(|_| std::array::from_fn(|_| rng.next_u64()))
         .collect_vec();
-    run_bitwise_keccakf::<E, Pcs>(setup_bitwise_keccak_gkr_circuit(), states, false, false);
+    run_bitwise_keccakf::<E, Pcs>(
+        setup_bitwise_keccak_gkr_circuit().expect("setup circuit error"),
+        states,
+        false,
+        false,
+    );
 }

gkr_iop/src/bin/lookup_keccak.rs → ceno_zkvm/src/bin/lookup_keccak.rs

@@ -1,6 +1,6 @@
+use ceno_zkvm::precompiles::{run_faster_keccakf, setup_lookup_keccak_gkr_circuit};
 use clap::{Parser, command};
 use ff_ext::GoldilocksExt2;
-use gkr_iop::precompiles::{run_faster_keccakf, setup_lookup_keccak_gkr_circuit};
 use itertools::Itertools;
 use mpcs::BasefoldDefault;
 use rand::{RngCore, SeedableRng};
@@ -71,7 +71,12 @@ fn main() {
         .map(|_| std::array::from_fn(|_| rng.next_u64()))
         .collect_vec();
     let circuit_setup = setup_lookup_keccak_gkr_circuit();
-    let proof =
-        run_faster_keccakf::<E, Pcs>(circuit_setup, states, true, true).expect("generate proof");
+    let proof = run_faster_keccakf::<E, Pcs>(
+        circuit_setup.expect("setup circuit error"),
+        states,
+        true,
+        true,
+    )
+    .expect("generate proof");
     tracing::info!("lookup keccak proof stat: {}", proof);
 }