diff --git a/src/config.rs b/src/config.rs index 4dd809d..1ff68c1 100644 --- a/src/config.rs +++ b/src/config.rs @@ -109,4 +109,15 @@ pub trait StarkGenericConfig { /// [`System::new`](crate::system::System::new) rejects circuits whose /// constraint degree requires a larger quotient degree. fn max_quotient_degree(&self) -> usize; + + /// Log2 of the blowup the PCS applies when committing: a degree-`N` + /// trace is stored as a low-degree extension with `2^log_blowup · N` + /// evaluations. + /// + /// This must be the blowup `Pcs::commit` ACTUALLY applies, not a bound: + /// the prover uses it to rebuild committed LDEs directly from + /// polynomial coefficients (see `lde_from_coefficients` in the prover + /// module), and a mismatch produces commitments to the wrong + /// evaluations. + fn log_blowup(&self) -> usize; } diff --git a/src/prover.rs b/src/prover.rs index 7ba4829..78cbe3f 100644 --- a/src/prover.rs +++ b/src/prover.rs @@ -84,7 +84,7 @@ //! construction over the LDE rows. All circuits are committed together. //! //! ```text -//! FFT work: Σ_i w_i · n_i · B · log₂(n_i · B) field multiplications +//! FFT work: Σ_i w_i · (B+1) · n_i · log₂(n_i) field multiplications //! Hashing: Σ_i n_i · B Merkle leaf hashes //! ``` //! @@ -107,7 +107,7 @@ //! n_i × w2_i · D base elements) are committed identically to stage 1. //! //! ```text -//! FFT work: Σ_i w2_i · D · n_i · B · log₂(n_i · B) field multiplications +//! FFT work: Σ_i w2_i · D · (B+1) · n_i · log₂(n_i) field multiplications //! Hashing: Σ_i n_i · B Merkle leaf hashes //! ``` //! @@ -122,11 +122,17 @@ //! The quotient polynomial is split into q_i coefficient slices (each of //! degree n_i) and committed as a single q_i·D-column matrix per circuit on //! the trace domain, so the opening phase pays its per-matrix costs once per -//! circuit rather than once per slice. +//! circuit rather than once per slice. The committed LDE is built directly +//! from the slice coefficients, skipping the trace-domain DFT whose output +//! `Pcs::commit` would immediately invert: one forward DFT recovers the +//! coefficients, a fused parallel gather slices them with the LDE's coset +//! shift folded in (`shifted_quotient_slices`), and one zero-padded DFT +//! produces the committed LDE (`lde_from_shifted_coefficients`). //! //! ```text //! Constraint eval: Σ_i n_i · q_i · eval_cost(k_i) field operations -//! Quotient FFT: Σ_i q_i · D · n_i · B · log₂(n_i · B) field multiplications +//! Quotient iFFT: Σ_i D · q_i · n_i · log₂(q_i · n_i) field multiplications +//! Quotient LDE: Σ_i q_i · D · B · n_i · log₂(B · n_i) field multiplications //! Hashing: Σ_i q_i · n_i · B Merkle leaf hashes //! ``` //! @@ -152,7 +158,7 @@ //! The total per-proof cost is approximately: //! //! ```text -//! C_prove ≈ Σ_i n_i · B · log₂(n_i · B) · W_i (FFT — all commit rounds) +//! C_prove ≈ Σ_i (B+1) · n_i · log₂(n_i) · W_i (FFT — all commit rounds) //! + Σ_i n_i · q_i · eval_cost(k_i) (constraint evaluation) //! + Σ_i n_i · B · W_i (barycentric interpolation) //! + Q · R · log₂ H (FRI query phase) @@ -187,9 +193,9 @@ use p3_dft::{Radix2DitParallel, TwoAdicSubgroupDft}; use p3_field::{ Algebra, BasedVectorSpace, Field, PackedValue, PrimeCharacteristicRing, TwoAdicField, }; -use p3_matrix::{Matrix, dense::RowMajorMatrix}; +use p3_matrix::{Matrix, bitrev::BitReversibleMatrix, dense::RowMajorMatrix}; use p3_maybe_rayon::prelude::*; -use p3_util::log2_strict_usize; +use p3_util::{log2_strict_usize, reverse_bits_len}; use serde::{Deserialize, Serialize}; /// Polynomial commitments included in the proof. @@ -251,8 +257,9 @@ impl Proof { impl System where SC: StarkGenericConfig, - // Two-adicity is needed to re-base the quotient's coefficient slices - // onto the trace domain; every FRI-based config is two-adic anyway. + // Two-adicity is needed to slice the quotient into coefficient slices + // and rebuild their committed LDE from those coefficients; every + // FRI-based config is two-adic anyway. Val: TwoAdicField + Ord, { /// Generates a STARK proof for the system with a single claim. @@ -324,7 +331,8 @@ where } // Cost: "Stage 1 commit" — coset LDE (FFT) of each trace from n_i to - // n_i·B rows, then Merkle tree. FFT work: Σ w_i · n_i · B · log₂(n_i·B). + // n_i·B rows (an iDFT plus B coset DFTs per column), then Merkle + // tree. FFT work: Σ w_i · (B+1) · n_i · log₂(n_i). let _g = tracing::info_span!("stark/stage1_commit").entered(); let mut log_degrees = vec![]; let evaluations = witness @@ -401,7 +409,7 @@ where drop(_g); // Cost: "Stage 2 commit" — LDE + Merkle for flattened extension traces. - // FFT work: Σ w2_i · D · n_i · B · log₂(n_i·B). + // FFT work: Σ w2_i · D · (B+1) · n_i · log₂(n_i). let _g = tracing::info_span!("stark/stage2_commit").entered(); let evaluations = stage_2_traces.into_iter().map(|trace| { let degree = trace.height(); @@ -422,14 +430,16 @@ where // Constraint challenge. let constraint_challenge: SC::Challenge = challenger.sample_algebra_element(); - // Cost: "Quotient computation and commit" — constraint evaluation on the - // quotient domain (Σ n_i·q_i·eval_cost(k_i)) plus LDE + Merkle of the - // quotient sub-polynomials (Σ q_i·D·n_i·B·log₂(n_i·B)). + // Cost: "Quotient computation and commit" — constraint evaluation on + // the quotient domain (Σ n_i·q_i·eval_cost(k_i)), the forward DFT of + // the flattened quotient (Σ D·q_i·n_i·log₂(q_i·n_i)), then LDE + + // Merkle of the sub-polynomials (Σ q_i·D·B·n_i·log₂(B·n_i)). let _g = tracing::info_span!("stark/quotient").entered(); debug_assert_eq!(intermediate_accumulators.len(), active_indices.len()); debug_assert_eq!(log_degrees.len(), active_indices.len()); let dft = Radix2DitParallel::>::default(); - let quotient_evaluations = active_indices + let log_blowup = self.config.log_blowup(); + let quotient_ldes: Vec<_> = active_indices .iter() .zip(log_degrees.iter()) .zip(intermediate_accumulators.iter()) @@ -490,31 +500,30 @@ where // commit all slices as ONE `q·D`-column matrix on the trace // domain — instead of one matrix per slice on the split // cosets — so the opening phase pays its per-matrix costs - // once per circuit rather than once per slice. - let coefficients = - dft.coset_idft_batch(quotient_flat, quotient_domain.first_point()); - let ext_degree = >>::DIMENSION; - let n = 1 << log_degree; - let width = quotient_degree * ext_degree; - let mut sliced = Vec::with_capacity(n * width); - for row in 0..n { - for chunk in 0..quotient_degree { - sliced.extend_from_slice( - &coefficients.values[(chunk * n + row) * ext_degree - ..(chunk * n + row + 1) * ext_degree], - ); - } - } - let quotient_chunks_evals = dft - .coset_dft_batch( - RowMajorMatrix::new(sliced, width), - trace_domain.first_point(), - ) - .to_row_major_matrix(); + // once per circuit rather than once per slice. The committed + // LDE is built straight from these coefficients: evaluating + // them onto the trace domain only for `Pcs::commit` to + // inverse-DFT that evaluation right back would waste two + // size-n transforms per column. The slicing itself is fused + // with the iDFT's scaling passes into one parallel gather + // (see `shifted_quotient_slices`). acc = *next_acc; - (trace_domain, quotient_chunks_evals) - }); - let (quotient_commit, quotient_data) = pcs.commit(quotient_evaluations); + let sliced = shifted_quotient_slices( + &dft, + quotient_flat, + quotient_domain.first_point(), + quotient_degree, + ); + lde_from_shifted_coefficients(&dft, sliced, log_blowup) + }) + .collect(); + // `commit_ldes` skips the randomization a hiding PCS applies inside + // `commit`; this prover targets non-hiding configurations only. + assert!( + !>::ZK, + "committing the quotient from coefficients bypasses hiding-PCS randomization" + ); + let (quotient_commit, quotient_data) = pcs.commit_ldes(quotient_ldes); challenger.observe(quotient_commit.clone()); drop(_g); @@ -594,6 +603,153 @@ where } } +/// From the quotient's evaluations on its disjoint domain — `q·n` rows of +/// `D` base columns on the coset `shift·H` with `|H| = q·n` — produce the +/// `n`-row, `q·D`-column matrix of slice coefficients with the committed +/// LDE's `GENERATOR` shift already folded in: the input +/// [`lde_from_shifted_coefficients`] expects. +/// +/// Semantically this is three steps: coset iDFT to coefficients, slicing +/// `Q(X) = Σₖ X^{k·n}·cₖ(X)` into rows `[c₀ | … | c_{q−1}]`, and +/// pre-scaling row `r` by `GENERATOR^r`. Executed literally (the library +/// entry points) those steps cost a bit-reversal materialization, a serial +/// row-swap pass, two serial full-matrix scaling passes, and a serial +/// gather. But the composition collapses: with `N = q·n` and `S` the raw +/// bit-reversed storage of the forward DFT, +/// +/// ```text +/// idft(f)ⱼ = N⁻¹ · dft(f)_{(N−j) mod N} = N⁻¹ · S[rev((N−j) mod N)] +/// ``` +/// +/// and the coset-unscale factor `shift^{−j}` at `j = k·n + r` splits into +/// `shift^{−k·n} · shift^{−r}`, whose row-dependent part cancels the LDE +/// pre-scale `GENERATOR^r` exactly, because the disjoint quotient domain's +/// shift IS the generator (asserted below; `create_disjoint_domain` on a +/// natural trace domain guarantees it). What survives is a single parallel +/// gather off the DFT storage with ONE constant weight per slice: +/// `wₖ = N⁻¹ · GENERATOR^{−k·n}`. +fn shifted_quotient_slices( + dft: &Radix2DitParallel, + quotient_evals: RowMajorMatrix, + domain_shift: F, + quotient_degree: usize, +) -> RowMajorMatrix { + assert_eq!( + domain_shift, + F::GENERATOR, + "quotient domain shift must equal the LDE shift for the scalings to cancel" + ); + let ext_degree = quotient_evals.width(); + let big_height = quotient_evals.height(); + let log_big_height = log2_strict_usize(big_height); + debug_assert_eq!(big_height % quotient_degree, 0); + let n = big_height / quotient_degree; + let width = quotient_degree * ext_degree; + // Raw storage of the forward DFT: natural index `k` lives at row + // `rev(k)`, and the unwrap out of the bit-reversed view is copy-free. + let storage = dft.dft_batch(quotient_evals).bit_reverse_rows(); + let n_inv = F::ONE.div_2exp_u64(log_big_height as u64); + let weight_step = F::GENERATOR.exp_u64(n as u64).inverse(); + let weights: Vec = weight_step + .powers() + .take(quotient_degree) + .map(|w| w * n_inv) + .collect(); + let mut values = F::zero_vec(n * width); + values + .par_chunks_mut(width) + .enumerate() + .for_each(|(row, out)| { + for (chunk, weight) in weights.iter().enumerate() { + let j = chunk * n + row; + let src = reverse_bits_len( + big_height.wrapping_sub(j) & (big_height - 1), + log_big_height, + ); + let src = &storage.values[src * ext_degree..(src + 1) * ext_degree]; + for (out, src) in out[chunk * ext_degree..(chunk + 1) * ext_degree] + .iter_mut() + .zip(src) + { + *out = *src * *weight; + } + } + }); + RowMajorMatrix::new(values, width) +} + +/// Low-degree extension of column polynomials given by their COEFFICIENTS +/// with the `GENERATOR` coset shift already folded in (row `j` +/// pre-multiplied by `GENERATOR^j`, which [`shifted_quotient_slices`] +/// produces for free), in the exact layout `Pcs::commit` stores for +/// evaluations on the natural domain: `2^log_blowup` row blocks, where +/// block `b` holds the evaluations on the coset `GENERATOR · w^rev(b) · H` +/// in bit-reversed row order (`H` is the size-`n` subgroup, `w` generates +/// the size-`2^log_blowup · n` subgroup, and `rev` reverses `log_blowup` +/// bits). Globally that is the bit-reversal of the natural order of the +/// whole blown-up coset — i.e. +/// `coset_lde_batch(evals, log_blowup, GENERATOR).bit_reverse_rows()`, +/// which is what `TwoAdicFriPcs::commit` computes. +/// +/// Committing the result via `Pcs::commit_ldes` is therefore bit-identical +/// to `Pcs::commit` on the columns' trace-domain evaluations — field +/// arithmetic is exact, so equal polynomials give equal evaluations no +/// matter which transform produced them — while skipping both that +/// evaluation DFT and the inverse DFT `commit` would open with. +/// +/// The whole extension is ONE size-`2^log_blowup · n` transform: zero-pad +/// the shifted coefficients to the LDE height (which leaves the column +/// polynomials unchanged) and DFT. That spends `log_blowup` more butterfly +/// layers than `2^log_blowup` separate size-`n` coset DFTs would, but one +/// batched transform is what the memory traffic wants: no per-coset matrix +/// clones, no per-coset serial shift-scaling passes inside +/// `coset_dft_batch`, no reassembly copies, and `Radix2DitParallel`'s +/// native output order is already the bit-reversed storage order, so the +/// final unwrap is copy-free. +fn lde_from_shifted_coefficients( + dft: &Radix2DitParallel, + mut coefficients: RowMajorMatrix, + log_blowup: usize, +) -> RowMajorMatrix { + let height = coefficients.height(); + coefficients.pad_to_height(height << log_blowup, F::ZERO); + dft.dft_batch(coefficients).bit_reverse_rows() +} + +/// Reference form of [`lde_from_shifted_coefficients`] taking PLAIN +/// coefficients: folds the `GENERATOR` shift in explicitly. Only the +/// pinning tests need it; the prover gets the shift for free inside +/// [`shifted_quotient_slices`]. +#[cfg(test)] +fn lde_from_coefficients( + dft: &Radix2DitParallel, + mut coefficients: RowMajorMatrix, + log_blowup: usize, +) -> RowMajorMatrix { + scale_rows_by_powers(&mut coefficients, F::GENERATOR); + lde_from_shifted_coefficients(dft, coefficients, log_blowup) +} + +/// Multiplies row `j` of `mat` by `base^j`, in parallel: each chunk of rows +/// pays one exponentiation and steps serially from there. +#[cfg(test)] +fn scale_rows_by_powers(mat: &mut RowMajorMatrix, base: F) { + const ROWS_PER_CHUNK: usize = 512; + let width = mat.width(); + mat.values + .par_chunks_mut(ROWS_PER_CHUNK * width) + .enumerate() + .for_each(|(chunk, rows)| { + let mut weight = base.exp_u64((chunk * ROWS_PER_CHUNK) as u64); + for row in rows.chunks_mut(width) { + for value in row { + *value *= weight; + } + weight *= base; + } + }); +} + /// Evaluates the folded constraints on the quotient domain and divides by /// the vanishing polynomial, producing the quotient values. #[allow(clippy::too_many_arguments)] @@ -804,3 +960,83 @@ where }) }) } + +#[cfg(test)] +mod tests { + use super::*; + use crate::types::Val; + use rand::{RngExt, SeedableRng, rngs::SmallRng}; + + /// `lde_from_coefficients` must reproduce, value for value, the matrix + /// `TwoAdicFriPcs::commit` stores for the same polynomials given as + /// trace-domain evaluations: `coset_lde_batch` with the generator shift + /// (the natural domain's shift is one), then a bit-reversal. This pins + /// the exact substitution the quotient commit relies on. + #[test] + fn lde_from_coefficients_matches_commit_transform() { + let mut rng = SmallRng::seed_from_u64(0); + let dft = Radix2DitParallel::::default(); + for log_height in [0usize, 1, 2, 5, 8] { + for log_blowup in [1usize, 2, 3] { + for width in [1usize, 2, 7] { + let height = 1 << log_height; + let coefficients = RowMajorMatrix::new( + (0..height * width).map(|_| rng.random()).collect(), + width, + ); + let evaluations = dft + .coset_dft_batch(coefficients.clone(), Val::ONE) + .to_row_major_matrix(); + let expected = dft + .coset_lde_batch(evaluations, log_blowup, Val::GENERATOR) + .bit_reverse_rows() + .to_row_major_matrix(); + let got = lde_from_coefficients(&dft, coefficients, log_blowup); + assert_eq!(got, expected, "h=2^{log_height} B=2^{log_blowup} w={width}"); + } + } + } + } + + /// `shifted_quotient_slices` must reproduce, value for value, the naive + /// composition it replaces: coset iDFT off the quotient domain, slicing + /// the coefficients into `q` chunks per row, and folding the committed + /// LDE's `GENERATOR` shift into the rows. This pins the scaling + /// cancellation the fused gather relies on. + #[test] + fn shifted_quotient_slices_matches_naive_composition() { + let mut rng = SmallRng::seed_from_u64(1); + let dft = Radix2DitParallel::::default(); + for log_n in [0usize, 1, 2, 5, 7] { + for quotient_degree in [1usize, 2, 4] { + for ext_degree in [1usize, 2] { + let n = 1 << log_n; + let big_height = n * quotient_degree; + let evals = RowMajorMatrix::new( + (0..big_height * ext_degree).map(|_| rng.random()).collect(), + ext_degree, + ); + // Naive path: coset iDFT, slice, fold the LDE shift in. + let coefficients = dft.coset_idft_batch(evals.clone(), Val::GENERATOR); + let width = quotient_degree * ext_degree; + let mut sliced = Vec::with_capacity(n * width); + for row in 0..n { + for chunk in 0..quotient_degree { + sliced.extend_from_slice( + &coefficients.values[(chunk * n + row) * ext_degree + ..(chunk * n + row + 1) * ext_degree], + ); + } + } + let mut expected = RowMajorMatrix::new(sliced, width); + scale_rows_by_powers(&mut expected, Val::GENERATOR); + let got = shifted_quotient_slices(&dft, evals, Val::GENERATOR, quotient_degree); + assert_eq!( + got, expected, + "n=2^{log_n} q={quotient_degree} D={ext_degree}" + ); + } + } + } + } +} diff --git a/src/test_circuits/baby_bear_config.rs b/src/test_circuits/baby_bear_config.rs index ccf019e..ca35637 100644 --- a/src/test_circuits/baby_bear_config.rs +++ b/src/test_circuits/baby_bear_config.rs @@ -46,6 +46,7 @@ struct BabyBearPoseidon2Config { challenger_seed: Vec, max_log_degree: usize, max_quotient_degree: usize, + log_blowup: usize, } impl BabyBearPoseidon2Config { @@ -90,6 +91,7 @@ impl BabyBearPoseidon2Config { challenger_seed, max_log_degree, max_quotient_degree, + log_blowup: commitment_parameters.log_blowup, } } } @@ -118,6 +120,10 @@ impl StarkGenericConfig for BabyBearPoseidon2Config { fn max_quotient_degree(&self) -> usize { self.max_quotient_degree } + + fn log_blowup(&self) -> usize { + self.log_blowup + } } /// A minimal AIR: enforces `a * b == c` per row, with a self-canceling diff --git a/src/types.rs b/src/types.rs index 236b9ca..ee42b65 100644 --- a/src/types.rs +++ b/src/types.rs @@ -103,6 +103,8 @@ pub struct GoldilocksBlake3Config { /// Largest quotient degree the PCS can serve trace evaluations for /// (the FRI blowup factor). max_quotient_degree: usize, + /// Log2 of the blowup the PCS applies when committing. + log_blowup: usize, } impl GoldilocksBlake3Config { @@ -133,6 +135,7 @@ impl GoldilocksBlake3Config { challenger_seed, max_log_degree, max_quotient_degree, + log_blowup: commitment_parameters.log_blowup, } } } @@ -157,6 +160,10 @@ impl StarkGenericConfig for GoldilocksBlake3Config { fn max_quotient_degree(&self) -> usize { self.max_quotient_degree } + + fn log_blowup(&self) -> usize { + self.log_blowup + } } /// Parameters of the polynomial commitment: Reed-Solomon rate and Merkle