[Experiment] ceno + babybear

Cargo.toml

@@ -43,21 +43,21 @@ num-bigint = { version = "0.4.6" }
 num-derive = "0.4"
 num-traits = "0.2"
 p3 = { path = "p3" }
-p3-baby-bear = { git = "https://github.com/Plonky3/plonky3", rev = "1ba4e5c" }
-p3-challenger = { git = "https://github.com/Plonky3/plonky3", rev = "1ba4e5c" }
-p3-commit = { git = "https://github.com/Plonky3/plonky3", rev = "1ba4e5c" }
-p3-dft = { git = "https://github.com/Plonky3/plonky3", rev = "1ba4e5c" }
-p3-field = { git = "https://github.com/Plonky3/plonky3", rev = "1ba4e5c" }
-p3-fri = { git = "https://github.com/Plonky3/plonky3", rev = "1ba4e5c" }
-p3-goldilocks = { git = "https://github.com/Plonky3/plonky3", rev = "1ba4e5c" }
-p3-matrix = { git = "https://github.com/Plonky3/plonky3", rev = "1ba4e5c" }
-p3-maybe-rayon = { git = "https://github.com/Plonky3/plonky3", rev = "1ba4e5c" }
-p3-mds = { git = "https://github.com/Plonky3/plonky3", rev = "1ba4e5c" }
-p3-merkle-tree = { git = "https://github.com/Plonky3/plonky3", rev = "1ba4e5c" }
-p3-poseidon = { git = "https://github.com/Plonky3/plonky3", rev = "1ba4e5c" }
-p3-poseidon2 = { git = "https://github.com/Plonky3/plonky3", rev = "1ba4e5c" }
-p3-symmetric = { git = "https://github.com/Plonky3/plonky3", rev = "1ba4e5c" }
-p3-util = { git = "https://github.com/Plonky3/plonky3", rev = "1ba4e5c" }
+p3-baby-bear = { git = "https://github.com/Plonky3/plonky3", rev = "539bbc84085efb609f4f62cb03cf49588388abdb" }
+p3-challenger = { git = "https://github.com/Plonky3/plonky3", rev = "539bbc84085efb609f4f62cb03cf49588388abdb" }
+p3-commit = { git = "https://github.com/Plonky3/plonky3", rev = "539bbc84085efb609f4f62cb03cf49588388abdb" }
+p3-dft = { git = "https://github.com/Plonky3/plonky3", rev = "539bbc84085efb609f4f62cb03cf49588388abdb" }
+p3-field = { git = "https://github.com/Plonky3/plonky3", rev = "539bbc84085efb609f4f62cb03cf49588388abdb" }
+p3-fri = { git = "https://github.com/Plonky3/plonky3", rev = "539bbc84085efb609f4f62cb03cf49588388abdb" }
+p3-goldilocks = { git = "https://github.com/Plonky3/plonky3", rev = "539bbc84085efb609f4f62cb03cf49588388abdb" }
+p3-matrix = { git = "https://github.com/Plonky3/plonky3", rev = "539bbc84085efb609f4f62cb03cf49588388abdb" }
+p3-maybe-rayon = { git = "https://github.com/Plonky3/plonky3", rev = "539bbc84085efb609f4f62cb03cf49588388abdb" }
+p3-mds = { git = "https://github.com/Plonky3/plonky3", rev = "539bbc84085efb609f4f62cb03cf49588388abdb" }
+p3-merkle-tree = { git = "https://github.com/Plonky3/plonky3", rev = "539bbc84085efb609f4f62cb03cf49588388abdb" }
+p3-poseidon = { git = "https://github.com/Plonky3/plonky3", rev = "539bbc84085efb609f4f62cb03cf49588388abdb" }
+p3-poseidon2 = { git = "https://github.com/Plonky3/plonky3", rev = "539bbc84085efb609f4f62cb03cf49588388abdb" }
+p3-symmetric = { git = "https://github.com/Plonky3/plonky3", rev = "539bbc84085efb609f4f62cb03cf49588388abdb" }
+p3-util = { git = "https://github.com/Plonky3/plonky3", rev = "539bbc84085efb609f4f62cb03cf49588388abdb" }
 paste = "1"
 poseidon = { path = "./poseidon" }
 pprof2 = { version = "0.13", features = ["flamegraph"] }

ceno_zkvm/src/bin/e2e.rs

@@ -343,13 +343,10 @@ fn run_inner<
     let proof_bytes = bincode::serialize(&zkvm_proof).unwrap();
     fs::write(&proof_file, proof_bytes).unwrap();
     let vk_bytes = bincode::serialize(&vk).unwrap();
-    fs::write(&vk_file, vk_bytes).unwrap();
+    std::fs::write(&vk_file, vk_bytes).unwrap();
 
-    if checkpoint > Checkpoint::PrepVerify {
-        let verifier = ZKVMVerifier::new(vk);
-        verify(&zkvm_proof, &verifier).expect("Verification failed");
-        soundness_test(zkvm_proof, &verifier);
-    }
+    let verifier = ZKVMVerifier::new(vk);
+    verify(&zkvm_proof, &verifier).expect("Verification failed");
 }
 
 fn soundness_test<E: ExtensionField, Pcs: PolynomialCommitmentScheme<E>>(

ceno_zkvm/src/e2e.rs

@@ -20,17 +20,19 @@ use ceno_emul::{
     StepRecord, Tracer, VMState, WORD_SIZE, WordAddr, host_utils::read_all_messages,
 };
 use clap::ValueEnum;
-use ff_ext::ExtensionField;
+use ff_ext::{BabyBearExt4, ExtensionField};
 #[cfg(debug_assertions)]
 use ff_ext::{Instrumented, PoseidonField};
 use gkr_iop::hal::ProverBackend;
 use itertools::{Itertools, MinMaxResult, chain};
-use mpcs::{PolynomialCommitmentScheme, SecurityLevel};
+use mpcs::{Basefold, BasefoldRSParams, PolynomialCommitmentScheme, SecurityLevel};
+use p3::{babybear::BabyBear, goldilocks::Goldilocks};
 use serde::Serialize;
 use std::{
     collections::{BTreeSet, HashMap, HashSet},
     sync::Arc,
 };
+use tracing::info;
 use transcript::BasicTranscript as Transcript;
 
 /// The polynomial commitment scheme kind
@@ -75,6 +77,14 @@ pub enum FieldType {
     BabyBear,
 }
 
+// pub type E = GoldilocksExt2;
+// pub type B = Goldilocks;
+// pub type Pcs = Basefold<GoldilocksExt2, BasefoldRSParams>;
+
+pub type E = BabyBearExt4;
+pub type B = BabyBear;
+pub type Pcs = Basefold<E, BasefoldRSParams>;
+
 pub struct FullMemState<Record> {
     pub mem: Vec<Record>,
     pub io: Vec<Record>,
@@ -507,6 +517,7 @@ pub fn generate_witness<E: ExtensionField>(
 pub enum Checkpoint {
     PrepE2EProving,
     PrepWitnessGen,
+    PrepProof,
     PrepVerify,
     #[default]
     Complete,

ceno_zkvm/src/instructions/riscv/branch/branch_circuit.rs

@@ -153,8 +153,8 @@ impl<E: ExtensionField, I: RIVInstruction> Instruction<E> for BranchCircuit<E, I
         if let Some(equal) = &config.is_equal {
             equal.assign_instance(
                 instance,
-                E::BaseField::from_canonical_u64(rs2.as_u64()),
-                E::BaseField::from_canonical_u64(rs1.as_u64()),
+                E::BaseField::from_wrapped_u64(rs2.as_u64()),
+                E::BaseField::from_wrapped_u64(rs1.as_u64()),
             )?;
         }
 

ceno_zkvm/src/instructions/riscv/jump/jalr.rs

@@ -1,6 +1,6 @@
-use std::marker::PhantomData;
-
+use either::Either;
 use ff_ext::ExtensionField;
+use std::marker::PhantomData;
 
 use crate::{
     Value,
@@ -81,10 +81,11 @@ impl<E: ExtensionField> Instruction<E> for JalrInstruction<E> {
             (overflow.expr(), Some((overflow, tmp)))
         };
 
+        let pow2_32 = E::BaseField::from_wrapped_u64(1u64 << 32).expr();
         circuit_builder.require_equal(
             || "rs1+imm = next_pc_unrounded + overflow*2^32",
             rs1_read.value() + imm.expr(),
-            next_pc_addr.expr_unaligned() + overflow_expr * (1u64 << 32),
+            next_pc_addr.expr_unaligned() + overflow_expr * pow2_32,
         )?;
 
         circuit_builder.require_equal(