wip convert local final ram circuit to gkr-iop circuit

Author: hero78119

ceno_zkvm/src/instructions/riscv/rv32im/mmu.rs

@@ -46,7 +46,9 @@ impl<E: ExtensionField> MmuConfig<'_, E> {
         let hints_config = cs.register_table_circuit::<HintsCircuit<E>>();
         let stack_init_config = cs.register_table_circuit::<StackInitCircuit<E>>();
         let heap_init_config = cs.register_table_circuit::<HeapInitCircuit<E>>();
+        println!("register LocalFinalCircuit");
         let local_final_circuit = cs.register_table_circuit::<LocalFinalCircuit<E>>();
+        println!("end register LocalFinalCircuit");
         let ram_bus_circuit = cs.register_table_circuit::<RBCircuit<E>>();
 
         Self {

ceno_zkvm/src/scheme/cpu/mod.rs

@@ -681,98 +681,19 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> MainSumcheckProver<C
                 Some(gkr_proof),
             ))
         } else {
-            let (wits_in_evals, fixed_in_evals, main_sumcheck_proof, rt) =
-                if next_pow2_instance_padding(num_instances) == num_instances {
-                    let span = entered_span!("fixed::evals + witin::evals");
-                    let mut evals = input
-                        .witness
-                        .par_iter()
-                        .chain(input.fixed.par_iter())
-                        .map(|poly| poly.evaluate(&rt_tower[..poly.num_vars()]))
-                        .collect::<Vec<_>>();
-                    let fixed_in_evals = evals.split_off(input.witness.len());
-                    let wits_in_evals = evals;
-                    exit_span!(span);
-                    (wits_in_evals, fixed_in_evals, None, rt_tower)
-                } else {
-                    assert!(cs.r_table_expressions.len() <= 1);
-                    assert!(cs.w_table_expressions.len() <= 1);
-
-                    let sel_type = SelectorType::Prefix(E::BaseField::ZERO, 0.into());
-                    let mut sel_mle = sel_type.compute(&rt_tower, num_instances).unwrap();
-
-                    // `wit` := witin ++ fixed
-                    // we concat eq in between `wit` := witin ++ eqs ++ fixed
-                    let all_witins = input
-                        .witness
-                        .iter()
-                        .map(|mle| Either::Left(mle.as_ref()))
-                        .chain(vec![Either::Right(&mut sel_mle)])
-                        .chain(input.fixed.iter().map(|mle| Either::Left(mle.as_ref())))
-                        .collect_vec();
-                    assert_eq!(
-                        all_witins.len() as WitnessId,
-                        cs.num_witin + cs.num_structural_witin + cs.num_fixed as WitnessId,
-                        "all_witins.len() {} != layer.n_witin {} + layer.n_structural_witin {} + layer.n_fixed {}",
-                        all_witins.len(),
-                        cs.num_witin,
-                        cs.num_structural_witin,
-                        cs.num_fixed,
-                    );
-                    let builder = VirtualPolynomialsBuilder::new_with_mles(
-                        num_threads,
-                        rt_tower.len(),
-                        all_witins,
-                    );
-
-                    let alpha_pows_expr = (2..)
-                        .take(cs.r_table_expressions.len() + cs.w_table_expressions.len())
-                        .map(|id| Expression::Challenge(id as ChallengeId, 1, E::ONE, E::ZERO))
-                        .collect_vec();
-                    let zero_check_expr: Expression<E> = cs
-                        .r_table_expressions
-                        .iter()
-                        .take(1)
-                        .chain(cs.w_table_expressions.iter().take(1))
-                        .zip_eq(&alpha_pows_expr)
-                        .map(|(expr, alpha)| alpha * expr.expr.expr())
-                        .sum();
-                    let zero_check_monomial = monomialize_expr_to_wit_terms(
-                        &zero_check_expr,
-                        cs.num_witin as WitnessId,
-                        cs.num_structural_witin as WitnessId,
-                        cs.num_fixed as WitnessId,
-                    );
-                    let main_sumcheck_challenges = chain!(
-                        challenges.iter().copied(),
-                        get_challenge_pows(
-                            cs.w_table_expressions.len() + cs.r_table_expressions.len(),
-                            transcript,
-                        )
-                    )
-                    .collect_vec();
-
-                    let span = entered_span!("IOPProverState::prove", profiling_4 = true);
-                    let (proof, prover_state) = IOPProverState::prove(
-                        builder.to_virtual_polys_with_monomial_terms(
-                            &zero_check_monomial,
-                            &[],
-                            &main_sumcheck_challenges,
-                        ),
-                        transcript,
-                    );
-                    exit_span!(span);
-                    let rt = prover_state
-                        .challenges
-                        .iter()
-                        .map(|c| c.elements)
-                        .collect_vec();
-                    let mut wits_in_evals = prover_state.get_mle_flatten_final_evaluations();
-                    let mut rest = wits_in_evals.split_off(cs.num_witin as usize);
-                    let rest = rest.split_off(cs.num_structural_witin as usize);
-                    let fixed_in_evals = rest;
-                    (wits_in_evals, fixed_in_evals, Some(proof.proofs), rt)
-                };
+            let (wits_in_evals, fixed_in_evals, main_sumcheck_proof, rt) = {
+                let span = entered_span!("fixed::evals + witin::evals");
+                let mut evals = input
+                    .witness
+                    .par_iter()
+                    .chain(input.fixed.par_iter())
+                    .map(|poly| poly.evaluate(&rt_tower[..poly.num_vars()]))
+                    .collect::<Vec<_>>();
+                let fixed_in_evals = evals.split_off(input.witness.len());
+                let wits_in_evals = evals;
+                exit_span!(span);
+                (wits_in_evals, fixed_in_evals, None, rt_tower)
+            };
 
             Ok((
                 rt,

ceno_zkvm/src/scheme/prover.rs

@@ -210,6 +210,7 @@ impl<
         let (points, evaluations) = self.pk.circuit_pks.iter().enumerate().try_fold(
             (vec![], vec![]),
             |(mut points, mut evaluations), (index, (circuit_name, pk))| {
+                println!("prove circuit_name {circuit_name}");
                 let num_instances = circuit_name_num_instances_mapping
                     .get(&circuit_name)
                     .copied()

ceno_zkvm/src/scheme/utils.rs

@@ -345,12 +345,16 @@ pub fn build_main_witness<
         }
 
         if let Some(gkr_circuit) = gkr_circuit {
-            // opcode must have at least one read/write/lookup
+            // circuit must have at least one read/write/lookup
             assert!(
-                cs.lk_expressions.is_empty()
-                    || !cs.r_expressions.is_empty()
-                    || !cs.w_expressions.is_empty(),
-                "assert opcode circuit"
+                cs.r_expressions.len()
+                    + cs.w_expressions.len()
+                    + cs.lk_expressions.len()
+                    + cs.r_table_expressions.len()
+                    + cs.w_table_expressions.len()
+                    + cs.lk_table_expressions.len()
+                    > 0,
+                "assert circuit"
             );
 
             let (_, gkr_circuit_out) = gkr_witness::<E, PCS, PB, PD>(

ceno_zkvm/src/scheme/verifier.rs

@@ -164,6 +164,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMVerifier<E, PCS>
         for (index, proof) in &vm_proof.chip_proofs {
             assert!(proof.num_instances > 0);
             let circuit_name = &self.vk.circuit_index_to_name[index];
+            println!("verify circuit_name {circuit_name}");
             let circuit_vk = &self.vk.circuit_vks[circuit_name];
 
             // check chip proof is well-formed
@@ -356,9 +357,9 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMVerifier<E, PCS>
         } = &composed_cs;
         let num_instances = proof.num_instances;
         let (r_counts_per_instance, w_counts_per_instance, lk_counts_per_instance) = (
-            cs.r_expressions.len(),
-            cs.w_expressions.len(),
-            cs.lk_expressions.len(),
+            cs.r_expressions.len() + cs.r_table_expressions.len(),
+            cs.w_expressions.len() + cs.w_table_expressions.len(),
+            cs.lk_expressions.len() + cs.lk_table_expressions.len() * 2,
         );
         let num_batched = r_counts_per_instance + w_counts_per_instance + lk_counts_per_instance;
 
@@ -529,9 +530,8 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMVerifier<E, PCS>
             "[prod_record] mismatch length"
         );
 
-        let input_opening_point = if next_pow2_instance_padding(proof.num_instances)
-            == proof.num_instances
-        {
+        let ram_bus_circuit = false;
+        let input_opening_point = if !ram_bus_circuit {
             // evaluate the evaluation of structural mles at input_opening_point by verifier
             let structural_evals = if with_rw {
                 // only iterate r set, as read/write set round should match
@@ -620,73 +620,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMVerifier<E, PCS>
             }
             rt_tower
         } else {
-            // TODO LocalFinalTable goes here, merge flow into gkr_iop
-            assert_eq!(cs.lk_table_expressions.len(), 0);
-            assert!(proof.main_sumcheck_proofs.is_some());
-            assert_eq!(cs.num_structural_witin, 1);
-            assert_eq!(prod_point_and_eval.len(), 1);
-
-            // verify opening same point layer sumcheck
-            let alpha_pow = get_challenge_pows(
-                cs.r_table_expressions.len() + cs.w_table_expressions.len(),
-                transcript,
-            );
-
-            //  \sum_i alpha_{i} * (out_r_eval{i} - ONE)
-            //  + \sum_i alpha_{i} * (out_w_eval{i} - ONE)
-            let claim_sum = prod_point_and_eval
-                .iter()
-                .zip_eq(alpha_pow.iter())
-                .map(|(point_and_eval, alpha)| *alpha * (point_and_eval.eval - E::ONE))
-                .sum::<E>();
-            let sel_subclaim = IOPVerifierState::verify(
-                claim_sum,
-                &IOPProof {
-                    proofs: proof.main_sumcheck_proofs.clone().unwrap(),
-                },
-                &VPAuxInfo {
-                    max_degree: SEL_DEGREE,
-                    max_num_variables: expected_max_rounds,
-                    phantom: PhantomData,
-                },
-                transcript,
-            );
-            let (input_opening_point, sumcheck_eval) = (
-                sel_subclaim.point.iter().map(|c| c.elements).collect_vec(),
-                sel_subclaim.expected_evaluation,
-            );
-            let structural_evals = vec![eq_eval_less_or_equal_than(
-                proof.num_instances - 1,
-                &prod_point_and_eval[0].point,
-                &input_opening_point,
-            )];
-
-            let expected_evals = interleave(
-                &cs.r_table_expressions, // r
-                &cs.w_table_expressions, // w
-            )
-            .map(|rw| &rw.expr)
-            .zip(alpha_pow.iter())
-            .map(|(expr, alpha)| {
-                *alpha
-                    * eval_by_expr_with_instance(
-                        &proof.fixed_in_evals,
-                        &proof.wits_in_evals,
-                        &structural_evals,
-                        pi,
-                        challenges,
-                        expr,
-                    )
-                    .right()
-                    .unwrap()
-            })
-            .sum::<E>();
-            if expected_evals != sumcheck_eval {
-                return Err(ZKVMError::VerifyError(
-                    "sel evaluation verify failed".into(),
-                ));
-            }
-            input_opening_point
+            unimplemented!("shard ram bus circuit go here");
         };
 
         // assume public io is tiny vector, so we evaluate it directly without PCS

ceno_zkvm/src/structs.rs

@@ -135,9 +135,7 @@ impl<E: ExtensionField> ComposedConstrainSystem<E> {
     }
 
     pub fn is_opcode_circuit(&self) -> bool {
-        self.zkvm_v1_css.lk_table_expressions.is_empty()
-            && self.zkvm_v1_css.r_table_expressions.is_empty()
-            && self.zkvm_v1_css.w_table_expressions.is_empty()
+        self.gkr_circuit.is_some()
     }
 
     /// return number of lookup operation
@@ -219,18 +217,13 @@ impl<E: ExtensionField> ZKVMConstraintSystem<E> {
     pub fn register_table_circuit<TC: TableCircuit<E>>(&mut self) -> TC::TableConfig {
         let mut cs = ConstraintSystem::new(|| format!("riscv_table/{}", TC::name()));
         let mut circuit_builder = CircuitBuilder::<E>::new(&mut cs);
-        let config = TC::construct_circuit(&mut circuit_builder, &self.params).unwrap();
-        assert!(
-            self.circuit_css
-                .insert(
-                    TC::name(),
-                    ComposedConstrainSystem {
-                        zkvm_v1_css: cs,
-                        gkr_circuit: None
-                    }
-                )
-                .is_none()
-        );
+        let (config, gkr_iop_circuit) =
+            TC::build_gkr_iop_circuit(&mut circuit_builder, &self.params).unwrap();
+        let cs = ComposedConstrainSystem {
+            zkvm_v1_css: cs,
+            gkr_circuit: gkr_iop_circuit,
+        };
+        assert!(self.circuit_css.insert(TC::name(), cs).is_none());
         config
     }
 

ceno_zkvm/src/tables/mod.rs

@@ -1,7 +1,14 @@
 use crate::{circuit_builder::CircuitBuilder, error::ZKVMError, structs::ProgramParams};
 use ff_ext::ExtensionField;
 use std::collections::HashMap;
+use itertools::Itertools;
+use multilinear_extensions::{StructuralWitInType, ToExpr};
 use witness::RowMajorMatrix;
+use gkr_iop::chip::Chip;
+use gkr_iop::gkr::GKRCircuit;
+use gkr_iop::gkr::layer::Layer;
+use gkr_iop::selector::SelectorType;
+
 mod range;
 pub use range::*;
 
@@ -29,6 +36,14 @@ pub trait TableCircuit<E: ExtensionField> {
         params: &ProgramParams,
     ) -> Result<Self::TableConfig, ZKVMError>;
 
+    fn build_gkr_iop_circuit(
+        cb: &mut CircuitBuilder<E>,
+        param: &ProgramParams,
+    ) -> Result<(Self::TableConfig, Option<GKRCircuit<E>>), ZKVMError> {
+        let config = Self::construct_circuit(cb, param)?;
+        Ok((config, None))
+    }
+
     fn generate_fixed_traces(
         config: &Self::TableConfig,
         num_fixed: usize,