bench: microbenches for the hot mixed base x extension field kernels (chunk 1 of #260)

crates/backend/poly/src/benchmark_kernels.rs

@@ -0,0 +1,120 @@
+//! Ignored-test microbenches for hot mixed base x extension kernels of this crate:
+//! [`eval_base_packed`] and [`finger_print_packed`].
+//!
+//! Part of the delayed modular reduction work tracked in
+//! https://github.com/leanEthereum/leanVM/issues/260. Each bench first asserts the kernel
+//! output against an independent reference computation, so the harness doubles as a
+//! regression test when the kernels are rewritten.
+
+use std::hint::black_box;
+use std::time::Instant;
+
+use field::{PackedValue, PrimeCharacteristicRing};
+use koala_bear::{KoalaBear, QuinticExtensionFieldKB};
+use rand::rngs::StdRng;
+use rand::{RngExt, SeedableRng};
+
+use crate::*;
+
+type F = KoalaBear;
+type EF = QuinticExtensionFieldKB;
+
+#[test]
+#[ignore]
+fn bench_eval_base_packed() {
+    // cargo test --release --package poly --lib -- benchmark_kernels::bench_eval_base_packed --exact --nocapture --ignored
+
+    let n_vars = 22;
+    let mut rng = StdRng::seed_from_u64(0);
+    let evals: Vec<F> = (0..1usize << n_vars).map(|_| rng.random()).collect();
+    let point: Vec<EF> = (0..n_vars).map(|_| rng.random()).collect();
+
+    // Reference: the recursive strategy is an independent code path.
+    let expected = evals.evaluate_sequential(&MultilinearPoint(point.clone()));
+    assert_eq!(eval_base_packed::<EF, true>(&evals, &point), expected);
+
+    // warming
+    for _ in 0..3 {
+        let _ = black_box(eval_base_packed::<EF, true>(&evals, &point));
+    }
+
+    let n_iters = 30;
+    let time = Instant::now();
+    let mut acc = EF::ZERO;
+    for _ in 0..n_iters {
+        acc += eval_base_packed::<EF, true>(&evals, &point);
+    }
+    let elapsed = time.elapsed();
+    let _ = black_box(acc);
+    println!(
+        "eval_base_packed ({} vars): {:.3} ms/call, {:.0} Melems/s",
+        n_vars,
+        elapsed.as_secs_f64() * 1e3 / n_iters as f64,
+        (n_iters as u64 * (1u64 << n_vars)) as f64 / elapsed.as_secs_f64() / 1e6
+    );
+}
+
+#[test]
+#[ignore]
+fn bench_finger_print_packed() {
+    // cargo test --release --package poly --lib -- benchmark_kernels::bench_finger_print_packed --exact --nocapture --ignored
+
+    let mut rng = StdRng::seed_from_u64(0);
+    // Memory-style tuples (address, value) and bytecode-style tuples (12 instruction
+    // columns + index), the narrowest and widest logup uses.
+    run_finger_print_packed::<2>("memory-style", &mut rng);
+    run_finger_print_packed::<13>("bytecode-style", &mut rng);
+}
+
+fn run_finger_print_packed<const N_DATA: usize>(label: &str, rng: &mut StdRng) {
+    const N_ALPHAS: usize = 16;
+    assert!(N_ALPHAS > N_DATA);
+    let n_rows = 1usize << 17; // packed rows
+    let width = packing_width::<EF>();
+
+    let alphas: Vec<EF> = (0..N_ALPHAS).map(|_| rng.random()).collect();
+    let alphas_packed: Vec<EFPacking<EF>> = alphas.iter().map(|a| EFPacking::<EF>::from(*a)).collect();
+    let domainsep: F = rng.random();
+    let domainsep_packed = PFPacking::<EF>::from(domainsep);
+    let rows: Vec<[PFPacking<EF>; N_DATA]> = (0..n_rows)
+        .map(|_| core::array::from_fn(|_| PFPacking::<EF>::from_fn(|_| rng.random())))
+        .collect();
+
+    // Reference: scalar finger_print on every lane; the packed kernel must match in total.
+    let mut total_ref = EF::ZERO;
+    for row in &rows {
+        for lane in 0..width {
+            let data: Vec<EF> = row.iter().map(|d| EF::from(d.as_slice()[lane])).collect();
+            total_ref += finger_print(EF::from(domainsep), &data, &alphas);
+        }
+    }
+    let total = rows.iter().fold(EFPacking::<EF>::ZERO, |acc, row| {
+        acc + finger_print_packed::<EF>(domainsep_packed, row, &alphas_packed)
+    });
+    assert_eq!(
+        unpack_extension::<EF, Vec<EF>>(&[total]).iter().copied().sum::<EF>(),
+        total_ref
+    );
+
+    // warming
+    let mut acc = EFPacking::<EF>::ZERO;
+    for row in &rows {
+        acc += finger_print_packed::<EF>(domainsep_packed, row, &alphas_packed);
+    }
+
+    let n_passes = 20;
+    let time = Instant::now();
+    for _ in 0..n_passes {
+        for row in &rows {
+            acc += finger_print_packed::<EF>(domainsep_packed, row, &alphas_packed);
+        }
+    }
+    let elapsed = time.elapsed();
+    let _ = black_box(acc);
+    let calls = (n_passes * n_rows) as f64;
+    println!(
+        "finger_print_packed ({label}, {N_DATA} data, {N_ALPHAS} alphas): {:.1} ns/call, {:.0}M scalar rows/s",
+        elapsed.as_secs_f64() * 1e9 / calls,
+        calls * width as f64 / elapsed.as_secs_f64() / 1e6
+    );
+}

crates/backend/poly/src/lib.rs

@@ -26,3 +26,6 @@ pub use wrappers::*;
 
 mod multilinear_utils;
 pub use multilinear_utils::*;
+
+#[cfg(test)]
+mod benchmark_kernels;

crates/backend/sumcheck/Cargo.toml

@@ -14,3 +14,4 @@ tracing.workspace = true
 
 [dev-dependencies]
 koala-bear = { path = "../koala-bear", package = "koala-bear" }
+rand.workspace = true

crates/backend/sumcheck/src/benchmark_product_sumcheck.rs

@@ -0,0 +1,71 @@
+//! Ignored-test microbench for the product-sumcheck quadratic round kernel
+//! ([`compute_product_sumcheck_polynomial`] over base-packed evals x extension-packed
+//! weights, i.e. the [`sumcheck_quadratic`] hot loop).
+//!
+//! Part of the delayed modular reduction work tracked in
+//! https://github.com/leanEthereum/leanVM/issues/260. The bench first asserts the kernel
+//! output against a scalar reference computation, so the harness doubles as a regression
+//! test when the kernel is rewritten.
+
+use std::hint::black_box;
+use std::time::Instant;
+
+use field::{PackedValue, PrimeCharacteristicRing};
+use koala_bear::{KoalaBear, QuinticExtensionFieldKB};
+use poly::*;
+use rand::rngs::StdRng;
+use rand::{RngExt, SeedableRng};
+
+use crate::compute_product_sumcheck_polynomial;
+
+type F = KoalaBear;
+type EF = QuinticExtensionFieldKB;
+
+#[test]
+#[ignore]
+fn bench_product_sumcheck_quadratic_round() {
+    // cargo test --release --package sumcheck --lib -- benchmark_product_sumcheck::bench_product_sumcheck_quadratic_round --exact --nocapture --ignored
+
+    let n_vars = 20;
+    let n = 1usize << n_vars;
+    let mut rng = StdRng::seed_from_u64(0);
+    let evals: Vec<F> = (0..n).map(|_| rng.random()).collect();
+    let weights: Vec<EF> = (0..n).map(|_| rng.random()).collect();
+
+    let evals_packed: &[PFPacking<EF>] = PFPacking::<EF>::pack_slice(&evals);
+    let weights_packed: Vec<EFPacking<EF>> = pack_extension(&weights);
+
+    // Scalar reference for the claimed sum and the two computed coefficients.
+    let half = n / 2;
+    let sum = weights.iter().zip(&evals).map(|(&w, &e)| w * e).sum::<EF>();
+    let c0_ref = (0..half).map(|i| weights[i] * evals[i]).sum::<EF>();
+    let c2_ref = (0..half)
+        .map(|i| (weights[half + i] - weights[i]) * (evals[half + i] - evals[i]))
+        .sum::<EF>();
+    let c1_ref = sum - c0_ref.double() - c2_ref;
+
+    let compute = || {
+        compute_product_sumcheck_polynomial(evals_packed, &weights_packed, sum, |e| {
+            unpack_extension::<EF, Vec<EF>>(&[e])
+        })
+    };
+    assert_eq!(compute().coeffs, vec![c0_ref, c1_ref, c2_ref]);
+
+    // warming
+    for _ in 0..3 {
+        black_box(compute());
+    }
+
+    let n_iters = 30;
+    let time = Instant::now();
+    for _ in 0..n_iters {
+        black_box(compute());
+    }
+    let elapsed = time.elapsed();
+    println!(
+        "product sumcheck quadratic round ({} vars): {:.3} ms/call, {:.0} Melems/s",
+        n_vars,
+        elapsed.as_secs_f64() * 1e3 / n_iters as f64,
+        (n_iters * n) as f64 / elapsed.as_secs_f64() / 1e6
+    );
+}

crates/backend/sumcheck/src/lib.rs

@@ -14,3 +14,6 @@ pub use sc_computation::*;
 
 mod product_computation;
 pub use product_computation::*;
+
+#[cfg(test)]
+mod benchmark_product_sumcheck;

crates/whir/src/benchmark_first_round.rs

@@ -0,0 +1,92 @@
+//! Ignored-test microbench for the WHIR round-0 combine kernel
+//! ([`combine_and_compute_first_round`]): one pass writing the rounds-1+ weight buffer
+//! while accumulating the round-0 quadratic.
+//!
+//! Part of the delayed modular reduction work tracked in
+//! https://github.com/leanEthereum/leanVM/issues/260. The bench first asserts the kernel
+//! output against a scalar reference computation, so the harness doubles as a regression
+//! test when the kernel is rewritten.
+
+use std::hint::black_box;
+use std::time::Instant;
+
+use field::{PackedValue, PrimeCharacteristicRing};
+use koala_bear::{KoalaBear, QuinticExtensionFieldKB};
+use rand::rngs::StdRng;
+use rand::{RngExt, SeedableRng};
+
+use poly::{EvaluationsList, MultilinearPoint, PFPacking, eval_eq_scaled, unpack_extension};
+
+use crate::SparseStatement;
+use crate::open::{build_lazy_combine_terms, combine_and_compute_first_round};
+
+type F = KoalaBear;
+type EF = QuinticExtensionFieldKB;
+
+#[test]
+#[ignore]
+fn bench_combine_and_compute_first_round() {
+    // cargo test --release --package whir --lib -- benchmark_first_round::bench_combine_and_compute_first_round --exact --nocapture --ignored
+
+    let n_vars = 20;
+    let n = 1usize << n_vars;
+    let n_statements = 4;
+    let mut rng = StdRng::seed_from_u64(0);
+    let evals: Vec<F> = (0..n).map(|_| rng.random()).collect();
+    let evals_packed: &[PFPacking<EF>] = PFPacking::<EF>::pack_slice(&evals);
+    let gamma: EF = rng.random();
+
+    // Dense equality statements, the hot path of the round-0 combine.
+    let statements: Vec<SparseStatement<EF>> = (0..n_statements)
+        .map(|_| {
+            let point = MultilinearPoint((0..n_vars).map(|_| rng.random()).collect::<Vec<EF>>());
+            let value = evals.evaluate(&point);
+            SparseStatement::dense(point, value)
+        })
+        .collect();
+    let terms = build_lazy_combine_terms::<EF>(&statements, gamma);
+
+    // Scalar reference: materialized weights, combined sum, and round-0 coefficients.
+    let mut w_ref = vec![EF::ZERO; n];
+    let mut combined_sum_ref = EF::ZERO;
+    let mut gamma_pow = EF::ONE;
+    for statement in &statements {
+        let eq = eval_eq_scaled(&statement.point.0, gamma_pow);
+        for (w, e) in w_ref.iter_mut().zip(eq.iter()) {
+            *w += *e;
+        }
+        combined_sum_ref += statement.values[0].value * gamma_pow;
+        gamma_pow *= gamma;
+    }
+    assert_eq!(terms.combined_sum, combined_sum_ref);
+    let half = n / 2;
+    let c0_ref = (0..half).map(|i| w_ref[i] * evals[i]).sum::<EF>();
+    let c2_ref = (0..half)
+        .map(|i| (w_ref[half + i] - w_ref[i]) * (evals[half + i] - evals[i]))
+        .sum::<EF>();
+    let c1_ref = combined_sum_ref - c0_ref.double() - c2_ref;
+
+    let (first_poly, weights_buf) = combine_and_compute_first_round(evals_packed, &terms);
+    assert_eq!(first_poly.coeffs, vec![c0_ref, c1_ref, c2_ref]);
+    let weights_unpacked: Vec<EF> = unpack_extension(&weights_buf);
+    assert_eq!(weights_unpacked, w_ref);
+
+    // warming
+    for _ in 0..2 {
+        black_box(combine_and_compute_first_round(evals_packed, &terms));
+    }
+
+    let n_iters = 10;
+    let time = Instant::now();
+    for _ in 0..n_iters {
+        black_box(combine_and_compute_first_round(evals_packed, &terms));
+    }
+    let elapsed = time.elapsed();
+    println!(
+        "WHIR round-0 combine ({} vars, {} statements): {:.3} ms/call, {:.0} Melems/s",
+        n_vars,
+        n_statements,
+        elapsed.as_secs_f64() * 1e3 / n_iters as f64,
+        (n_iters * n) as f64 / elapsed.as_secs_f64() / 1e6
+    );
+}

crates/whir/src/lib.rs

@@ -27,6 +27,9 @@ pub(crate) use utils::*;
 mod matrix;
 pub(crate) use matrix::*;
 
+#[cfg(test)]
+mod benchmark_first_round;
+
 #[derive(Clone, Debug)]
 pub struct SparseStatement<EF> {
     pub total_num_variables: usize,

crates/whir/src/open.rs

@@ -739,7 +739,7 @@ fn gather_run_terms<'a, EF: ExtensionField<PF<EF>>>(
 }
 
 /// One parallel pass: write the rounds-1+ weight buffer and accumulate the round-0 quadratic.
-fn combine_and_compute_first_round<EF>(
+pub(crate) fn combine_and_compute_first_round<EF>(
     evals: &[PFPacking<EF>],
     terms: &LazyCombineTerms<EF>,
 ) -> (DensePolynomial<EF>, ArenaVec<EFPacking<EF>>)