[precompile] part1 integrate keccak precompile into e2e flow (#976)

This raised PR is not fully sound yet, however it serve as a step
milestone for easier to review

### change scope
- define new `ComposedConstrainSystem` to wrap `original constrain
system + gkr-iop circuit`. I decide for our first version, tower
argument will remain in original constrain system because its dynamism
layers not fit well into gkr-iop design. So I tend to keep both
constrain system separately, then invoke css in proving flow
accordingly.
- define `Keccak` as a new (ecall-)opcode followed our conventions, and
clean up previous workaround as `LargeDummy<KccakSpec>` which is not
used for this purpose
- refactor lookup tables utilization from zkvm crate to gkr-iop so both
crate could be shared.

### TODOs in followup PR(s)
- build correct circuit of `keccak` ecall opcode to deal with
state(memory) consistency => existing ceno zkvm circuit builder has few
good utilization to wrap boilerplate code. But gkr-iop define it's own
constrain system so it's hard to leverage existing utilization.
- assign witness mechanism
- integrate into main prover flow + mpcs

### benchmarks

On fibonacci e2e against master branch

| Benchmark | Median Time (s) | Median Change (%) |

|----------------------------------|------------------|----------------------------------------------|
| fibonacci_max_steps_1048576 | 2.0902 | -1.1437% (No change in
performance detected) |
| fibonacci_max_steps_2097152 | 3.5378 | -0.5016% (No change in
performance detected) |
| fibonacci_max_steps_4194304 | 6.8585 | -1.9079% (No change in
performance detected) |

and lookup keccak

| Benchmark | Median Time (s) | Median Change (%) |

|----------------------------------|------------------|----------------------------------------------|
| keccak_lookup_f_4096 | 0.88786 | +1.8583% (No change in performance
detected) |
| keccak_lookup_f_8192 | 1.7728 | +0.9121% (No change in performance
detected) |

both benchmark remain unchanged

Author: hero78119

Cargo.lock

@@ -77,6 +77,7 @@ strum_macros = "0.26"
 substrate-bn = { version = "0.6.0" }
 sumcheck = { path = "sumcheck" }
 thiserror = "1" # do we need this?
+thread_local = "1.1"
 tiny-keccak = { version = "2.0.2", features = ["keccak"] }
 tracing = { version = "0.1", features = [
   "attributes",

Cargo.toml

@@ -47,7 +47,7 @@ generic_static = "0.2"
 parse-size = "1.1"
 rand.workspace = true
 tempfile = "3.14"
-thread_local = "1.1"
+thread_local.workspace = true
 tiny-keccak.workspace = true
 
 [target.'cfg(unix)'.dependencies]

ceno_zkvm/Cargo.toml

@@ -61,7 +61,7 @@ fn bench_add(c: &mut Criterion) {
         .circuit_pks
         .get(&AddInstruction::<E>::name())
         .unwrap();
-    let num_witin = circuit_pk.get_cs().num_witin;
+    let num_witin = circuit_pk.get_cs().num_witin();
 
     for instance_num_vars in 20..22 {
         // expand more input size once runtime is acceptable
@@ -79,7 +79,7 @@ fn bench_add(c: &mut Criterion) {
                         let num_instances = 1 << instance_num_vars;
                         let rmms = BTreeMap::from([(
                             0,
-                            RowMajorMatrix::rand(&mut OsRng, num_instances, num_witin as usize),
+                            RowMajorMatrix::rand(&mut OsRng, num_instances, num_witin),
                         )]);
 
                         let instant = std::time::Instant::now();

ceno_zkvm/benches/riscv_add.rs

@@ -10,7 +10,7 @@ use ff_ext::ExtensionField;
 use crate::{
     ROMType,
     error::ZKVMError,
-    structs::{ProgramParams, ProvingKey, RAMType, VerifyingKey},
+    structs::{ProgramParams, RAMType},
 };
 use multilinear_extensions::monomial::Term;
 use p3::field::FieldAlgebra;
@@ -189,12 +189,6 @@ impl<E: ExtensionField> ConstraintSystem<E> {
         }
     }
 
-    pub fn key_gen(self) -> ProvingKey<E> {
-        ProvingKey {
-            vk: VerifyingKey { cs: self },
-        }
-    }
-
     pub fn create_witin<NR: Into<String>, N: FnOnce() -> NR>(&mut self, n: N) -> WitIn {
         let wit_in = WitIn { id: self.num_witin };
         self.num_witin = self.num_witin.strict_add(1);

ceno_zkvm/src/circuit_builder.rs

@@ -1,10 +1,6 @@
 use ceno_emul::StepRecord;
 use ff_ext::ExtensionField;
-use gkr_iop::{
-    ProtocolBuilder, ProtocolWitnessGenerator,
-    gkr::{GKRCircuit, GKRCircuitOutput, GKRCircuitWitness},
-    hal::ProverBackend,
-};
+use gkr_iop::gkr::GKRCircuit;
 use multilinear_extensions::util::max_usable_threads;
 use rayon::{
     iter::{IndexedParallelIterator, ParallelIterator},
@@ -25,10 +21,19 @@ pub trait Instruction<E: ExtensionField> {
     }
 
     fn name() -> String;
+
+    /// construct circuit and manipulate circuit builder, then return the respective config
     fn construct_circuit(
         circuit_builder: &mut CircuitBuilder<E>,
     ) -> Result<Self::InstructionConfig, ZKVMError>;
 
+    /// giving config, extract optional gkr circuit
+    fn extract_gkr_iop_circuit(
+        _config: &mut Self::InstructionConfig,
+    ) -> Result<Option<GKRCircuit<E>>, ZKVMError> {
+        Ok(None)
+    }
+
     // assign single instance giving step from trace
     fn assign_instance(
         config: &Self::InstructionConfig,
@@ -72,178 +77,3 @@ pub trait Instruction<E: ExtensionField> {
         Ok((raw_witin, lk_multiplicity))
     }
 }
-
-pub struct GKRinfo {
-    pub and_lookups: usize,
-    pub xor_lookups: usize,
-    pub range_lookups: usize,
-    pub aux_wits: usize,
-}
-
-impl GKRinfo {
-    #[allow(dead_code)]
-    fn lookup_total(&self) -> usize {
-        self.and_lookups + self.xor_lookups + self.range_lookups
-    }
-}
-
-// Trait that should be implemented by opcodes which use the
-// GKR-IOP prover. Presently, such opcodes should also implement
-// the Instruction trait and in general methods of GKRIOPInstruction
-// will call corresponding methods of Instruction and do something extra
-// with respect to syncing state between GKR-IOP and old-style circuits/witnesses
-pub trait GKRIOPInstruction<E: ExtensionField>
-where
-    Self: Instruction<E>,
-{
-    type Layout: ProtocolWitnessGenerator<E> + ProtocolBuilder<E>;
-
-    /// Similar to Instruction::construct_circuit; generally
-    /// meant to extend InstructionConfig with GKR-specific
-    /// fields
-    #[allow(unused_variables)]
-    fn construct_circuit_with_gkr_iop(
-        cb: &mut CircuitBuilder<E>,
-    ) -> Result<Self::InstructionConfig, ZKVMError> {
-        unimplemented!();
-    }
-
-    /// Should generate phase1 witness for GKR from step records
-    fn phase1_witness_from_steps(
-        layout: &Self::Layout,
-        steps: &[StepRecord],
-    ) -> RowMajorMatrix<E::BaseField>;
-
-    /// Similar to Instruction::assign_instance, but with access to GKR lookups and wits
-    fn assign_instance_with_gkr_iop(
-        config: &Self::InstructionConfig,
-        instance: &mut [E::BaseField],
-        lk_multiplicity: &mut LkMultiplicity,
-        step: &StepRecord,
-        lookups: &[E::BaseField],
-        aux_wits: &[E::BaseField],
-    ) -> Result<(), ZKVMError>;
-
-    /// Similar to Instruction::assign_instances, but with access to the GKR layout.
-    #[allow(clippy::type_complexity)]
-    fn assign_instances_with_gkr_iop<'a, PB: ProverBackend<E = E>>(
-        _config: &Self::InstructionConfig,
-        _num_witin: usize,
-        _steps: Vec<StepRecord>,
-        _gkr_circuit: &GKRCircuit<E>,
-        _gkr_layout: &Self::Layout,
-    ) -> Result<
-        (
-            RowMajorMatrix<E::BaseField>,
-            GKRCircuitWitness<'a, PB>,
-            GKRCircuitOutput<'a, PB>,
-            LkMultiplicity,
-        ),
-        ZKVMError,
-    > {
-        todo!()
-        // let nthreads = max_usable_threads();
-        // let num_instance_per_batch = if steps.len() > 256 {
-        //     steps.len().div_ceil(nthreads)
-        // } else {
-        //     steps.len()
-        // }
-        // .max(1);
-        // let lk_multiplicity = LkMultiplicity::default();
-        // let mut raw_witin =
-        //     RowMajorMatrix::<E::BaseField>::new(steps.len(), num_witin, Self::padding_strategy());
-
-        // let raw_witin_iter = raw_witin.par_batch_iter_mut(num_instance_per_batch);
-
-        // let (gkr_witness, gkr_output) = gkr_layout.gkr_witness(
-        //     gkr_circuit,
-        //     &Self::phase1_witness_from_steps(gkr_layout, &steps),
-        //     &[],
-        // );
-
-        // let (lookups, aux_wits) = {
-        //     // Extract lookups and auxiliary witnesses from GKR protocol
-        //     // Here we assume that the gkr_witness's last layer is a convenience layer which holds
-        //     // the output records for all instances; further, we assume that the last ```Self::lookup_count()```
-        //     // elements of this layer are the lookup arguments.
-        //     let mut lookups = vec![vec![]; steps.len()];
-        //     let mut aux_wits: Vec<Vec<E::BaseField>> = vec![vec![]; steps.len()];
-        //     let last_layer = gkr_witness.layers.last().unwrap().bases.clone();
-        //     let len = last_layer.len();
-        //     let lookup_total = Self::gkr_info().lookup_total();
-
-        //     let non_lookups = if len == 0 {
-        //         0
-        //     } else {
-        //         assert!(len >= lookup_total);
-        //         len - lookup_total
-        //     };
-
-        //     for witness in last_layer.iter().skip(non_lookups) {
-        //         for i in 0..witness.len() {
-        //             lookups[i].push(witness[i]);
-        //         }
-        //     }
-
-        //     let n_layers = gkr_witness.layers.len();
-
-        //     for i in 0..steps.len() {
-        //         // Omit last layer, which stores outputs and not real witnesses
-        //         for layer in gkr_witness.layers[..n_layers - 1].iter() {
-        //             for base in layer.bases.iter() {
-        //                 aux_wits[i].push(base[i]);
-        //             }
-        //         }
-        //     }
-
-        //     (lookups, aux_wits)
-        // };
-
-        // raw_witin_iter
-        //     .zip(
-        //         steps
-        //             .iter()
-        //             .enumerate()
-        //             .collect_vec()
-        //             .par_chunks(num_instance_per_batch),
-        //     )
-        //     .flat_map(|(instances, steps)| {
-        //         let mut lk_multiplicity = lk_multiplicity.clone();
-        //         instances
-        //             .chunks_mut(num_witin)
-        //             .zip(steps)
-        //             .map(|(instance, (i, step))| {
-        //                 Self::assign_instance_with_gkr_iop(
-        //                     config,
-        //                     instance,
-        //                     &mut lk_multiplicity,
-        //                     step,
-        //                     &lookups[*i],
-        //                     &aux_wits[*i],
-        //                 )
-        //             })
-        //             .collect::<Vec<_>>()
-        //     })
-        //     .collect::<Result<(), ZKVMError>>()?;
-
-        // raw_witin.padding_by_strategy();
-        // Ok((raw_witin, gkr_witness, gkr_output, lk_multiplicity))
-    }
-
-    /// Lookup and witness counts used by GKR proof
-    fn gkr_info() -> GKRinfo;
-
-    /// Returns corresponding column in RMM for the i-th
-    /// output evaluation of the GKR proof
-    #[allow(unused_variables)]
-    fn output_evals_map(i: usize) -> usize {
-        unimplemented!();
-    }
-
-    /// Returns corresponding column in RMM for the i-th
-    /// witness of the GKR proof
-    #[allow(unused_variables)]
-    fn witness_map(i: usize) -> usize {
-        unimplemented!();
-    }
-}

ceno_zkvm/src/instructions.rs

@@ -90,7 +90,7 @@ pub struct DummyConfig<E: ExtensionField> {
 
 impl<E: ExtensionField> DummyConfig<E> {
     #[allow(clippy::too_many_arguments)]
-    pub(super) fn construct_circuit(
+    pub fn construct_circuit(
         circuit_builder: &mut CircuitBuilder<E>,
         kind: InsnKind,
         with_rs1: bool,