wip

Author: mcalancea

ceno_host/tests/test_elf.rs

@@ -231,7 +231,7 @@ fn test_ceno_rt_keccak() -> Result<()> {
     let steps = run(&mut state)?;
 
     // Expect the program to have written successive states between Keccak permutations.
-    const ITERATIONS: usize = 3;
+    const ITERATIONS: usize = 4;
     let keccak_outs = sample_keccak_f(ITERATIONS);
 
     let all_messages = read_all_messages(&state);

ceno_zkvm/Cargo.toml

@@ -23,7 +23,7 @@ gkr_iop = { path = "../gkr_iop" }
 mpcs = { path = "../mpcs" }
 multilinear_extensions = { version = "0", path = "../multilinear_extensions" }
 p3 = { path = "../p3" }
-subprotocols = {path = "../subprotocols"}
+subprotocols = { path = "../subprotocols" }
 sumcheck = { version = "0", path = "../sumcheck" }
 transcript = { path = "../transcript" }
 witness = { path = "../witness" }

ceno_zkvm/src/instructions.rs

@@ -1,10 +1,6 @@
 use ceno_emul::StepRecord;
 use ff_ext::ExtensionField;
-use gkr_iop::{
-    ProtocolBuilder, ProtocolWitnessGenerator,
-    gkr::{GKRCircuitWitness, layer::LayerWitness},
-    precompiles::KeccakLayout,
-};
+use gkr_iop::{ProtocolBuilder, ProtocolWitnessGenerator, gkr::GKRCircuitWitness};
 use itertools::Itertools;
 use multilinear_extensions::util::max_usable_threads;
 use rayon::{
@@ -87,31 +83,34 @@ impl GKRinfo {
     }
 }
 
+// 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;
 
+    /// 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!();
     }
 
-    // Returns index of `i-th` GKR-IOP output-eval in PCS
-    fn output_map(i: usize) -> usize {
-        unimplemented!();
-    }
-
+    /// Should generate phase1 witness for GKR from step records
     fn phase1_witness_from_steps(
         layout: &Self::Layout,
         steps: &[StepRecord],
     ) -> Vec<Vec<E::BaseField>>;
 
-    // Number of lookup arguments used by this GKR proof
-    fn gkr_info() -> GKRinfo;
-
+    /// 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],
@@ -121,6 +120,8 @@ where
         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(
         config: &Self::InstructionConfig,
         num_witin: usize,
@@ -146,10 +147,8 @@ where
             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_layout.gkr_witness(
-            &Self::phase1_witness_from_steps(gkr_layout, &steps),
-            &vec![],
-        );
+        let gkr_witness =
+            gkr_layout.gkr_witness(&Self::phase1_witness_from_steps(gkr_layout, &steps), &[]);
 
         let (lookups, aux_wits) = {
             // Extract lookups and auxiliary witnesses from GKR protocol
@@ -159,7 +158,10 @@ where
             let mut lookups = vec![vec![]; steps.len()];
             let last_layer = gkr_witness.layers.last().unwrap().bases.clone();
             let len = last_layer.len();
-            for witness in last_layer[len - Self::gkr_info().lookup_total()..].iter() {
+            for witness in last_layer
+                .iter()
+                .skip(len - Self::gkr_info().lookup_total())
+            {
                 for i in 0..witness.len() {
                     lookups[i].push(witness[i]);
                 }
@@ -194,7 +196,6 @@ where
                     .chunks_mut(num_witin)
                     .zip(steps)
                     .map(|(instance, (i, step))| {
-                        // dbg!(i, step);
                         Self::assign_instance_with_gkr_iop(
                             config,
                             instance,
@@ -211,4 +212,21 @@ where
         raw_witin.padding_by_strategy();
         Ok((raw_witin, gkr_witness, 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/riscv/dummy/dummy_ecall.rs

@@ -112,7 +112,6 @@ impl<E: ExtensionField, S: SyscallSpec> Instruction<E> for LargeEcallDummy<E, S>
 
         // Assign memory.
         for ((addr, value, writer), op) in config.mem_writes.iter().zip_eq(&ops.mem_ops) {
-            dbg!(&op.value.before);
             set_val!(instance, value.before, op.value.before as u64);
             set_val!(instance, value.after, op.value.after as u64);
             set_val!(instance, addr, u64::from(op.addr));
@@ -174,14 +173,10 @@ impl<E: ExtensionField> GKRIOPInstruction<E> for LargeEcallDummy<E, KeccakSpec>
         partial_config.lookups = lookups;
         partial_config.aux_wits = aux_wits;
 
-        dbg!(&partial_config.mem_writes[0].1.after);
-        dbg!(&partial_config.mem_writes[1].1.before);
-        dbg!(&partial_config.lookups[0]);
-
         Ok(partial_config)
     }
 
-    fn output_map(i: usize) -> usize {
+    fn output_evals_map(i: usize) -> usize {
         if i < 50 {
             27 + 6 * i
         } else if i < 100 {
@@ -225,9 +220,6 @@ impl<E: ExtensionField> GKRIOPInstruction<E> for LargeEcallDummy<E, KeccakSpec>
         let mut wit_iter = lookups.iter().map(|f| f.to_canonical_u64());
         let mut var_iter = config.lookups.iter();
 
-        dbg!(wit_iter.clone().count());
-        dbg!(var_iter.clone().count());
-
         let mut pop_arg = || -> u64 {
             let wit = wit_iter.next().unwrap();
             let var = var_iter.next().unwrap();
@@ -247,7 +239,6 @@ impl<E: ExtensionField> GKRIOPInstruction<E> for LargeEcallDummy<E, KeccakSpec>
             lk_multiplicity.assert_ux::<16>(pop_arg());
         }
 
-        dbg!(aux_wits.len());
         for (aux_wit_var, aux_wit) in zip_eq(config.aux_wits.iter(), aux_wits) {
             set_val!(instance, aux_wit_var, (aux_wit.to_canonical_u64()));
         }

ceno_zkvm/src/instructions/riscv/rv32im.rs

@@ -466,7 +466,6 @@ pub struct DummyExtraConfig<E: ExtensionField> {
 impl<E: ExtensionField> DummyExtraConfig<E> {
     pub fn construct_circuits(cs: &mut ZKVMConstraintSystem<E>) -> Self {
         let ecall_config = cs.register_opcode_circuit::<EcallDummy<E>>();
-        // let keccak_config = cs.register_opcode_circuit::<LargeEcallDummy<E, KeccakSpec>>();
         let keccak_config = cs.register_keccakf_circuit();
         let secp256k1_add_config =
             cs.register_opcode_circuit::<LargeEcallDummy<E, Secp256k1AddSpec>>();
@@ -511,7 +510,6 @@ impl<E: ExtensionField> DummyExtraConfig<E> {
     ) {
         fixed.register_opcode_circuit::<EcallDummy<E>>(cs);
         fixed.register_keccakf_circuit(cs);
-        // fixed.register_opcode_circuit::<LargeEcallDummy<E, KeccakSpec>>(cs);
         fixed.register_opcode_circuit::<LargeEcallDummy<E, Secp256k1AddSpec>>(cs);
         fixed.register_opcode_circuit::<LargeEcallDummy<E, Secp256k1DecompressSpec>>(cs);
         fixed.register_opcode_circuit::<LargeEcallDummy<E, Secp256k1DoubleSpec>>(cs);
@@ -566,12 +564,6 @@ impl<E: ExtensionField> DummyExtraConfig<E> {
 
         witness.assign_keccakf_circuit(cs, &self.keccak_config, keccak_steps)?;
 
-        // witness.assign_opcode_circuit::<LargeEcallDummy<E, KeccakSpec>>(
-        //    cs,
-        //    &self.keccak_config,
-        //    keccak_steps,
-        //)?;
-
         witness.assign_opcode_circuit::<LargeEcallDummy<E, Secp256k1AddSpec>>(
             cs,
             &self.secp256k1_add_config,

ceno_zkvm/src/keygen.rs

@@ -1,6 +1,5 @@
 use crate::{
     error::ZKVMError,
-    instructions::riscv::dummy::LargeEcallDummy,
     structs::{ZKVMConstraintSystem, ZKVMFixedTraces, ZKVMProvingKey},
 };
 use ff_ext::ExtensionField;

ceno_zkvm/src/scheme.rs

@@ -57,6 +57,7 @@ pub struct ZKVMOpcodeProof<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>
 }
 
 #[derive(Clone, Serialize)]
+// WARN/TODO: depends on serde's `arc` feature which might not behave correctly
 pub struct GKROpcodeProof<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> {
     output_evals: Vec<E>,
     prover_output: GKRProverOutput<E, Evaluation<E>>,

ceno_zkvm/src/scheme/prover.rs

@@ -7,7 +7,7 @@ use std::{
     sync::Arc,
 };
 
-use itertools::{Itertools, chain, enumerate, izip};
+use itertools::{Itertools, enumerate, izip};
 use mpcs::PolynomialCommitmentScheme;
 use multilinear_extensions::{
     mle::{IntoMLE, MultilinearExtension},
@@ -27,10 +27,7 @@ use witness::{RowMajorMatrix, next_pow2_instance_padding};
 use crate::{
     error::ZKVMError,
     expression::Instance,
-    instructions::{
-        Instruction,
-        riscv::{dummy::LargeEcallDummy, ecall::EcallDummy},
-    },
+    instructions::{Instruction, riscv::dummy::LargeEcallDummy},
     scheme::{
         GKROpcodeProof,
         constants::{MAINCONSTRAIN_SUMCHECK_BATCH_SIZE, NUM_FANIN, NUM_FANIN_LOGUP},
@@ -261,6 +258,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
     /// 1: witness layer inferring from input -> output
     /// 2: proof (sumcheck reduce) from output to input
     #[allow(clippy::too_many_arguments)]
+    #[allow(clippy::type_complexity)]
     #[tracing::instrument(skip_all, name = "create_opcode_proof", fields(circuit_name=name,profiling_2), level="trace")]
     pub fn create_opcode_proof(
         &self,
@@ -695,6 +693,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
             opening_dur.elapsed(),
         );
         exit_span!(pcs_open_span);
+
         let wits_commit = PCS::get_pure_commitment(&wits_commit);
 
         let gkr_opcode_proof = if let Some((gkr_iop_pk, gkr_wit)) = gkr_iop_pk {
@@ -711,25 +710,26 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
                 .iter()
                 .map(|base| PointAndEval {
                     point: point.clone(),
-                    eval: subprotocols::utils::evaluate_mle_ext(&base, &point),
+                    eval: subprotocols::utils::evaluate_mle_ext(base, &point),
                 })
                 .collect_vec();
 
-            // out_evals from point and output polynomials instead of *_records which is combined
-
             let prover_output = gkr_circuit
-                .prove(gkr_wit, &out_evals, &vec![], transcript)
+                .prove(gkr_wit, &out_evals, &[], transcript)
                 .expect("Failed to prove phase");
             // unimplemented!("cannot fully handle GKRIOP component yet")
 
             dbg!(&prover_output.opening_evaluations.len());
+            let _gkr_open_point = prover_output.opening_evaluations[0].point.clone();
+            // TODO: open polynomials for GKR proof
+
             let output_evals = out_evals.into_iter().map(|pae| pae.eval).collect_vec();
 
             Some(GKROpcodeProof {
                 output_evals,
                 prover_output,
                 circuit: gkr_circuit,
-                _marker: PhantomData::default(),
+                _marker: PhantomData,
             })
         } else {
             None

ceno_zkvm/src/scheme/verifier.rs

@@ -276,8 +276,6 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMVerifier<E, PCS>
             cs.lk_expressions.len(),
         );
 
-        // dbg!(&cs.r_expressions);
-
         let (log2_r_count, log2_w_count, log2_lk_count) = (
             ceil_log2(r_counts_per_instance),
             ceil_log2(w_counts_per_instance),
@@ -332,8 +330,6 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMVerifier<E, PCS>
             logup_q_evals[0].point.clone(),
         );
 
-        // dbg!(&rt_r, &rt_w, &rt_lk);
-
         let alpha_pow = get_challenge_pows(
             MAINCONSTRAIN_SUMCHECK_BATCH_SIZE + cs.assert_zero_sumcheck_expressions.len(),
             transcript,
@@ -422,26 +418,18 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMVerifier<E, PCS>
             for (i, gkr_out_eval) in gkr_iop.output_evals.iter().enumerate() {
                 assert_eq!(
                     *gkr_out_eval,
-                    proof.wits_in_evals[LargeEcallDummy::<E, KeccakSpec>::output_map(i)],
+                    proof.wits_in_evals[LargeEcallDummy::<E, KeccakSpec>::output_evals_map(i)],
                     "{i}"
                 );
-                if *gkr_out_eval
-                    == proof.wits_in_evals[LargeEcallDummy::<E, KeccakSpec>::output_map(i)]
-                {
-                    matches += 1;
-                } else {
-                    dbg!(i);
-                }
             }
-            dbg!(matches);
             // Verify GKR proof
             let point = Arc::new(input_opening_point.clone());
             let out_evals = gkr_iop
                 .output_evals
                 .iter()
                 .map(|eval| gkr_iop::evaluation::PointAndEval {
                     point: point.clone(),
-                    eval: eval.clone(),
+                    eval: *eval,
                 })
                 .collect_vec();
 
@@ -450,7 +438,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMVerifier<E, PCS>
                 .verify(
                     gkr_iop.prover_output.gkr_proof.clone(),
                     &out_evals,
-                    &vec![],
+                    &[],
                     transcript,
                 )
                 .expect("GKR-IOP verify failure");

ceno_zkvm/src/structs.rs

@@ -374,12 +374,10 @@ impl<E: ExtensionField> ZKVMWitnesses<E> {
 
     pub fn assign_keccakf_circuit(
         &mut self,
-        // TODO: do without mutability requirement
         css: &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();
@@ -392,12 +390,6 @@ impl<E: ExtensionField> ZKVMWitnesses<E> {
             )?;
         self.keccak_gkr_wit = gkr_witness;
 
-        // // Intercept row-major matrix, convert into KeccakTrace and obtain phase1_wit
-        // self.keccak_phase1wit = css
-        //     .keccak_gkr_iop
-        //     .layout
-        //     .phase1_witness(KeccakTrace::from(witness.clone()));
-
         assert!(
             self.witnesses_opcodes
                 .insert(LargeEcallDummy::<E, KeccakSpec>::name(), witness)