key gen with entry pc

Author: hero78119

ceno_zkvm/benches/riscv_add.rs

@@ -55,7 +55,7 @@ fn bench_add(c: &mut Criterion) {
 
     let pk = zkvm_cs
         .clone()
-        .key_gen::<Pcs>(pp, vp, zkvm_fixed_traces)
+        .key_gen::<Pcs>(pp, vp, 0, zkvm_fixed_traces)
         .expect("keygen failed");
 
     let (max_num_variables, security_level) = default_backend_config();

ceno_zkvm/src/e2e.rs

@@ -895,7 +895,12 @@ impl<E: ExtensionField> E2EProgramCtx<'_, E> {
             .system_config
             .zkvm_cs
             .clone()
-            .key_gen::<PCS>(pp.clone(), vp.clone(), self.zkvm_fixed_traces.clone())
+            .key_gen::<PCS>(
+                pp.clone(),
+                vp.clone(),
+                self.program.entry,
+                self.zkvm_fixed_traces.clone(),
+            )
             .expect("keygen failed");
         let vk = pk.get_vk_slow();
         (pk, vk)
@@ -915,6 +920,7 @@ impl<E: ExtensionField> E2EProgramCtx<'_, E> {
             .key_gen::<PCS>(
                 pb.get_pp().clone(),
                 pb.get_vp().clone(),
+                self.program.entry,
                 self.zkvm_fixed_traces.clone(),
             )
             .expect("keygen failed");

ceno_zkvm/src/keygen.rs

@@ -12,6 +12,7 @@ impl<E: ExtensionField> ZKVMConstraintSystem<E> {
         self,
         pp: PCS::ProverParam,
         vp: PCS::VerifierParam,
+        entry_pc: u32,
         mut vm_fixed_traces: ZKVMFixedTraces<E>,
     ) -> Result<ZKVMProvingKey<E, PCS>, ZKVMError> {
         let mut vm_pk = ZKVMProvingKey::new(pp.clone(), vp);
@@ -41,6 +42,8 @@ impl<E: ExtensionField> ZKVMConstraintSystem<E> {
         vm_pk.initial_global_state_expr = self.initial_global_state_expr;
         vm_pk.finalize_global_state_expr = self.finalize_global_state_expr;
 
+        vm_pk.set_program_entry_pc(entry_pc);
+
         Ok(vm_pk)
     }
 }

ceno_zkvm/src/scheme/tests.rs

@@ -132,6 +132,7 @@ fn test_rw_lk_expression_combination() {
             .key_gen::<Pcs>(
                 device.backend.pp.clone(),
                 device.backend.vp.clone(),
+                0,
                 zkvm_fixed_traces,
             )
             .unwrap();
@@ -308,7 +309,7 @@ fn test_single_add_instance_e2e() {
 
     let pk = zkvm_cs
         .clone()
-        .key_gen::<Pcs>(pp, vp, zkvm_fixed_traces)
+        .key_gen::<Pcs>(pp, vp, program.entry, zkvm_fixed_traces)
         .expect("keygen failed");
     let vk = pk.get_vk_slow();
 

ceno_zkvm/src/scheme/verifier.rs

@@ -8,7 +8,7 @@ use ff_ext::{Instrumented, PoseidonField};
 use super::{ZKVMChipProof, ZKVMProof};
 use crate::{
     error::ZKVMError,
-    instructions::riscv::constants::{INIT_PC_IDX, SHARD_ID_IDX},
+    instructions::riscv::constants::{END_PC_IDX, INIT_CYCLE_IDX, INIT_PC_IDX, SHARD_ID_IDX},
     scheme::constants::NUM_FANIN,
     structs::{ComposedConstrainSystem, PointAndEval, TowerProofs, VerifyingKey, ZKVMVerifyingKey},
 };
@@ -84,13 +84,13 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMVerifier<E, PCS>
         let _logup_sum = E::ZERO;
         assert!(!vm_proofs.is_empty());
         let num_proofs = vm_proofs.len();
-        vm_proofs
+        let _end_pc = vm_proofs
             .into_iter()
             .zip_eq(transcripts)
             // optionally halt on last chunk
             .zip_eq(iter::repeat_n(false, num_proofs - 1).chain(iter::once(expect_halt)))
             .enumerate()
-            .try_for_each(|(shard_id, ((vm_proof, transcript), expect_halt))| {
+            .try_fold(None, |prev_pc, (shard_id, ((vm_proof, transcript), expect_halt))| {
                 // require ecall/halt proof to exist, depend on whether we expect a halt.
                 let has_halt = vm_proof.has_halt(&self.vk);
                 if has_halt != expect_halt {
@@ -101,12 +101,19 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMVerifier<E, PCS>
                             .into(),
                     ));
                 }
-                // TODO check init ts = 4
-                // TODO check init_pc match prev end_pc
+                // check init cycle = 4
+                assert_eq!(vm_proof.pi_evals[INIT_CYCLE_IDX], E::from_canonical_usize(4));
+                // check init_pc match prev end_pc
+                if let Some(prev_pc) = prev_pc {
+                    assert_eq!(vm_proof.pi_evals[INIT_PC_IDX], prev_pc);
+                } else {
+                    // first chunk, check program entry
+                    assert_eq!(vm_proof.pi_evals[INIT_PC_IDX], E::from_canonical_u32(self.vk.entry_pc));
+                }
+                let end_pc = vm_proof.pi_evals[END_PC_IDX];
                 // TODO add to global ec
                 let _shard_ram_bus = self.verify_proof_validity(shard_id, vm_proof, transcript)?;
-
-                Ok(())
+                Ok(Some(end_pc))
             })?;
         // TODO check global ec_sum == 0
         Ok(true)

ceno_zkvm/src/structs.rs

@@ -409,6 +409,8 @@ impl<E: ExtensionField> ZKVMWitnesses<E> {
 pub struct ZKVMProvingKey<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> {
     pub pp: PCS::ProverParam,
     pub vp: PCS::VerifierParam,
+    // entry program counter
+    pub entry_pc: u32,
     // pk for opcode and table circuits
     pub circuit_pks: BTreeMap<String, ProvingKey<E>>,
     pub fixed_commit_wd: Option<Arc<<PCS as PolynomialCommitmentScheme<E>>::CommitmentWithWitness>>,
@@ -425,6 +427,7 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProvingKey<E, PC
         Self {
             pp,
             vp,
+            entry_pc: 0,
             circuit_pks: BTreeMap::new(),
             initial_global_state_expr: Expression::ZERO,
             finalize_global_state_expr: Expression::ZERO,
@@ -451,12 +454,17 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProvingKey<E, PC
         }
         Ok(())
     }
+
+    pub(crate) fn set_program_entry_pc(&mut self, entry_pc: u32) {
+        self.entry_pc = entry_pc;
+    }
 }
 
 impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProvingKey<E, PCS> {
     pub fn get_vk_slow(&self) -> ZKVMVerifyingKey<E, PCS> {
         ZKVMVerifyingKey {
             vp: self.vp.clone(),
+            entry_pc: self.entry_pc,
             circuit_vks: self
                 .circuit_pks
                 .iter()
@@ -483,6 +491,8 @@ impl<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> ZKVMProvingKey<E, PC
 ))]
 pub struct ZKVMVerifyingKey<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> {
     pub vp: PCS::VerifierParam,
+    // entry program counter
+    pub entry_pc: u32,
     // vk for opcode and table circuits
     pub circuit_vks: BTreeMap<String, VerifyingKey<E>>,
     pub fixed_commit: Option<<PCS as PolynomialCommitmentScheme<E>>::Commitment>,