diff --git a/Cargo.lock b/Cargo.lock index 09dbfb301..0f6cf8756 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -103,6 +103,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" dependencies = [ "cfg-if", + "getrandom 0.3.4", "once_cell", "version_check", "zerocopy", @@ -2517,6 +2518,22 @@ dependencies = [ "serde_core", ] +[[package]] +name = "inferno" +version = "0.12.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "90807d610575744524d9bdc69f3885d96f0e6c3354565b0828354a7ff2a262b8" +dependencies = [ + "ahash", + "itoa", + "log", + "num-format", + "once_cell", + "quick-xml", + "rgb", + "str_stack", +] + [[package]] name = "inotify" version = "0.11.2" @@ -3692,6 +3709,16 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "521739c6d2bac4aa25192232afe6841231376b2b26d4d9fae5ecf8ca5772e441" +[[package]] +name = "num-format" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a652d9771a63711fd3c3deb670acfbe5c30a4072e664d7a3bf5a9e1056ac72c3" +dependencies = [ + "arrayvec", + "itoa", +] + [[package]] name = "num-integer" version = "0.1.46" @@ -4506,11 +4533,15 @@ dependencies = [ "anyhow", "ark-ff", "ark-std", + "inferno", "provekit-backend-bn254", "provekit-backend-goldilocks", "provekit-common", "provekit-prover", "provekit-verifier", + "tracing", + "tracing-flame", + "tracing-subscriber", "whir", ] @@ -4605,6 +4636,15 @@ version = "1.2.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a1d01941d82fa2ab50be1e79e6714289dd7cde78eba4c074bc5a4374f650dfe0" +[[package]] +name = "quick-xml" +version = "0.39.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cdcc8dd4e2f670d309a5f0e83fe36dfdc05af317008fea29144da1a2ac858e5e" +dependencies = [ + "memchr", +] + [[package]] name = "quickcheck" version = "1.1.0" @@ -4990,6 +5030,15 @@ dependencies = [ "hmac", ] +[[package]] +name = "rgb" +version = "0.8.53" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47b34b781b31e5d73e9fbc8689c70551fd1ade9a19e3e28cfec8580a79290cc4" +dependencies = [ + "bytemuck", +] + [[package]] name = "ring" version = "0.17.14" @@ -5901,6 +5950,12 @@ version = "1.0.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9e08d8363704e6c71fc928674353e6b7c23dcea9d82d7012c8faf2a3a025f8d0" +[[package]] +name = "str_stack" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f446288b699d66d0fd2e30d1cfe7869194312524b3b9252594868ed26ef056a" + [[package]] name = "strength_reduce" version = "0.2.4" @@ -6476,6 +6531,17 @@ dependencies = [ "tracing-subscriber", ] +[[package]] +name = "tracing-flame" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bae117ee14789185e129aaee5d93750abe67fdc5a9a62650452bfe4e122a3a9" +dependencies = [ + "lazy_static", + "tracing", + "tracing-subscriber", +] + [[package]] name = "tracing-log" version = "0.2.0" diff --git a/Cargo.toml b/Cargo.toml index b184cc15b..821309f82 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -123,6 +123,7 @@ chrono = "0.4.41" # This is a workaround because different package selection based on target does not mix well with workspace dependencies. divan = "0.1.21" hex = "0.4.3" +inferno = { version = "0.12", default-features = false } itertools = "0.14.0" num-bigint = "0.4" paste = "1.0.15" @@ -153,6 +154,7 @@ tokio-util = "0.7.13" tower = "0.5.2" tower-http = { version = "0.6.6", features = ["cors", "timeout", "trace"] } tracing = "0.1.41" +tracing-flame = "0.2" tracing-subscriber = { version = "0.3.18", features = ["env-filter", "ansi"] } tracing-tracy = "=0.11.4" tracy-client = "=0.18.0" diff --git a/provekit/backend/bn254/src/mavros_prove.rs b/provekit/backend/bn254/src/mavros_prove.rs index 79de3ea2e..bf177b7fb 100644 --- a/provekit/backend/bn254/src/mavros_prove.rs +++ b/provekit/backend/bn254/src/mavros_prove.rs @@ -18,7 +18,7 @@ use { }, provekit_prover::{prove_from_alphas, run_zk_sumcheck_prover, WhirR1CSCommitment}, tracing::instrument, - whir::transcript::ProverState, + whir::{algebra::embedding::Identity, transcript::ProverState}, }; pub trait MavrosR1CSProver { @@ -53,13 +53,11 @@ impl MavrosR1CSProver for WhirR1CSScheme { .as_ref() .ok_or_else(|| anyhow::anyhow!("c1 must carry blinding state"))?; - let [a, b, c] = [witgen.out_a, witgen.out_b, witgen.out_c]; // `g` is committed separately in the extension field; open `blinding_eval` // against the blinding vector (g flattened at offset 0). let (alpha, blinding_eval) = run_zk_sumcheck_prover( - a, - b, - c, + &Identity::new(), + [witgen.out_a, witgen.out_b, witgen.out_c], &mut merlin, self.m_0, &blinding.polynomial, diff --git a/provekit/common/src/utils/sumcheck.rs b/provekit/common/src/utils/sumcheck.rs index 93a7cc525..43bc70e56 100644 --- a/provekit/common/src/utils/sumcheck.rs +++ b/provekit/common/src/utils/sumcheck.rs @@ -5,6 +5,7 @@ use { R1CS, }, ark_ff::Field, + rayon::iter::{IndexedParallelIterator, IntoParallelRefIterator, ParallelIterator}, std::array, tracing::instrument, whir::algebra::embedding::Embedding, @@ -102,6 +103,73 @@ fn sumcheck_fold_map_reduce_inner( } } +/// Compute the sum of a vector valued function over the boolean hypercube in +/// the leading variable, with base-field `mles` against an extension-field +/// `ext_mle`. +pub fn mixed_sumcheck_map_reduce( + mles: [&[F]; N], + ext_mle: &[G], + map: impl Fn([(F, F); N], (G, G)) -> [G; M] + Send + Sync + Copy, +) -> [G; M] { + let size = ext_mle.len(); + assert!(size.is_power_of_two()); + assert!(size >= 2); + assert!(mles.iter().all(|mle| mle.len() == size)); + + let mles = mles.map(|mle| mle.split_at(size / 2)); + mixed_map_reduce_inner::(mles, ext_mle.split_at(size / 2), map) +} + +fn mixed_map_reduce_inner( + mles: [(&[F], &[F]); N], + ext_mle: (&[G], &[G]), + map: impl Fn([(F, F); N], (G, G)) -> [G; M] + Send + Sync + Copy, +) -> [G; M] { + let size = ext_mle.0.len(); + if size * (N + 1) * 2 > workload_size::() { + // Split slices + let pairs = mles.map(|(p0, p1)| (p0.split_at(size / 2), p1.split_at(size / 2))); + let left = pairs.map(|((l0, _), (l1, _))| (l0, l1)); + let right = pairs.map(|((_, r0), (_, r1))| (r0, r1)); + let (ext_p0, ext_p1) = (ext_mle.0.split_at(size / 2), ext_mle.1.split_at(size / 2)); + + // Parallel recurse + let (l, r) = rayon::join( + || mixed_map_reduce_inner(left, (ext_p0.0, ext_p1.0), map), + || mixed_map_reduce_inner(right, (ext_p0.1, ext_p1.1), map), + ); + + // Combine results + array::from_fn(|i| l[i] + r[i]) + } else { + let mut result = [G::zero(); M]; + for i in 0..size { + let e = mles.map(|(p0, p1)| (p0[i], p1[i])); + let local = map(e, (ext_mle.0[i], ext_mle.1[i])); + result.iter_mut().zip(local).for_each(|(r, l)| *r += l); + } + result + } +} + +/// Fold the leading variable of a base-field mle at an extension-field point, +/// lifting it into the extension: `out[i] = mle[i] + point * (mle[i + n/2] - +/// mle[i])`. +pub fn mixed_fold( + embedding: &M, + mle: &[M::Source], + point: M::Target, +) -> Vec { + assert!(mle.len().is_power_of_two()); + assert!(mle.len() >= 2); + let (p0, p1) = mle.split_at(mle.len() / 2); + p0.par_iter() + .zip(p1.par_iter()) + .with_min_len(workload_size::()) + .map(|(&e0, &e1)| embedding.mixed_add(embedding.mixed_mul(point, e1 - e0), e0)) + .collect() +} + /// List of evaluations for eq(r, x) over the boolean hypercube, truncated to /// `num_entries` elements. When `num_entries < 2^r.len()`, avoids allocating /// the full hypercube. @@ -230,7 +298,10 @@ mod tests { use { super::*, ark_std::{One, Zero}, - whir::algebra::fields::Field64 as FieldElement, + whir::algebra::{ + embedding::{Basefield, Identity}, + fields::{Field64 as FieldElement, Field64_3}, + }, }; fn fe(v: i64) -> FieldElement { @@ -241,6 +312,10 @@ mod tests { } } + fn fe3(a: i64, b: i64, c: i64) -> Field64_3 { + Field64_3::from_base_prime_field_elems(vec![fe(a), fe(b), fe(c)]).unwrap() + } + /// Build a small 3×4 R1CS for matrix tests. /// /// A = [[1, 2, 0, 0], B = [[0, 1, 0, 0], C = [[0, 0, 1, 0], @@ -360,6 +435,131 @@ mod tests { assert_eq!(result, vec![calculate_eq(&r, &[fe(0), fe(0)])]); } + // mixed_sumcheck_map_reduce / mixed_fold + + /// The cubic sumcheck round evaluations, base mles against ext eq. + fn mixed_round_map( + embedding: &Basefield, + [a, b, c]: [(FieldElement, FieldElement); 3], + eq: (Field64_3, Field64_3), + ) -> [Field64_3; 3] { + let f0 = embedding.mixed_mul(eq.0, a.0 * b.0 - c.0); + let f_em1 = embedding.mixed_mul( + eq.0 + eq.0 - eq.1, + (a.0 + a.0 - a.1) * (b.0 + b.0 - b.1) - (c.0 + c.0 - c.1), + ); + let f_inf = embedding.mixed_mul(eq.1 - eq.0, (a.1 - a.0) * (b.1 - b.0)); + [f0, f_em1, f_inf] + } + + /// The same round evaluations, everything lifted to the extension. + fn ext_round_map([a, b, c, eq]: [(Field64_3, Field64_3); 4]) -> [Field64_3; 3] { + let f0 = eq.0 * (a.0 * b.0 - c.0); + let f_em1 = + (eq.0 + eq.0 - eq.1) * ((a.0 + a.0 - a.1) * (b.0 + b.0 - b.1) - (c.0 + c.0 - c.1)); + let f_inf = (eq.1 - eq.0) * (a.1 - a.0) * (b.1 - b.0); + [f0, f_em1, f_inf] + } + + fn mixed_test_mles(n: usize) -> ([Vec; 3], Vec) { + let a = (0..n).map(|i| fe(i as i64 + 1)).collect(); + let b = (0..n).map(|i| fe(2 * i as i64 + 3)).collect(); + let c = (0..n).map(|i| fe(7 * i as i64 + 5)).collect(); + let eq = (0..n) + .map(|i| fe3(i as i64 + 2, 3 * i as i64 + 1, i as i64)) + .collect(); + ([a, b, c], eq) + } + + #[test] + fn test_mixed_map_reduce_matches_lifted() { + let embedding = Basefield::::new(); + let ([a, b, c], eq) = mixed_test_mles(16); + + let mixed = mixed_sumcheck_map_reduce([&a[..], &b[..], &c[..]], &eq, |mles, eq| { + mixed_round_map(&embedding, mles, eq) + }); + + let (mut la, mut lb, mut lc, mut leq) = ( + embedding.map_vec(a), + embedding.map_vec(b), + embedding.map_vec(c), + eq, + ); + let lifted = + sumcheck_fold_map_reduce([&mut la, &mut lb, &mut lc, &mut leq], None, ext_round_map); + + assert_eq!(mixed, lifted); + } + + #[test] + fn test_mixed_fold_matches_lifted() { + let embedding = Basefield::::new(); + let n = 8; + let mle: Vec = (0..n).map(|i| fe(3 * i as i64 + 2)).collect(); + let point = fe3(5, 7, 11); + + let folded = mixed_fold(&embedding, &mle, point); + + let lifted = embedding.map_vec(mle); + let (p0, p1) = lifted.split_at(n / 2); + let expected: Vec = p0 + .iter() + .zip(p1) + .map(|(&l, &h)| l + point * (h - l)) + .collect(); + assert_eq!(folded, expected); + } + + #[test] + fn test_mixed_fold_identity() { + let n = 4; + let mle: Vec = (0..n) + .map(|i| fe3(i as i64 + 1, i as i64, 2 * i as i64)) + .collect(); + let point = fe3(9, 4, 6); + + let folded = mixed_fold(&Identity::new(), &mle, point); + + let expected: Vec = (0..n / 2) + .map(|i| mle[i] + point * (mle[i + n / 2] - mle[i])) + .collect(); + assert_eq!(folded, expected); + } + + /// Explicit `mixed_fold` + fold-free round must match the fused in-place + /// fold inside `sumcheck_fold_map_reduce` — the parity the prover's + /// base-field first round relies on. + #[test] + fn test_explicit_mixed_fold_matches_fused_fold() { + let embedding = Basefield::::new(); + let ([a, b, c], eq) = mixed_test_mles(8); + let alpha = fe3(3, 8, 2); + + let (mut la, mut lb, mut lc, mut leq) = ( + embedding.map_vec(a.clone()), + embedding.map_vec(b.clone()), + embedding.map_vec(c.clone()), + eq.clone(), + ); + let fused = sumcheck_fold_map_reduce( + [&mut la, &mut lb, &mut lc, &mut leq], + Some(alpha), + ext_round_map, + ); + + let (mut fa, mut fb, mut fc, mut feq) = ( + mixed_fold(&embedding, &a, alpha), + mixed_fold(&embedding, &b, alpha), + mixed_fold(&embedding, &c, alpha), + mixed_fold(&Identity::new(), &eq, alpha), + ); + let explicit = + sumcheck_fold_map_reduce([&mut fa, &mut fb, &mut fc, &mut feq], None, ext_round_map); + + assert_eq!(fused, explicit); + } + /// transpose_r1cs_matrices #[test] diff --git a/provekit/prover/src/whir_r1cs.rs b/provekit/prover/src/whir_r1cs.rs index 8456b813e..885082372 100644 --- a/provekit/prover/src/whir_r1cs.rs +++ b/provekit/prover/src/whir_r1cs.rs @@ -1,7 +1,7 @@ use { ::tracing::instrument, anyhow::{ensure, Result}, - ark_ff::Field, + ark_ff::{Field, One}, ark_std::{ rand::distributions::{Distribution, Standard}, Zero, @@ -16,8 +16,8 @@ use { pad_to_power_of_two, sumcheck::{ calculate_evaluations_over_boolean_hypercube_for_eq, calculate_witness_bounds, - eval_cubic_poly, multiply_transposed_by_eq_alpha, sumcheck_fold_map_reduce, - transpose_r1cs_matrices, + eval_cubic_poly, mixed_fold, mixed_sumcheck_map_reduce, + multiply_transposed_by_eq_alpha, sumcheck_fold_map_reduce, transpose_r1cs_matrices, }, }, Base, Ext, FieldHash, PrefixCovector, ProofField, PublicInputs, WhirR1CSProof, @@ -156,28 +156,21 @@ where let (a, b, c) = calculate_witness_bounds(&r1cs, &full_witness); drop(full_witness); - // Witness bounds are computed in the base field; lift them to the - // extension for the sumcheck (a no-op under the `Identity` embedding, - // where base == ext). let embedding = ::default(); - let (a, b, c) = ( - embedding.map_vec(a), - embedding.map_vec(b), - embedding.map_vec(c), - ); let blinding = commitments[0] .blinding .as_ref() .ok_or_else(|| anyhow::anyhow!("c1 must carry blinding state"))?; - // The Spartan sumcheck runs entirely in the extension field; `g` is - // native ext, committed separately. `blinding_eval` opens the ext - // blinding vector (which holds `g` flattened at offset 0). + // The witness bounds stay in the base field for the first sumcheck + // round; the fold at the first challenge lifts them into the + // extension. `g` is native ext, committed separately; `blinding_eval` + // opens the ext blinding vector (which holds `g` flattened at + // offset 0). let (alpha, blinding_eval) = run_zk_sumcheck_prover( - a, - b, - c, + &embedding, + [a, b, c], &mut merlin, self.m_0, &blinding.polynomial, @@ -495,17 +488,25 @@ pub fn pad_to_pow2_len_min2(v: &mut Vec) { } } +/// Run the blinded Spartan sumcheck. The first round consumes the base-field +/// mles directly; the fold at the first challenge lifts them into the +/// extension, where the remaining rounds run. #[instrument(skip_all)] -pub fn run_zk_sumcheck_prover>( - mut a: Vec, - mut b: Vec, - mut c: Vec, +pub fn run_zk_sumcheck_prover( + embedding: &M, + mles: [Vec; 3], merlin: &mut ProverState, m_0: usize, - blinding_polynomial: &[[F; 4]], - blinding_vector: &[F], -) -> (Vec, F) { - let r: Vec = merlin.verifier_message_vec(m_0); + blinding_polynomial: &[[M::Target; 4]], + blinding_vector: &[M::Target], +) -> (Vec, M::Target) +where + M: Embedding, + M::Target: Codec, + S: DuplexSpongeInterface, +{ + let [mut a, mut b, mut c] = mles; + let r: Vec = merlin.verifier_message_vec(m_0); let mut eq = calculate_evaluations_over_boolean_hypercube_for_eq(&r, 1 << r.len()); pad_to_pow2_len_min2(&mut a); @@ -513,26 +514,57 @@ pub fn run_zk_sumcheck_prover pad_to_pow2_len_min2(&mut c); pad_to_pow2_len_min2(&mut eq); - let mut alpha = Vec::::with_capacity(m_0); + let mut alpha = Vec::::with_capacity(m_0); let sum_g_reduce = sum_over_hypercube(blinding_polynomial); merlin.prover_message(&sum_g_reduce); - let rho: F = merlin.verifier_message(); + let rho: M::Target = merlin.verifier_message(); // Prove that sum of F + ρ·G over the boolean hypercube equals ρ·Σ(G). let mut saved_val_for_sumcheck_equality_assertion = rho * sum_g_reduce; - let mut fold = None; + // 2 is invertible in any field of odd characteristic; precompute 1/2 once + // and reuse it for each round's cubic-coefficient solve. + let half = M::Target::one() / (M::Target::one() + M::Target::one()); + + // First round: a, b, c are base-field, so the constraint products stay in + // the base field with one mixed mul per point against the ext eq. + let hhat = mixed_sumcheck_map_reduce([&a[..], &b[..], &c[..]], &eq, |[a, b, c], eq| { + let f0 = embedding.mixed_mul(eq.0, a.0 * b.0 - c.0); + let f_em1 = embedding.mixed_mul( + eq.0 + eq.0 - eq.1, + (a.0 + a.0 - a.1) * (b.0 + b.0 - b.1) - (c.0 + c.0 - c.1), + ); + let f_inf = embedding.mixed_mul(eq.1 - eq.0, (a.1 - a.0) * (b.1 - b.0)); + + [f0, f_em1, f_inf] + }); + let g_poly = compute_blinding_coefficients_for_round(blinding_polynomial, 0, &[]); + let (alpha_0, saved_val) = combined_round_message( + merlin, + hhat, + g_poly, + rho, + half, + saved_val_for_sumcheck_equality_assertion, + ); + saved_val_for_sumcheck_equality_assertion = saved_val; + alpha.push(alpha_0); + + // Fold at the first challenge, lifting the mles into the extension. + let folded_a = mixed_fold(embedding, &a, alpha_0); + let folded_b = mixed_fold(embedding, &b, alpha_0); + let folded_c = mixed_fold(embedding, &c, alpha_0); + let folded_eq = mixed_fold(&Identity::new(), &eq, alpha_0); + drop((a, b, c, eq)); + let (mut a, mut b, mut c, mut eq) = (folded_a, folded_b, folded_c, folded_eq); - // 2 is invertible in any field of odd characteristic. - let half = (F::one() + F::one()) - .inverse() - .expect("field characteristic is not 2"); + let mut fold = None; - for idx in 0..m_0 { - let [hhat_i_at_0, hhat_i_at_em1, hhat_i_at_inf_over_x_cube] = + for idx in 1..m_0 { + let hhat = sumcheck_fold_map_reduce([&mut a, &mut b, &mut c, &mut eq], fold, |[a, b, c, eq]| { let f0 = eq.0 * (a.0 * b.0 - c.0); let f_em1 = (eq.0 + eq.0 - eq.1) @@ -550,47 +582,18 @@ pub fn run_zk_sumcheck_prover let g_poly = compute_blinding_coefficients_for_round(blinding_polynomial, idx, alpha.as_slice()); - - let mut combined_hhat_i_coeffs = [F::zero(); 4]; - - combined_hhat_i_coeffs[0] = hhat_i_at_0 + rho * g_poly[0]; - - let g_at_minus_one = g_poly[0] - g_poly[1] + g_poly[2] - g_poly[3]; - let combined_at_em1 = hhat_i_at_em1 + rho * g_at_minus_one; - - combined_hhat_i_coeffs[2] = (saved_val_for_sumcheck_equality_assertion + combined_at_em1 - - combined_hhat_i_coeffs[0] - - combined_hhat_i_coeffs[0] - - combined_hhat_i_coeffs[0]) - * half; - - combined_hhat_i_coeffs[3] = hhat_i_at_inf_over_x_cube + rho * g_poly[3]; - - combined_hhat_i_coeffs[1] = saved_val_for_sumcheck_equality_assertion - - combined_hhat_i_coeffs[0] - - combined_hhat_i_coeffs[0] - - combined_hhat_i_coeffs[3] - - combined_hhat_i_coeffs[2]; - - assert_eq!( + let (alpha_i, saved_val) = combined_round_message( + merlin, + hhat, + g_poly, + rho, + half, saved_val_for_sumcheck_equality_assertion, - combined_hhat_i_coeffs[0] - + combined_hhat_i_coeffs[0] - + combined_hhat_i_coeffs[1] - + combined_hhat_i_coeffs[2] - + combined_hhat_i_coeffs[3] ); - - for coeff in &combined_hhat_i_coeffs { - merlin.prover_message(coeff); - } - let alpha_i: F = merlin.verifier_message(); + saved_val_for_sumcheck_equality_assertion = saved_val; alpha.push(alpha_i); fold = Some(alpha_i); - - saved_val_for_sumcheck_equality_assertion = - eval_cubic_poly(combined_hhat_i_coeffs, alpha_i); } drop((a, b, c, eq)); @@ -601,6 +604,57 @@ pub fn run_zk_sumcheck_prover (alpha, blinding_eval) } +/// Absorb the combined round polynomial `hhat_i + ρ·g_i` into the transcript +/// and draw the round challenge; returns it with the sumcheck equality value +/// for the next round. +fn combined_round_message>( + merlin: &mut ProverState, + hhat: [F; 3], + g_poly: [F; 4], + rho: F, + half: F, + saved_val_for_sumcheck_equality_assertion: F, +) -> (F, F) { + let [hhat_i_at_0, hhat_i_at_em1, hhat_i_at_inf_over_x_cube] = hhat; + + let mut combined_hhat_i_coeffs = [F::zero(); 4]; + + combined_hhat_i_coeffs[0] = hhat_i_at_0 + rho * g_poly[0]; + + let g_at_minus_one = g_poly[0] - g_poly[1] + g_poly[2] - g_poly[3]; + let combined_at_em1 = hhat_i_at_em1 + rho * g_at_minus_one; + + combined_hhat_i_coeffs[2] = (saved_val_for_sumcheck_equality_assertion + combined_at_em1 + - combined_hhat_i_coeffs[0] + - combined_hhat_i_coeffs[0] + - combined_hhat_i_coeffs[0]) + * half; + + combined_hhat_i_coeffs[3] = hhat_i_at_inf_over_x_cube + rho * g_poly[3]; + + combined_hhat_i_coeffs[1] = saved_val_for_sumcheck_equality_assertion + - combined_hhat_i_coeffs[0] + - combined_hhat_i_coeffs[0] + - combined_hhat_i_coeffs[3] + - combined_hhat_i_coeffs[2]; + + assert_eq!( + saved_val_for_sumcheck_equality_assertion, + combined_hhat_i_coeffs[0] + + combined_hhat_i_coeffs[0] + + combined_hhat_i_coeffs[1] + + combined_hhat_i_coeffs[2] + + combined_hhat_i_coeffs[3] + ); + + for coeff in &combined_hhat_i_coeffs { + merlin.prover_message(coeff); + } + let alpha_i: F = merlin.verifier_message(); + + (alpha_i, eval_cubic_poly(combined_hhat_i_coeffs, alpha_i)) +} + fn create_weights_and_evaluations( embedding: &M, m: usize, diff --git a/tooling/provekit-fixtures/Cargo.toml b/tooling/provekit-fixtures/Cargo.toml index 3e9db1523..fd6add37d 100644 --- a/tooling/provekit-fixtures/Cargo.toml +++ b/tooling/provekit-fixtures/Cargo.toml @@ -24,6 +24,10 @@ whir.workspace = true # Concrete proof fields to instantiate the generic harness. provekit-backend-bn254.workspace = true provekit-backend-goldilocks.workspace = true +inferno.workspace = true +tracing.workspace = true +tracing-flame.workspace = true +tracing-subscriber.workspace = true [lints] workspace = true diff --git a/tooling/provekit-fixtures/src/harness.rs b/tooling/provekit-fixtures/src/harness.rs index 879efa613..8aec25db9 100644 --- a/tooling/provekit-fixtures/src/harness.rs +++ b/tooling/provekit-fixtures/src/harness.rs @@ -287,3 +287,67 @@ where scheme.prove_noir(merlin, r1cs_owned, vec![c1, c2], full_witness, &public)?; Ok(start.elapsed()) } + +/// Pre-built inputs for a timed single-commit prove, constructed outside the +/// timer/flamegraph so scheme construction and the witness clones are not +/// measured. Opaque: build with [`prove_setup`], consume with +/// [`time_prove_core`]. +pub struct ProveInputs { + scheme: WhirR1CSScheme

, + r1cs: R1CS>, + witness_commit: Vec>, + witness_prove: Vec>, +} + +/// Build the scheme and owned witness copies for a single-commit prove. Kept +/// separate from [`time_prove_core`] so a profiler can install its subscriber +/// between setup and the timed region (only the latter is captured). +pub fn prove_setup

(r1cs: &R1CS>, witness: &[Base

]) -> ProveInputs

+where + P: FieldHash, + Standard: Distribution> + Distribution>, +{ + let scheme = WhirR1CSScheme::

::new_for_r1cs(r1cs, witness.len(), 0, Vec::new(), true, HASH); + ProveInputs { + scheme, + r1cs: r1cs.clone(), + witness_commit: witness.to_vec(), + witness_prove: witness.to_vec(), + } +} + +/// Time the single-commit proving core — `commit` + `prove_noir` — from +/// pre-built inputs, returning the elapsed time alongside the proof and scheme. +pub fn time_prove_core

( + inp: ProveInputs

, + public_inputs: &PublicInputs>, +) -> Result<(Duration, WhirR1CSProof, WhirR1CSScheme

)> +where + P: FieldHash, + Standard: Distribution> + Distribution>, +{ + let ProveInputs { + scheme, + r1cs, + witness_commit, + witness_prove, + } = inp; + + // Transcript binding derives from the scheme; built before the timer starts. + let instance = public_inputs.hash_bytes::

(HASH); + let ds = scheme.create_domain_separator().instance(&instance); + let num_witnesses = r1cs.num_witnesses(); + let num_constraints = r1cs.num_constraints(); + + let start = Instant::now(); + let mut merlin = ProverState::new(&ds, P::transcript_sponge(HASH)); + let commitment = scheme.commit( + &mut merlin, + num_witnesses, + num_constraints, + witness_commit, + true, + )?; + let proof = scheme.prove_noir(merlin, r1cs, vec![commitment], witness_prove, public_inputs)?; + Ok((start.elapsed(), proof, scheme)) +} diff --git a/tooling/provekit-fixtures/tests/profile.rs b/tooling/provekit-fixtures/tests/profile.rs new file mode 100644 index 000000000..70f333452 --- /dev/null +++ b/tooling/provekit-fixtures/tests/profile.rs @@ -0,0 +1,226 @@ +//! `#[ignore]` prover profiling harness. +//! +//! Two probes (run with `--release`): +//! +//! ```text +//! # Size sweep: BF vs EF vs bn254 at 2^14..2^20 +//! cargo test -p provekit-fixtures --release --test profile size_sweep -- --ignored --nocapture +//! +//! # Flamegraph + stage breakdown for one 2^20 goldilocks-BF prove +//! cargo test -p provekit-fixtures --release --test profile flamegraph_2pow20_bf -- --ignored --nocapture +//! ``` + +use { + ark_std::rand::distributions::{Distribution, Standard}, + provekit_common::{Base, Ext, FieldHash, PublicInputs}, + provekit_fixtures::{ + builders::squaring_chain, + harness::{prove_setup, time_prove_core}, + }, + provekit_verifier::WhirR1CSVerifier, + std::{ + collections::HashMap, + path::PathBuf, + time::{Duration, Instant}, + }, +}; + +/// Log2 witness sizes to sweep. +const SIZES: [u32; 4] = [14, 16, 18, 20]; + +/// Build a `2^log_size`-witness squaring chain, prove, verify, and return +/// `(prove, verify, narg_bytes, hints_bytes)`. +fn run_one

(log_size: u32) -> (Duration, Duration, usize, usize) +where + P: FieldHash, + Standard: Distribution> + Distribution>, +{ + let depth = (1usize << log_size) - 2; // num_witnesses = depth + 2 = 2^log_size + let (r1cs, w) = squaring_chain::>(2, depth); + let public_inputs = PublicInputs::from_vec(vec![w[1]]); + + // Only commit + prove_noir are timed (setup excluded above). + let inp = prove_setup::

(&r1cs, &w); + let (prove_t, proof, scheme) = time_prove_core::

(inp, &public_inputs).expect("prove"); + + let t = Instant::now(); + scheme + .verify(&proof, &public_inputs, &r1cs) + .expect("verify"); + let verify_t = t.elapsed(); + + ( + prove_t, + verify_t, + proof.narg_string.len(), + proof.hints.len(), + ) +} + +/// Run one (field, size) cell and print its row immediately (so partial results +/// survive a later OOM). Wrapped in `catch_unwind` so a panic in one cell does +/// not abort the rest of the sweep. +fn row

(label: &str, log_size: u32) +where + P: FieldHash, + Standard: Distribution> + Distribution>, +{ + let res = std::panic::catch_unwind(|| run_one::

(log_size)); + match res { + Ok((p, v, narg, hints)) => println!( + "{:>5} | {:<13} | {:>10.3?} | {:>9.3?} | {:>8} | {:>9}", + format!("2^{log_size}"), + label, + p, + v, + narg, + hints + ), + Err(_) => println!( + "{:>5} | {:<13} | {:>10} | {:>9} | {:>8} | {:>9}", + format!("2^{log_size}"), + label, + "FAILED", + "-", + "-", + "-" + ), + } +} + +#[test] +#[ignore = "profiling probe; run --release -- --ignored --nocapture"] +fn size_sweep() { + provekit_backend_goldilocks::register(); + provekit_backend_bn254::register(); + + println!("\n size | field | prove | verify | narg | hints"); + println!("------+---------------+------------+-----------+----------+----------"); + for &s in &SIZES { + row::("goldilocks-BF", s); + } + for &s in &SIZES { + row::("goldilocks-EF", s); + } + for &s in &SIZES { + row::("bn254", s); + } + println!(); +} + +/// Workspace `target/profile/` directory (created if absent). Lives under the +/// gitignored build dir so the artifacts are not committed and the probe is +/// portable across machines. +fn profile_dir() -> PathBuf { + let dir = PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .join("..") + .join("..") + .join("target") + .join("profile"); + std::fs::create_dir_all(&dir).expect("create profile dir"); + dir +} + +#[test] +#[ignore = "profiling probe; run --release -- --ignored --nocapture"] +fn flamegraph_2pow20_bf() { + use { + tracing_flame::FlameLayer, + tracing_subscriber::{prelude::*, registry::Registry}, + }; + + provekit_backend_goldilocks::register(); + + let dir = profile_dir(); + let folded_path = dir.join("goldilocks_bf_2pow20.folded"); + let svg_path = dir.join("goldilocks_bf_2pow20.svg"); + + // Fixture-gen + prove setup happen BEFORE the flame layer installs, so only + // the prove core (commit + prove_noir) is captured. + let log_size = 20u32; + let depth = (1usize << log_size) - 2; + let (r1cs, w) = squaring_chain::>(2, depth); + let public_inputs = PublicInputs::from_vec(vec![w[1]]); + let inp = prove_setup::(&r1cs, &w); + + // Global subscriber required: prove() uses rayon, and only the global + // default captures spans created on rayon worker threads. + let (flame_layer, guard) = + FlameLayer::with_file(&folded_path).expect("create flame layer file"); + let subscriber = Registry::default().with(flame_layer); + tracing::subscriber::set_global_default(subscriber).expect("set global subscriber"); + + // Timed + captured: the true prove core only (setup/clones already done). + let (prove_t, proof, _scheme) = + time_prove_core::(inp, &public_inputs) + .expect("prove"); + println!( + "\n[flamegraph] 2^20 goldilocks-BF prove={prove_t:.3?} narg={} hints={}", + proof.narg_string.len(), + proof.hints.len() + ); + + // Flush folded stacks. + drop(guard); + + // Render SVG via inferno. + let folded = std::fs::read_to_string(&folded_path).expect("read folded"); + { + let svg = std::fs::File::create(&svg_path).expect("create svg"); + let mut opts = inferno::flamegraph::Options::default(); + opts.title = "provekit 2^20 goldilocks-BF prove".to_string(); + inferno::flamegraph::from_reader(&mut opts, folded.as_bytes(), svg) + .expect("render flamegraph"); + } + + println!("[flamegraph] folded: {}", folded_path.display()); + println!("[flamegraph] svg: {}", svg_path.display()); + + analyze_folded(&folded); +} + +/// Parse tracing-flame folded stacks and print (1) per-leaf self-time table and +/// (2) the top folded stacks by total weight. +fn analyze_folded(folded: &str) { + // Each line: `frame1;frame2;...;leaf weight`. In folded format the weight is + // the self-time of that exact stack, so summing by leaf gives per-span self + // time, and aggregating identical stacks gives total stack time. + let mut leaf_self: HashMap<&str, u64> = HashMap::new(); + let mut stack_total: HashMap<&str, u64> = HashMap::new(); + let mut total: u64 = 0; + + for line in folded.lines() { + let Some((stack, weight_str)) = line.rsplit_once(' ') else { + continue; + }; + let Ok(weight) = weight_str.trim().parse::() else { + continue; + }; + total += weight; + *stack_total.entry(stack).or_default() += weight; + let leaf = stack.rsplit(';').next().unwrap_or(stack); + *leaf_self.entry(leaf).or_default() += weight; + } + + if total == 0 { + println!("[flamegraph] no folded samples captured"); + return; + } + let pct = |w: u64| 100.0 * (w as f64) / (total as f64); + + let mut leaves: Vec<(&str, u64)> = leaf_self.into_iter().collect(); + leaves.sort_by_key(|leaf| std::cmp::Reverse(leaf.1)); + println!("\n=== 2^20 goldilocks-BF: per-span SELF time (top 25) ==="); + println!(" self% | weight | span"); + for (leaf, w) in leaves.iter().take(25) { + println!(" {:>5.1}% | {:>12} | {}", pct(*w), w, leaf); + } + + let mut stacks: Vec<(&str, u64)> = stack_total.into_iter().collect(); + stacks.sort_by_key(|stack| std::cmp::Reverse(stack.1)); + println!("\n=== 2^20 goldilocks-BF: top 20 folded stacks by total time ==="); + for (stack, w) in stacks.iter().take(20) { + println!(" {:>5.1}% | {:>12} | {}", pct(*w), w, stack); + } + println!("\n[flamegraph] total folded weight = {total}"); +}