diff --git a/Cargo.lock b/Cargo.lock index 082feb8d..ddeb67c4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -251,6 +251,7 @@ dependencies = [ "cfg-if", "constant_time_eq", "cpufeatures 0.3.0", + "rayon-core", ] [[package]] @@ -1119,6 +1120,7 @@ dependencies = [ "field", "flock-core", "num-traits", + "post_gkr", "proptest", "tracing", "transcript", @@ -1142,6 +1144,21 @@ dependencies = [ "rayon", ] +[[package]] +name = "post_gkr" +version = "0.1.0" +dependencies = [ + "common", + "field", + "num-traits", + "poly", + "rand_core 0.10.1", + "rand_pcg 0.10.2", + "rayon", + "tracing", + "transcript", +] + [[package]] name = "ppv-lite86" version = "0.2.21" diff --git a/Cargo.toml b/Cargo.toml index 7ff204d4..30097fe5 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -20,6 +20,7 @@ spartan = { path = "crates/spartan" } transcript = { path = "crates/transcript" } verifier = { path = "crates/verifier" } gkr = { path = "crates/gkr"} +post_gkr = { path = "crates/post_gkr" } aes = "0.9.2" anyhow = "1.0.93" diff --git a/crates/common/src/table.rs b/crates/common/src/table.rs index c2e63051..61196980 100644 --- a/crates/common/src/table.rs +++ b/crates/common/src/table.rs @@ -346,7 +346,7 @@ mod tests { for (index, element) in packed.iter().enumerate() { let column = index / groups_per_column; let i_hi = index % groups_per_column; - let bits = u128::from(element.lo) | (u128::from(element.hi) << 64); + let bits = element.to_u128(); for v in 0..PACKED_BITS { assert_eq!( (bits >> v) & 1 == 1, @@ -401,7 +401,7 @@ mod tests { let table = BitTable::new(shape, &packed).unwrap(); for (index, element) in table.column(2).iter().enumerate() { - let bits = u128::from(element.lo) | (u128::from(element.hi) << 64); + let bits = element.to_u128(); for offset in 0..PACKED_BITS { let row = index * PACKED_BITS + offset; assert_eq!(table.bit(2, row), (bits >> offset) & 1 == 1, "row {row}"); diff --git a/crates/common/src/virtual_map.rs b/crates/common/src/virtual_map.rs index 0f1dc9fd..a59e2d03 100644 --- a/crates/common/src/virtual_map.rs +++ b/crates/common/src/virtual_map.rs @@ -13,6 +13,8 @@ use field::{F128, Fq}; use num_traits::ConstZero; +#[cfg(feature = "parallel")] +use rayon::prelude::*; use crate::{ BitZParams, ClaimError, LinearClaim, OpeningQuery, Shape, VirtualParams, VirtualParamsError, @@ -47,6 +49,8 @@ pub enum VirtualMapError { ClaimWeightCountMismatch, /// The input omits coordinates of `h`, or the output does not have one weight per bit of `f`. WeightCountMismatch, + /// The map's weights on `f` outnumber the committed bits. + CommittedShapeTooSmall, } /// What the protocol asks of `M`. @@ -153,8 +157,9 @@ impl<'a, const Q: u128, M: VirtualMap> VirtualStatement<'a, Q, M> { /// /// Coefficients use `column * row_count + row` order. The map drops weights on /// virtual padding. The opening subtracts the constant-column weight from the - /// target and zero-pads the remaining weights to the commitment size. A - /// single-column `InnerProduct` holds the dense weights with column weight one. + /// target and zero-pads the remaining weights to the commitment size, refusing + /// more weights than it has. A single-column `InnerProduct` holds the dense + /// weights with column weight one. pub fn transpose_query(&self, query: OpeningQuery) -> Result { let shape = self.params.claim().shape(); let (weights, target) = match query { @@ -170,12 +175,10 @@ impl<'a, const Q: u128, M: VirtualMap> VirtualStatement<'a, Q, M> { { return Err(VirtualMapError::ClaimWeightCountMismatch); } - let weights = claim - .column_weights() - .iter() - .flat_map(|column| claim.row_weights().iter().map(move |row| *row * *column)) - .collect(); - (weights, claim.target()) + ( + flatten(claim.row_weights(), claim.column_weights()), + claim.target(), + ) } }; let transposed = self.map.transpose(&weights)?; @@ -185,7 +188,11 @@ impl<'a, const Q: u128, M: VirtualMap> VirtualStatement<'a, Q, M> { } let target = transposed.adjusted_target(target); let mut weights = transposed.into_weights(); - weights.resize(1 << self.params.committed_shape().log_bits(), F128::ZERO); + let committed_bits = 1 << self.params.committed_shape().log_bits(); + if weights.len() > committed_bits { + return Err(VirtualMapError::CommittedShapeTooSmall); + } + weights.resize(committed_bits, F128::ZERO); let shape = Shape::new(self.params.committed_shape().log_bits(), 0) .expect("a valid committed bit count permits a single-column shape"); let claim = LinearClaim::from_shape(&shape, weights, vec![F128::from(1u64)], target) @@ -194,6 +201,28 @@ impl<'a, const Q: u128, M: VirtualMap> VirtualStatement<'a, Q, M> { } } +/// `row_weights (x) column_weights` written out per bit of `h`, one column +/// after another. +fn flatten(row_weights: &[F128], column_weights: &[F128]) -> Vec { + let mut weights = vec![F128::ZERO; row_weights.len() * column_weights.len()]; + let process_column = |(slot, &scale): (&mut [F128], &F128)| { + for (weight, &row) in slot.iter_mut().zip(row_weights) { + *weight = scale * row; + } + }; + #[cfg(feature = "parallel")] + weights + .par_chunks_mut(row_weights.len()) + .zip(column_weights) + .for_each(process_column); + #[cfg(not(feature = "parallel"))] + weights + .chunks_mut(row_weights.len()) + .zip(column_weights) + .for_each(process_column); + weights +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/field/benches/binius64.rs b/crates/field/benches/binius64.rs index b67ff43e..6350bc68 100644 --- a/crates/field/benches/binius64.rs +++ b/crates/field/benches/binius64.rs @@ -63,7 +63,7 @@ fn operands(count: usize, seed: u64) -> Vec { } fn to_binius(a: F128) -> Ghash { - Ghash::from((a.hi as u128) << 64 | a.lo as u128) + Ghash::from(a.to_u128()) } fn from_binius(a: Ghash) -> F128 { @@ -117,7 +117,7 @@ fn check_agreement(xs: &[F128], ys: &[F128]) { ); // The timed exponents are the operands read as integers, so raising the // generator to `x` here is the same call the `pow` row makes. - let e = (x.hi as u128) << 64 | x.lo as u128; + let e = x.to_u128(); assert_eq!( F128::GENERATOR.pow(e), from_binius(to_binius(F128::GENERATOR).pow([e as u64, (e >> 64) as u64])), @@ -342,10 +342,7 @@ fn main() { // Full-width exponents, square-and-multiply on both sides. This crate's // fixed-base comb has no counterpart in binius64, so it is left out. - let exps: Vec = xs - .iter() - .map(|x| (x.hi as u128) << 64 | x.lo as u128) - .collect(); + let exps: Vec = xs.iter().map(|x| x.to_u128()).collect(); let bgen = to_binius(F128::GENERATOR); c.run( "pow/square-and-multiply", diff --git a/crates/field/benches/gf128.rs b/crates/field/benches/gf128.rs index 46a3b828..e0051247 100644 --- a/crates/field/benches/gf128.rs +++ b/crates/field/benches/gf128.rs @@ -113,10 +113,7 @@ fn inverse(bencher: Bencher) { /// Full-width exponents: what the fold values are. #[divan::bench] fn pow_square_and_multiply(bencher: Bencher) { - let exps: Vec = xs() - .iter() - .map(|x| (x.hi as u128) << 64 | x.lo as u128) - .collect(); + let exps: Vec = xs().iter().map(|x| x.to_u128()).collect(); bencher.counter(ItemsCount::new(N)).bench_local(|| { for &e in &exps { black_box(F128::GENERATOR.pow(e)); @@ -127,10 +124,7 @@ fn pow_square_and_multiply(bencher: Bencher) { #[divan::bench] fn pow_comb_w8(bencher: Bencher) { let comb = FixedBasePow::new(F128::GENERATOR, 8); - let exps: Vec = xs() - .iter() - .map(|x| (x.hi as u128) << 64 | x.lo as u128) - .collect(); + let exps: Vec = xs().iter().map(|x| x.to_u128()).collect(); bencher.counter(ItemsCount::new(N)).bench_local(|| { for &e in &exps { black_box(comb.pow(e)); diff --git a/crates/field/src/gf128.rs b/crates/field/src/gf128.rs index 53c3bce2..c5e44e08 100644 --- a/crates/field/src/gf128.rs +++ b/crates/field/src/gf128.rs @@ -118,6 +118,11 @@ impl F128 { const fn words(self) -> [u64; 2] { [self.lo, self.hi] } + + /// `lo || hi` as one `u128`. + pub const fn to_u128(&self) -> u128 { + self.lo as u128 | ((self.hi as u128) << 64) + } } /// The bit pattern as 32 hex digits, `hi` first — how the polynomial reads on diff --git a/crates/pcs/Cargo.toml b/crates/pcs/Cargo.toml index ad2b2f8f..b627754a 100644 --- a/crates/pcs/Cargo.toml +++ b/crates/pcs/Cargo.toml @@ -7,10 +7,11 @@ license.workspace = true [dependencies] bincode = { workspace = true } -blake3 = { workspace = true } +blake3 = { workspace = true, features = ["rayon"] } common = { workspace = true } field = { workspace = true, features = ["spongefish"] } flock-core = { workspace = true } +post_gkr = { workspace = true } transcript = { workspace = true } tracing = { workspace = true } num-traits = { workspace = true } diff --git a/crates/pcs/src/lib.rs b/crates/pcs/src/lib.rs index 2c636d7c..68ab2c40 100644 --- a/crates/pcs/src/lib.rs +++ b/crates/pcs/src/lib.rs @@ -22,7 +22,7 @@ //! Ring-switching transposes `(s_v)` into `(s_u)` and samples `batching_point`. //! It sets `packed_target = Σ_u eq(batching_point, u) · s_u`. //! Recursive Ligerito proves `Σ_y B(y) · q_pkd(y) = packed_target` against the committed root. -//! Quadratic sumcheck reduces factored inner-product claims to MLE claims before this opening protocol. +//! The post-GKR sumcheck (`post_gkr`) reduces factored inner-product claims to MLE claims before this opening protocol. //! //! # Interface //! @@ -95,7 +95,6 @@ mod ligerito; mod mle; mod opening; mod profiles; -mod sumcheck; mod transpose; #[cfg(test)] @@ -130,7 +129,7 @@ pub enum StatementBinding { /// `q̂(r) = Σ_{b ∈ {0,1}^m} q(b) · eq(b, r) = target`, where /// `eq(b, r) = ∏_i (b_i · r_i + (1 - b_i) · (1 - r_i))`. /// [`OpeningQuery::InnerProduct`] accepts row weights, column weights, and a target over `F128`. -/// Quadratic sumcheck reduces this claim to an MLE claim before the opening protocol. +/// The post-GKR sumcheck reduces this claim to an MLE claim before the opening protocol. pub trait CommitScheme { /// The public commitment. type Commitment; @@ -147,7 +146,7 @@ pub trait CommitScheme { /// Consumes the exact packed witness and proves either opening query. /// - /// Inner-product claims first pass through quadratic sumcheck and then the MLE opening protocol. + /// Inner-product claims first pass through the post-GKR sumcheck and then the MLE opening protocol. fn prove_lin( &self, data: &Self::ProverData, @@ -159,7 +158,7 @@ pub trait CommitScheme { /// Verifies either opening query against `commitment`. /// - /// Inner-product claims first pass through quadratic sumcheck and then the MLE opening protocol. + /// Inner-product claims first pass through the post-GKR sumcheck and then the MLE opening protocol. fn verify_lin( &self, commitment: &Self::Commitment, diff --git a/crates/pcs/src/opening.rs b/crates/pcs/src/opening.rs index f2683320..1bdf5d54 100644 --- a/crates/pcs/src/opening.rs +++ b/crates/pcs/src/opening.rs @@ -7,12 +7,13 @@ use flock_core::pcs::pack::PACKING_WIDTH as CLAIM_COUNT; use transcript::{ProverState, PublicTranscript, VerifierState}; use crate::bridge::{as_flock_f128, as_flock_f128s, from_flock_f128}; -use crate::ligerito::{self, ReducedProver}; -use crate::{OpeningQuery, Pcs, ProverData, Root, StatementBinding, mle, sumcheck}; +use crate::ligerito::{self, ReducedProver, validate_prover_data}; +use crate::{OpeningQuery, Pcs, ProverData, Root, StatementBinding, mle}; const MLE_STATEMENT_LABEL: &[u8] = b"bitz/pcs/mle-opening/v1"; -const INNER_PRODUCT_STATEMENT_LABEL: &[u8] = b"bitz/pcs/bit-inner-product/v2"; -const SUMCHECK_LABEL: &[u8] = b"bitz/pcs/inner-product-sumcheck/v1"; +const INNER_PRODUCT_STATEMENT_LABEL: &[u8] = b"bitz/pcs/bit-inner-product/v3"; +const INNER_PRODUCT_DIGEST_CONTEXT: &str = "bitz/pcs/bit-inner-product-weights/v1"; +const SUMCHECK_LABEL: &[u8] = b"bitz/pcs/inner-product-sumcheck/v2"; const MLE_CLAIMS_LABEL: &[u8] = b"bitz/pcs/mle-claims/v1"; const CHALLENGES_LABEL: &[u8] = b"bitz/pcs/ring-switch-challenges/v1"; @@ -82,6 +83,24 @@ impl From for VerifyError { } } +impl From for ProveError { + fn from(error: post_gkr::ProveError) -> Self { + match error { + post_gkr::ProveError::WitnessLengthMismatch => Self::PackedWitnessLengthMismatch, + post_gkr::ProveError::ClaimDoesNotHold => Self::InvalidClaim, + } + } +} + +impl From for VerifyError { + fn from(error: post_gkr::VerifyError) -> Self { + match error { + post_gkr::VerifyError::MalformedProof => Self::MalformedProof, + post_gkr::VerifyError::EvaluationMismatch => Self::VerificationFailed, + } + } +} + #[tracing::instrument(name = "Prove PCS opening", skip_all)] pub(crate) fn prove( pcs: &Pcs, @@ -102,22 +121,23 @@ pub(crate) fn prove( } OpeningQuery::InnerProduct { claim } => { validate_inner_product_claim(pcs, claim)?; - let prover = ReducedProver::new(pcs, data, packed_witness)?; + validate_prover_data(pcs, data)?; if statement_binding == StatementBinding::Bind { bind_inner_product_statement(pcs, &data.commitment().root, claim, transcript); } transcript.public_message(SUMCHECK_LABEL); - let reduced = sumcheck::prove(claim, prover.witness(), transcript)?; - let ring_switch = mle::RingSwitch::new(&reduced.point, pcs.params().m)?; - // AlreadyBound covers the original claim, before the reduction produces this MLE claim. - bind_mle_statement( + let reduced = post_gkr::prove(claim, &packed_witness, transcript)?; + // The evaluation claim the sumcheck leaves is opened like any + // other, and bound whatever the caller's mode: `AlreadyBound` + // covers the original claim only. + prove( pcs, - &data.commitment().root, - &reduced.point, - reduced.target, + data, + packed_witness, + &reduced, + StatementBinding::Bind, transcript, - ); - prove_mle(prover, ring_switch, reduced.target, transcript) + ) } } } @@ -144,16 +164,14 @@ pub(crate) fn verify( bind_inner_product_statement(pcs, &commitment.0, claim, transcript); } transcript.public_message(SUMCHECK_LABEL); - let reduced = sumcheck::verify(claim, transcript)?; - let ring_switch = mle::RingSwitch::new(&reduced.point, pcs.params().m)?; - bind_mle_statement( + let reduced = post_gkr::verify(claim, transcript)?; + verify( pcs, - &commitment.0, - &reduced.point, - reduced.target, + commitment, + &reduced, + StatementBinding::Bind, transcript, - ); - verify_mle(pcs, commitment, ring_switch, reduced.target, transcript) + ) } } } @@ -260,7 +278,8 @@ fn bind_mle_statement( transcript.public_message(&target); } -/// Binds both tensor factors before the first sumcheck challenge. +/// Binds both tensor factors before the first sumcheck challenge: their +/// lengths, a digest of their weights, and the target. fn bind_inner_product_statement( pcs: &Pcs, root: &[u8; 32], @@ -270,7 +289,28 @@ fn bind_inner_product_statement( transcript.public_message(INNER_PRODUCT_STATEMENT_LABEL); transcript.public_message(root); transcript.public_message(pcs); - transcript.public_message(claim); + transcript.public_message(&(claim.row_weights().len() as u64)); + transcript.public_message(&(claim.column_weights().len() as u64)); + + // Weights per digest update: `2^16` elements, one megabyte. + const INNER_PRODUCT_DIGEST_CHUNK: usize = 1 << 16; + + // Hashing this way is much faster that going through `public_message` directly. + // Does not depend on the chunking. + let mut hasher = blake3::Hasher::new_derive_key(INNER_PRODUCT_DIGEST_CONTEXT); + let mut buffer = Vec::with_capacity(INNER_PRODUCT_DIGEST_CHUNK * 16); + for factor in [claim.row_weights(), claim.column_weights()] { + for chunk in factor.chunks(INNER_PRODUCT_DIGEST_CHUNK) { + buffer.clear(); + for weight in chunk { + buffer.extend_from_slice(&weight.to_bytes()); + } + hasher.update_rayon(&buffer); + } + } + transcript.public_message(hasher.finalize().as_bytes()); + + transcript.public_message(&claim.target()); } #[cfg(test)] diff --git a/crates/pcs/src/opening/tests.rs b/crates/pcs/src/opening/tests.rs index a704f6f8..e4c03182 100644 --- a/crates/pcs/src/opening/tests.rs +++ b/crates/pcs/src/opening/tests.rs @@ -69,14 +69,11 @@ fn inner_product_proof_composes_sumcheck_with_a_bound_mle_opening() { let mut verifier = build_verifier(SESSION, INSTANCE, &proof); bind_inner_product_statement(&fixture.pcs, &fixture.root.0, &fixture.claim, &mut verifier); verifier.public_message(SUMCHECK_LABEL); - let reduced = sumcheck::verify(&fixture.claim, &mut verifier).unwrap(); + let reduced = post_gkr::verify(&fixture.claim, &mut verifier).unwrap(); verify( &fixture.pcs, &fixture.root, - &OpeningQuery::Mle { - point: reduced.point, - target: reduced.target, - }, + &reduced, StatementBinding::Bind, &mut verifier, ) diff --git a/crates/pcs/src/sumcheck.rs b/crates/pcs/src/sumcheck.rs deleted file mode 100644 index c8f6c6d6..00000000 --- a/crates/pcs/src/sumcheck.rs +++ /dev/null @@ -1,157 +0,0 @@ -//! The input contains arbitrary row and column weights over `F128`. -//! Bit `column * row_weights.len() + row` has weight `row_weights[row] * column_weights[column]`. -//! Each packed witness element stores bits 0 through 63 in `lo`, then bits 64 through 127 in `hi`. -//! Packed element `i`, local bit `v`, supplies logical bit `128 * i + v`. -//! The target is the claimed weighted sum of the original committed bits. -//! -//! `LinearClaim` ensures nonempty factors whose lengths are powers of two. -//! Opening code checks the total weight count, packed witness length, and retained prover parameters. -//! Opening code or the outer caller binds the original statement before this stage. -//! Opening code always adds the sumcheck domain before calling `prove` or `verify`. -//! -//! This module owns the round messages, challenges, and terminal product check in the supplied transcript. -//! It returns the witness MLE claim and leaves the remaining transcript for the opening protocol. -//! Opening code binds that claim and performs the final PCS opening, which authenticates the witness evaluation. - -use common::LinearClaim; -use field::F128; -use num_traits::ConstZero; -use transcript::{ProverState, VerifierState}; - -use crate::{ProveError, VerifyError}; - -/// A pending evaluation claim over the original committed bit polynomial. -pub(super) struct MleClaim { - /// The evaluation point, with row coordinates before column coordinates, in low-bit-first order. - /// Its length is `log2(row_weights.len()) + log2(column_weights.len())`. - pub(super) point: Vec, - /// The witness MLE evaluation, rather than the terminal product evaluation. - pub(super) target: F128, -} - -/// Proves the reduction using a temporary dense table of witness evaluations. -#[tracing::instrument(name = "Prove inner-product sumcheck", skip_all)] -pub(super) fn prove( - claim: &LinearClaim, - packed_witness: &[F128], - transcript: &mut ProverState, -) -> Result { - let mut rows = claim.row_weights().to_vec(); - let mut columns = claim.column_weights().to_vec(); - let bit_len = rows.len() * columns.len(); - if packed_witness.len().checked_mul(128) != Some(bit_len) { - return Err(ProveError::PackedWitnessLengthMismatch); - } - let mut witness = Vec::with_capacity(bit_len); - for packed in packed_witness { - let bits = u128::from(packed.lo) | (u128::from(packed.hi) << 64); - for bit in 0..128 { - witness.push(F128::from(((bits >> bit) & 1) as u64)); - } - } - - let mut target = claim.target(); - let mut point = Vec::with_capacity(bit_len.ilog2() as usize); - while witness.len() > 1 { - let coefficients = round_polynomial(&witness, &rows, &columns); - // In characteristic two, g(0) + g(1) = a1 + a2. - if coefficients[1] + coefficients[2] != target { - return Err(ProveError::InvalidClaim); - } - transcript.prover_message(&coefficients); - let challenge = transcript.verifier_message::(); - target = evaluate_round(coefficients, challenge); - point.push(challenge); - fold(&mut witness, challenge); - if rows.len() > 1 { - fold(&mut rows, challenge); - } else { - fold(&mut columns, challenge); - } - } - - let evaluation = witness[0]; - if target != evaluation * rows[0] * columns[0] { - return Err(ProveError::InvalidClaim); - } - transcript.prover_message(&evaluation); - Ok(MleClaim { - point, - target: evaluation, - }) -} - -/// Verifies sumcheck and returns a witness claim that still requires the final PCS opening. -#[tracing::instrument(name = "Verify inner-product sumcheck", skip_all)] -pub(super) fn verify( - claim: &LinearClaim, - transcript: &mut VerifierState<'_>, -) -> Result { - let mut rows = claim.row_weights().to_vec(); - let mut columns = claim.column_weights().to_vec(); - let rounds = rows.len().ilog2() as usize + columns.len().ilog2() as usize; - let mut point = Vec::with_capacity(rounds); - let mut target = claim.target(); - for _ in 0..rounds { - let coefficients = transcript - .prover_message::<[F128; 3]>() - .map_err(|_| VerifyError::MalformedProof)?; - if coefficients[1] + coefficients[2] != target { - return Err(VerifyError::VerificationFailed); - } - let challenge = transcript.verifier_message::(); - target = evaluate_round(coefficients, challenge); - point.push(challenge); - if rows.len() > 1 { - fold(&mut rows, challenge); - } else { - fold(&mut columns, challenge); - } - } - - let evaluation = transcript - .prover_message::() - .map_err(|_| VerifyError::MalformedProof)?; - // Multiplication also handles zero weights; the MLE opening authenticates the witness value. - if target != evaluation * rows[0] * columns[0] { - return Err(VerifyError::VerificationFailed); - } - Ok(MleClaim { - point, - target: evaluation, - }) -} - -/// Returns the coefficients of the next degree-two round polynomial. -fn round_polynomial(witness: &[F128], rows: &[F128], columns: &[F128]) -> [F128; 3] { - let mut at_zero = F128::ZERO; - let mut at_one = F128::ZERO; - let mut quadratic = F128::ZERO; - for (pair_index, pair) in witness.chunks_exact(2).enumerate() { - let index = 2 * pair_index; - let weight_zero = rows[index % rows.len()] * columns[index / rows.len()]; - let weight_one = rows[(index + 1) % rows.len()] * columns[(index + 1) / rows.len()]; - at_zero += pair[0] * weight_zero; - at_one += pair[1] * weight_one; - quadratic += (pair[0] + pair[1]) * (weight_zero + weight_one); - } - [at_zero, at_zero + at_one + quadratic, quadratic] -} - -fn evaluate_round([constant, linear, quadratic]: [F128; 3], challenge: F128) -> F128 { - constant + challenge * (linear + challenge * quadratic) -} - -/// Fixes the lowest remaining coordinate without changing the order of the remaining coordinates. -fn fold(values: &mut Vec, challenge: F128) { - let len = values.len() / 2; - for index in 0..len { - let zero = values[2 * index]; - let one = values[2 * index + 1]; - values[index] = zero + challenge * (zero + one); - } - values.truncate(len); -} - -#[cfg(test)] -mod tests; diff --git a/crates/pcs/src/sumcheck/tests.rs b/crates/pcs/src/sumcheck/tests.rs deleted file mode 100644 index 198888a1..00000000 --- a/crates/pcs/src/sumcheck/tests.rs +++ /dev/null @@ -1,213 +0,0 @@ -use std::sync::OnceLock; - -use common::{LinearClaim, Shape}; -use field::F128; -use num_traits::ConstZero; -use transcript::{Proof, PublicTranscript, VerifierState, build_prover, build_verifier}; - -use super::{prove, verify}; -use crate::{ - CommitScheme, HashKind, LigeritoProfile, OpeningQuery, Pcs, StatementBinding, VerifyError, -}; - -const M: usize = 22; -const SESSION: &[u8] = b"pcs-sumcheck-format-test"; -const INSTANCE: &[u8] = b"factored-inner-product"; -const SET_BITS: [usize; 8] = [0, 63, 64, 127, 128, 255, 256, (1 << M) - 1]; -const ROUND_BYTES: usize = 3 * 16; -const EVALUATION_OFFSET: usize = M * ROUND_BYTES; - -fn shape() -> Shape { - Shape::new(8, M - 8).unwrap() -} - -fn factor_weight(index: usize) -> F128 { - let index = index as u64; - F128::new( - index.wrapping_mul(0x9e37_79b9_7f4a_7c15) ^ 0x0123_4567_89ab_cdef, - index.rotate_left(29) ^ 0xa5a5_5a5a_f0f0_0f0f, - ) -} - -fn sparse_witness(packed_len: usize) -> Vec { - let mut witness = vec![F128::ZERO; packed_len]; - for index in SET_BITS { - if index % 128 < 64 { - witness[index / 128].lo |= 1 << (index % 128); - } else { - witness[index / 128].hi |= 1 << (index % 128 - 64); - } - } - witness -} - -fn bind_claim(transcript: &mut impl PublicTranscript, claim: &LinearClaim) { - transcript.public_message(b"sumcheck-test/claim/v1" as &[u8]); - transcript.public_message(claim); -} - -struct Fixture { - claim: LinearClaim, - proof: Proof, - point: Vec, - evaluation: F128, -} - -impl Fixture { - fn build() -> Self { - // Shape requires at least 2^22 bits, so this is the smallest valid commitment size. - let shape = shape(); - let target = SET_BITS - .into_iter() - .map(|index| { - factor_weight(index / shape.rows() + shape.rows()) - * factor_weight(index % shape.rows()) - }) - .sum(); - let claim = LinearClaim::from_shape( - &shape, - (0..shape.rows()).map(factor_weight).collect(), - (0..shape.columns()) - .map(|column| factor_weight(column + shape.rows())) - .collect(), - target, - ) - .unwrap(); - let witness = sparse_witness(1 << shape.log_packed_len()); - let mut prover = build_prover(SESSION, INSTANCE); - bind_claim(&mut prover, &claim); - let reduced = prove(&claim, &witness, &mut prover).unwrap(); - Self { - claim, - proof: prover.finish(), - point: reduced.point, - evaluation: reduced.target, - } - } - - fn verifier<'proof>(&self, proof: &'proof Proof) -> VerifierState<'proof> { - let mut verifier = build_verifier(SESSION, INSTANCE, proof); - bind_claim(&mut verifier, &self.claim); - verifier - } -} - -fn fixture() -> &'static Fixture { - static FIXTURE: OnceLock = OnceLock::new(); - FIXTURE.get_or_init(Fixture::build) -} - -#[test] -fn standalone_proof_returns_the_witness_mle_evaluation() { - let fixture = fixture(); - let mut verifier = fixture.verifier(&fixture.proof); - let reduced = verify(&fixture.claim, &mut verifier).unwrap(); - assert_eq!(reduced.point, fixture.point); - assert_eq!(reduced.target, fixture.evaluation); - let expected = SET_BITS - .into_iter() - .map(|index| { - reduced.point.iter().copied().enumerate().fold( - F128::from(1u64), - |product, (coordinate, value)| { - product - * if (index >> coordinate) & 1 == 1 { - value - } else { - F128::from(1u64) + value - } - }, - ) - }) - .sum::(); - assert_eq!(reduced.target, expected); - verifier.check_eof().unwrap(); -} - -#[test] -fn rejects_truncated_rounds_and_witness_evaluation() { - let fixture = fixture(); - for length in [ - 0, - 16, - ROUND_BYTES - 1, - EVALUATION_OFFSET - 1, - EVALUATION_OFFSET + 15, - ] { - let mut proof = fixture.proof.clone(); - proof.narg_string.truncate(length); - let mut verifier = fixture.verifier(&proof); - assert_eq!( - verify(&fixture.claim, &mut verifier).err(), - Some(VerifyError::MalformedProof), - ); - } -} - -#[test] -fn rejects_changed_round_coefficients_and_witness_evaluation() { - let fixture = fixture(); - for offset in [ - 0, - 16, - 32, - ROUND_BYTES * (M / 2) + 16, - ROUND_BYTES * (M - 1) + 32, - EVALUATION_OFFSET, - ] { - let mut proof = fixture.proof.clone(); - proof.narg_string[offset] ^= 1; - let mut verifier = fixture.verifier(&proof); - assert_eq!( - verify(&fixture.claim, &mut verifier).err(), - Some(VerifyError::VerificationFailed), - ); - } -} - -#[test] -fn zero_weight_factor_still_requires_the_correct_pcs_witness_evaluation() { - let shape = shape(); - let pcs = Pcs::new(&shape, LigeritoProfile::Fast, HashKind::Blake3).unwrap(); - let witness = sparse_witness(pcs.packed_len()); - let (commitment, data) = pcs.commit(&witness).unwrap(); - - for zero_rows in [true, false] { - let mut row_weights = (0..shape.rows()).map(factor_weight).collect::>(); - let mut column_weights = (0..shape.columns()) - .map(|column| factor_weight(column + shape.rows())) - .collect::>(); - if zero_rows { - row_weights.fill(F128::ZERO); - } else { - column_weights.fill(F128::ZERO); - } - let query = OpeningQuery::InnerProduct { - claim: LinearClaim::from_shape(&shape, row_weights, column_weights, F128::ZERO) - .unwrap(), - }; - let mut prover = build_prover(SESSION, b"zero-inner-product-factor"); - pcs.prove_lin( - &data, - witness.clone(), - &query, - StatementBinding::Bind, - &mut prover, - ) - .unwrap(); - let proof = prover.finish(); - let mut verifier = build_verifier(SESSION, b"zero-inner-product-factor", &proof); - pcs.verify_lin(&commitment, &query, StatementBinding::Bind, &mut verifier) - .unwrap(); - verifier.check_eof().unwrap(); - - // A zero factor leaves this value unconstrained until the full PCS checks the MLE opening. - let mut changed_proof = proof; - changed_proof.narg_string[EVALUATION_OFFSET] ^= 1; - let mut verifier = build_verifier(SESSION, b"zero-inner-product-factor", &changed_proof); - assert_eq!( - pcs.verify_lin(&commitment, &query, StatementBinding::Bind, &mut verifier), - Err(VerifyError::VerificationFailed), - ); - } -} diff --git a/crates/post_gkr/Cargo.toml b/crates/post_gkr/Cargo.toml new file mode 100644 index 00000000..ab5c7b4a --- /dev/null +++ b/crates/post_gkr/Cargo.toml @@ -0,0 +1,23 @@ +[package] +name = "post_gkr" +version.workspace = true +edition.workspace = true +rust-version.workspace = true +license.workspace = true + +[features] +default = ["parallel"] +parallel = ["dep:rayon"] + +[dependencies] +common = { workspace = true } +field = { workspace = true, features = ["spongefish"] } +num-traits = { workspace = true } +poly = { workspace = true } +rayon = { workspace = true, optional = true } +tracing = { workspace = true } +transcript = { workspace = true } + +[dev-dependencies] +rand_core = { workspace = true } +rand_pcg = { workspace = true } diff --git a/crates/post_gkr/src/lib.rs b/crates/post_gkr/src/lib.rs new file mode 100644 index 00000000..d4ba69de --- /dev/null +++ b/crates/post_gkr/src/lib.rs @@ -0,0 +1,465 @@ +//! Step 6 of 2.1. "The prime field case": from the GKR's inner product +//! claim on the committed bits to the MLE evaluation claim the opening takes. +//! +//! The GKR protocol ends on ` = C - 1` +//! over the `2^t x 2^s` committed bits `f`, a `LinearClaim` with row +//! factor `omega` and column factor `eq(., r_c)`. A virtualization `h = M f` +//! (2.4. "Virtual F_2-linear transforms, NP-complete multi-domain linear +//! relations, and hybrid proof systems") moves it onto `f` as ``, +//! one weight per bit: the same type with a single column of weight one. +//! +//! The opening scheme (`pcs`) proves evaluation claims `MLE[f](r) = v` and +//! does the ring switch from the bits to the packed vector itself (Binius's, +//! as Step 6 cites it). This crate's sumcheck turns the linear claim into +//! such an evaluation claim: all `m = t + s` variables are bound against the +//! bits, the row ones first as the bits are indexed. The verifier's closing +//! weight `MLE[rows](rho_b) MLE[columns](rho_c)` costs `2^t + 2^s` +//! multiplications, so it is linear in the bits only when a factor is. +//! +//! The first round is taken off the packed bits: a bit is zero or one, so +//! the round's coefficients are sums of weights with no multiplication, and +//! the tables the prover folds to have one entry per two bits. The weights +//! are never written out in full: their folded table is the folded row +//! factor tensored with the column factor. Proof: `2m + 1` elements, `rho` +//! low coordinate first ([`OpeningQuery::Mle`]). +//! +//! # Transcript +//! +//! Each round's coefficients are written and absorbed before its challenge +//! is squeezed; the closing evaluation follows the last round. The claim is +//! not re-absorbed here: the opening scheme binds its statement before it +//! calls the sumcheck. The round count derives from the claim and is not +//! absorbed either. + +mod sumcheck; +#[cfg(test)] +mod test_util; + +use crate::sumcheck::{Pair, RoundMessage}; +use common::shape::PACK_BITS; +use common::{LinearClaim, OpeningQuery}; +use field::F128; +use num_traits::{ConstOne, ConstZero}; +#[cfg(feature = "parallel")] +use poly::parallel::workload_size; +#[cfg(feature = "parallel")] +use rayon::prelude::*; +use transcript::{ProverState, VerifierState}; + +/// A reduction the prover cannot run. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum ProveError { + /// The packed witness is not one element per 128 bits of the claim. + WitnessLengthMismatch, + /// The claim's weights against the witness do not give the target. + ClaimDoesNotHold, +} + +/// A reduction the verifier rejects. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum VerifyError { + /// A record is missing or does not decode. + MalformedProof, + /// The weights at the challenges times `v` is not the running claim. + EvaluationMismatch, +} + +/// The bits one packed element carries. +const ELEMENT_BITS: usize = 1 << PACK_BITS; + +/// The even bit positions of a packed element, `0b01010101...` +const EVEN: u128 = 0x55555555555555555555555555555555; + +/// Packed elements per parallel task, enough to fill the cache. +#[cfg(feature = "parallel")] +const ELEMENTS_PER_TASK: usize = workload_size::() / ELEMENT_BITS; + +/// Runs the sumcheck for `claim` on `packed`, the committed `f` column +/// major, and returns the evaluation claim it leaves: `MLE[f](rho) = v` as +/// an [`OpeningQuery::Mle`], `rho` the challenges low coordinate first and +/// `v` the prover's closing evaluation. +#[tracing::instrument(name = "Prove inner-product sumcheck", skip_all)] +pub fn prove( + claim: &LinearClaim, + packed: &[F128], + transcript: &mut ProverState, +) -> Result { + prove_factors( + claim.row_weights(), + claim.column_weights(), + claim.target(), + packed, + transcript, + ) +} + +/// Replays the sumcheck for `claim`, checks that the closing evaluation +/// fits, and returns the evaluation claim for the opening scheme to verify +/// against the commitment. +#[tracing::instrument(name = "Verify inner-product sumcheck", skip_all)] +pub fn verify( + claim: &LinearClaim, + transcript: &mut VerifierState<'_>, +) -> Result { + verify_factors( + claim.row_weights(), + claim.column_weights(), + claim.target(), + transcript, + ) +} + +/// [`prove`] on the claim's parts. `rows` has `2^t` weights with `t >= 7` +/// and `columns` `2^s`, as `LinearClaim` guarantees. +fn prove_factors( + rows: &[F128], + columns: &[F128], + target: F128, + packed: &[F128], + transcript: &mut ProverState, +) -> Result { + let factors = Factors { rows, columns }; + if packed.len() != factors.packed_len() { + return Err(ProveError::WitnessLengthMismatch); + } + if factors.weighted_sum(packed) != target { + return Err(ProveError::ClaimDoesNotHold); + } + + let message = factors.first_message(packed); + transcript.prover_message(&message); + let challenge: F128 = transcript.verifier_message(); + let running = advance(target, message, challenge); + let mut pair = Pair::new(factors.fold(challenge), bind_first(packed, challenge)); + let (rest, target) = sumcheck::prove(&mut pair, running, transcript); + let mut point = vec![challenge]; + point.extend(rest); + Ok(OpeningQuery::Mle { point, target }) +} + +/// [`verify`] on the claim's parts. +fn verify_factors( + rows: &[F128], + columns: &[F128], + target: F128, + transcript: &mut VerifierState<'_>, +) -> Result { + let log_rows = rows.len().trailing_zeros() as usize; + let rounds = log_rows + columns.len().trailing_zeros() as usize; + // `MLE[rows (x) columns](rho) = MLE[rows](rho_b) MLE[columns](rho_c)`. + let weight = |point: &[F128]| { + sumcheck::evaluate(rows, &point[..log_rows]) + * sumcheck::evaluate(columns, &point[log_rows..]) + }; + let (point, target) = sumcheck::verify(rounds, target, weight, transcript)?; + Ok(OpeningQuery::Mle { point, target }) +} + +/// The weights `rows (x) columns` over the bits, read without being written +/// out: bit `(c << t) | b` weighs `rows[b] columns[c]`, so packed element +/// `e` covers rows `128 (e mod 2^(t-7)) ..` of column `e / 2^(t-7)`. +struct Factors<'a> { + rows: &'a [F128], + columns: &'a [F128], +} + +impl Factors<'_> { + fn packed_len(&self) -> usize { + (self.rows.len() >> PACK_BITS) * self.columns.len() + } + + /// The row weights element `index` covers, and its column's weight. + fn element(&self, index: usize) -> (&[F128], F128) { + let per_column = self.rows.len() >> PACK_BITS; + let rows = &self.rows[(index % per_column) << PACK_BITS..][..ELEMENT_BITS]; + (rows, self.columns[index / per_column]) + } + + /// ``: the weights at the set bits. + fn weighted_sum(&self, packed: &[F128]) -> F128 { + let element = |(index, &element): (usize, &F128)| -> F128 { + let (rows, column) = self.element(index); + column + * iter_over_set_bits(element.to_u128()) + .map(|v| rows[v]) + .sum::() + }; + #[cfg(feature = "parallel")] + return packed + .par_iter() + .enumerate() + .with_min_len(ELEMENTS_PER_TASK) + .map(element) + .sum(); + #[cfg(not(feature = "parallel"))] + packed.iter().enumerate().map(element).sum() + } + + /// `(a_0, a_2)` of the first round off the bits: `a_0` sums the weights + /// at even positions whose bit is set, `a_2` the pair sums `w_0 + w_1` + /// where the two bits differ, since `(w_0 + w_1)(f_0 + f_1)` is that + /// exactly then. + fn first_message(&self, packed: &[F128]) -> RoundMessage { + let element = |(index, &element): (usize, &F128)| -> (F128, F128) { + let (rows, column) = self.element(index); + let bits = element.to_u128(); + let a0: F128 = iter_over_set_bits(bits & EVEN).map(|v| rows[v]).sum(); + let a2: F128 = iter_over_set_bits((bits ^ (bits >> 1)) & EVEN) + .map(|v| rows[v] + rows[v + 1]) + .sum(); + (column * a0, column * a2) + }; + let add = |(a0, a2): (F128, F128), (b0, b2): (F128, F128)| (a0 + b0, a2 + b2); + #[cfg(feature = "parallel")] + let (a0, a2) = packed + .par_iter() + .enumerate() + .with_min_len(ELEMENTS_PER_TASK) + .map(element) + .reduce(|| (F128::ZERO, F128::ZERO), add); + #[cfg(not(feature = "parallel"))] + let (a0, a2) = packed + .iter() + .enumerate() + .map(element) + .fold((F128::ZERO, F128::ZERO), add); + [a0, a2] + } + + /// The weights with their first variable bound: the folded row factor + /// tensored with the column factor, one entry per two bits. + fn fold(self, challenge: F128) -> Vec { + let rows = sumcheck::folded(self.rows, challenge); + let mut table = Vec::with_capacity(rows.len() * self.columns.len()); + for &column in self.columns { + table.extend(rows.iter().map(|&row| column * row)); + } + table + } +} + +/// The positions of the set bits of `word`, ascending. +fn iter_over_set_bits(mut word: u128) -> impl Iterator { + std::iter::from_fn(move || { + (word != 0).then(|| { + let position = word.trailing_zeros() as usize; + word &= word - 1; + position + }) + }) +} + +/// `MLE[f]` with its first variable bound: over each pair of bits +/// `(f_0, f_1)`, `f_0 + rho (f_0 + f_1)`, one of `0`, `1`, `rho` and +/// `1 + rho`. +fn bind_first(packed: &[F128], challenge: F128) -> Vec { + let challenge_plus_one = challenge + F128::ONE; + let values = [F128::ZERO, challenge_plus_one, challenge, F128::ONE]; + let element = |(slot, &element): (&mut [F128], &F128)| { + let bits = element.to_u128(); + for (pair, value) in slot.iter_mut().enumerate() { + *value = values[((bits >> (2 * pair)) & 3) as usize]; + } + }; + let mut table = vec![F128::ZERO; packed.len() * ELEMENT_BITS / 2]; + #[cfg(feature = "parallel")] + table + .par_chunks_mut(ELEMENT_BITS / 2) + .zip(packed) + .with_min_len(ELEMENTS_PER_TASK) + .for_each(element); + #[cfg(not(feature = "parallel"))] + table + .chunks_mut(ELEMENT_BITS / 2) + .zip(packed) + .for_each(element); + table +} + +/// `h' = p(rho)` with `a_1 = h + a_2`. +fn advance(claim: F128, [a0, a2]: RoundMessage, challenge: F128) -> F128 { + a0 + challenge * (claim + a2 + challenge * a2) +} + +#[cfg(test)] +mod tests { + use common::Shape; + use poly::DenseMultilinearExtension; + use transcript::{Proof, build_prover, build_verifier}; + + use super::*; + use crate::sumcheck::inner_product; + use crate::test_util::{Leaf, random, rng}; + + /// The evaluation claim's parts. + fn mle(query: &OpeningQuery) -> (&[F128], F128) { + let OpeningQuery::Mle { point, target } = query else { + panic!("the sumcheck leaves an evaluation claim"); + }; + (point, *target) + } + + fn reduced(leaf: &Leaf) -> (OpeningQuery, Proof) { + let mut prover = build_prover("post_gkr-tests", "reduce"); + let sent = prove_factors( + &leaf.rows, + &leaf.columns, + leaf.target, + &leaf.packed, + &mut prover, + ) + .unwrap(); + (sent, prover.finish()) + } + + fn verified(leaf: &Leaf, proof: &Proof) -> Result { + let mut verifier = build_verifier("post_gkr-tests", "reduce", proof); + let received = verify_factors(&leaf.rows, &leaf.columns, leaf.target, &mut verifier)?; + assert!(verifier.check_eof().is_ok()); + Ok(received) + } + + #[test] + fn set_bits_are_read_in_ascending_order() { + let element = F128::new(0b1011, 1 << 63); + assert_eq!( + iter_over_set_bits(element.to_u128()).collect::>(), + vec![0, 1, 3, 127] + ); + assert_eq!(iter_over_set_bits(0).count(), 0); + assert_eq!(iter_over_set_bits(u128::MAX).count(), 128); + } + + /// Factored over several columns, and a single column of weight one: + /// the shape the transposition leaves. + #[test] + fn the_two_sides_agree_on_a_true_claim() { + for (log_rows, log_columns, seed) in [(8, 2, 63), (10, 0, 64)] { + let leaf = Leaf::random(log_rows, log_columns, seed); + let (sent, proof) = reduced(&leaf); + assert_eq!(proof.narg_string.len(), (2 * 10 + 1) * 16); + assert!(proof.hints.is_empty()); + + let received = verified(&leaf, &proof).unwrap(); + assert_eq!(sent, received); + let (point, target) = mle(&received); + assert_eq!(point.len(), 10); + assert_eq!( + leaf.bits_extension().evaluate(point).unwrap(), + target, + "{log_rows} x {log_columns}" + ); + } + } + + /// Through `LinearClaim`, at the smallest committed shape. + #[test] + fn a_linear_claim_reduces_to_an_evaluation_the_witness_satisfies() { + let leaf = Leaf::sparse(7, 15, 65); + let claim = LinearClaim::from_shape( + &Shape::new(7, 15).unwrap(), + leaf.rows.clone(), + leaf.columns.clone(), + leaf.target, + ) + .unwrap(); + let mut prover = build_prover("post_gkr-tests", "reduce"); + let sent = prove(&claim, &leaf.packed, &mut prover).unwrap(); + let proof = prover.finish(); + let mut verifier = build_verifier("post_gkr-tests", "reduce", &proof); + assert_eq!(verify(&claim, &mut verifier), Ok(sent.clone())); + let (point, target) = mle(&sent); + assert_eq!(leaf.evaluate(point), target); + } + + /// The first round off the bits sends what the generic round over the + /// weights and bits written out would, and folds to the same tables. + #[test] + fn the_first_round_off_the_bits_is_the_generic_round() { + let leaf = Leaf::random(7, 2, 67); + let factors = Factors { + rows: &leaf.rows, + columns: &leaf.columns, + }; + let weights = leaf.weights(); + let written_out: Vec = (0..1 << 9) + .map(|index| F128::from(leaf.bit(index))) + .collect(); + assert_eq!(factors.weighted_sum(&leaf.packed), leaf.target); + assert_eq!(inner_product(&weights, &written_out), leaf.target); + + let mut generic = Pair::new(weights.clone(), written_out.clone()); + let mut prover = build_prover("post_gkr-tests", "reduce"); + let (point, _) = sumcheck::prove(&mut generic, leaf.target, &mut prover); + let proof = prover.finish(); + let message = factors.first_message(&leaf.packed); + assert_eq!( + proof.narg_string[..32], + [message[0].to_bytes(), message[1].to_bytes()].concat() + ); + + let fold = |table: Vec| { + let mut folded = DenseMultilinearExtension { evaluations: table }; + folded.fold(&point[..1]).unwrap(); + folded.evaluations + }; + assert_eq!(bind_first(&leaf.packed, point[0]), fold(written_out)); + assert_eq!(factors.fold(point[0]), fold(weights)); + } + + /// Every record off by one bit fails the closing check; a different + /// target does too; a proof cut short does not decode. + #[test] + fn a_reduction_altered_in_transit_is_caught() { + let mut leaf = Leaf::random(7, 1, 66); + let (_, proof) = reduced(&leaf); + let records = 2 * 8 + 1; + assert_eq!(proof.narg_string.len(), records * 16); + for record in 0..records { + let mut altered = proof.clone(); + altered.narg_string[record * 16] ^= 1; + assert_eq!( + verified(&leaf, &altered), + Err(VerifyError::EvaluationMismatch), + "record {record}" + ); + } + let mut short = proof.clone(); + short.narg_string.truncate(proof.narg_string.len() - 16); + assert_eq!(verified(&leaf, &short), Err(VerifyError::MalformedProof)); + + leaf.target += F128::ONE; + assert_eq!( + verified(&leaf, &proof), + Err(VerifyError::EvaluationMismatch) + ); + } + + #[test] + fn the_reduction_refuses_what_it_cannot_prove() { + let leaf = Leaf::random(7, 1, 56); + let mut transcript = build_prover("post_gkr-tests", "reduce"); + assert_eq!( + prove_factors( + &leaf.rows, + &leaf.columns, + leaf.target + random(&mut rng(57)), + &leaf.packed, + &mut transcript + ), + Err(ProveError::ClaimDoesNotHold) + ); + assert_eq!( + prove_factors( + &leaf.rows, + &leaf.columns, + leaf.target, + &leaf.packed[1..], + &mut transcript + ), + Err(ProveError::WitnessLengthMismatch) + ); + // Nothing was written before the refusals. + assert!(transcript.finish().narg_string.is_empty()); + } +} diff --git a/crates/post_gkr/src/sumcheck.rs b/crates/post_gkr/src/sumcheck.rs new file mode 100644 index 00000000..004b2d4e --- /dev/null +++ b/crates/post_gkr/src/sumcheck.rs @@ -0,0 +1,356 @@ +//! The degree-two sumcheck: from `sum_x W(x) V(x) = h_0` over two tables to +//! `MLE[V](rho) = v`, the verifier checking `MLE[W](rho) * v = h_n` with the +//! `MLE[W](rho)` its caller computes. +//! +//! Each round splits off the lowest remaining variable of both tables and +//! sends the round polynomial +//! +//! ```text +//! p(X) = sum_{x'} MLE[W](X, x') MLE[V](X, x') = a_0 + a_1 X + a_2 X^2. +//! ``` +//! +//! In characteristic two `p(0) + p(1) = a_1 + a_2`, so the running claim +//! `h = p(0) + p(1)` fixes `a_1 = h + a_2` and only `(a_0, a_2)` is sent +//! ([`RoundMessage`]). Both sides continue with `h' = p(rho)` and the prover +//! folds its tables at `rho`. There is no per-round check: a wrong message +//! in any round surfaces in the closing check. Soundness error at most +//! `2n / |E|` for `n` rounds, plus the probability that `MLE[W](rho) = 0`. + +use crate::VerifyError; +use common::shape::PACK_BITS; +use field::{F128, Wide256}; +use poly::eq_table; +#[cfg(feature = "parallel")] +use poly::parallel::workload_size; +#[cfg(feature = "parallel")] +use rayon::prelude::*; +use transcript::{ProverState, VerifierState}; + +/// `(a_0, a_2)` of `p(X) = a_0 + a_1 X + a_2 X^2`; `a_1` the running claim +/// implies. +pub(crate) type RoundMessage = [F128; 2]; + +/// Entries below which [`evaluate`] and [`folded`] run on one thread: the +/// pool's overhead is that of a few hundred thousand multiplications. +#[cfg(feature = "parallel")] +const PARALLEL_EVALUATE_MIN: usize = 1 << 18; + +/// The public weights `W` and the prover's values `V` over the same +/// variables, folded together as the rounds bind them. +#[derive(Debug, Clone)] +pub(crate) struct Pair { + weights: Vec, + values: Vec, +} + +impl Pair { + /// Two tables of the same power-of-two length. + pub(crate) fn new(weights: Vec, values: Vec) -> Self { + debug_assert_eq!(weights.len(), values.len()); + debug_assert!(weights.len().is_power_of_two()); + Self { weights, values } + } + + /// `(MLE[W](rho), MLE[V](rho))` once every variable is bound. + fn bound(&self) -> (F128, F128) { + debug_assert_eq!(self.weights.len(), 1); + (self.weights[0], self.values[0]) + } + + /// `(a_0, a_2)`: `a_0 = sum w_0 v_0` and `a_2 = sum (w_0 + w_1)(v_0 + v_1)` + /// over adjacent entries, since `MLE[W](X, x') = w_0 + X (w_0 + w_1)` + /// and likewise for `MLE[V]`. + fn message(&self) -> RoundMessage { + fn coeffs(this: &Pair) -> (Wide256, Wide256) { + // Large tables are split into cache-sized chunks summed on the Rayon pool. + #[cfg(feature = "parallel")] + { + // An even chunk length keeps every pair inside one chunk. + let chunk = workload_size::() & !1; + if this.weights.len() > chunk { + return this + .weights + .par_chunks(chunk) + .zip(this.values.par_chunks(chunk)) + .map(|(w, v)| coefficients_serial(w, v)) + .reduce( + || (Wide256::zero(), Wide256::zero()), + |(a0, a2), (b0, b2)| (a0 + b0, a2 + b2), + ); + } + } + coefficients_serial(&this.weights, &this.values) + } + + let (a0, a2) = coeffs(self); + [a0.reduce(), a2.reduce()] + } + + fn fold(&mut self, challenge: F128) { + fold(&mut self.weights, challenge); + fold(&mut self.values, challenge); + } +} + +/// Runs the sumcheck over `pair`, folding it in place, and writes the +/// closing evaluation. Returns the challenges in the order they were drawn +/// and `v = MLE[V](rho)`. +pub(crate) fn prove( + pair: &mut Pair, + claim: F128, + transcript: &mut ProverState, +) -> (Vec, F128) { + let rounds = pair.weights.len().trailing_zeros() as usize; + let (point, running) = prove_rounds(pair, rounds, claim, transcript); + let evaluation = prove_evaluation(pair, running, transcript); + (point, evaluation) +} + +/// Replays the sumcheck from the records and checks the closing evaluation +/// against `weight(rho) = MLE[W](rho)`. Returns the challenges and `v`. +pub(crate) fn verify( + rounds: usize, + claim: F128, + weight: impl FnOnce(&[F128]) -> F128, + transcript: &mut VerifierState<'_>, +) -> Result<(Vec, F128), VerifyError> { + let (point, running) = verify_rounds(rounds, claim, transcript)?; + let evaluation = verify_evaluation(weight(&point), running, transcript)?; + Ok((point, evaluation)) +} + +/// Runs `rounds` rounds over `pair`, folding it in place. Returns the +/// challenges in the order they were drawn and the running claim. +fn prove_rounds( + pair: &mut Pair, + rounds: usize, + mut claim: F128, + transcript: &mut ProverState, +) -> (Vec, F128) { + let mut point = Vec::with_capacity(rounds); + for _ in 0..rounds { + let message = pair.message(); + transcript.prover_message(&message); + let challenge: F128 = transcript.verifier_message(); + claim = super::advance(claim, message, challenge); + point.push(challenge); + pair.fold(challenge); + } + (point, claim) +} + +/// Replays `rounds` rounds from the records alone. Returns the challenges +/// and the running claim. +fn verify_rounds( + rounds: usize, + mut claim: F128, + transcript: &mut VerifierState<'_>, +) -> Result<(Vec, F128), VerifyError> { + let mut point = Vec::with_capacity(rounds); + for _ in 0..rounds { + let message: RoundMessage = transcript + .prover_message() + .map_err(|_| VerifyError::MalformedProof)?; + let challenge: F128 = transcript.verifier_message(); + claim = super::advance(claim, message, challenge); + point.push(challenge); + } + Ok((point, claim)) +} + +/// Writes `v = MLE[V](rho)`, the one entry left in the folded pair. +fn prove_evaluation(pair: &Pair, claim: F128, transcript: &mut ProverState) -> F128 { + let (weight, evaluation) = pair.bound(); + debug_assert_eq!(weight * evaluation, claim); + transcript.prover_message(&evaluation); + evaluation +} + +/// Reads `v` and checks `MLE[W](rho) * v = h` for the caller's `MLE[W](rho)`. +fn verify_evaluation( + bound_weight: F128, + claim: F128, + transcript: &mut VerifierState<'_>, +) -> Result { + let evaluation: F128 = transcript + .prover_message() + .map_err(|_| VerifyError::MalformedProof)?; + if bound_weight * evaluation != claim { + return Err(VerifyError::EvaluationMismatch); + } + Ok(evaluation) +} + +/// Fixes the lowest remaining variable of `table` at `challenge`. +pub(crate) fn fold(table: &mut Vec, challenge: F128) { + *table = folded(table, challenge) +} + +/// [`fold`] into a fresh table, for a table that is only borrowed. +pub(crate) fn folded(table: &[F128], challenge: F128) -> Vec { + let entry = |pair: &[F128]| pair[0] + challenge * (pair[0] + pair[1]); + #[cfg(feature = "parallel")] + if table.len() >= PARALLEL_EVALUATE_MIN { + return table.par_chunks_exact(2).map(entry).collect(); + } + table.chunks_exact(2).map(entry).collect() +} + +/// `MLE[weights](point)`, one multiplication per weight: the weights have +/// no succinct form. The equality table is factored at the pack width, so +/// the larger factor is one element per 128 weights. +pub(crate) fn evaluate(weights: &[F128], point: &[F128]) -> F128 { + debug_assert_eq!(weights.len(), 1 << point.len()); + let (low, high) = point.split_at(point.len().min(PACK_BITS as usize)); + let eq_low = eq_table(low); + let eq_high = eq_table(high); + let term = |(chunk, weight): (&[F128], &F128)| *weight * inner_product(chunk, &eq_low); + #[cfg(feature = "parallel")] + if weights.len() >= PARALLEL_EVALUATE_MIN { + return weights + .par_chunks_exact(eq_low.len()) + .zip(&eq_high) + .map(term) + .sum(); + } + weights + .chunks_exact(eq_low.len()) + .zip(&eq_high) + .map(term) + .sum() +} + +pub(crate) fn inner_product(a: &[F128], b: &[F128]) -> F128 { + debug_assert_eq!(a.len(), b.len()); + a.iter() + .zip(b) + .fold(Wide256::zero(), |sum, (x, y)| sum + Wide256::mul(*x, *y)) + .reduce() +} + +fn coefficients_serial(weights: &[F128], values: &[F128]) -> (Wide256, Wide256) { + let mut a0 = Wide256::zero(); + let mut a2 = Wide256::zero(); + for (w, v) in weights.chunks_exact(2).zip(values.chunks_exact(2)) { + a0 += Wide256::mul(w[0], v[0]); + a2 += Wide256::mul(w[0] + w[1], v[0] + v[1]); + } + (a0, a2) +} + +#[cfg(test)] +mod tests { + use num_traits::{ConstOne, ConstZero}; + use poly::DenseMultilinearExtension; + use transcript::{build_prover, build_verifier}; + + use super::*; + use crate::test_util::{random, random_elements, rng}; + + fn pair(n: usize, rng: &mut rand_pcg::Pcg64) -> Pair { + Pair::new(random_elements(rng, 1 << n), random_elements(rng, 1 << n)) + } + + /// `sum_x MLE[W](x, x') MLE[V](x, x')` at `x`, by interpolating each pair. + fn round_polynomial(pair: &Pair, x: F128) -> F128 { + let interpolate = |entries: &[F128]| entries[0] + x * (entries[0] + entries[1]); + pair.weights + .chunks_exact(2) + .zip(pair.values.chunks_exact(2)) + .map(|(w, v)| interpolate(w) * interpolate(v)) + .sum() + } + + #[test] + fn the_message_is_the_round_polynomial() { + let mut rng = rng(1); + let pair = pair(5, &mut rng); + let claim = inner_product(&pair.weights, &pair.values); + let message = pair.message(); + assert_eq!( + round_polynomial(&pair, F128::ZERO) + round_polynomial(&pair, F128::ONE), + claim + ); + assert_eq!(message[0], round_polynomial(&pair, F128::ZERO)); + let rho = random(&mut rng); + assert_eq!( + crate::advance(claim, message, rho), + round_polynomial(&pair, rho) + ); + } + + #[test] + fn an_honest_run_closes_and_a_false_claim_does_not() { + let mut rng = rng(3); + let mut pair = pair(6, &mut rng); + let (weights, values) = (pair.weights.clone(), pair.values.clone()); + let claim = inner_product(&weights, &values); + + let mut prover = build_prover("post_gkr-tests", "sumcheck"); + let (point, running) = prove_rounds(&mut pair, 6, claim, &mut prover); + let evaluation = prove_evaluation(&pair, running, &mut prover); + let proof = prover.finish(); + assert_eq!(proof.narg_string.len(), (2 * 6 + 1) * 16); + + let extension = |table: &[F128]| { + DenseMultilinearExtension { + evaluations: table.to_vec(), + } + .evaluate(&point) + .unwrap() + }; + assert_eq!(evaluation, extension(&values)); + assert_eq!(evaluate(&weights, &point), extension(&weights)); + + let mut verifier = build_verifier("post_gkr-tests", "sumcheck", &proof); + let (same_point, same_running) = verify_rounds(6, claim, &mut verifier).unwrap(); + assert_eq!(same_point, point); + assert_eq!(same_running, running); + assert_eq!( + verify_evaluation(evaluate(&weights, &point), running, &mut verifier), + Ok(evaluation) + ); + assert!(verifier.check_eof().is_ok()); + + // The same records against a claim one off: the gap survives every + // round and the closing check catches it. + let mut verifier = build_verifier("post_gkr-tests", "sumcheck", &proof); + let (_, running) = verify_rounds(6, claim + F128::ONE, &mut verifier).unwrap(); + assert_eq!( + verify_evaluation(evaluate(&weights, &point), running, &mut verifier), + Err(VerifyError::EvaluationMismatch) + ); + } + + #[test] + fn the_evaluation_and_the_fold_agree_with_the_extension() { + let mut rng = rng(9); + for n in [3, 7, 9, 13] { + let weights = random_elements(&mut rng, 1 << n); + let point = random_elements(&mut rng, n); + assert_eq!( + evaluate(&weights, &point), + inner_product(&weights, &eq_table(&point)), + "{n} variables" + ); + let mut table = weights.clone(); + fold(&mut table, point[0]); + assert_eq!(folded(&weights, point[0]), table); + } + } + + #[test] + fn a_truncated_reduction_is_malformed() { + let mut pair = pair(3, &mut rng(10)); + let claim = inner_product(&pair.weights, &pair.values); + let mut prover = build_prover("post_gkr-tests", "sumcheck"); + prove_rounds(&mut pair, 3, claim, &mut prover); + let mut proof = prover.finish(); + proof.narg_string.truncate(proof.narg_string.len() - 16); + let mut verifier = build_verifier("post_gkr-tests", "sumcheck", &proof); + assert_eq!( + verify_rounds(3, claim, &mut verifier).err(), + Some(VerifyError::MalformedProof) + ); + } +} diff --git a/crates/post_gkr/src/test_util.rs b/crates/post_gkr/src/test_util.rs new file mode 100644 index 00000000..8dfa3eaa --- /dev/null +++ b/crates/post_gkr/src/test_util.rs @@ -0,0 +1,134 @@ +//! Random instances whose claims hold by construction. + +use field::F128; +use num_traits::ConstZero; +use poly::{DenseMultilinearExtension, eq_table}; +use rand_core::{Rng, SeedableRng}; +use rand_pcg::Pcg64; + +pub fn rng(seed: u64) -> Pcg64 { + Pcg64::seed_from_u64(seed) +} + +pub fn random(rng: &mut Pcg64) -> F128 { + F128::new(rng.next_u64(), rng.next_u64()) +} + +pub fn random_elements(rng: &mut Pcg64, count: usize) -> Vec { + (0..count).map(|_| random(rng)).collect() +} + +/// A claim ` = target` with the witness it is about, +/// `f` packed column major with `2^log_rows` rows. +pub struct Leaf { + pub log_rows: usize, + pub log_columns: usize, + pub rows: Vec, + pub columns: Vec, + /// From the definition. + pub target: F128, + pub packed: Vec, +} + +impl Leaf { + /// Random weights over a random witness. + pub fn random(log_rows: usize, log_columns: usize, seed: u64) -> Self { + let mut rng = rng(seed); + let packed = random_elements(&mut rng, 1 << (log_rows + log_columns - 7)); + Self::with_witness(log_rows, log_columns, packed, &mut rng) + } + + /// Random weights over a witness of a few set bits, so that the + /// extension stays cheap to evaluate at large shapes. + pub fn sparse(log_rows: usize, log_columns: usize, seed: u64) -> Self { + let mut rng = rng(seed); + let mut packed = vec![F128::ZERO; 1 << (log_rows + log_columns - 7)]; + for _ in 0..64 { + let index = rng.next_u64() as usize % packed.len(); + packed[index] = random(&mut rng); + } + Self::with_witness(log_rows, log_columns, packed, &mut rng) + } + + fn with_witness( + log_rows: usize, + log_columns: usize, + packed: Vec, + rng: &mut Pcg64, + ) -> Self { + let rows = random_elements(rng, 1 << log_rows); + let columns = random_elements(rng, 1 << log_columns); + let mut leaf = Self { + log_rows, + log_columns, + rows, + columns, + target: F128::ZERO, + packed, + }; + leaf.target = leaf + .set_bits() + .map(|index| leaf.rows[index % (1 << log_rows)] * leaf.columns[index >> log_rows]) + .sum(); + leaf + } + + /// Bit `index` of the packed witness. + pub fn bit(&self, index: usize) -> bool { + let element = self.packed[index >> 7]; + let half = if index % 128 < 64 { + element.lo + } else { + element.hi + }; + (half >> (index % 64)) & 1 == 1 + } + + fn set_bits(&self) -> impl Iterator + '_ { + (0..self.packed.len() << 7).filter(|&index| self.bit(index)) + } + + /// `rows (x) columns` written out per bit. + pub fn weights(&self) -> Vec { + (0..self.packed.len() << 7) + .map(|index| { + self.rows[index % (1 << self.log_rows)] * self.columns[index >> self.log_rows] + }) + .collect() + } + + /// `MLE[f]` over all `t + s` variables. Dense, so small instances only. + pub fn bits_extension(&self) -> DenseMultilinearExtension { + let log_bits = self.log_rows + self.log_columns; + let table = (0..1usize << log_bits) + .map(|index| F128::from(self.bit(index))) + .collect(); + DenseMultilinearExtension::from_evaluations(log_bits, table).unwrap() + } + + /// `MLE[f](point)` from the set bits alone. + pub fn evaluate(&self, point: &[F128]) -> F128 { + let (low, high) = point.split_at(7); + let eq_low = eq_table(low); + let eq_high = eq_table(high); + self.set_bits() + .map(|index| eq_high[index >> 7] * eq_low[index % 128]) + .sum() + } +} + +#[test] +fn the_leaf_is_its_own_extension_on_the_cube() { + let leaf = Leaf::random(7, 2, 98); + let extension = leaf.bits_extension(); + assert_eq!(extension.num_vars(), 9); + assert_eq!(extension[(3 << 7) | 5], F128::from(leaf.bit((3 << 7) | 5))); + let written_out: F128 = extension + .iter() + .zip(leaf.weights()) + .map(|(f, weight)| weight * *f) + .sum(); + assert_eq!(written_out, leaf.target); + let point: Vec = (0..9).map(|i| F128::new(i + 3, 7)).collect(); + assert_eq!(leaf.evaluate(&point), extension.evaluate(&point).unwrap()); +} diff --git a/crates/prover/src/lib.rs b/crates/prover/src/lib.rs index a341e91a..dec471bf 100644 --- a/crates/prover/src/lib.rs +++ b/crates/prover/src/lib.rs @@ -7,5 +7,5 @@ pub mod setup; pub use fold::SendError; pub use prove::{ProveError, VirtualWitness}; -pub use reduce::gkr_reduce; +pub use reduce::{ReduceError, gkr_reduce}; pub use setup::BitZProver; diff --git a/crates/prover/src/prove.rs b/crates/prover/src/prove.rs index 87c62705..28cbc197 100644 --- a/crates/prover/src/prove.rs +++ b/crates/prover/src/prove.rs @@ -1,14 +1,13 @@ //! `ProveBitZ`. use common::{ - BitTable, ClaimError, LinearClaim, OpeningQuery, TableError, VirtualMap, VirtualMapError, - VirtualStatement, + BitTable, LinearClaim, OpeningQuery, TableError, VirtualMap, VirtualMapError, VirtualStatement, }; use field::{F128, Fq}; use pcs::{CommitScheme, Pcs, ProveError as OpeningProveError, ProverData, StatementBinding}; use transcript::ProverState; -use crate::{BitZProver, SendError, reduce::gkr_reduce}; +use crate::{BitZProver, SendError, reduce::ReduceError, reduce::gkr_reduce}; /// A proof the prover cannot produce. #[derive(Debug, Clone, PartialEq, Eq)] @@ -21,8 +20,8 @@ pub enum ProveError { Witness(TableError), /// The fold round failed. Fold(SendError), - /// The derived GKR weight counts do not match the table shape. - Reduction(ClaimError), + /// The GKR left no claim. + Reduction(ReduceError), /// The opening failed, so the reduction's claim was never discharged. Opening(OpeningProveError), } @@ -132,7 +131,7 @@ impl BitZProver { .transpose_query(query) .map_err(ProveError::VirtualMap)?; - // Step 6: run PCS sumcheck, ring switching, and opening on committed bits. + // Step 6: the post-GKR sumcheck, ring switching, and opening on committed bits. // Bind the PCS parameters and transposed query before its challenges. pcs.prove_lin( data, diff --git a/crates/prover/src/reduce.rs b/crates/prover/src/reduce.rs index 4364cf26..fc726004 100644 --- a/crates/prover/src/reduce.rs +++ b/crates/prover/src/reduce.rs @@ -8,10 +8,21 @@ use common::{BitTable, ClaimError, Fold, LinearClaim, OpeningQuery, TransposeError}; use field::F128; use gkr::{GrandProductCircuit, gpgkr_prove}; -use num_traits::ConstOne; +use num_traits::{ConstOne, ConstZero}; use poly::eq_table; use transcript::ProverState; +/// A reduction that produced no claim. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum ReduceError { + /// The fold's batch point has a zero coordinate, which the verifier's + /// GKR rounds divide by. A `2^-128` accident of the transcript, not + /// something a prover can arrange. + DegenerateChallenge, + /// The derived weight counts do not match the table shape. + Claim(ClaimError), +} + #[inline(never)] #[tracing::instrument(name = "Build grand-product circuit", level = "debug", skip_all)] fn init_circuit(table: &BitTable, fold: &Fold) -> GrandProductCircuit { @@ -56,7 +67,10 @@ pub fn gkr_reduce( transcript: &mut ProverState, fold: &Fold, table: &BitTable, -) -> Result { +) -> Result { + if fold.zeta.contains(&F128::ZERO) { + return Err(ReduceError::DegenerateChallenge); + } let circuit = init_circuit(table, fold); let (_last_value, witnesses) = circuit.batched_eval(table.shape().columns()); @@ -80,7 +94,8 @@ pub fn gkr_reduce( .collect(); let u2 = eq_table(&alfa_c); - let claim = LinearClaim::from_shape(table.shape(), u1, u2, inner_product_claim)?; + let claim = LinearClaim::from_shape(table.shape(), u1, u2, inner_product_claim) + .map_err(ReduceError::Claim)?; Ok(OpeningQuery::InnerProduct { claim }) } diff --git a/crates/verifier/src/reduce.rs b/crates/verifier/src/reduce.rs index eec8a599..31f5a574 100644 --- a/crates/verifier/src/reduce.rs +++ b/crates/verifier/src/reduce.rs @@ -9,13 +9,16 @@ use common::{ClaimError, Fold, LinearClaim, OpeningQuery, Shape}; use field::F128; -use num_traits::ConstOne; +use num_traits::{ConstOne, ConstZero}; use transcript::VerifierState; #[derive(Debug, Clone, PartialEq, Eq)] pub enum ReduceError { /// The GKR transcript does not satisfy a sumcheck or child-product relation. GKR, + /// The fold's batch point has a zero coordinate, which the GKR rounds + /// divide by. + DegenerateChallenge, /// The derived weight counts do not match the configured shape. Claim(ClaimError), } @@ -26,6 +29,9 @@ pub(crate) fn gkr_reduce( fold: &Fold, shape: &Shape, ) -> Result { + if fold.zeta.contains(&F128::ZERO) { + return Err(ReduceError::DegenerateChallenge); + } // Each layer halves the row count, leaving one product per column. let r1 = fold.row_images.len().max(1).ilog2();