optimize witness cloning (#89)

* optimize sumcheck algo

circuit witness: direct witness on mle

devirgo style on phase1_output

* temporarily exclude singer-pro from default workspace members

* fux build error in mpcs

* [Experiment] new zkVM design (#91)

* optimize sumcheck algo

circuit witness: direct witness on mle

devirgo style on phase1_output

* initial version for new zkVM design

* riscv add prototype implementation

* add new zkVM prover

* new package ceno_zkvm

* record witness generation

* add transcript

* add verifier

* code cleanup

* rename expression

* prover record_r/record_w sumcheck

* main sel sumcheck proof/verify

* tower product witness inference

* tower product sumcheck prove/verify

* chores: fix tower sumcheck witness length error and clean up

* verify record and zero expression

* tower sumcheck prove/verify pass

* WIP test main sel prove/verify

* add benchmark

* chores: interleaving with default value

* main constraint sumcheck prove/verify pass

* chores: mock witness

* main constraint sumcheck verify final claim assertion pass

* restructure ceno_zkvm package

* refine expression format

* wip lookup

* lookup in logup implemetation with integration test pass

* chores: code cosmetics

* optimize with 2-stage sumcheck #103

* chores: refine virtual polys naming

* fix proper ts and pc counting

* try sumcheck bench

* refine global state in riscv

* degree > 1 main constraint sumcheck implementation #107 (#108)

* monomial expression to virtual poly

* degree > 1 sumcheck batched with main constraint

* succint selector evaluation

* refine succint selector evaluation formula and documentation

* wip fix interleaving edge case

* deal with interleaving_mles instance = 1 case

* chores: code cosmetics and address review comments

* fix logup padding with chip record challenge

* riscv opcode type & combine add/sub opcode & dependency trim

* ci whitelist ceno_zkvm lint/clippy

* address review comments and naming cosmetics

* remove unnessesary to_vec operation

* tower verifier logup p(x) constant check

* cleanup and hide thread-based logic

* soundness fix: derive new sumcheck batched challenge for each round

* fix sel evaluation point and add TODO check

* fix sumcheck batched challenge deriving order

* chore: rename pc step size & fine tune project structure

* fix lint error

.github/workflows/lints.yml

@@ -57,11 +57,11 @@ jobs:
         uses: actions-rs/cargo@v1
         with:
           command: make
-          args: fmt-check
+          args: fmt-check-selected-packages
 
       - name: Run clippy
         uses: actions-rs/cargo@v1
         with:
           command: make
-          args: clippy
+          args: clippy-check-selected-packages
 

Cargo.toml

@@ -10,6 +10,7 @@ members = [
     "singer-utils",
     "sumcheck",
     "transcript",
+    "ceno_zkvm"
 ]
 
 [workspace.package]

Makefile.toml

@@ -5,7 +5,7 @@ RAYON_NUM_THREADS = "${CORE}"
 
 [tasks.tests]
 command = "cargo"
-args = ["test", "--lib", "--release", "--all"]
+args = ["test", "--lib", "--release", "--workspace", "--exclude", "singer-pro"]
 
 [tasks.fmt-check]
 command = "cargo"
@@ -18,3 +18,11 @@ args = ["fmt", "--all"]
 [tasks.clippy]
 command = "cargo"
 args = ["clippy", "--all-features", "--all-targets", "--", "-D", "warnings"]
+
+[tasks.fmt-check-selected-packages]
+command = "cargo"
+args = ["fmt", "-p", "ceno_zkvm", "--", "--check"]
+
+[tasks.clippy-check-selected-packages]
+command = "cargo"
+args = ["clippy", "-p", "ceno_zkvm", "--", "-D", "warnings"]

ceno_zkvm/Cargo.toml

@@ -0,0 +1,42 @@
+[package]
+name = "ceno_zkvm"
+version.workspace = true
+edition.workspace = true
+license.workspace = true
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+ark-std.workspace = true
+ff.workspace = true
+goldilocks.workspace = true
+rayon.workspace = true
+serde.workspace = true
+
+transcript = { path = "../transcript" }
+sumcheck = { version = "0.1.0", path = "../sumcheck" }
+multilinear_extensions = { version = "0.1.0", path = "../multilinear_extensions" }
+ff_ext = { path = "../ff_ext" }
+
+itertools = "0.12.0"
+strum = "0.25.0"
+strum_macros = "0.25.3"
+paste = "1.0.14"
+tracing-subscriber = { version = "0.3.17", features = ["env-filter"] }
+tracing-flame = "0.2.0"
+tracing = "0.1.40"
+
+[dev-dependencies]
+pprof = { version = "0.13", features = ["flamegraph"]}
+criterion = { version = "0.5", features = ["html_reports"] }
+cfg-if = "1.0.0"
+const_env = "0.1.2"
+
+[features]
+
+[profile.bench]
+opt-level = 0
+
+[[bench]]
+name = "riscv_add"
+harness = false

ceno_zkvm/benches/riscv_add.rs

@@ -0,0 +1,127 @@
+#![allow(clippy::manual_memcpy)]
+#![allow(clippy::needless_range_loop)]
+
+use std::time::{Duration, Instant};
+
+use ark_std::test_rng;
+use ceno_zkvm::{
+    circuit_builder::CircuitBuilder,
+    instructions::{riscv::addsub::AddInstruction, Instruction},
+    scheme::prover::ZKVMProver,
+};
+use const_env::from_env;
+use criterion::*;
+
+use ff_ext::ff::Field;
+use goldilocks::{Goldilocks, GoldilocksExt2};
+use itertools::Itertools;
+use multilinear_extensions::mle::IntoMLE;
+use transcript::Transcript;
+
+cfg_if::cfg_if! {
+  if #[cfg(feature = "flamegraph")] {
+    criterion_group! {
+      name = op_add;
+      config = Criterion::default().warm_up_time(Duration::from_millis(3000)).with_profiler(pprof::criterion::PProfProfiler::new(100, pprof::criterion::Output::Flamegraph(None)));
+      targets = bench_add
+    }
+  } else {
+    criterion_group! {
+      name = op_add;
+      config = Criterion::default().warm_up_time(Duration::from_millis(3000));
+      targets = bench_add
+    }
+  }
+}
+
+criterion_main!(op_add);
+
+const NUM_SAMPLES: usize = 10;
+#[from_env]
+const RAYON_NUM_THREADS: usize = 8;
+
+pub fn is_power_of_2(x: usize) -> bool {
+    (x != 0) && ((x & (x - 1)) == 0)
+}
+
+fn bench_add(c: &mut Criterion) {
+    let max_threads = {
+        if !is_power_of_2(RAYON_NUM_THREADS) {
+            #[cfg(not(feature = "non_pow2_rayon_thread"))]
+            {
+                panic!(
+                    "add --features non_pow2_rayon_thread to enable unsafe feature which support non pow of 2 rayon thread pool"
+                );
+            }
+
+            #[cfg(feature = "non_pow2_rayon_thread")]
+            {
+                use sumcheck::{local_thread_pool::create_local_pool_once, util::ceil_log2};
+                let max_thread_id = 1 << ceil_log2(RAYON_NUM_THREADS);
+                create_local_pool_once(1 << ceil_log2(RAYON_NUM_THREADS), true);
+                max_thread_id
+            }
+        } else {
+            RAYON_NUM_THREADS
+        }
+    };
+    let mut circuit_builder = CircuitBuilder::<GoldilocksExt2>::new();
+    let _ = AddInstruction::construct_circuit(&mut circuit_builder);
+    let circuit = circuit_builder.finalize_circuit();
+    let num_witin = circuit.num_witin;
+
+    let prover = ZKVMProver::new(circuit); // circuit clone due to verifier alos need circuit reference
+    let mut transcript = Transcript::new(b"riscv");
+
+    for instance_num_vars in 20..22 {
+        // expand more input size once runtime is acceptable
+        let mut group = c.benchmark_group(format!("add_op_{}", instance_num_vars));
+        group.sample_size(NUM_SAMPLES);
+
+        // Benchmark the proving time
+        group.bench_function(
+            BenchmarkId::new("prove_add", format!("prove_add_log2_{}", instance_num_vars)),
+            |b| {
+                b.iter_with_setup(
+                    || {
+                        let mut rng = test_rng();
+                        let real_challenges = [E::random(&mut rng), E::random(&mut rng)];
+                        (rng, real_challenges)
+                    },
+                    |(mut rng, real_challenges)| {
+                        // generate mock witness
+                        let num_instances = 1 << instance_num_vars;
+                        let wits_in = (0..num_witin as usize)
+                            .map(|_| {
+                                (0..num_instances)
+                                    .map(|_| Goldilocks::random(&mut rng))
+                                    .collect::<Vec<Goldilocks>>()
+                                    .into_mle()
+                                    .into()
+                            })
+                            .collect_vec();
+                        let timer = Instant::now();
+                        let _ = prover
+                            .create_proof(
+                                wits_in,
+                                num_instances,
+                                max_threads,
+                                &mut transcript,
+                                &real_challenges,
+                            )
+                            .expect("create_proof failed");
+                        println!(
+                            "AddInstruction::create_proof, instance_num_vars = {}, time = {}",
+                            instance_num_vars,
+                            timer.elapsed().as_secs_f64()
+                        );
+                    },
+                );
+            },
+        );
+
+        group.finish();
+    }
+
+    type E = GoldilocksExt2;
+}

ceno_zkvm/src/chip_handler.rs

@@ -0,0 +1,36 @@
+use ff_ext::ExtensionField;
+
+use crate::{
+    error::ZKVMError,
+    expression::WitIn,
+    structs::{PCUInt, TSUInt, UInt64},
+};
+
+pub mod general;
+pub mod global_state;
+pub mod register;
+
+pub trait GlobalStateRegisterMachineChipOperations<E: ExtensionField> {
+    fn state_in(&mut self, pc: &PCUInt, ts: &TSUInt) -> Result<(), ZKVMError>;
+
+    fn state_out(&mut self, pc: &PCUInt, ts: &TSUInt) -> Result<(), ZKVMError>;
+}
+
+pub trait RegisterChipOperations<E: ExtensionField> {
+    fn register_read(
+        &mut self,
+        register_id: &WitIn,
+        prev_ts: &mut TSUInt,
+        ts: &mut TSUInt,
+        values: &UInt64,
+    ) -> Result<TSUInt, ZKVMError>;
+
+    fn register_write(
+        &mut self,
+        register_id: &WitIn,
+        prev_ts: &mut TSUInt,
+        ts: &mut TSUInt,
+        prev_values: &UInt64,
+        values: &UInt64,
+    ) -> Result<TSUInt, ZKVMError>;
+}