From e102acf3200fce4a2a40afeb92c626bfa01a47d3 Mon Sep 17 00:00:00 2001 From: Koko Bhadra Date: Wed, 22 Jul 2026 12:40:45 -0400 Subject: [PATCH] bench: microbenches for the hot mixed base x extension field kernels --- Cargo.lock | 1 + crates/backend/poly/src/benchmark_kernels.rs | 120 ++++++++++++++++++ crates/backend/poly/src/lib.rs | 3 + crates/backend/sumcheck/Cargo.toml | 1 + .../src/benchmark_product_sumcheck.rs | 71 +++++++++++ crates/backend/sumcheck/src/lib.rs | 3 + crates/whir/src/benchmark_first_round.rs | 92 ++++++++++++++ crates/whir/src/lib.rs | 3 + crates/whir/src/open.rs | 2 +- 9 files changed, 295 insertions(+), 1 deletion(-) create mode 100644 crates/backend/poly/src/benchmark_kernels.rs create mode 100644 crates/backend/sumcheck/src/benchmark_product_sumcheck.rs create mode 100644 crates/whir/src/benchmark_first_round.rs diff --git a/Cargo.lock b/Cargo.lock index 15cb51f3a..9820fbea8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -878,6 +878,7 @@ dependencies = [ "koala-bear", "parallel", "poly", + "rand", "tracing", "zk-alloc", ] diff --git a/crates/backend/poly/src/benchmark_kernels.rs b/crates/backend/poly/src/benchmark_kernels.rs new file mode 100644 index 000000000..ec114a6cb --- /dev/null +++ b/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 = (0..1usize << n_vars).map(|_| rng.random()).collect(); + let point: Vec = (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::(&evals, &point), expected); + + // warming + for _ in 0..3 { + let _ = black_box(eval_base_packed::(&evals, &point)); + } + + let n_iters = 30; + let time = Instant::now(); + let mut acc = EF::ZERO; + for _ in 0..n_iters { + acc += eval_base_packed::(&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(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::(); + + let alphas: Vec = (0..N_ALPHAS).map(|_| rng.random()).collect(); + let alphas_packed: Vec> = alphas.iter().map(|a| EFPacking::::from(*a)).collect(); + let domainsep: F = rng.random(); + let domainsep_packed = PFPacking::::from(domainsep); + let rows: Vec<[PFPacking; N_DATA]> = (0..n_rows) + .map(|_| core::array::from_fn(|_| PFPacking::::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 = 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::::ZERO, |acc, row| { + acc + finger_print_packed::(domainsep_packed, row, &alphas_packed) + }); + assert_eq!( + unpack_extension::>(&[total]).iter().copied().sum::(), + total_ref + ); + + // warming + let mut acc = EFPacking::::ZERO; + for row in &rows { + acc += finger_print_packed::(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::(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 + ); +} diff --git a/crates/backend/poly/src/lib.rs b/crates/backend/poly/src/lib.rs index c72266ca0..96056e197 100644 --- a/crates/backend/poly/src/lib.rs +++ b/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; diff --git a/crates/backend/sumcheck/Cargo.toml b/crates/backend/sumcheck/Cargo.toml index 47e7a2c7e..9aeac4d9a 100644 --- a/crates/backend/sumcheck/Cargo.toml +++ b/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 diff --git a/crates/backend/sumcheck/src/benchmark_product_sumcheck.rs b/crates/backend/sumcheck/src/benchmark_product_sumcheck.rs new file mode 100644 index 000000000..cd90ec2e9 --- /dev/null +++ b/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 = (0..n).map(|_| rng.random()).collect(); + let weights: Vec = (0..n).map(|_| rng.random()).collect(); + + let evals_packed: &[PFPacking] = PFPacking::::pack_slice(&evals); + let weights_packed: Vec> = 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::(); + let c0_ref = (0..half).map(|i| weights[i] * evals[i]).sum::(); + let c2_ref = (0..half) + .map(|i| (weights[half + i] - weights[i]) * (evals[half + i] - evals[i])) + .sum::(); + let c1_ref = sum - c0_ref.double() - c2_ref; + + let compute = || { + compute_product_sumcheck_polynomial(evals_packed, &weights_packed, sum, |e| { + unpack_extension::>(&[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 + ); +} diff --git a/crates/backend/sumcheck/src/lib.rs b/crates/backend/sumcheck/src/lib.rs index 5ec3c06ba..aa1da9df1 100644 --- a/crates/backend/sumcheck/src/lib.rs +++ b/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; diff --git a/crates/whir/src/benchmark_first_round.rs b/crates/whir/src/benchmark_first_round.rs new file mode 100644 index 000000000..1339598dc --- /dev/null +++ b/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 = (0..n).map(|_| rng.random()).collect(); + let evals_packed: &[PFPacking] = PFPacking::::pack_slice(&evals); + let gamma: EF = rng.random(); + + // Dense equality statements, the hot path of the round-0 combine. + let statements: Vec> = (0..n_statements) + .map(|_| { + let point = MultilinearPoint((0..n_vars).map(|_| rng.random()).collect::>()); + let value = evals.evaluate(&point); + SparseStatement::dense(point, value) + }) + .collect(); + let terms = build_lazy_combine_terms::(&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::(); + let c2_ref = (0..half) + .map(|i| (w_ref[half + i] - w_ref[i]) * (evals[half + i] - evals[i])) + .sum::(); + 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 = 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 + ); +} diff --git a/crates/whir/src/lib.rs b/crates/whir/src/lib.rs index a3e84f9cf..ef6f5aeaf 100644 --- a/crates/whir/src/lib.rs +++ b/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 { pub total_num_variables: usize, diff --git a/crates/whir/src/open.rs b/crates/whir/src/open.rs index cfafa784a..6076d4607 100644 --- a/crates/whir/src/open.rs +++ b/crates/whir/src/open.rs @@ -739,7 +739,7 @@ fn gather_run_terms<'a, EF: ExtensionField>>( } /// One parallel pass: write the rounds-1+ weight buffer and accumulate the round-0 quadratic. -fn combine_and_compute_first_round( +pub(crate) fn combine_and_compute_first_round( evals: &[PFPacking], terms: &LazyCombineTerms, ) -> (DensePolynomial, ArenaVec>)