WIP refactor #799 (#952)

### Change scope
- [x] unify `Expression` with ceno
- [x] unify sumcheck with ceno
- [ ] WIP GKR witness generation, take bit benchmark as example

---------

Co-authored-by: Zhang Zhuo <[email protected]>

Author: hero78119

.github/workflows/integration.yml

@@ -58,3 +58,11 @@ jobs:
         env:
           RUSTFLAGS: "-C opt-level=3"
         run: cargo run --release --package ceno_zkvm --bin e2e -- --platform=ceno --hints=10 --public-io=4191 examples/target/riscv32im-ceno-zkvm-elf/release/examples/fibonacci
+
+      - name: Install cargo make
+        run: |
+          cargo make --version || cargo install cargo-make
+
+      - name: Test install Ceno cli
+        run: |
+          cargo make cli

.github/workflows/lints.yml

@@ -43,6 +43,7 @@ jobs:
       - name: Install cargo make
         run: |
           cargo make --version || cargo install cargo-make
+
       - name: Check code format
         run: cargo fmt --all --check
 

Cargo.lock

@@ -12,7 +12,6 @@ members = [
   "sumcheck_macro",
   "poseidon",
   "gkr_iop",
-  "subprotocols",
   "sumcheck",
   "transcript",
   "whir",
@@ -75,9 +74,9 @@ serde = { version = "1.0", features = ["derive", "rc"] }
 serde_json = "1.0"
 strum = "0.26"
 strum_macros = "0.26"
-subprotocols = { path = "subprotocols" }
 substrate-bn = { version = "0.6.0" }
 sumcheck = { path = "sumcheck" }
+thiserror = "1" # do we need this?
 tiny-keccak = { version = "2.0.2", features = ["keccak"] }
 tracing = { version = "0.1", features = [
   "attributes",

Cargo.toml

@@ -30,3 +30,17 @@ args = [
 ]
 command = "cargo"
 workspace = false
+
+[tasks.cli]
+args = [
+  "install",
+  "--features",
+  "jemalloc",
+  "--features",
+  "nightly-features",
+  "--path",
+  "./ceno_cli",
+]
+command = "cargo"
+env = { "JEMALLOC_SYS_WITH_MALLOC_CONF" = "retain:true,metadata_thp:always,thp:always,dirty_decay_ms:-1,muzzy_decay_ms:-1,abort_conf:true" }
+workspace = false

Makefile.toml

@@ -23,6 +23,9 @@ tracing.workspace = true
 tracing-forest.workspace = true
 tracing-subscriber.workspace = true
 
+[target.'cfg(unix)'.dependencies]
+tikv-jemallocator = { version = "0.6", optional = true }
+
 ceno_emul = { path = "../ceno_emul" }
 ceno_host = { path = "../ceno_host" }
 ceno_zkvm = { path = "../ceno_zkvm" }
@@ -33,6 +36,8 @@ mpcs = { path = "../mpcs" }
 vergen-git2 = { version = "1", features = ["build", "cargo", "rustc", "emit_and_set"] }
 
 [features]
+jemalloc = ["dep:tikv-jemallocator", "ceno_zkvm/jemalloc"]
+jemalloc-prof = ["jemalloc", "tikv-jemallocator?/profiling"]
 nightly-features = [
   "ceno_zkvm/nightly-features",
   "ff_ext/nightly-features",

ceno_cli/Cargo.toml

@@ -1,10 +1,17 @@
 use crate::{commands::*, utils::*};
 use anyhow::Context;
+#[cfg(all(feature = "jemalloc", unix, not(test)))]
+use ceno_zkvm::print_allocated_bytes;
 use clap::{Args, Parser, Subcommand};
 
 mod commands;
 mod utils;
 
+// Use jemalloc as global allocator for performance
+#[cfg(all(feature = "jemalloc", unix, not(test)))]
+#[global_allocator]
+static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
 const CENO_VERSION: &str = env!("CENO_VERSION");
 
 #[derive(Parser)]
@@ -86,4 +93,8 @@ fn main() {
         print_error(e);
         std::process::exit(1);
     }
+    #[cfg(all(feature = "jemalloc", unix, not(test)))]
+    {
+        print_allocated_bytes();
+    }
 }

ceno_cli/src/main.rs

@@ -24,7 +24,6 @@ rand_chacha.workspace = true
 rayon.workspace = true
 serde.workspace = true
 serde_json.workspace = true
-subprotocols.workspace = true
 sumcheck.workspace = true
 transcript = { path = "../transcript" }
 witness = { path = "../witness" }
@@ -51,6 +50,10 @@ tempfile = "3.14"
 thread_local = "1.1"
 tiny-keccak.workspace = true
 
+[target.'cfg(unix)'.dependencies]
+tikv-jemalloc-ctl = { version = "0.6", features = ["stats"], optional = true }
+tikv-jemallocator = { version = "0.6", optional = true }
+
 [dev-dependencies]
 cfg-if.workspace = true
 criterion.workspace = true
@@ -65,6 +68,8 @@ glob = "0.3"
 default = ["forbid_overflow"]
 flamegraph = ["pprof2/flamegraph", "pprof2/criterion"]
 forbid_overflow = []
+jemalloc = ["dep:tikv-jemallocator", "dep:tikv-jemalloc-ctl"]
+jemalloc-prof = ["jemalloc", "tikv-jemallocator?/profiling"]
 nightly-features = [
   "p3/nightly-features",
   "ff_ext/nightly-features",

ceno_zkvm/Cargo.toml

@@ -0,0 +1,4 @@
+// Use jemalloc as global allocator for performance
+#[cfg(all(feature = "jemalloc", unix))]
+#[global_allocator]
+static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;

ceno_zkvm/benches/alloc.rs

@@ -7,6 +7,7 @@ use ceno_zkvm::{
     e2e::{Checkpoint, Preset, run_e2e_with_checkpoint, setup_platform},
     scheme::{constants::MAX_NUM_VARIABLES, verifier::ZKVMVerifier},
 };
+mod alloc;
 use criterion::*;
 
 use ff_ext::GoldilocksExt2;
@@ -31,7 +32,7 @@ fn setup() -> (Program, Platform) {
     let stack_size = 32768;
     let heap_size = 2097152;
     let pub_io_size = 16;
-    let program = Program::load_elf(ceno_examples::fibonacci, u32::MAX).unwrap();
+    let program = Program::load_elf(ceno_examples::guest_keccak, u32::MAX).unwrap();
     let platform = setup_platform(Preset::Ceno, &program, stack_size, heap_size, pub_io_size);
     (program, platform)
 }

ceno_zkvm/benches/fibonacci.rs

@@ -7,6 +7,7 @@ use ceno_zkvm::{
     e2e::{Checkpoint, Preset, run_e2e_with_checkpoint, setup_platform},
     scheme::constants::MAX_NUM_VARIABLES,
 };
+mod alloc;
 use criterion::*;
 use ff_ext::GoldilocksExt2;
 use mpcs::{BasefoldDefault, SecurityLevel};

ceno_zkvm/benches/fibonacci_witness.rs

@@ -7,6 +7,7 @@ use ceno_zkvm::{
     e2e::{Checkpoint, Preset, run_e2e_with_checkpoint, setup_platform},
     scheme::constants::MAX_NUM_VARIABLES,
 };
+mod alloc;
 use criterion::*;
 use ff_ext::GoldilocksExt2;
 use mpcs::{BasefoldDefault, SecurityLevel};

ceno_zkvm/benches/is_prime.rs

@@ -7,6 +7,7 @@ use ceno_zkvm::{
     e2e::{Checkpoint, Preset, run_e2e_with_checkpoint, setup_platform},
     scheme::constants::MAX_NUM_VARIABLES,
 };
+mod alloc;
 use criterion::*;
 use ff_ext::GoldilocksExt2;
 use mpcs::{BasefoldDefault, SecurityLevel};

ceno_zkvm/benches/quadratic_sorting.rs

@@ -6,6 +6,7 @@ use ceno_zkvm::{
     scheme::prover::ZKVMProver,
     structs::{ZKVMConstraintSystem, ZKVMFixedTraces},
 };
+mod alloc;
 use criterion::*;
 
 use ceno_zkvm::scheme::constants::MAX_NUM_VARIABLES;

ceno_zkvm/benches/riscv_add.rs

@@ -1,5 +1,7 @@
 use ceno_emul::{IterAddresses, Platform, Program, WORD_SIZE, Word};
 use ceno_host::{CenoStdin, memory_from_file};
+#[cfg(all(feature = "jemalloc", unix, not(test)))]
+use ceno_zkvm::print_allocated_bytes;
 use ceno_zkvm::{
     e2e::{
         Checkpoint, FieldType, PcsKind, Preset, run_e2e_with_checkpoint, setup_platform,
@@ -26,6 +28,11 @@ use tracing_subscriber::{
 };
 use transcript::BasicTranscript as Transcript;
 
+// Use jemalloc as global allocator for performance
+#[cfg(all(feature = "jemalloc", unix, not(test)))]
+#[global_allocator]
+static GLOBAL: tikv_jemallocator::Jemalloc = tikv_jemallocator::Jemalloc;
+
 fn parse_size(s: &str) -> Result<u32, parse_size::Error> {
     parse_size::Config::new()
         .with_binary()
@@ -272,6 +279,11 @@ fn main() {
                 Checkpoint::PrepVerify, // FIXME: when whir and babybear is ready
             )
         }
+    };
+
+    #[cfg(all(feature = "jemalloc", unix, not(test)))]
+    {
+        print_allocated_bytes();
     }
 }
 

ceno_zkvm/src/bin/e2e.rs

@@ -209,7 +209,7 @@ impl<E: ExtensionField> GKRIOPInstruction<E> for LargeEcallDummy<E, KeccakSpec>
             })
             .collect_vec();
 
-        layout.phase1_witness(KeccakTrace { instances })
+        layout.phase1_witness_group(KeccakTrace { instances })
     }
 
     fn assign_instance_with_gkr_iop(

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

@@ -20,6 +20,8 @@ pub mod stats;
 pub mod structs;
 mod uint;
 mod utils;
+#[cfg(all(feature = "jemalloc", unix, not(test)))]
+pub use utils::print_allocated_bytes;
 mod witness;
 
 pub use structs::ROMType;

ceno_zkvm/src/lib.rs

@@ -530,7 +530,7 @@ impl<'a, E: ExtensionField + Hash> MockProver<E> {
             {
                 let right = -right.as_ref();
 
-                let left_evaluated = wit_infer_by_expr(&[], wits_in, &[], pi, &challenge, left);
+                let left_evaluated = wit_infer_by_expr(&[], wits_in, &[], pi, &challenge, &left);
                 let left_evaluated = left_evaluated.get_base_field_vec();
 
                 let right_evaluated = wit_infer_by_expr(&[], wits_in, &[], pi, &challenge, &right);