quadratic_sorting and is_prime benchmarks

ceno_zkvm/Cargo.toml

@@ -19,7 +19,9 @@ serde.workspace = true
 serde_json.workspace = true
 
 base64 = "0.22"
+ceno-examples = { path = "../examples-builder" }
 ceno_emul = { path = "../ceno_emul" }
+ceno_host = { path = "../ceno_host" }
 ff_ext = { path = "../ff_ext" }
 mpcs = { path = "../mpcs" }
 multilinear_extensions = { version = "0", path = "../multilinear_extensions" }
@@ -51,6 +53,7 @@ criterion.workspace = true
 pprof2.workspace = true
 proptest.workspace = true
 
+
 [build-dependencies]
 glob = "0.3"
 
@@ -71,3 +74,11 @@ name = "fibonacci"
 [[bench]]
 harness = false
 name = "fibonacci_witness"
+
+[[bench]]
+harness = false
+name = "quadratic_sorting"
+
+[[bench]]
+harness = false
+name = "is_prime"

ceno_zkvm/benches/is_prime.rs

@@ -0,0 +1,73 @@
+use std::time::Duration;
+
+use ceno_emul::{Platform, Program};
+use ceno_host::CenoStdin;
+use ceno_zkvm::{
+    self,
+    e2e::{Checkpoint, Preset, run_e2e_with_checkpoint, setup_platform},
+};
+use criterion::*;
+
+use goldilocks::GoldilocksExt2;
+use mpcs::BasefoldDefault;
+
+criterion_group! {
+  name = is_prime;
+  config = Criterion::default().warm_up_time(Duration::from_millis(5000));
+  targets = is_prime_1
+}
+
+criterion_main!(is_prime);
+
+const NUM_SAMPLES: usize = 10;
+type Pcs = BasefoldDefault<E>;
+type E = GoldilocksExt2;
+
+// Relevant init data for fibonacci run
+fn setup() -> (Program, Platform) {
+    let stack_size = 32 << 10;
+    let heap_size = 2 << 20;
+    let pub_io_size = 16;
+    let program = Program::load_elf(ceno_examples::is_prime, u32::MAX).unwrap();
+    let platform = setup_platform(Preset::Ceno, &program, stack_size, heap_size, pub_io_size);
+    (program, platform)
+}
+
+fn is_prime_1(c: &mut Criterion) {
+    let (program, platform) = setup();
+
+    for n in [100u32, 10000u32] {
+        let max_steps = usize::MAX;
+        let mut hints = CenoStdin::default();
+        hints.write(&n).unwrap();
+        let hints: Vec<u32> = (&hints).into();
+
+        let mut group = c.benchmark_group("is_prime".to_string());
+        group.sample_size(NUM_SAMPLES);
+
+        // Benchmark the proving time
+        group.bench_function(BenchmarkId::new("is_prime", format!("n = {}", n)), |b| {
+            b.iter_custom(|iters| {
+                let mut time = Duration::new(0, 0);
+
+                for _ in 0..iters {
+                    let (_, prove) = run_e2e_with_checkpoint::<E, Pcs>(
+                        program.clone(),
+                        platform.clone(),
+                        hints.clone(),
+                        max_steps,
+                        Checkpoint::PrepE2EProving,
+                    );
+                    let instant = std::time::Instant::now();
+                    prove();
+                    time += instant.elapsed();
+                }
+                time
+            });
+        });
+
+        group.finish();
+    }
+
+    type E = GoldilocksExt2;
+}

ceno_zkvm/benches/quadratic_sorting.rs

@@ -0,0 +1,80 @@
+use std::time::Duration;
+
+use ceno_emul::{Platform, Program};
+use ceno_host::CenoStdin;
+use ceno_zkvm::{
+    self,
+    e2e::{Checkpoint, Preset, run_e2e_with_checkpoint, setup_platform},
+};
+use criterion::*;
+
+use goldilocks::GoldilocksExt2;
+use mpcs::BasefoldDefault;
+
+criterion_group! {
+  name = quadratic_sorting;
+  config = Criterion::default().warm_up_time(Duration::from_millis(5000));
+  targets = quadratic_sorting_1
+}
+
+criterion_main!(quadratic_sorting);
+
+const NUM_SAMPLES: usize = 10;
+type Pcs = BasefoldDefault<E>;
+type E = GoldilocksExt2;
+
+// Relevant init data for fibonacci run
+fn setup() -> (Program, Platform) {
+    let stack_size = 32 << 10;
+    let heap_size = 2 << 20;
+    let pub_io_size = 16;
+
+    let program = Program::load_elf(ceno_examples::quadratic_sorting, u32::MAX).unwrap();
+    let platform = setup_platform(Preset::Ceno, &program, stack_size, heap_size, pub_io_size);
+    (program, platform)
+}
+
+fn quadratic_sorting_1(c: &mut Criterion) {
+    let (program, platform) = setup();
+    use rand::{Rng, SeedableRng};
+    let mut rng = rand::rngs::StdRng::seed_from_u64(42);
+
+    for n in [100, 500] {
+        let max_steps = usize::MAX;
+        let mut hints = CenoStdin::default();
+        hints
+            .write(&(0..n).map(|_| rng.gen::<u32>()).collect::<Vec<_>>())
+            .unwrap();
+        let hints: Vec<u32> = (&hints).into();
+
+        let mut group = c.benchmark_group("quadratic_sorting".to_string());
+        group.sample_size(NUM_SAMPLES);
+
+        // Benchmark the proving time
+        group.bench_function(
+            BenchmarkId::new("quadratic_sorting", format!("n = {}", n)),
+            |b| {
+                b.iter_custom(|iters| {
+                    let mut time = Duration::new(0, 0);
+                    for _ in 0..iters {
+                        let (_, prove) = run_e2e_with_checkpoint::<E, Pcs>(
+                            program.clone(),
+                            platform.clone(),
+                            hints.clone(),
+                            max_steps,
+                            Checkpoint::PrepE2EProving,
+                        );
+                        let instant = std::time::Instant::now();
+                        prove();
+                        time += instant.elapsed();
+                    }
+                    time
+                });
+            },
+        );
+
+        group.finish();
+    }
+
+    type E = GoldilocksExt2;
+}

ceno_zkvm/src/e2e.rs

@@ -14,18 +14,14 @@ use crate::{
     tables::{MemFinalRecord, MemInitRecord, ProgramTableCircuit, ProgramTableConfig},
 };
 use ceno_emul::{
-    ByteAddr, CENO_PLATFORM, EmuContext, InsnKind, IterAddresses, Platform, Program, StepRecord,
-    Tracer, VMState, WORD_SIZE, WordAddr,
+    CENO_PLATFORM, EmuContext, InsnKind, IterAddresses, Platform, Program, StepRecord, Tracer,
+    VMState, WORD_SIZE, WordAddr,
 };
 use clap::ValueEnum;
 use ff_ext::ExtensionField;
-use itertools::{Itertools, MinMaxResult, chain};
+use itertools::{Itertools, chain};
 use mpcs::PolynomialCommitmentScheme;
-use std::{
-    collections::{BTreeSet, HashMap, HashSet},
-    iter::zip,
-    sync::Arc,
-};
+use std::{collections::BTreeSet, iter::zip, sync::Arc};
 use transcript::BasicTranscript as Transcript;
 
 pub struct FullMemState<Record> {
@@ -129,7 +125,6 @@ fn emulate_program(
             }
         })
         .collect_vec();
-    debug_memory_ranges(&vm, &mem_final);
 
     // Find the final public IO cycles.
     let io_final = io_init
@@ -548,50 +543,3 @@ pub fn run_e2e_verify<E: ExtensionField, PCS: PolynomialCommitmentScheme<E>>(
         None => tracing::error!("Unfinished execution. max_steps={:?}.", max_steps),
     }
 }
-
-fn debug_memory_ranges(vm: &VMState, mem_final: &[MemFinalRecord]) {
-    let accessed_addrs = vm
-        .tracer()
-        .final_accesses()
-        .iter()
-        .filter(|(_, &cycle)| (cycle != 0))
-        .map(|(&addr, _)| addr.baddr())
-        .filter(|addr| vm.platform().can_read(addr.0))
-        .collect_vec();
-
-    let handled_addrs = mem_final
-        .iter()
-        .filter(|rec| rec.cycle != 0)
-        .map(|rec| ByteAddr(rec.addr))
-        .collect::<HashSet<_>>();
-
-    tracing::debug!(
-        "Memory range (accessed): {:?}",
-        format_segments(vm.platform(), accessed_addrs.iter().copied())
-    );
-    tracing::debug!(
-        "Memory range (handled):  {:?}",
-        format_segments(vm.platform(), handled_addrs.iter().copied())
-    );
-
-    for addr in &accessed_addrs {
-        assert!(handled_addrs.contains(addr), "unhandled addr: {:?}", addr);
-    }
-}
-
-fn format_segments(
-    platform: &Platform,
-    addrs: impl Iterator<Item = ByteAddr>,
-) -> HashMap<String, MinMaxResult<ByteAddr>> {
-    addrs
-        .into_grouping_map_by(|addr| format_segment(platform, addr.0))
-        .minmax()
-}
-
-fn format_segment(platform: &Platform, addr: u32) -> String {
-    format!(
-        "{}{}",
-        if platform.can_read(addr) { "R" } else { "-" },
-        if platform.can_write(addr) { "W" } else { "-" },
-    )
-}

examples/examples/is_prime.rs

@@ -0,0 +1,30 @@
+extern crate ceno_rt;
+use rkyv::Archived;
+
+fn is_prime(n: u32) -> bool {
+    if n < 2 {
+        return false;
+    }
+    let mut i = 2;
+    while i * i <= n {
+        if n % i == 0 {
+            return false;
+        }
+        i += 1;
+    }
+
+    true
+}
+
+fn main() {
+    let n: &Archived<u32> = ceno_rt::read();
+    let mut cnt_primes = 0;
+
+    for i in 0..=n.into() {
+        cnt_primes += is_prime(i) as u32;
+    }
+
+    if cnt_primes > 1000 * 1000 {
+        panic!();
+    }
+}

examples/examples/quadratic_sorting.rs

@@ -1,6 +1,4 @@
 extern crate ceno_rt;
-use ceno_rt::println;
-use core::fmt::Write;
 use rkyv::vec::ArchivedVec;
 
 fn sort<T: Ord>(slice: &mut [T]) {
@@ -18,6 +16,4 @@ fn main() {
     let input: &ArchivedVec<u32> = ceno_rt::read();
     let mut scratch = input.to_vec();
     sort(&mut scratch);
-    // Print any output you feel like, eg the first element of the sorted vector:
-    println!("{}", scratch[0]);
 }