refactor & clean-up

Author: mcalancea

ceno_emul/src/syscalls.rs

@@ -22,6 +22,8 @@ pub trait SyscallSpec {
     const CODE: u32;
 
     const HAS_LOOKUPS: bool = false;
+
+    const GKR_OUTPUTS: usize = 0;
 }
 
 /// Trace the inputs and effects of a syscall.

ceno_zkvm/src/e2e.rs

@@ -295,7 +295,7 @@ fn generate_fixed_traces<E: ExtensionField>(
 }
 
 pub fn generate_witness<E: ExtensionField>(
-    system_config: &mut ConstraintSystemConfig<E>,
+    system_config: &ConstraintSystemConfig<E>,
     emul_result: EmulationResult,
     program: &Program,
     is_mock_proving: bool,
@@ -312,7 +312,7 @@ pub fn generate_witness<E: ExtensionField>(
         .unwrap();
     system_config
         .dummy_config
-        .assign_opcode_circuit(&mut system_config.zkvm_cs, &mut zkvm_witness, dummy_records)
+        .assign_opcode_circuit(&system_config.zkvm_cs, &mut zkvm_witness, dummy_records)
         .unwrap();
     zkvm_witness.finalize_lk_multiplicities(is_mock_proving);
 
@@ -394,7 +394,7 @@ pub fn run_e2e_with_checkpoint<
     };
 
     let program = Arc::new(program);
-    let mut system_config = construct_configs::<E>(program_params);
+    let system_config = construct_configs::<E>(program_params);
     let reg_init = system_config.mmu_config.initial_registers();
 
     // IO is not used in this program, but it must have a particular size at the moment.
@@ -432,7 +432,7 @@ pub fn run_e2e_with_checkpoint<
                     init_full_mem,
                     platform,
                     hints,
-                    &mut system_config,
+                    &system_config,
                     pk,
                     zkvm_fixed_traces,
                     is_mock_proving,
@@ -451,13 +451,11 @@ pub fn run_e2e_with_checkpoint<
     if let Checkpoint::PrepWitnessGen = checkpoint {
         return (
             None,
-            Box::new(move || {
-                _ = generate_witness(&mut system_config, emul_result, &program, false)
-            }),
+            Box::new(move || _ = generate_witness(&system_config, emul_result, &program, false)),
         );
     }
 
-    let zkvm_witness = generate_witness(&mut system_config, emul_result, &program, is_mock_proving);
+    let zkvm_witness = generate_witness(&system_config, emul_result, &program, is_mock_proving);
 
     // proving
     let prover = ZKVMProver::new(pk);
@@ -498,7 +496,7 @@ pub fn run_e2e_proof<E: ExtensionField + LkMultiplicityKey, PCS: PolynomialCommi
     init_full_mem: InitMemState,
     platform: Platform,
     hints: Vec<u32>,
-    system_config: &mut ConstraintSystemConfig<E>,
+    system_config: &ConstraintSystemConfig<E>,
     pk: ZKVMProvingKey<E, PCS>,
     zkvm_fixed_traces: ZKVMFixedTraces<E>,
     is_mock_proving: bool,

ceno_zkvm/src/instructions.rs

@@ -74,11 +74,39 @@ pub trait Instruction<E: ExtensionField> {
     }
 }
 
+pub struct GKRinfo {
+    pub and_lookups: usize,
+    pub xor_lookups: usize,
+    pub range_lookups: usize,
+    pub aux_wits: usize,
+}
+
+impl GKRinfo {
+    fn lookup_total(&self) -> usize {
+        self.and_lookups + self.xor_lookups + self.range_lookups
+    }
+}
+
 pub trait GKRIOPInstruction<E: ExtensionField>
 where
     Self: Instruction<E>,
 {
     type Layout: ProtocolWitnessGenerator<E> + ProtocolBuilder;
+
+    fn construct_circuit_with_gkr_iop(
+        cb: &mut CircuitBuilder<E>,
+    ) -> Result<Self::InstructionConfig, ZKVMError> {
+        unimplemented!();
+    }
+
+    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;
+
     fn assign_instance_with_gkr_iop(
         config: &Self::InstructionConfig,
         instance: &mut [E::BaseField],
@@ -88,16 +116,11 @@ where
         aux_wits: &[E::BaseField],
     ) -> Result<(), ZKVMError>;
 
-    fn phase1_witness_from_steps(
-        layout: &Self::Layout,
-        steps: &[StepRecord],
-    ) -> Vec<Vec<E::BaseField>>;
-
     fn assign_instances_with_gkr_iop(
         config: &Self::InstructionConfig,
         num_witin: usize,
         steps: Vec<StepRecord>,
-        gkr_layout: &mut Self::Layout,
+        gkr_layout: &Self::Layout,
     ) -> Result<
         (
             RowMajorMatrix<E::BaseField>,
@@ -124,17 +147,14 @@ where
         );
 
         let (lookups, aux_wits) = {
-            // Keccak-specific
+            // 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()];
-            for witness in gkr_witness
-                .layers
-                .last()
-                .unwrap()
-                .bases
-                .iter()
-                .skip(100)
-                .cloned()
-            {
+            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 i in 0..witness.len() {
                     lookups[i].push(witness[i]);
                 }
@@ -144,6 +164,7 @@ where
             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]);
@@ -154,8 +175,6 @@ where
             (lookups, aux_wits)
         };
 
-        // dbg!(&lookups);
-
         raw_witin_iter
             .zip(
                 steps

ceno_zkvm/src/instructions/riscv/dummy/dummy_ecall.rs

@@ -11,7 +11,7 @@ use crate::{
     error::ZKVMError,
     expression::{ToExpr, WitIn},
     instructions::{
-        GKRIOPInstruction, Instruction,
+        GKRIOPInstruction, GKRinfo, Instruction,
         riscv::{constants::UInt, insn_base::WriteRD},
     },
     set_val,
@@ -21,10 +21,7 @@ use ff_ext::FieldInto;
 
 use gkr_iop::{
     ProtocolWitnessGenerator,
-    precompiles::{
-        AND_LOOKUPS_PER_ROUND, KeccakLayout, KeccakTrace, RANGE_LOOKUPS_PER_ROUND,
-        XOR_LOOKUPS_PER_ROUND,
-    },
+    precompiles::{AND_LOOKUPS, KeccakLayout, KeccakTrace, RANGE_LOOKUPS, XOR_LOOKUPS},
 };
 
 /// LargeEcallDummy can handle any instruction and produce its effects,
@@ -37,7 +34,7 @@ impl<E: ExtensionField, S: SyscallSpec> Instruction<E> for LargeEcallDummy<E, S>
     type InstructionConfig = LargeEcallConfig<E>;
 
     fn name() -> String {
-        format!("{}_DUMMY", S::NAME)
+        S::NAME.to_owned()
     }
     fn construct_circuit(cb: &mut CircuitBuilder<E>) -> Result<Self::InstructionConfig, ZKVMError> {
         let dummy_insn = DummyConfig::construct_circuit(
@@ -78,45 +75,9 @@ impl<E: ExtensionField, S: SyscallSpec> Instruction<E> for LargeEcallDummy<E, S>
             })
             .collect::<Result<Vec<_>, _>>()?;
 
-        // Temporarily set this to < 24 to avoid cb.num_witin overflow
-        let active_rounds = 24;
-
-        let mut lookups = Vec::with_capacity(
-            active_rounds
-                * (3 * AND_LOOKUPS_PER_ROUND + 3 * XOR_LOOKUPS_PER_ROUND + RANGE_LOOKUPS_PER_ROUND),
-        );
-
-        let mut aux_wits = vec![];
-
-        if S::HAS_LOOKUPS {
-            dbg!(lookups.capacity());
-
-            for round in 0..active_rounds {
-                for i in 0..AND_LOOKUPS_PER_ROUND {
-                    let a = cb.create_witin(|| format!("and_lookup_{round}_{i}_a"));
-                    let b = cb.create_witin(|| format!("and_lookup_{round}_{i}_b"));
-                    let c = cb.create_witin(|| format!("and_lookup_{round}_{i}_c"));
-                    cb.lookup_and_byte(a.into(), b.into(), c.into())?;
-                    lookups.extend(vec![a, b, c]);
-                }
-                for i in 0..XOR_LOOKUPS_PER_ROUND {
-                    let a = cb.create_witin(|| format!("xor_lookup_{round}_{i}_a"));
-                    let b = cb.create_witin(|| format!("xor_lookup_{round}_{i}_b"));
-                    let c = cb.create_witin(|| format!("xor_lookup_{round}_{i}_c"));
-                    cb.lookup_xor_byte(a.into(), b.into(), c.into())?;
-                    lookups.extend(vec![a, b, c]);
-                }
-                for i in 0..RANGE_LOOKUPS_PER_ROUND {
-                    let wit = cb.create_witin(|| format!("range_lookup_{round}_{i}"));
-                    cb.assert_ux::<_, _, 16>(|| "nada", wit.into())?;
-                    lookups.push(wit);
-                }
-            }
-
-            for i in 0..40144 {
-                aux_wits.push(cb.create_witin(|| format!("aux_wit{i}")));
-            }
-        }
+        // Will be filled in by GKR Instruction trait
+        let lookups = vec![];
+        let aux_wits = vec![];
 
         Ok(LargeEcallConfig {
             dummy_insn,
@@ -164,6 +125,57 @@ impl<E: ExtensionField, S: SyscallSpec> Instruction<E> for LargeEcallDummy<E, S>
 impl<E: ExtensionField> GKRIOPInstruction<E> for LargeEcallDummy<E, KeccakSpec> {
     type Layout = KeccakLayout<E>;
 
+    fn gkr_info() -> crate::instructions::GKRinfo {
+        GKRinfo {
+            and_lookups: 3 * AND_LOOKUPS,
+            xor_lookups: 3 * XOR_LOOKUPS,
+            range_lookups: RANGE_LOOKUPS,
+            aux_wits: 40144,
+        }
+    }
+
+    fn construct_circuit_with_gkr_iop(
+        cb: &mut CircuitBuilder<E>,
+    ) -> Result<Self::InstructionConfig, ZKVMError> {
+        let mut partial_config = Self::construct_circuit(cb)?;
+
+        assert!(partial_config.lookups.is_empty());
+        assert!(partial_config.aux_wits.is_empty());
+
+        // TODO: capacity
+        let mut lookups = vec![];
+        let mut aux_wits = vec![];
+
+        for i in 0..AND_LOOKUPS {
+            let a = cb.create_witin(|| format!("and_lookup_{i}_a"));
+            let b = cb.create_witin(|| format!("and_lookup_{i}_b"));
+            let c = cb.create_witin(|| format!("and_lookup_{i}_c"));
+            cb.lookup_and_byte(a.into(), b.into(), c.into())?;
+            lookups.extend(vec![a, b, c]);
+        }
+        for i in 0..XOR_LOOKUPS {
+            let a = cb.create_witin(|| format!("xor_lookup_{i}_a"));
+            let b = cb.create_witin(|| format!("xor_lookup_{i}_b"));
+            let c = cb.create_witin(|| format!("xor_lookup_{i}_c"));
+            cb.lookup_xor_byte(a.into(), b.into(), c.into())?;
+            lookups.extend(vec![a, b, c]);
+        }
+        for i in 0..RANGE_LOOKUPS {
+            let wit = cb.create_witin(|| format!("range_lookup_{i}"));
+            cb.assert_ux::<_, _, 16>(|| "nada", wit.into())?;
+            lookups.push(wit);
+        }
+
+        for i in 0..40144 {
+            aux_wits.push(cb.create_witin(|| format!("aux_wit{i}")));
+        }
+
+        partial_config.lookups = lookups;
+        partial_config.aux_wits = aux_wits;
+
+        Ok(partial_config)
+    }
+
     fn phase1_witness_from_steps(
         layout: &Self::Layout,
         steps: &[StepRecord],
@@ -195,7 +207,6 @@ impl<E: ExtensionField> GKRIOPInstruction<E> for LargeEcallDummy<E, KeccakSpec>
     ) -> Result<(), ZKVMError> {
         Self::assign_instance(config, instance, lk_multiplicity, step)?;
 
-        let active_rounds = 24;
         let mut wit_iter = lookups.iter().map(|f| f.to_canonical_u64());
         let mut var_iter = config.lookups.iter();
 
@@ -206,29 +217,19 @@ impl<E: ExtensionField> GKRIOPInstruction<E> for LargeEcallDummy<E, KeccakSpec>
             let wit = wit_iter.next().unwrap();
             let var = var_iter.next().unwrap();
             set_val!(instance, var, wit);
-            // set_val!(instance, var, 0);
             wit
         };
 
-        for _round in 0..active_rounds {
-            for _i in 0..AND_LOOKUPS_PER_ROUND {
-                lk_multiplicity.lookup_and_byte(pop_arg(), pop_arg());
-                // lk_multiplicity.lookup_and_byte(0, 0);
-                // pop_arg();
-                // pop_arg();
-                pop_arg();
-            }
-            for _i in 0..XOR_LOOKUPS_PER_ROUND {
-                lk_multiplicity.lookup_xor_byte(pop_arg(), pop_arg());
-                // lk_multiplicity.lookup_xor_byte(0, 0);
-                // pop_arg();
-                // pop_arg();
-                pop_arg();
-            }
-            for _i in 0..RANGE_LOOKUPS_PER_ROUND {
-                lk_multiplicity.assert_ux::<16>(pop_arg());
-                // pop_arg();
-            }
+        for _i in 0..AND_LOOKUPS {
+            lk_multiplicity.lookup_and_byte(pop_arg(), pop_arg());
+            pop_arg();
+        }
+        for _i in 0..XOR_LOOKUPS {
+            lk_multiplicity.lookup_xor_byte(pop_arg(), pop_arg());
+            pop_arg();
+        }
+        for _i in 0..RANGE_LOOKUPS {
+            lk_multiplicity.assert_ux::<16>(pop_arg());
         }
 
         dbg!(aux_wits.len());

ceno_zkvm/src/instructions/riscv/rv32im.rs

@@ -526,7 +526,7 @@ impl<E: ExtensionField> DummyExtraConfig<E> {
 
     pub fn assign_opcode_circuit(
         &self,
-        cs: &mut ZKVMConstraintSystem<E>,
+        cs: &ZKVMConstraintSystem<E>,
         witness: &mut ZKVMWitnesses<E>,
         steps: GroupedSteps,
     ) -> Result<(), ZKVMError> {

ceno_zkvm/src/scheme/prover.rs

@@ -481,7 +481,8 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProver<E, PCS> {
                 );
             }
 
-            dbg!(sel_r.len());
+            // dbg!(sel_r);
+            // panic!();
 
             let mut sel_w = build_eq_x_r_vec(&rt_w[log2_w_count..]);
             if num_instances < sel_w.len() {