From d4eb2e44e09c40325e3be273bfbb2375c43f328b Mon Sep 17 00:00:00 2001
From: Andrew Milson <andrew.j.milson@gmail.com>
Date: Sun, 23 Jun 2024 22:17:51 -0400
Subject: [PATCH] Add benchmarks for GKR lookups

---
 crates/prover/Cargo.toml                     |   4 +
 crates/prover/benches/fft.rs                 |   2 +-
 crates/prover/benches/lookups.rs             | 113 +++++++++++++++++++
 crates/prover/src/core/lookups/gkr_prover.rs |   2 +
 4 files changed, 120 insertions(+), 1 deletion(-)
 create mode 100644 crates/prover/benches/lookups.rs

diff --git a/crates/prover/Cargo.toml b/crates/prover/Cargo.toml
index bc018e896..de8b54981 100644
--- a/crates/prover/Cargo.toml
+++ b/crates/prover/Cargo.toml
@@ -63,6 +63,10 @@ harness = false
 harness = false
 name = "field"
 
+[[bench]]
+harness = false
+name = "lookups"
+
 [[bench]]
 harness = false
 name = "matrix"
diff --git a/crates/prover/benches/fft.rs b/crates/prover/benches/fft.rs
index a8dd61ff7..b5c7a6e47 100644
--- a/crates/prover/benches/fft.rs
+++ b/crates/prover/benches/fft.rs
@@ -18,7 +18,7 @@ use stwo_prover::core::poly::circle::CanonicCoset;
 pub fn simd_ifft(c: &mut Criterion) {
     let mut group = c.benchmark_group("iffts");
 
-    for log_size in 16..=28 {
+    for log_size in 16..=16 {
         let domain = CanonicCoset::new(log_size).circle_domain();
         let twiddle_dbls = get_itwiddle_dbls(domain.half_coset);
         let twiddle_dbls_refs = twiddle_dbls.iter().map(|x| x.as_slice()).collect_vec();
diff --git a/crates/prover/benches/lookups.rs b/crates/prover/benches/lookups.rs
new file mode 100644
index 000000000..fff8fee48
--- /dev/null
+++ b/crates/prover/benches/lookups.rs
@@ -0,0 +1,113 @@
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
+use rand::distributions::{Distribution, Standard};
+use rand::rngs::SmallRng;
+use rand::{Rng, SeedableRng};
+use stwo_prover::core::backend::simd::SimdBackend;
+use stwo_prover::core::backend::CpuBackend;
+use stwo_prover::core::channel::{Blake2sChannel, Channel};
+use stwo_prover::core::fields::Field;
+use stwo_prover::core::lookups::gkr_prover::{prove_batch, GkrOps, Layer};
+use stwo_prover::core::lookups::mle::{Mle, MleOps};
+use stwo_prover::core::vcs::blake2_hash::Blake2sHash;
+
+const LOG_N_ROWS: u32 = 16;
+
+fn bench_gkr_grand_product<B: GkrOps>(c: &mut Criterion, id: &str) {
+    let mut rng = SmallRng::seed_from_u64(0);
+    let layer = Layer::<B>::GrandProduct(gen_random_mle(&mut rng, LOG_N_ROWS));
+    c.bench_function(&format!("{id} grand product lookup 2^{LOG_N_ROWS}"), |b| {
+        b.iter_batched(
+            || layer.clone(),
+            |layer| prove_batch(&mut test_channel(), vec![layer]),
+            BatchSize::LargeInput,
+        )
+    });
+    c.bench_function(
+        &format!("{id} grand product lookup batch 4x 2^{LOG_N_ROWS}"),
+        |b| {
+            b.iter_batched(
+                || vec![layer.clone(), layer.clone(), layer.clone(), layer.clone()],
+                |layers| prove_batch(&mut test_channel(), layers),
+                BatchSize::LargeInput,
+            )
+        },
+    );
+}
+
+fn bench_gkr_logup_generic<B: GkrOps>(c: &mut Criterion, id: &str) {
+    let mut rng = SmallRng::seed_from_u64(0);
+    let generic_layer = Layer::<B>::LogUpGeneric {
+        numerators: gen_random_mle(&mut rng, LOG_N_ROWS),
+        denominators: gen_random_mle(&mut rng, LOG_N_ROWS),
+    };
+    c.bench_function(&format!("{id} generic logup lookup 2^{LOG_N_ROWS}"), |b| {
+        b.iter_batched(
+            || generic_layer.clone(),
+            |layer| prove_batch(&mut test_channel(), vec![layer]),
+            BatchSize::LargeInput,
+        )
+    });
+}
+
+fn bench_gkr_logup_multiplicities<B: GkrOps>(c: &mut Criterion, id: &str) {
+    let mut rng = SmallRng::seed_from_u64(0);
+    let multiplicities_layer = Layer::<B>::LogUpMultiplicities {
+        numerators: gen_random_mle(&mut rng, LOG_N_ROWS),
+        denominators: gen_random_mle(&mut rng, LOG_N_ROWS),
+    };
+    c.bench_function(
+        &format!("{id} multiplicities logup lookup 2^{LOG_N_ROWS}"),
+        |b| {
+            b.iter_batched(
+                || multiplicities_layer.clone(),
+                |layer| prove_batch(&mut test_channel(), vec![layer]),
+                BatchSize::LargeInput,
+            )
+        },
+    );
+}
+
+fn bench_gkr_logup_singles<B: GkrOps>(c: &mut Criterion, id: &str) {
+    let mut rng = SmallRng::seed_from_u64(0);
+    let singles_layer = Layer::<B>::LogUpSingles {
+        denominators: gen_random_mle(&mut rng, LOG_N_ROWS),
+    };
+    c.bench_function(&format!("{id} singles logup lookup 2^{LOG_N_ROWS}"), |b| {
+        b.iter_batched(
+            || singles_layer.clone(),
+            |layer| prove_batch(&mut test_channel(), vec![layer]),
+            BatchSize::LargeInput,
+        )
+    });
+}
+
+/// Generates a random multilinear polynomial.
+fn gen_random_mle<B: MleOps<F>, F: Field>(rng: &mut impl Rng, n_variables: u32) -> Mle<B, F>
+where
+    Standard: Distribution<F>,
+{
+    Mle::new((0..1 << n_variables).map(|_| rng.gen()).collect())
+}
+
+fn gkr_lookup_benches(c: &mut Criterion) {
+    bench_gkr_grand_product::<SimdBackend>(c, "simd");
+    bench_gkr_logup_generic::<SimdBackend>(c, "simd");
+    bench_gkr_logup_multiplicities::<SimdBackend>(c, "simd");
+    bench_gkr_logup_singles::<SimdBackend>(c, "simd");
+
+    bench_gkr_grand_product::<CpuBackend>(c, "cpu");
+    bench_gkr_logup_generic::<CpuBackend>(c, "cpu");
+    bench_gkr_logup_multiplicities::<CpuBackend>(c, "cpu");
+    bench_gkr_logup_singles::<CpuBackend>(c, "cpu");
+}
+
+pub fn test_channel() -> Blake2sChannel {
+    let seed = Blake2sHash::from(vec![0; 32]);
+    Blake2sChannel::new(seed)
+}
+
+criterion_group!(
+    name = benches;
+    config = Criterion::default().sample_size(100);
+    targets = gkr_lookup_benches);
+criterion_main!(benches);
diff --git a/crates/prover/src/core/lookups/gkr_prover.rs b/crates/prover/src/core/lookups/gkr_prover.rs
index 4b58915f2..b2a27f454 100644
--- a/crates/prover/src/core/lookups/gkr_prover.rs
+++ b/crates/prover/src/core/lookups/gkr_prover.rs
@@ -89,6 +89,8 @@ impl<B: ColumnOps<SecureField>> Deref for EqEvals<B> {
 /// numerators and denominators.
 ///
 /// [LogUp]: https://eprint.iacr.org/2023/1284.pdf
+#[derive(Educe)]
+#[educe(Debug, Clone)]
 pub enum Layer<B: GkrOps> {
     GrandProduct(Mle<B, SecureField>),
     LogUpGeneric {