Merge remote-tracking branch 'origin/06-23-Add_benchmarks_for_GKR_loo…

…kups' into andrew/dev/components/eq_evals

* origin/06-23-Add_benchmarks_for_GKR_lookups:
  Add benchmarks for GKR lookups
  Implement LogupOps for SIMD backend
  Implement GrandProductOps for SIMD backend
  Implement GkrOps for SIMD backend
  Implement MleOps for SIMD backend
  Add GKR implementation of Logup lookups
  Add GKR implementation of Grand Product lookups

crates/prover/Cargo.toml

@@ -63,6 +63,10 @@ harness = false
 harness = false
 name = "field"
 
+[[bench]]
+harness = false
+name = "lookups"
+
 [[bench]]
 harness = false
 name = "matrix"

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();

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);