[Experiment] ceno + babybear

ceno_zkvm/src/e2e.rs

@@ -20,20 +20,24 @@ use ceno_emul::{
     Tracer, VMState, WORD_SIZE, WordAddr,
 };
 use clap::ValueEnum;
-use ff_ext::{ExtensionField, GoldilocksExt2};
+use ff_ext::{BabyBearExt4, ExtensionField, GoldilocksExt2};
 use itertools::{Itertools, MinMaxResult, chain};
 use mpcs::{Basefold, BasefoldRSParams, PolynomialCommitmentScheme};
-use p3::goldilocks::Goldilocks;
+use p3::{babybear::BabyBear, goldilocks::Goldilocks};
 use std::{
     collections::{BTreeSet, HashMap, HashSet},
     sync::Arc,
 };
 use tracing::info;
 use transcript::{BasicTranscript as Transcript, BasicTranscriptWithStat, StatisticRecorder};
 
-pub type E = GoldilocksExt2;
-pub type B = Goldilocks;
-pub type Pcs = Basefold<GoldilocksExt2, BasefoldRSParams>;
+// 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> {
     mem: Vec<Record>,
@@ -417,6 +421,7 @@ pub enum Checkpoint {
     Keygen,
     PrepE2EProving,
     PrepWitnessGen,
+    PrepProof,
     PrepSanityCheck,
     Complete,
 }
@@ -553,6 +558,13 @@ pub fn run_e2e_with_checkpoint<
 
     let verifier = ZKVMVerifier::new(vk.clone());
 
+    if let Checkpoint::PrepProof = checkpoint {
+        return (
+            (Some(zkvm_proof.clone()), Some(vk)),
+            Box::new(move || run_e2e_verify::<E, _>(&verifier, zkvm_proof, exit_code, max_steps)),
+        );
+    }
+
     run_e2e_verify::<E, _>(&verifier, zkvm_proof.clone(), exit_code, max_steps);
 
     if let Checkpoint::PrepSanityCheck = checkpoint {

ceno_zkvm/src/instructions/riscv/rv32im.rs

@@ -63,14 +63,14 @@ pub struct Rv32imConfig<E: ExtensionField> {
     pub sra_config: <SraInstruction<E> as Instruction<E>>::InstructionConfig,
     pub slt_config: <SltInstruction<E> as Instruction<E>>::InstructionConfig,
     pub sltu_config: <SltuInstruction<E> as Instruction<E>>::InstructionConfig,
-    pub mul_config: <MulInstruction<E> as Instruction<E>>::InstructionConfig,
-    pub mulh_config: <MulhInstruction<E> as Instruction<E>>::InstructionConfig,
-    pub mulhsu_config: <MulhsuInstruction<E> as Instruction<E>>::InstructionConfig,
-    pub mulhu_config: <MulhuInstruction<E> as Instruction<E>>::InstructionConfig,
-    pub divu_config: <DivuInstruction<E> as Instruction<E>>::InstructionConfig,
-    pub remu_config: <RemuInstruction<E> as Instruction<E>>::InstructionConfig,
-    pub div_config: <DivInstruction<E> as Instruction<E>>::InstructionConfig,
-    pub rem_config: <RemInstruction<E> as Instruction<E>>::InstructionConfig,
+    // pub mul_config: <MulInstruction<E> as Instruction<E>>::InstructionConfig,
+    // pub mulh_config: <MulhInstruction<E> as Instruction<E>>::InstructionConfig,
+    // pub mulhsu_config: <MulhsuInstruction<E> as Instruction<E>>::InstructionConfig,
+    // pub mulhu_config: <MulhuInstruction<E> as Instruction<E>>::InstructionConfig,
+    // pub divu_config: <DivuInstruction<E> as Instruction<E>>::InstructionConfig,
+    // pub remu_config: <RemuInstruction<E> as Instruction<E>>::InstructionConfig,
+    // pub div_config: <DivInstruction<E> as Instruction<E>>::InstructionConfig,
+    // pub rem_config: <RemInstruction<E> as Instruction<E>>::InstructionConfig,
 
     // ALU with imm
     pub addi_config: <AddiInstruction<E> as Instruction<E>>::InstructionConfig,
@@ -133,14 +133,14 @@ impl<E: ExtensionField> Rv32imConfig<E> {
         let sra_config = cs.register_opcode_circuit::<SraInstruction<E>>();
         let slt_config = cs.register_opcode_circuit::<SltInstruction<E>>();
         let sltu_config = cs.register_opcode_circuit::<SltuInstruction<E>>();
-        let mul_config = cs.register_opcode_circuit::<MulInstruction<E>>();
-        let mulh_config = cs.register_opcode_circuit::<MulhInstruction<E>>();
-        let mulhsu_config = cs.register_opcode_circuit::<MulhsuInstruction<E>>();
-        let mulhu_config = cs.register_opcode_circuit::<MulhuInstruction<E>>();
-        let divu_config = cs.register_opcode_circuit::<DivuInstruction<E>>();
-        let remu_config = cs.register_opcode_circuit::<RemuInstruction<E>>();
-        let div_config = cs.register_opcode_circuit::<DivInstruction<E>>();
-        let rem_config = cs.register_opcode_circuit::<RemInstruction<E>>();
+        // let mul_config = cs.register_opcode_circuit::<MulInstruction<E>>();
+        // let mulh_config = cs.register_opcode_circuit::<MulhInstruction<E>>();
+        // let mulhsu_config = cs.register_opcode_circuit::<MulhsuInstruction<E>>();
+        // let mulhu_config = cs.register_opcode_circuit::<MulhuInstruction<E>>();
+        // let divu_config = cs.register_opcode_circuit::<DivuInstruction<E>>();
+        // let remu_config = cs.register_opcode_circuit::<RemuInstruction<E>>();
+        // let div_config = cs.register_opcode_circuit::<DivInstruction<E>>();
+        // let rem_config = cs.register_opcode_circuit::<RemInstruction<E>>();
 
         // alu with imm opcodes
         let addi_config = cs.register_opcode_circuit::<AddiInstruction<E>>();
@@ -200,14 +200,14 @@ impl<E: ExtensionField> Rv32imConfig<E> {
             sra_config,
             slt_config,
             sltu_config,
-            mul_config,
-            mulh_config,
-            mulhsu_config,
-            mulhu_config,
-            divu_config,
-            remu_config,
-            div_config,
-            rem_config,
+            // mul_config,
+            // mulh_config,
+            // mulhsu_config,
+            // mulhu_config,
+            // divu_config,
+            // remu_config,
+            // div_config,
+            // rem_config,
             // alu with imm
             addi_config,
             andi_config,
@@ -268,14 +268,14 @@ impl<E: ExtensionField> Rv32imConfig<E> {
         fixed.register_opcode_circuit::<SraInstruction<E>>(cs);
         fixed.register_opcode_circuit::<SltInstruction<E>>(cs);
         fixed.register_opcode_circuit::<SltuInstruction<E>>(cs);
-        fixed.register_opcode_circuit::<MulInstruction<E>>(cs);
-        fixed.register_opcode_circuit::<MulhInstruction<E>>(cs);
-        fixed.register_opcode_circuit::<MulhsuInstruction<E>>(cs);
-        fixed.register_opcode_circuit::<MulhuInstruction<E>>(cs);
-        fixed.register_opcode_circuit::<DivuInstruction<E>>(cs);
-        fixed.register_opcode_circuit::<RemuInstruction<E>>(cs);
-        fixed.register_opcode_circuit::<DivInstruction<E>>(cs);
-        fixed.register_opcode_circuit::<RemInstruction<E>>(cs);
+        // fixed.register_opcode_circuit::<MulInstruction<E>>(cs);
+        // fixed.register_opcode_circuit::<MulhInstruction<E>>(cs);
+        // fixed.register_opcode_circuit::<MulhsuInstruction<E>>(cs);
+        // fixed.register_opcode_circuit::<MulhuInstruction<E>>(cs);
+        // fixed.register_opcode_circuit::<DivuInstruction<E>>(cs);
+        // fixed.register_opcode_circuit::<RemuInstruction<E>>(cs);
+        // fixed.register_opcode_circuit::<DivInstruction<E>>(cs);
+        // fixed.register_opcode_circuit::<RemInstruction<E>>(cs);
         // alu with imm
         fixed.register_opcode_circuit::<AddiInstruction<E>>(cs);
         fixed.register_opcode_circuit::<AndiInstruction<E>>(cs);
@@ -370,14 +370,14 @@ impl<E: ExtensionField> Rv32imConfig<E> {
         assign_opcode!(SRA, SraInstruction<E>, sra_config);
         assign_opcode!(SLT, SltInstruction<E>, slt_config);
         assign_opcode!(SLTU, SltuInstruction<E>, sltu_config);
-        assign_opcode!(MUL, MulInstruction<E>, mul_config);
-        assign_opcode!(MULH, MulhInstruction<E>, mulh_config);
-        assign_opcode!(MULHSU, MulhsuInstruction<E>, mulhsu_config);
-        assign_opcode!(MULHU, MulhuInstruction<E>, mulhu_config);
-        assign_opcode!(DIVU, DivuInstruction<E>, divu_config);
-        assign_opcode!(REMU, RemuInstruction<E>, remu_config);
-        assign_opcode!(DIV, DivInstruction<E>, div_config);
-        assign_opcode!(REM, RemInstruction<E>, rem_config);
+        // assign_opcode!(MUL, MulInstruction<E>, mul_config);
+        // assign_opcode!(MULH, MulhInstruction<E>, mulh_config);
+        // assign_opcode!(MULHSU, MulhsuInstruction<E>, mulhsu_config);
+        // assign_opcode!(MULHU, MulhuInstruction<E>, mulhu_config);
+        // assign_opcode!(DIVU, DivuInstruction<E>, divu_config);
+        // assign_opcode!(REMU, RemuInstruction<E>, remu_config);
+        // assign_opcode!(DIV, DivInstruction<E>, div_config);
+        // assign_opcode!(REM, RemInstruction<E>, rem_config);
         // alu with imm
         assign_opcode!(ADDI, AddiInstruction<E>, addi_config);
         assign_opcode!(ANDI, AndiInstruction<E>, andi_config);
@@ -411,11 +411,11 @@ impl<E: ExtensionField> Rv32imConfig<E> {
         // ecall / halt
         witness.assign_opcode_circuit::<HaltInstruction<E>>(cs, &self.halt_config, halt_records)?;
 
-        assert_eq!(
-            all_records.keys().cloned().collect::<BTreeSet<_>>(),
-            // these are opcodes that haven't been implemented
-            [INVALID, ECALL].into_iter().collect::<BTreeSet<_>>(),
-        );
+        // assert_eq!(
+        //     all_records.keys().cloned().collect::<BTreeSet<_>>(),
+        //     // these are opcodes that haven't been implemented
+        //     [INVALID, ECALL].into_iter().collect::<BTreeSet<_>>(),
+        // );
         Ok(GroupedSteps(all_records))
     }
 
@@ -621,7 +621,7 @@ impl<E: ExtensionField> DummyExtraConfig<E> {
 
         let _ = steps.remove(&INVALID);
         let keys: Vec<&InsnKind> = steps.keys().collect::<Vec<_>>();
-        assert!(steps.is_empty(), "unimplemented opcodes: {:?}", keys);
+        // assert!(steps.is_empty(), "unimplemented opcodes: {:?}", keys);
         Ok(())
     }
 }

ceno_zkvm/src/scheme/verifier.rs

@@ -30,6 +30,7 @@ use super::{
     ZKVMOpcodeProof, ZKVMProof, ZKVMTableProof, constants::MAINCONSTRAIN_SUMCHECK_BATCH_SIZE,
 };
 
+#[derive(Clone)]
 pub struct ZKVMVerifier<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>> {
     pub(crate) vk: ZKVMVerifyingKey<E, PCS>,
 }

poseidon/Cargo.toml

@@ -19,3 +19,8 @@ unroll = "0.1"
 [dev-dependencies]
 ark-std.workspace = true
 rand.workspace = true
+
+[features]
+default = ["babybear"]
+babybear = []
+goldilocks = []

poseidon/src/constants.rs

@@ -1 +1,5 @@
+#[cfg(not(feature = "babybear"))]
 pub const DIGEST_WIDTH: usize = 4;
+
+#[cfg(feature = "babybear")]
+pub const DIGEST_WIDTH: usize = 8;