From 892a3da2ed02396ae558bbcc7842fd3b2fc00bc4 Mon Sep 17 00:00:00 2001 From: Joaquin Carletti <56092489+ColoCarletti@users.noreply.github.com> Date: Fri, 18 Sep 2026 16:44:28 -0300 Subject: [PATCH] feat(multilinear): the sumcheck, and the two arguments built on it MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit One brick and two uses. The sumcheck reduces a sum over `2^n` points to a single evaluation: round by round the prover sends a one-variable polynomial, the verifier draws, and the cube halves. The form here does not send `g(0)` — it is recovered from the claim the round entered with, since `g(0) + g(1)` is that claim — which saves a whole pass over the cube and one element per round, and has a consequence worth stating plainly: **a round never rejects.** Verification recovers `g(0)`, absorbs, draws and interpolates; the only rejection a sumcheck has is its final claim, and discharging that is the caller's. On top of it, two arguments. The zerocheck says a polynomial vanishes on the whole cube, by weighting it with `eq(r, ·)` at the cost of one degree. LogUp-GKR says a bus balances, by reducing a tree of fractions layer by layer. Neither is a separate protocol: both are a weight and a claim handed to the same rounds. `claim_reduce` ties the shifted views a table reads back to the columns they read, so a shifted read is not a second commitment. `batch` is what makes that one pass instead of several: statements over one cube share their factors and their rounds, combined with powers of a challenge. The claims are absorbed **before** that challenge is drawn — which is the only thing stopping a prover from picking a statement after seeing it. `gpu` is the seam to the device, and this is the half that declines: every entry point returns `None`, so the whole path is host-only and can be read without a card. The half that does the work lands last. `zerocheck` is the reference and says so: a real table's constraints go in as a compiled `batch::Rule` beside its two bus claims, so all three share one pass over one factor list. Nothing outside tests calls it. It stays because the tests that do call it are `stark`'s, they run the production polynomial through it, and three of them are negative. What pinches it: 128 tests, in debug and parallel. The one to know about is `batch::the_claims_are_bound_to_the_batching_challenge`, which is not a test of a proof — pull the three claim absorptions out of both sides and every other test still passes, because the two sides still agree and the proof still verifies. What is lost is soundness, and the only thing that sees it is the transcript: two different claims have to give different challenges. Checked that it is the one test that fails. Not wired to anything. Nothing in the tree calls this crate yet. --- Cargo.lock | 12 + Cargo.toml | 1 + crypto/multilinear/Cargo.toml | 29 + crypto/multilinear/src/batch.rs | 762 ++++++++++++++++++++++++ crypto/multilinear/src/claim_reduce.rs | 647 ++++++++++++++++++++ crypto/multilinear/src/eq.rs | 541 +++++++++++++++++ crypto/multilinear/src/gkr.rs | 781 +++++++++++++++++++++++++ crypto/multilinear/src/gpu.rs | 549 +++++++++++++++++ crypto/multilinear/src/lib.rs | 120 ++++ crypto/multilinear/src/logup.rs | 551 +++++++++++++++++ crypto/multilinear/src/mle.rs | 417 +++++++++++++ crypto/multilinear/src/poly.rs | 428 ++++++++++++++ crypto/multilinear/src/program.rs | 567 ++++++++++++++++++ crypto/multilinear/src/selector.rs | 210 +++++++ crypto/multilinear/src/sumcheck.rs | 624 ++++++++++++++++++++ crypto/multilinear/src/zerocheck.rs | 276 +++++++++ 16 files changed, 6515 insertions(+) create mode 100644 crypto/multilinear/Cargo.toml create mode 100644 crypto/multilinear/src/batch.rs create mode 100644 crypto/multilinear/src/claim_reduce.rs create mode 100644 crypto/multilinear/src/eq.rs create mode 100644 crypto/multilinear/src/gkr.rs create mode 100644 crypto/multilinear/src/gpu.rs create mode 100644 crypto/multilinear/src/lib.rs create mode 100644 crypto/multilinear/src/logup.rs create mode 100644 crypto/multilinear/src/mle.rs create mode 100644 crypto/multilinear/src/poly.rs create mode 100644 crypto/multilinear/src/program.rs create mode 100644 crypto/multilinear/src/selector.rs create mode 100644 crypto/multilinear/src/sumcheck.rs create mode 100644 crypto/multilinear/src/zerocheck.rs diff --git a/Cargo.lock b/Cargo.lock index 5f46638a9..40c46e3e4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -933,6 +933,18 @@ dependencies = [ "libc", ] +[[package]] +name = "multilinear" +version = "0.1.0" +dependencies = [ + "crypto", + "math", + "rayon", + "rkyv", + "serde", + "thiserror", +] + [[package]] name = "munge" version = "0.4.7" diff --git a/Cargo.toml b/Cargo.toml index 8f9bbe7d3..26071b181 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -7,6 +7,7 @@ members = [ "crypto/ecsm", "crypto/math", "crypto/math-cuda", + "crypto/multilinear", "bin/cli", ] # Riscv-only bare-metal crate, path-dependent from crypto/crypto (target-gated), diff --git a/crypto/multilinear/Cargo.toml b/crypto/multilinear/Cargo.toml new file mode 100644 index 000000000..7237ebe8e --- /dev/null +++ b/crypto/multilinear/Cargo.toml @@ -0,0 +1,29 @@ +[package] +name = "multilinear" +description = "Multilinear extensions, sumcheck and zerocheck over the Boolean hypercube" +version = "0.1.0" +edition = "2024" +license.workspace = true + +[dependencies] +# Same feature set as `stark`: the proofs this crate produces have to serialize +# in the same format the univariate ones do. +math = { path = "../math", features = [ + "std", + "lambdaworks-serde-binary", + "rkyv", +] } +crypto = { path = "../crypto", features = ["std", "serde", "rkyv"] } +thiserror = "1.0.38" +serde = { version = "1.0", features = ["derive"] } +# pointer_width_64: proof-format pointer width — see prover/Cargo.toml. +rkyv = { version = "0.8.10", default-features = false, features = [ + "alloc", + "bytecheck", + "aligned", + "pointer_width_64", +] } +rayon = { version = "1.8.0", optional = true } + +[features] +parallel = ["dep:rayon", "math/parallel", "crypto/parallel"] diff --git a/crypto/multilinear/src/batch.rs b/crypto/multilinear/src/batch.rs new file mode 100644 index 000000000..c3a7240bc --- /dev/null +++ b/crypto/multilinear/src/batch.rs @@ -0,0 +1,762 @@ +//! Several sumcheck statements over one cube, proved in a single pass. +//! +//! A zerocheck is `Σ_x eq(r,x)·C(x) = 0`. A LogUp-GKR input-layer claim is +//! `Σ_x eq(z,x)·P(x) = p(z)`. Both have the shape `Σ_x weight(x)·poly(x) = +//! claimed` over the same trace, so batching them with powers of a challenge +//! costs one pass instead of one per argument. +//! +//! The factors are **one list, shared by every statement**: a column read by a +//! constraint and by a bus fingerprint is folded once, not twice. Each +//! statement is a rule that indexes that list, so a statement's weight table is +//! just another factor it multiplies in. + +use crypto::fiat_shamir::is_transcript::IsTranscript; +use math::field::{element::FieldElement, traits::IsField}; + +use crate::{ + Error, challenge_powers, + mle::Mle, + poly::SumcheckPolynomial, + program::{self, Program}, + sumcheck::{self, SumcheckProof}, +}; + +/// One statement's rule: what it makes of the batch's factors. +/// +/// Reads the **whole** factor list, so a sub-argument that already combines a +/// prefix of it can be used unchanged. +pub struct Rule<'a, F: IsField> { + body: Body<'a, F>, + degree: usize, +} + +/// A rule is either a closure or the program it is. A statement that can say +/// which program it is gets the sumcheck's device path; the closure form is for +/// the ones that cannot. +enum Body<'a, F: IsField> { + Closure(RuleFn<'a, F>), + Compiled(Program), +} + +/// A statement's value from the batch's factor values. +type RuleFn<'a, F> = Box]) -> FieldElement + Sync + 'a>; + +impl<'a, F: IsField> Rule<'a, F> { + /// `degree` must upper-bound the rule's total degree in the factors, its + /// weight table included. + pub fn new( + degree: usize, + eval: impl Fn(&[FieldElement]) -> FieldElement + Sync + 'a, + ) -> Self { + Self { + body: Body::Closure(Box::new(eval)), + degree, + } + } + + /// The same rule as straight-line code over the factor values. + pub fn compiled(degree: usize, program: Program) -> Self { + Self { + body: Body::Compiled(program), + degree, + } + } + + pub fn degree(&self) -> usize { + self.degree + } + + /// The program this rule is, when it has one. + pub fn program(&self) -> Option<&Program> { + match &self.body { + Body::Compiled(program) => Some(program), + Body::Closure(_) => None, + } + } + + pub fn apply(&self, values: &[FieldElement]) -> FieldElement { + self.apply_in(values, &mut Vec::new()) + } + + /// The same, reusing the caller's scratch for a compiled rule's steps. + pub fn apply_in( + &self, + values: &[FieldElement], + scratch: &mut Vec>, + ) -> FieldElement { + match &self.body { + Body::Closure(eval) => eval(values), + Body::Compiled(program) => program.eval(values, scratch), + } + } +} + +/// The batch as one polynomial: `Σ_i lambda^i · rule_i`. +pub struct Batched<'a, F: IsField> { + polys: Vec>, + rules: Vec>, + lambdas: Vec>, + num_vars: usize, + degree: usize, + /// The whole batch as one program, when every rule is compiled. The round + /// loop runs this instead of the rules, and it is what a device gets. + program: Option>, + /// Whether the rounds may still go to a device. False once one has handed + /// its factors back: what it stopped for is that the cube got too small to + /// be worth sending anywhere. + dispatch: bool, +} + +impl<'a, F: IsField + 'static> Batched<'a, F> { + /// `lambdas` weights the statements; there must be one per rule. + pub fn new( + polys: Vec>, + rules: Vec>, + lambdas: Vec>, + ) -> Result { + if rules.is_empty() { + return Err(Error::EmptyPolynomial); + } + if lambdas.len() != rules.len() { + return Err(Error::VariableCountMismatch { + expected: rules.len(), + got: lambdas.len(), + }); + } + let num_vars = polys.first().map(|p| p.num_vars()).unwrap_or(0); + for p in &polys { + if p.num_vars() != num_vars { + return Err(Error::VariableCountMismatch { + expected: num_vars, + got: p.num_vars(), + }); + } + } + let degree = rules.iter().map(Rule::degree).max().unwrap_or(0); + let programs: Option>> = rules.iter().map(Rule::program).collect(); + let program = match programs { + Some(programs) => Some(program::combine(&programs, &lambdas)?), + None => None, + }; + Ok(Self { + polys, + rules, + lambdas, + num_vars, + degree, + program, + dispatch: true, + }) + } + + /// Puts `polys` in front of the ones already here. + /// + /// A batch whose first factors a device holds is built over the rest + /// alone; this is what brings them back when the device turns the rounds + /// down and the host has to run them. + pub fn prepend(&mut self, mut polys: Vec>) -> Result<(), Error> { + if polys.is_empty() { + return Ok(()); + } + let num_vars = polys[0].num_vars(); + for p in polys.iter().chain(&self.polys) { + if p.num_vars() != num_vars { + return Err(Error::VariableCountMismatch { + expected: num_vars, + got: p.num_vars(), + }); + } + } + polys.append(&mut self.polys); + self.polys = polys; + self.num_vars = num_vars; + Ok(()) + } + + /// The batch's factors, replaced whole by the ones a device folded. + /// + /// A batch whose leading factors a device holds is built over the rest + /// alone; when the device stops part-way it hands back **every** factor, + /// resident or not, as its rounds left them — so this replaces the list + /// rather than adding to it. The rounds that follow stay here: the cube + /// the device stopped at is the one it was no longer worth sending. + pub fn adopt(&mut self, polys: Vec>) -> Result<(), Error> { + let num_vars = polys.first().map(Mle::num_vars).unwrap_or(0); + for p in &polys { + if p.num_vars() != num_vars { + return Err(Error::VariableCountMismatch { + expected: num_vars, + got: p.num_vars(), + }); + } + } + self.polys = polys; + self.num_vars = num_vars; + self.dispatch = false; + Ok(()) + } + + /// The batch as one program, when it has one. + pub fn program(&self) -> Option<&Program> { + self.program.as_ref() + } +} + +impl SumcheckPolynomial for Batched<'_, F> { + fn num_vars(&self) -> usize { + self.num_vars + } + + fn degree(&self) -> usize { + self.degree + } + + fn polys(&self) -> &[Mle] { + &self.polys + } + + fn combine(&self, values: &[FieldElement]) -> FieldElement { + self.combine_in(values, &mut Vec::new()) + } + + fn combine_in( + &self, + values: &[FieldElement], + scratch: &mut Vec>, + ) -> FieldElement { + if let Some(program) = &self.program { + return program.eval(values, scratch); + } + self.rules + .iter() + .zip(&self.lambdas) + .fold(FieldElement::zero(), |acc, (rule, lambda)| { + acc + lambda * rule.apply_in(values, scratch) + }) + } + + fn fix_first_variable(&mut self, r: &FieldElement) -> Result<(), Error> { + for p in &mut self.polys { + p.fix_first_variable_in_place(r)?; + } + self.num_vars -= 1; + Ok(()) + } + + fn program(&self) -> Option<&Program> { + self.program.as_ref().filter(|_| self.dispatch) + } + + fn accept_folded(&mut self, polys: Vec>) -> Result<(), Error> { + if polys.len() != self.polys.len() { + return Err(Error::VariableCountMismatch { + expected: self.polys.len(), + got: polys.len(), + }); + } + self.num_vars = polys.first().map(Mle::num_vars).unwrap_or(0); + self.polys = polys; + Ok(()) + } +} + +/// The degree the batched sumcheck runs at: the worst statement's. +pub fn degree_of(rules: &[Rule<'_, F>]) -> usize { + rules.iter().map(Rule::degree).max().unwrap_or(0) +} + +/// Proves every statement in one sumcheck, returning the proof and its point. +/// +/// `claims[i]` is what `Σ_x rule_i(x)` must come to. They are absorbed before +/// the batching challenge, so the prover cannot pick a statement after seeing +/// it. +pub fn prove( + polys: Vec>, + rules: Vec>, + claims: &[FieldElement], + transcript: &mut T, +) -> Result<(SumcheckProof, Vec>), Error> +where + F: IsField + 'static, + T: IsTranscript, +{ + if claims.len() != rules.len() { + return Err(Error::VariableCountMismatch { + expected: rules.len(), + got: claims.len(), + }); + } + for claim in claims { + transcript.append_field_element(claim); + } + let lambdas = challenge_powers(&transcript.sample_field_element(), rules.len()); + sumcheck::prove(Batched::new(polys, rules, lambdas)?, transcript) +} + +/// A batched sumcheck's proof, the point its rounds drew, and what every +/// factor slot was bound to there — empty when the factors were the host's own +/// and folded away as they went. +type ResidentProof = (SumcheckProof, Vec>, Vec>); + +/// The same, over factors a device already holds — the batch's first ones, in +/// the order the rules read them. +/// +/// Returns the proof, the point the rounds drew, and **what every slot was +/// bound to there**, in slot order. That last one is free on the device path — +/// the rounds fold the factors where they lie, so the values are already up +/// there — and empty when the host ran the rounds, which drops each factor as +/// it folds. +/// +/// `extra` is what the device does not have (the weight tables), and `absent` +/// makes what it does. That second one is a builder and not a value because on +/// the device path it is never called: building a table's factors here is most +/// of what the argument used to spend on the host, and the whole point of them +/// being up there is not to. +/// +/// It *is* called when the device turns the rounds down. Nothing has been said +/// about the factors to the transcript by then — only the claims and the +/// batching challenge, which are the same either way — so the host path +/// continues from where the device left off and produces the same proof. +pub fn prove_resident( + extra: Vec>, + device: Option>, + absent: B, + rules: Vec>, + claims: &[FieldElement], + transcript: &mut T, +) -> Result, Error> +where + F: IsField + 'static, + T: IsTranscript, + B: Fn() -> Result>, Error>, +{ + if claims.len() != rules.len() { + return Err(Error::VariableCountMismatch { + expected: rules.len(), + got: claims.len(), + }); + } + let Some(device) = device else { + let mut polys = absent()?; + polys.extend(extra); + let (proof, point) = prove(polys, rules, claims, transcript)?; + return Ok((proof, point, Vec::new())); + }; + + for claim in claims { + transcript.append_field_element(claim); + } + let lambdas = challenge_powers(&transcript.sample_field_element(), rules.len()); + let mut batched = Batched::new(extra, rules, lambdas)?; + let degree = batched.degree().max(1); + if let Some(program) = batched.program() { + let attempt = crate::gpu::prove_sumcheck_resident( + &device, + batched.polys(), + program, + degree, + |evaluations| { + for e in evaluations { + transcript.append_field_element(e); + } + transcript.sample_field_element() + }, + ); + if let Some(outcome) = attempt { + let (mut rounds, mut point, factors) = outcome?; + // The device stopped where the cube stopped being worth sending; + // the rest of the rounds run over the factors it folded. + batched.adopt(factors)?; + let left = batched.num_vars(); + let (tail, tail_point) = sumcheck::prove_rounds(&mut batched, left, transcript)?; + rounds.extend(tail); + point.extend(tail_point); + // Binding every variable is evaluating at the point, so the factor + // values the caller needs are the factors themselves by now. + let bound: Option>> = batched + .polys() + .iter() + .map(|factor| factor.as_constant().cloned()) + .collect(); + let bound = bound.ok_or(Error::NoVariablesLeft)?; + return Ok((SumcheckProof { rounds }, point, bound)); + } + } + // Declined before the first round: the host runs them, and for that the + // factors have to be here after all. + batched.prepend(absent()?)?; + let (proof, point) = sumcheck::prove(batched, transcript)?; + Ok((proof, point, Vec::new())) +} + +/// Checks the batched sumcheck against the factor values it reduces to. +/// +/// `values_at` is handed the sumcheck point and returns every factor's value +/// there. A weight table is not committed, so the verifier computes it in +/// closed form; the trace factors come from the proof, and binding *those* to +/// the committed columns is the caller's next step. +/// +/// Returns the point. +pub fn verify( + proof: &SumcheckProof, + rules: &[Rule<'_, F>], + claims: &[FieldElement], + values_at: V, + num_vars: usize, + transcript: &mut T, +) -> Result>, Error> +where + F: IsField + 'static, + T: IsTranscript, + V: FnOnce(&[FieldElement]) -> Result>, Error>, +{ + if rules.is_empty() { + return Err(Error::EmptyPolynomial); + } + if claims.len() != rules.len() { + return Err(Error::VariableCountMismatch { + expected: rules.len(), + got: claims.len(), + }); + } + for claim in claims { + transcript.append_field_element(claim); + } + let lambdas = challenge_powers(&transcript.sample_field_element(), rules.len()); + + let claimed = lambdas + .iter() + .zip(claims) + .fold(FieldElement::::zero(), |acc, (l, c)| acc + l * c); + let claim = sumcheck::verify(proof, claimed, num_vars, degree_of(rules), transcript)?; + + let values = values_at(&claim.point)?; + let rebuilt = rules + .iter() + .zip(&lambdas) + .fold(FieldElement::::zero(), |acc, (rule, lambda)| { + acc + lambda * rule.apply(&values) + }); + if rebuilt != claim.expected_evaluation { + return Err(Error::BatchMismatch); + } + + Ok(claim.point) +} + +#[cfg(test)] +mod tests { + use super::*; + use crypto::fiat_shamir::default_transcript::DefaultTranscript; + use math::field::goldilocks::GoldilocksField as F; + + use crate::{ + eq::{eq_eval, eq_mle}, + gkr::{self, FractionLayer, FractionTree}, + }; + + type FE = FieldElement; + + fn transcript() -> DefaultTranscript { + DefaultTranscript::::new(b"batch-test") + } + + /// Factor layout shared by every statement below. + const MULT: usize = 0; + const VALUE: usize = 1; + const SQUARE: usize = 2; + const EQ_R: usize = 3; + const EQ_Z: usize = 4; + + /// A bus and a constraint over the same two columns: `mult` is the signed + /// multiplicity, `value` the fingerprint, `square` its square. Each + /// fingerprint is sent once and received once, so the bus balances. + fn columns(num_vars: usize) -> [Mle; 3] { + let size = 1usize << num_vars; + let half = size / 2; + let value: Vec = (0..size).map(|i| FE::from((i % half) as u64 + 1)).collect(); + let mult: Vec = (0..size) + .map(|i| if i < half { FE::one() } else { -FE::one() }) + .collect(); + let square: Vec = value.iter().map(|x| x * x).collect(); + [ + Mle::new(mult).unwrap(), + Mle::new(value).unwrap(), + Mle::new(square).unwrap(), + ] + } + + /// The three statements, all indexing the one factor list: + /// + /// - the zerocheck, `Σ_x eq(r,x)·(value² − square) = 0`; + /// - the bus's numerator claim, `Σ_x eq(z,x)·mult = p(z)`; + /// - its denominator claim, `Σ_x eq(z,x)·(alpha − value) = q(z)`. + /// + /// The denominator is never a column: it is `alpha − value`, read off the + /// committed `value`. + fn statements(alpha: FE) -> Vec> { + vec![ + Rule::new(3, |v: &[FE]| v[EQ_R] * (v[VALUE] * v[VALUE] - v[SQUARE])), + Rule::new(2, |v: &[FE]| v[EQ_Z] * v[MULT]), + Rule::new(2, move |v: &[FE]| v[EQ_Z] * (alpha - v[VALUE])), + ] + } + + /// Runs GKR over the bus, then settles its input-layer claim in the same + /// sumcheck as the constraint's zerocheck. + /// + /// `batch_mult` overrides the multiplicity column the *batch* reads, while + /// the tree keeps the original — a prover arguing about a different table + /// than the one it ran GKR over. + fn fuse( + cols: [Mle; 3], + num_vars: usize, + batch_mult: Option>, + ) -> Result { + let alpha = FE::from(97); + let [mult, value, square] = cols; + let batch_mult = batch_mult.unwrap_or_else(|| mult.clone()); + let denominator = Mle::new(value.evals().iter().map(|x| alpha - x).collect())?; + let tree = FractionTree::build(FractionLayer::new(mult.clone(), denominator)?)?; + let output = tree.output(); + + let mut prover = transcript(); + let gkr_out = gkr::prove(&tree, &mut prover)?; + // The zerocheck challenge, drawn from the same transcript. + let r: Vec = (0..num_vars) + .map(|_| prover.sample_field_element()) + .collect(); + let z = gkr_out.claim.point.clone(); + + let factors = || -> Result>, Error> { + Ok(vec![ + batch_mult.clone(), + value.clone(), + square.clone(), + eq_mle(&r)?, + eq_mle(&z)?, + ]) + }; + let claims = [FE::zero(), gkr_out.claim.p, gkr_out.claim.q]; + + let (proof, point) = prove(factors()?, statements(alpha), &claims, &mut prover)?; + let trace_values: Vec = factors()?[..=SQUARE] + .iter() + .map(|f| f.evaluate(&point)) + .collect::>()?; + + let mut verifier = transcript(); + let gkr_claim = gkr::verify(&gkr_out.proof, output, &mut verifier)?; + let vr: Vec = (0..num_vars) + .map(|_| verifier.sample_field_element()) + .collect(); + + let checked = verify( + &proof, + &statements(alpha), + &claims, + // The weights are not committed: the verifier computes them. + |at: &[FE]| { + let mut values = trace_values.clone(); + values.push(eq_eval(&vr, at)?); + values.push(eq_eval(&gkr_claim.point, at)?); + Ok(values) + }, + num_vars, + &mut verifier, + )?; + assert_eq!(checked, point); + Ok(proof.rounds.len()) + } + + /// The design decision this module exists for: the constraint's zerocheck + /// and the bus's input-layer claim share one sumcheck, so every column is + /// folded once. + /// **The claims are bound to the batching challenge.** + /// + /// Absorbing them before drawing `lambda` is what stops a prover from + /// picking a statement after seeing it. **No proof can show that**: take + /// the absorption out of both sides and everything still verifies, because + /// both sides stay in step — what is gone is the soundness, not the + /// agreement. The transcript is what shows it, so that is what this tests: + /// two batches alike but for a claim must not draw the same challenge, and + /// a proof drawn from a different challenge is a different proof. + #[test] + fn the_claims_are_bound_to_the_batching_challenge() { + let polys = columns(4).to_vec(); + let alpha = FE::from(97); + let eq_r = eq_mle(&[FE::from(3), FE::from(5), FE::from(7), FE::from(11)]).unwrap(); + let eq_z = eq_mle(&[FE::from(13), FE::from(17), FE::from(19), FE::from(23)]).unwrap(); + let all = || { + let mut v = polys.clone(); + v.push(eq_r.clone()); + v.push(eq_z.clone()); + v + }; + let claims = [FE::zero(), FE::from(41), FE::from(43)]; + let (first, _) = prove(all(), statements(alpha), &claims, &mut transcript()).unwrap(); + + // One claim moved, nothing else. + let moved = [FE::zero(), FE::from(41) + FE::one(), FE::from(43)]; + let (second, _) = prove(all(), statements(alpha), &moved, &mut transcript()).unwrap(); + + let rounds = |p: &SumcheckProof| { + p.rounds + .iter() + .map(|r| r.evaluations.clone()) + .collect::>() + }; + assert_ne!( + rounds(&first), + rounds(&second), + "the claims are not bound to the challenge: a prover could pick one \ + after seeing it" + ); + } + + #[test] + fn a_bus_claim_settles_in_the_constraint_s_sumcheck() { + for num_vars in 1..=4usize { + let rounds = fuse(columns(num_vars), num_vars, None) + .unwrap_or_else(|e| panic!("num_vars={num_vars}: {e:?}")); + // One pass over the cube for all three statements, not one each. + assert_eq!(rounds, num_vars); + } + } + + #[test] + fn a_constraint_broken_in_one_row_is_rejected() { + let mut cols = columns(3); + let mut square = cols[2].evals().to_vec(); + square[5] += FE::one(); + cols[2] = Mle::new(square).unwrap(); + + assert!(fuse(cols, 3, None).is_err()); + } + + /// The GKR's input-layer claim is only worth anything if it lands on the + /// same table the rest of the argument reads. + #[test] + fn a_bus_table_the_tree_was_not_built_on_is_rejected() { + let cols = columns(3); + let mut mult = cols[0].evals().to_vec(); + mult[2] += FE::one(); + + assert!(fuse(cols, 3, Some(Mle::new(mult).unwrap())).is_err()); + } + + // --------------------------------------------------------------- + // The batching itself. + // --------------------------------------------------------------- + + fn mle(vals: &[u64]) -> Mle { + Mle::new(vals.iter().map(|v| FE::from(*v)).collect()).unwrap() + } + + /// Two statements over one factor list: `Σ eq(r,x)·a(x) = ã(r)` and + /// `Σ eq(r,x)·b(x) = b̃(r)`, sharing the weight table. + fn two_evaluation_claims(r: &[FE]) -> (Vec>, Vec>, [FE; 2]) { + let a = mle(&[3, 5, 8, 13]); + let b = mle(&[21, 34, 55, 89]); + let claims = [a.evaluate(r).unwrap(), b.evaluate(r).unwrap()]; + let polys = vec![a, b, eq_mle(r).unwrap()]; + let rules = vec![ + Rule::new(2, |v: &[FE]| v[2] * v[0]), + Rule::new(2, |v: &[FE]| v[2] * v[1]), + ]; + (polys, rules, claims) + } + + fn run_two(r: &[FE], claims: [FE; 2]) -> Result<(), Error> { + let (polys, rules, _) = two_evaluation_claims(r); + let values = polys.clone(); + let (proof, _) = prove(polys, rules, &claims, &mut transcript())?; + + let (_, rules, _) = two_evaluation_claims(r); + verify( + &proof, + &rules, + &claims, + |at: &[FE]| values.iter().map(|p| p.evaluate(at)).collect(), + r.len(), + &mut transcript(), + ) + .map(|_| ()) + } + + #[test] + fn statements_sharing_a_weight_table_hold_it_once() { + let r = [FE::from(11), FE::from(13)]; + let (polys, _, claims) = two_evaluation_claims(&r); + // Three factors for two statements: the weight is shared. + assert_eq!(polys.len(), 3); + run_two(&r, claims).unwrap(); + } + + #[test] + fn a_false_claim_in_one_statement_is_rejected() { + let r = [FE::from(11), FE::from(13)]; + let (_, _, claims) = two_evaluation_claims(&r); + let lying = [claims[0], claims[1] + FE::one()]; + assert!(run_two(&r, lying).is_err()); + } + + #[test] + fn the_batch_degree_is_the_worst_statement_s() { + let rules: Vec> = vec![ + Rule::new(2, |v: &[FE]| v[0]), + Rule::new(5, |v: &[FE]| v[0]), + Rule::new(3, |v: &[FE]| v[0]), + ]; + assert_eq!(degree_of(&rules), 5); + } + + #[test] + fn a_claim_per_statement_is_required() { + let r = [FE::from(11), FE::from(13)]; + let (polys, rules, _) = two_evaluation_claims(&r); + assert_eq!( + prove(polys, rules, &[FE::zero()], &mut transcript()).unwrap_err(), + Error::VariableCountMismatch { + expected: 2, + got: 1 + } + ); + } + + #[test] + fn factors_of_differing_heights_are_rejected() { + let rules: Vec> = vec![Rule::new(1, |v: &[FE]| v[0])]; + let result = Batched::new( + vec![mle(&[1, 2]), mle(&[1, 2, 3, 4])], + rules, + vec![FE::one()], + ); + assert!(matches!( + result.err(), + Some(Error::VariableCountMismatch { .. }) + )); + } + + #[test] + fn a_proof_replayed_under_another_transcript_is_rejected() { + let r = [FE::from(11), FE::from(13)]; + let (polys, rules, claims) = two_evaluation_claims(&r); + let values = polys.clone(); + let (proof, _) = prove(polys, rules, &claims, &mut transcript()).unwrap(); + + let (_, rules, _) = two_evaluation_claims(&r); + let mut other = DefaultTranscript::::new(b"a-different-statement"); + assert!( + verify( + &proof, + &rules, + &claims, + |at: &[FE]| values.iter().map(|p| p.evaluate(at)).collect(), + r.len(), + &mut other, + ) + .is_err() + ); + } +} diff --git a/crypto/multilinear/src/claim_reduce.rs b/crypto/multilinear/src/claim_reduce.rs new file mode 100644 index 000000000..1713ca665 --- /dev/null +++ b/crypto/multilinear/src/claim_reduce.rs @@ -0,0 +1,647 @@ +//! Binding a shifted read to the column it shifts. +//! +//! A constraint that reads the next step becomes a sumcheck factor that is the +//! cyclic shift of a committed column. The zerocheck does not care where a +//! factor came from, so a prover free to pick both the column and its "shift" +//! could satisfy a constraint with unrelated tables. The shift kernel closes +//! that: `f_shift_k(alpha) = Σ_y shift_k(alpha, y)·f(y)`, which is a claim about +//! the column itself. +//! +//! Every factor's claim is batched into one degree-2 sumcheck, so the whole +//! trace costs a single pass and leaves **one evaluation claim per committed +//! column, all at the same point** — one commitment opening each. +//! +//! The guarantee is conditional, and the caller supplies the other half: *if* +//! the column values at the reduced point are the committed columns' true +//! values, then the factor values were the true shifted values at `alpha`. +//! Pinning those is the commitment scheme's job. +//! +//! The columns are base-field, like the trace; the challenges and the claims +//! are not, so the batching reads them through mixed products. + +use crypto::fiat_shamir::is_transcript::IsTranscript; +use math::field::{ + element::FieldElement, + traits::{IsField, IsSubFieldOf}, +}; + +#[cfg(feature = "parallel")] +use rayon::prelude::*; + +use crate::{ + Error, challenge_powers, + eq::{shift_eval, shift_mle}, + mle::Mle, + poly::Composed, + program::{Builder, Program}, + sumcheck::{self, SumcheckProof}, +}; + +/// Which committed column a sumcheck factor reads, and how many steps ahead. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct FactorSource { + pub column: usize, + /// Frame-step offset; on the cube it is a cyclic shift by that many rows. + pub offset: usize, +} + +impl FactorSource { + /// The column itself. + pub const fn direct(column: usize) -> Self { + Self { column, offset: 0 } + } + + pub const fn shifted(column: usize, offset: usize) -> Self { + Self { column, offset } + } +} + +/// The batched reduction. +#[derive( + Clone, + Debug, + serde::Serialize, + serde::Deserialize, + rkyv::Archive, + rkyv::Serialize, + rkyv::Deserialize, +)] +#[serde(bound = "")] +pub struct ReduceProof { + pub sumcheck: SumcheckProof, + /// Every committed column's value at the reduced point. + pub column_values: Vec>, +} + +/// One evaluation claim per committed column, all at the same point. +#[derive(Clone, Debug)] +pub struct ReducedClaim { + pub point: Vec>, + pub column_values: Vec>, +} + +/// Distinct offsets in ascending order: the grouping both sides must agree on. +fn offsets(sources: &[FactorSource]) -> Vec { + let mut all: Vec = sources.iter().map(|s| s.offset).collect(); + all.sort_unstable(); + all.dedup(); + all +} + +/// `Σ_o K_o(y)·B_o(y)`, the factors laid out in kernel/column pairs. +fn pair_products(values: &[FieldElement]) -> FieldElement { + values + .chunks(2) + .fold(FieldElement::zero(), |acc, pair| acc + &pair[0] * &pair[1]) +} + +/// The same sum as straight-line code, over `pairs` kernel/column pairs. +fn pair_products_program(pairs: usize) -> Result, Error> { + let mut b = Builder::::new(); + let terms: Vec = (0..pairs) + .map(|pair| { + let kernel = b.var(2 * pair); + let column = b.var(2 * pair + 1); + b.mul(kernel, column) + }) + .collect(); + let root = b.sum(&terms); + b.finish(root) +} + +fn check_shape( + sources: &[FactorSource], + factor_values: &[FieldElement], + num_columns: usize, +) -> Result<(), Error> { + if sources.is_empty() { + return Err(Error::EmptyPolynomial); + } + if sources.len() != factor_values.len() { + return Err(Error::VariableCountMismatch { + expected: sources.len(), + got: factor_values.len(), + }); + } + if let Some(bad) = sources.iter().find(|s| s.column >= num_columns) { + return Err(Error::UnknownPolynomial { + index: bad.column, + len: num_columns, + }); + } + Ok(()) +} + +/// Accumulator cells a worker takes at a time. Small enough that a short table +/// still spreads, large enough that the columns streaming through one slab pay +/// for the handoff. +const ACCUMULATOR_CHUNK: usize = 1 << 12; + +/// `Σ_{i reads at `offset`} gamma^i · column_i`, the one polynomial that offset's +/// kernel multiplies. +fn batched_column( + columns: &[Mle], + sources: &[FactorSource], + weights: &[FieldElement], + offset: usize, + num_vars: usize, +) -> Result, Error> +where + F: IsField + IsSubFieldOf + 'static, + E: IsField + 'static, +{ + // Which columns read at this offset, so the walk below carries its whole + // list and the accumulator is written once rather than once per column. + let members: Vec<(usize, usize)> = sources + .iter() + .enumerate() + .filter(|(_, source)| source.offset == offset) + .map(|(i, source)| (i, source.column)) + .collect(); + let mut acc = vec![FieldElement::::zero(); 1usize << num_vars]; + // A slab of the accumulator, with every column streamed through it: a + // table brings dozens of columns, and this way its cells are touched once. + let fill = |(index, slab): (usize, &mut [FieldElement])| { + let at = index * ACCUMULATOR_CHUNK; + for (weight, column) in &members { + let values = &columns[*column].evals()[at..at + slab.len()]; + for (slot, value) in slab.iter_mut().zip(values) { + // The base element on the left: the only direction the tower + // gives. + *slot += value * &weights[*weight]; + } + } + }; + #[cfg(feature = "parallel")] + acc.par_chunks_mut(ACCUMULATOR_CHUNK) + .enumerate() + .for_each(fill); + #[cfg(not(feature = "parallel"))] + acc.chunks_mut(ACCUMULATOR_CHUNK).enumerate().for_each(fill); + Mle::new(acc) +} + +/// Reduces every factor's claimed value at `alpha` to one claim per column. +/// +/// Returns the proof and the reduced point. `factor_values[i]` must be the +/// value of the factor `sources[i]` describes — the column shifted by its +/// offset, evaluated at `alpha`. +pub fn prove( + columns: &[Mle], + sources: &[FactorSource], + factor_values: &[FieldElement], + alpha: &[FieldElement], + resident: Option<(&crate::gpu::ResidentColumns, usize)>, + transcript: &mut T, +) -> Result<(ReduceProof, Vec>), Error> +where + F: IsField + IsSubFieldOf + 'static, + E: IsField + 'static, + T: IsTranscript, +{ + check_shape(sources, factor_values, columns.len())?; + for column in columns { + if column.num_vars() != alpha.len() { + return Err(Error::VariableCountMismatch { + expected: alpha.len(), + got: column.num_vars(), + }); + } + } + + // The claims are what is being reduced, so the batching challenge must come + // after them. + for value in factor_values { + transcript.append_field_element(value); + } + let weights = challenge_powers(&transcript.sample_field_element(), sources.len()); + + let mut polys = Vec::with_capacity(2 * offsets(sources).len()); + for offset in offsets(sources) { + polys.push(shift_mle(alpha, offset)?); + polys.push(batched_column::( + columns, + sources, + &weights, + offset, + alpha.len(), + )?); + } + + let pairs = polys.len() / 2; + let (sumcheck, point) = sumcheck::prove( + Composed::new(polys, pair_products::, 2)? + .with_program(pair_products_program::(pairs)?), + transcript, + )?; + + // All at the same point, so they fold together: one upload and one launch + // per level for the table instead of per column. + let column_values = match crate::gpu::evaluate_many_base(columns, &point, resident) { + Some(values) => values, + None => columns + .iter() + .map(|c| c.evaluate_in(&point)) + .collect::, _>>()?, + }; + for value in &column_values { + transcript.append_field_element(value); + } + + Ok(( + ReduceProof { + sumcheck, + column_values, + }, + point, + )) +} + +/// Checks the reduction and returns the claims the commitment scheme must +/// settle. +pub fn verify( + proof: &ReduceProof, + sources: &[FactorSource], + factor_values: &[FieldElement], + alpha: &[FieldElement], + num_columns: usize, + transcript: &mut T, +) -> Result, Error> +where + E: IsField + 'static, + T: IsTranscript, +{ + check_shape(sources, factor_values, num_columns)?; + if proof.column_values.len() != num_columns { + return Err(Error::QueryCountMismatch { + expected: num_columns, + got: proof.column_values.len(), + }); + } + + for value in factor_values { + transcript.append_field_element(value); + } + let weights = challenge_powers(&transcript.sample_field_element(), sources.len()); + + let claimed = weights + .iter() + .zip(factor_values) + .fold(FieldElement::::zero(), |acc, (w, v)| acc + w * v); + let claim = sumcheck::verify(&proof.sumcheck, claimed, alpha.len(), 2, transcript)?; + + // The kernels are closed forms, so the residual is entirely about the + // columns — and the columns are what the commitments answer for. + let mut rebuilt = FieldElement::::zero(); + for offset in offsets(sources) { + let batched = + sources + .iter() + .enumerate() + .fold(FieldElement::::zero(), |acc, (i, source)| { + if source.offset == offset { + acc + &weights[i] * &proof.column_values[source.column] + } else { + acc + } + }); + rebuilt += shift_eval(alpha, &claim.point, offset)? * batched; + } + if rebuilt != claim.expected_evaluation { + return Err(Error::ShiftedReadMismatch); + } + + for value in &proof.column_values { + transcript.append_field_element(value); + } + + Ok(ReducedClaim { + point: claim.point, + column_values: proof.column_values.clone(), + }) +} + +/// The factor `source` describes, materialized: the column shifted cyclically +/// by its offset. +/// A factor's value at `point`, without building the view when there is none +/// to build. +/// +/// An unshifted source *is* its column, and a column is the biggest thing the +/// argument holds: copying one to read a single value out of it is the whole +/// trace copied again, once per factor. +pub fn evaluate_source( + columns: &[Mle], + source: &FactorSource, + point: &[FieldElement], +) -> Result, Error> +where + F: IsField + IsSubFieldOf + 'static, + E: IsField + 'static, +{ + let column = columns.get(source.column).ok_or(Error::UnknownPolynomial { + index: source.column, + len: columns.len(), + })?; + if source.offset.is_multiple_of(column.len()) { + return column.evaluate_in(point); + } + materialize(columns, source)?.evaluate_in(point) +} + +pub fn materialize( + columns: &[Mle], + source: &FactorSource, +) -> Result, Error> { + let column = columns.get(source.column).ok_or(Error::UnknownPolynomial { + index: source.column, + len: columns.len(), + })?; + let size = column.len(); + let shift = source.offset % size; + if shift == 0 { + return Ok(column.clone()); + } + Mle::new( + (0..size) + .map(|i| column.evals()[(i + shift) % size].clone()) + .collect(), + ) +} + +#[cfg(test)] +mod tests { + use super::*; + use crypto::fiat_shamir::default_transcript::DefaultTranscript; + use math::field::goldilocks::GoldilocksField as F; + + type FE = FieldElement; + + fn transcript() -> DefaultTranscript { + DefaultTranscript::::new(b"claim-reduce-test") + } + + fn column(num_vars: usize, seed: u64) -> Mle { + let vals: Vec = (0..(1u64 << num_vars)) + .map(|i| FE::from(i.wrapping_mul(seed).wrapping_add(seed * 7 + 1))) + .collect(); + Mle::new(vals).unwrap() + } + + fn point(num_vars: usize) -> Vec { + (0..num_vars).map(|i| FE::from(101 + i as u64)).collect() + } + + /// What an honest prover claims: each factor is its column shifted by its + /// offset, evaluated at `alpha`. + fn honest_values(columns: &[Mle], sources: &[FactorSource], alpha: &[FE]) -> Vec { + sources + .iter() + .map(|s| materialize(columns, s).unwrap().evaluate(alpha).unwrap()) + .collect() + } + + /// Proves honestly, then verifies whatever `factor_values` the caller wants + /// to present. + fn run( + columns: &[Mle], + sources: &[FactorSource], + prover_values: &[FE], + verifier_values: &[FE], + alpha: &[FE], + ) -> Result, Error> { + let (proof, _) = prove( + columns, + sources, + prover_values, + alpha, + None, + &mut transcript(), + )?; + verify( + &proof, + sources, + verifier_values, + alpha, + columns.len(), + &mut transcript(), + ) + } + + #[test] + fn honest_claims_reduce_to_the_columns() { + for num_vars in 1..=4usize { + let columns = vec![column(num_vars, 3), column(num_vars, 5)]; + let sources = [ + FactorSource::direct(0), + FactorSource::shifted(0, 1), + FactorSource::direct(1), + ]; + let alpha = point(num_vars); + let values = honest_values(&columns, &sources, &alpha); + + let claim = run(&columns, &sources, &values, &values, &alpha) + .unwrap_or_else(|e| panic!("num_vars={num_vars}: {e:?}")); + + // The claims handed on must be the columns' real values there. + for (c, value) in columns.iter().zip(&claim.column_values) { + assert_eq!(&c.evaluate(&claim.point).unwrap(), value); + } + } + } + + /// The reason this module exists: a "next step" factor that is not the shift + /// of the column it names cannot be passed off as one. + #[test] + fn an_unrelated_table_cannot_pass_as_a_shift() { + let num_vars = 3; + let columns = vec![column(num_vars, 3)]; + let sources = [FactorSource::direct(0), FactorSource::shifted(0, 1)]; + let alpha = point(num_vars); + + let mut lying = honest_values(&columns, &sources, &alpha); + lying[1] = column(num_vars, 11).evaluate(&alpha).unwrap(); + assert_ne!(lying[1], honest_values(&columns, &sources, &alpha)[1]); + + assert!(run(&columns, &sources, &lying, &lying, &alpha).is_err()); + } + + /// Same lie, but from a prover that proves the relation it wants: the + /// sumcheck is then internally consistent and the rebuild is what rejects. + #[test] + fn a_forged_column_value_is_rejected() { + let num_vars = 3; + let columns = vec![column(num_vars, 3), column(num_vars, 5)]; + let sources = [FactorSource::direct(0), FactorSource::shifted(1, 1)]; + let alpha = point(num_vars); + let values = honest_values(&columns, &sources, &alpha); + + let (mut proof, _) = + prove(&columns, &sources, &values, &alpha, None, &mut transcript()).unwrap(); + proof.column_values[1] += FE::one(); + + let err = verify( + &proof, + &sources, + &values, + &alpha, + columns.len(), + &mut transcript(), + ) + .unwrap_err(); + assert_eq!(err, Error::ShiftedReadMismatch); + } + + #[test] + fn swapping_two_column_values_is_rejected() { + let num_vars = 3; + let columns = vec![column(num_vars, 3), column(num_vars, 5)]; + let sources = [FactorSource::direct(0), FactorSource::direct(1)]; + let alpha = point(num_vars); + let values = honest_values(&columns, &sources, &alpha); + + let (mut proof, _) = + prove(&columns, &sources, &values, &alpha, None, &mut transcript()).unwrap(); + proof.column_values.swap(0, 1); + + assert_eq!( + verify( + &proof, + &sources, + &values, + &alpha, + columns.len(), + &mut transcript() + ) + .unwrap_err(), + Error::ShiftedReadMismatch + ); + } + + #[test] + fn a_forged_factor_value_is_rejected() { + let num_vars = 3; + let columns = vec![column(num_vars, 7)]; + let sources = [FactorSource::direct(0), FactorSource::shifted(0, 1)]; + let alpha = point(num_vars); + let honest = honest_values(&columns, &sources, &alpha); + let mut forged = honest.clone(); + forged[0] += FE::one(); + + assert!(run(&columns, &sources, &honest, &forged, &alpha).is_err()); + } + + #[test] + fn offsets_beyond_one_reduce_too() { + // The example AIRs in this repo read two steps ahead, so the kernel is + // not specialized to the rotation. + let num_vars = 3; + let columns = vec![column(num_vars, 3)]; + let sources = [ + FactorSource::direct(0), + FactorSource::shifted(0, 1), + FactorSource::shifted(0, 2), + ]; + let alpha = point(num_vars); + let values = honest_values(&columns, &sources, &alpha); + + run(&columns, &sources, &values, &values, &alpha).unwrap(); + } + + #[test] + fn one_sumcheck_covers_every_factor() { + // The cost of the reduction is one pass over the cube, not one per + // factor. + let num_vars = 4; + let columns = vec![column(num_vars, 3), column(num_vars, 5)]; + let sources: Vec = (0..2) + .flat_map(|c| { + [ + FactorSource::direct(c), + FactorSource::shifted(c, 1), + FactorSource::shifted(c, 2), + ] + }) + .collect(); + let alpha = point(num_vars); + let values = honest_values(&columns, &sources, &alpha); + + let (proof, _) = + prove(&columns, &sources, &values, &alpha, None, &mut transcript()).unwrap(); + assert_eq!(proof.sumcheck.rounds.len(), num_vars); + assert_eq!(proof.column_values.len(), 2); + } + + #[test] + fn a_proof_replayed_under_another_transcript_is_rejected() { + let num_vars = 3; + let columns = vec![column(num_vars, 3)]; + let sources = [FactorSource::direct(0), FactorSource::shifted(0, 1)]; + let alpha = point(num_vars); + let values = honest_values(&columns, &sources, &alpha); + + let (proof, _) = + prove(&columns, &sources, &values, &alpha, None, &mut transcript()).unwrap(); + let mut other = DefaultTranscript::::new(b"a-different-statement"); + assert!(verify(&proof, &sources, &values, &alpha, columns.len(), &mut other).is_err()); + } + + #[test] + fn a_source_naming_a_column_that_is_not_there_is_rejected() { + let columns = vec![column(3, 3)]; + let sources = [FactorSource::direct(1)]; + let alpha = point(3); + assert_eq!( + prove( + &columns, + &sources, + &[FE::zero()], + &alpha, + None, + &mut transcript() + ) + .unwrap_err(), + Error::UnknownPolynomial { index: 1, len: 1 } + ); + } + + #[test] + fn a_claim_per_factor_is_required() { + let columns = vec![column(3, 3)]; + let sources = [FactorSource::direct(0), FactorSource::shifted(0, 1)]; + let alpha = point(3); + assert_eq!( + prove( + &columns, + &sources, + &[FE::zero()], + &alpha, + None, + &mut transcript() + ) + .unwrap_err(), + Error::VariableCountMismatch { + expected: 2, + got: 1 + } + ); + } + + #[test] + fn materializing_offset_zero_is_the_column_itself() { + let columns = vec![column(3, 3)]; + assert_eq!( + materialize(&columns, &FactorSource::direct(0)).unwrap(), + columns[0] + ); + } + + #[test] + fn materializing_a_shift_moves_every_row_up() { + let columns = vec![column(3, 3)]; + let shifted = materialize(&columns, &FactorSource::shifted(0, 1)).unwrap(); + let size = columns[0].len(); + for i in 0..size { + assert_eq!(shifted.evals()[i], columns[0].evals()[(i + 1) % size]); + } + } +} diff --git a/crypto/multilinear/src/eq.rs b/crypto/multilinear/src/eq.rs new file mode 100644 index 000000000..8abfb3e50 --- /dev/null +++ b/crypto/multilinear/src/eq.rs @@ -0,0 +1,541 @@ +//! The equality kernel `eq(r, x) = ∏_i (r_i·x_i + (1 - r_i)(1 - x_i))` and the +//! shift kernel, `shift_k(x, y) = 1` iff `index(y) = index(x) + k mod 2^n`, of +//! which the cyclic rotation `rot` is the `k = 1` case. + +use math::field::{element::FieldElement, traits::IsField}; + +#[cfg(feature = "parallel")] +use rayon::prelude::*; + +use crate::{Error, mle::Mle}; + +/// Builds the table of `eq(r, x)` for every `x` in `{0,1}^n`, in `O(2^n)`. +/// +/// Doubles the table one variable at a time. Each step prepends its variable as +/// the new most significant bit, so the variables are consumed **back to +/// front**: that leaves variable 0 in the high bit, which is the indexing +/// convention [`Mle`](crate::mle::Mle) folds on. +pub fn eq_evals(r: &[FieldElement]) -> Vec> +where + FieldElement: Send + Sync, +{ + let mut table = vec![FieldElement::::zero(); 1usize << r.len()]; + // The length is `2^r.len()` by construction. + eq_evals_into(r, &FieldElement::::one(), &mut table).expect("the table is sized for r"); + table +} + +/// The same table, scaled by `by`, written into `dst`. +/// +/// The scale rides the seed rather than a pass over the result: the table is a +/// product over the variables, so one more factor at the start scales every +/// entry. And the doubling happens in place — a caller that wants the table +/// inside a bigger one (a stacked polynomial's weight, say) hands over that +/// range and pays no copy. +pub fn eq_evals_into( + r: &[FieldElement], + by: &FieldElement, + dst: &mut [FieldElement], +) -> Result<(), Error> +where + FieldElement: Send + Sync, +{ + if dst.len() != 1usize << r.len() { + return Err(Error::VariableCountMismatch { + expected: 1usize << r.len(), + got: dst.len(), + }); + } + dst[0] = by.clone(); + for (level, r_i) in r.iter().rev().enumerate() { + let one_minus = FieldElement::::one() - r_i; + // Each half of the doubled table is an independent scaling of the + // current one, and the halves are disjoint — so this is one pass over + // the level, and the last levels are the whole cube, which is where the + // pool is worth it. + let half = 1usize << level; + let (lo, hi) = dst[..half * 2].split_at_mut(half); + let scale = |(l, h): (&mut FieldElement, &mut FieldElement)| { + *h = &*l * r_i; + *l = &*l * &one_minus; + }; + #[cfg(feature = "parallel")] + if half >= crate::SERIAL_BELOW { + lo.par_iter_mut().zip(hi.par_iter_mut()).for_each(scale); + } else { + lo.iter_mut().zip(hi.iter_mut()).for_each(scale); + } + #[cfg(not(feature = "parallel"))] + lo.iter_mut().zip(hi.iter_mut()).for_each(scale); + } + Ok(()) +} + +/// The multilinear extension of `eq(r, ·)`. +pub fn eq_mle(r: &[FieldElement]) -> Result, Error> { + Mle::new(eq_evals(r)) +} + +/// Evaluates `eq(r, x)` directly, without materializing the table. +pub fn eq_eval( + r: &[FieldElement], + x: &[FieldElement], +) -> Result, Error> { + if r.len() != x.len() { + return Err(Error::VariableCountMismatch { + expected: r.len(), + got: x.len(), + }); + } + let one = FieldElement::::one(); + Ok(r.iter().zip(x).fold(one.clone(), |acc, (r_i, x_i)| { + acc * (r_i * x_i + (&one - r_i) * (&one - x_i)) + })) +} + +/// `rot` and `eq` together; the recursion needs both. +/// +/// `next_f(z) = Σ_y rot(z, y)·f(y)`, which is how a rotated table gets bound to +/// the committed column. +pub fn eq_and_rot_eval( + x: &[FieldElement], + y: &[FieldElement], +) -> Result<(FieldElement, FieldElement), Error> { + if x.len() != y.len() { + return Err(Error::VariableCountMismatch { + expected: x.len(), + got: y.len(), + }); + } + let one = FieldElement::::one(); + let mut eq = one.clone(); + let mut rot = one.clone(); + + // Recursion: rot(x, y) = y₀(1−x₀)·eq(rest) + (1−y₀)x₀·rot(rest), unrolled + // with the least significant variable outermost. `rot` on the tail carries + // "the low bits wrapped", which is what makes the next bit up increment. + // + // Our variable 0 is the most significant bit, so the fold runs front to + // back — the opposite of the usual least-significant-first presentation. + for (x_i, y_i) in x.iter().zip(y) { + rot = y_i * (&one - x_i) * &eq + (&one - y_i) * x_i * &rot; + eq *= x_i * y_i + (&one - x_i) * (&one - y_i); + } + Ok((eq, rot)) +} + +/// Just the rotation kernel. See [`eq_and_rot_eval`]. +pub fn rot_eval( + x: &[FieldElement], + y: &[FieldElement], +) -> Result, Error> { + Ok(eq_and_rot_eval(x, y)?.1) +} + +/// `shift_k(x, y) = 1` iff `index(y) = index(x) + k mod 2^n`, extended +/// multilinearly. `k = 0` is [`eq_eval`], `k = 1` is [`rot_eval`]. +/// +/// Adds the constant `k` bit by bit from the least significant end, which is +/// variable `n − 1`. `state[c]` is the weight of the bits handled so far having +/// produced carry `c`; the shift wraps, so both carries are accepted at the end. +pub fn shift_eval( + x: &[FieldElement], + y: &[FieldElement], + k: usize, +) -> Result, Error> { + if x.len() != y.len() { + return Err(Error::VariableCountMismatch { + expected: x.len(), + got: y.len(), + }); + } + let one = FieldElement::::one(); + let zero = FieldElement::::zero(); + let mut state = [one.clone(), zero.clone()]; + + for (t, (x_j, y_j)) in x.iter().zip(y).rev().enumerate() { + let eq_j = x_j * y_j + (&one - x_j) * (&one - y_j); + // y_j = 1 − x_j: x_j = 1 carries out, x_j = 0 does not. + let carry = x_j * (&one - y_j); + let no_carry = (&one - x_j) * y_j; + + let mut next = [zero.clone(), zero.clone()]; + for (c, weight) in state.iter().enumerate() { + match ((k >> t) & 1) + c { + 0 => next[0] += weight * &eq_j, + 1 => { + next[0] += weight * &no_carry; + next[1] += weight * &carry; + } + _ => next[1] += weight * &eq_j, + } + } + state = next; + } + + Ok(&state[0] + &state[1]) +} + +/// The table of `shift_k(x, ·)` over the cube, in `O(2^n)`. +/// +/// Same carry recursion as [`shift_eval`], but each step doubles the tables +/// instead of multiplying `y_j`'s weight out — prepending `y_j` as the new most +/// significant bit, which is the indexing [`Mle`] folds on. +/// +/// The two tables ride buffers the levels swap between rather than being +/// allocated per level: the last level alone is the whole cube, so allocating +/// per level is allocating the answer twice over. +pub fn shift_evals(x: &[FieldElement], k: usize) -> Vec> +where + FieldElement: Send + Sync, +{ + let size = 1usize << x.len(); + // Only the low `n` bits of `k` are ever read, so a shift by a multiple of + // the cube is no shift — and no shift is `eq`, which carries one table + // instead of two and is one pass per level instead of three. + if k.is_multiple_of(size) { + return eq_evals(x); + } + + let one = FieldElement::::one(); + let zero = FieldElement::::zero(); + // "No carry out" and "carry out". Only the seed has to be set: every cell + // a level reads is one that level or an earlier one wrote. + let mut cur = [vec![zero.clone(); size], vec![zero.clone(); size]]; + let mut next = [vec![zero.clone(); size], vec![zero; size]]; + cur[0][0] = one.clone(); + + let mut len = 1usize; + for (t, x_j) in x.iter().rev().enumerate() { + let one_minus = &one - x_j; + { + let [s0, s1] = &cur; + let (s0, s1) = (&s0[..len], &s1[..len]); + let [n0, n1] = &mut next; + let (n0_lo, n0_rest) = n0.split_at_mut(len); + let (n1_lo, n1_rest) = n1.split_at_mut(len); + let n0_hi = &mut n0_rest[..len]; + let n1_hi = &mut n1_rest[..len]; + // Every level is made of the same three blocks and one of zeros: + // b = (1 − x)·s0, m = x·s0 + (1 − x)·s1, h = x·s1 + // and all the shift's bit decides is where each one goes. + let (b, m, h, blank) = if (k >> t) & 1 == 0 { + (n0_lo, n0_hi, n1_lo, n1_hi) + } else { + (n0_hi, n1_lo, n1_hi, n0_lo) + }; + let scale_lo = |(o, v): (&mut FieldElement, &FieldElement)| *o = &one_minus * v; + let scale_hi = |(o, v): (&mut FieldElement, &FieldElement)| *o = x_j * v; + let mix = + |((o, a), c): ((&mut FieldElement, &FieldElement), &FieldElement)| { + *o = x_j * a + &one_minus * c + }; + let clear = |o: &mut FieldElement| *o = FieldElement::::zero(); + #[cfg(feature = "parallel")] + { + b.par_iter_mut().zip(s0.par_iter()).for_each(scale_lo); + h.par_iter_mut().zip(s1.par_iter()).for_each(scale_hi); + m.par_iter_mut() + .zip(s0.par_iter()) + .zip(s1.par_iter()) + .for_each(mix); + blank.par_iter_mut().for_each(clear); + } + #[cfg(not(feature = "parallel"))] + { + b.iter_mut().zip(s0.iter()).for_each(scale_lo); + h.iter_mut().zip(s1.iter()).for_each(scale_hi); + m.iter_mut().zip(s0.iter()).zip(s1.iter()).for_each(mix); + blank.iter_mut().for_each(clear); + } + } + len *= 2; + core::mem::swap(&mut cur, &mut next); + } + + let [no_carry, carried] = cur; + let add = |(a, b): (FieldElement, FieldElement)| a + b; + #[cfg(feature = "parallel")] + return no_carry.into_par_iter().zip(carried).map(add).collect(); + #[cfg(not(feature = "parallel"))] + return no_carry.into_iter().zip(carried).map(add).collect(); +} + +/// The multilinear extension of `shift_k(x, ·)`. +pub fn shift_mle(x: &[FieldElement], k: usize) -> Result, Error> +where + FieldElement: Send + Sync, +{ + Mle::new(shift_evals(x, k)) +} + +#[cfg(test)] +mod tests { + use super::*; + use math::field::goldilocks::GoldilocksField as F; + + type FE = FieldElement; + + fn point(vals: &[u64]) -> Vec { + vals.iter().map(|v| FE::from(*v)).collect() + } + + /// The hypercube corner for `index`, variable 0 most significant. + fn corner(index: usize, num_vars: usize) -> Vec { + (0..num_vars) + .map(|i| FE::from(((index >> (num_vars - 1 - i)) & 1) as u64)) + .collect() + } + + #[test] + fn is_the_indicator_on_the_hypercube() { + let n = 3; + for r_idx in 0..(1u64 << n) { + let r = point(&[(r_idx >> 2) & 1, (r_idx >> 1) & 1, r_idx & 1]); + let table = eq_evals(&r); + assert_eq!(table.len(), 1 << n); + for (x_idx, value) in table.iter().enumerate() { + let expected = if x_idx as u64 == r_idx { + FE::one() + } else { + FE::zero() + }; + assert_eq!(value, &expected, "eq(r={r_idx}, x={x_idx})"); + } + } + } + + #[test] + fn table_and_direct_evaluation_agree_on_corners() { + let r = point(&[5, 9, 2]); + for (x_idx, entry) in eq_evals(&r).into_iter().enumerate() { + let x = point(&[ + ((x_idx >> 2) & 1) as u64, + ((x_idx >> 1) & 1) as u64, + (x_idx & 1) as u64, + ]); + assert_eq!(entry, eq_eval(&r, &x).unwrap()); + } + } + + #[test] + fn table_is_the_multilinear_extension() { + // Evaluating the eq table as an MLE off the cube must match the + // product formula. + let r = point(&[3, 11]); + let mle = eq_mle(&r).unwrap(); + let x = point(&[7, 13]); + assert_eq!(mle.evaluate(&x).unwrap(), eq_eval(&r, &x).unwrap()); + } + + #[test] + fn sums_to_one_over_the_cube() { + // Σ_x eq(r, x) = 1 for any r, since eq interpolates a single corner. + let r = point(&[4, 6, 8]); + let total = eq_evals(&r).into_iter().fold(FE::zero(), |acc, v| acc + v); + assert_eq!(total, FE::one()); + } + + #[test] + fn is_symmetric_in_its_arguments() { + let a = point(&[2, 3]); + let b = point(&[5, 7]); + assert_eq!(eq_eval(&a, &b).unwrap(), eq_eval(&b, &a).unwrap()); + } + + #[test] + fn zero_variables_gives_the_empty_product() { + assert_eq!(eq_evals::(&[]), vec![FE::one()]); + assert_eq!(eq_eval::(&[], &[]).unwrap(), FE::one()); + } + + #[test] + fn rejects_mismatched_arity() { + assert_eq!( + eq_eval(&point(&[1, 2]), &point(&[1])).unwrap_err(), + Error::VariableCountMismatch { + expected: 2, + got: 1 + } + ); + } + + #[test] + fn rot_is_the_successor_indicator_on_the_hypercube() { + for num_vars in 1..=4usize { + let size = 1usize << num_vars; + for xi in 0..size { + for yi in 0..size { + let x = corner(xi, num_vars); + let y = corner(yi, num_vars); + let expected = if yi == (xi + 1) % size { + FE::one() + } else { + FE::zero() + }; + assert_eq!( + rot_eval(&x, &y).unwrap(), + expected, + "n={num_vars}, x={xi}, y={yi}" + ); + } + } + } + } + + #[test] + fn rot_wraps_the_last_index_to_the_first() { + let n = 3; + let last = corner(7, n); + let first = corner(0, n); + assert_eq!(rot_eval(&last, &first).unwrap(), FE::one()); + } + + #[test] + fn rot_reproduces_a_shifted_table() { + // The identity the commitment scheme will lean on: + // next_f(z) = Σ_y rot(z, y)·f(y), for z on the cube. + let n = 3; + let size = 1usize << n; + let f: Vec = (0..size as u64).map(|i| FE::from(i * 7 + 5)).collect(); + + for zi in 0..size { + let z = corner(zi, n); + let summed = (0..size).fold(FE::zero(), |acc, yi| { + acc + rot_eval(&z, &corner(yi, n)).unwrap() * f[yi] + }); + assert_eq!(summed, f[(zi + 1) % size], "z={zi}"); + } + } + + #[test] + fn rot_matches_its_table_off_the_cube() { + // The closed form must be the multilinear extension of the table, so + // the verifier can evaluate it at a random point. + let n = 3; + let size = 1usize << n; + let z = point(&[5, 9, 2]); + + // Brute-force the MLE of rot(·, y) for a fixed y by extending the table. + for yi in 0..size { + let table: Vec = (0..size) + .map(|xi| rot_eval(&corner(xi, n), &corner(yi, n)).unwrap()) + .collect(); + let mle = Mle::new(table).unwrap(); + assert_eq!( + mle.evaluate(&z).unwrap(), + rot_eval(&z, &corner(yi, n)).unwrap(), + "y={yi}" + ); + } + } + + #[test] + fn eq_and_rot_agree_with_the_standalone_helpers() { + let x = point(&[3, 11, 4]); + let y = point(&[7, 13, 2]); + let (eq, rot) = eq_and_rot_eval(&x, &y).unwrap(); + assert_eq!(eq, eq_eval(&x, &y).unwrap()); + assert_eq!(rot, rot_eval(&x, &y).unwrap()); + } + + #[test] + fn rot_rejects_mismatched_arity() { + assert!(rot_eval(&point(&[1, 2]), &point(&[1])).is_err()); + } + + #[test] + fn shift_is_the_offset_indicator_on_the_hypercube() { + for num_vars in 1..=4usize { + let size = 1usize << num_vars; + // Past `size` too, so wrapping is exercised as its own case. + for k in 0..(size + 3) { + for xi in 0..size { + for yi in 0..size { + let expected = if yi == (xi + k) % size { + FE::one() + } else { + FE::zero() + }; + assert_eq!( + shift_eval(&corner(xi, num_vars), &corner(yi, num_vars), k).unwrap(), + expected, + "n={num_vars}, k={k}, x={xi}, y={yi}" + ); + } + } + } + } + } + + #[test] + fn shift_zero_is_eq_and_shift_one_is_rot() { + // Off the cube, where agreeing on corners would not be enough. + let x = point(&[3, 11, 4]); + let y = point(&[7, 13, 2]); + assert_eq!(shift_eval(&x, &y, 0).unwrap(), eq_eval(&x, &y).unwrap()); + assert_eq!(shift_eval(&x, &y, 1).unwrap(), rot_eval(&x, &y).unwrap()); + } + + #[test] + fn shift_table_matches_the_closed_form() { + let x = point(&[5, 9, 2]); + for k in 0..10usize { + for (yi, entry) in shift_evals(&x, k).into_iter().enumerate() { + assert_eq!(entry, shift_eval(&x, &corner(yi, 3), k).unwrap(), "k={k}"); + } + } + } + + #[test] + fn shift_table_is_the_multilinear_extension() { + // The verifier evaluates the kernel at a random point while the prover + // folds the table, so the two must be the same polynomial. + let x = point(&[3, 11, 4]); + let z = point(&[7, 13, 2]); + for k in 0..8usize { + assert_eq!( + shift_mle(&x, k).unwrap().evaluate(&z).unwrap(), + shift_eval(&x, &z, k).unwrap(), + "k={k}" + ); + } + } + + #[test] + fn shift_reproduces_a_shifted_column() { + // The identity the reduction leans on: + // f_shift_k(z) = Σ_y shift_k(z, y)·f(y). + let n = 3; + let size = 1usize << n; + let f: Vec = (0..size as u64).map(|i| FE::from(i * 7 + 5)).collect(); + let z = point(&[5, 9, 2]); + + for k in 0..6usize { + let shifted = Mle::new((0..size).map(|i| f[(i + k) % size]).collect()).unwrap(); + let summed = shift_evals(&z, k) + .into_iter() + .zip(&f) + .fold(FE::zero(), |acc, (w, v)| acc + w * v); + assert_eq!(summed, shifted.evaluate(&z).unwrap(), "k={k}"); + } + } + + #[test] + fn shift_sums_to_one_over_the_cube() { + // The kernel picks out one corner, so its extension sums to one at any + // point — a cheap check that the carry recursion loses no weight. + let x = point(&[4, 6, 8]); + for k in 0..10usize { + let total = shift_evals(&x, k) + .into_iter() + .fold(FE::zero(), |acc, v| acc + v); + assert_eq!(total, FE::one(), "k={k}"); + } + } + + #[test] + fn shift_rejects_mismatched_arity() { + assert!(shift_eval(&point(&[1, 2]), &point(&[1]), 1).is_err()); + } +} diff --git a/crypto/multilinear/src/gkr.rs b/crypto/multilinear/src/gkr.rs new file mode 100644 index 000000000..130653275 --- /dev/null +++ b/crypto/multilinear/src/gkr.rs @@ -0,0 +1,781 @@ +//! LogUp as a tree of fractions: `p₁/q₁ + p₂/q₂ = (p₁q₂ + p₂q₁)/(q₁q₂)`, added +//! pairwise up a binary tree with GKR proving each layer against the one below. +//! Only the input layer is ever committed. +//! +//! Proves the sum is whatever the output claims. Checking that the output +//! numerator is zero — the bus balance — is the caller's. + +use crypto::fiat_shamir::is_transcript::IsTranscript; +use math::field::{element::FieldElement, traits::IsField}; + +#[cfg(feature = "parallel")] +use rayon::prelude::*; + +use crate::{ + Error, + eq::{eq_eval, eq_mle}, + mle::Mle, + poly::SumcheckPolynomial, + program::{Builder, Program}, + sumcheck::{self, SumcheckProof}, +}; + +/// One level of the tree: numerators and denominators over the same cube. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct FractionLayer { + pub p: Mle, + pub q: Mle, +} + +impl FractionLayer { + pub fn new(p: Mle, q: Mle) -> Result { + if p.num_vars() != q.num_vars() { + return Err(Error::VariableCountMismatch { + expected: p.num_vars(), + got: q.num_vars(), + }); + } + Ok(Self { p, q }) + } + + pub fn num_vars(&self) -> usize { + self.p.num_vars() + } + + /// Adds the two halves pointwise, giving the layer one level up. + pub fn fold(&self) -> Result + where + FieldElement: Send + Sync, + { + if self.num_vars() == 0 { + return Err(Error::NoVariablesLeft); + } + let half = self.p.len() / 2; + let (p_lo, p_hi) = self.p.evals().split_at(half); + let (q_lo, q_hi) = self.q.evals().split_at(half); + + // Every index folds on its own, and building the tree is a pass over + // the input layer at every level, so this is worth the pool. + let both = |i: usize| { + ( + &p_lo[i] * &q_hi[i] + &p_hi[i] * &q_lo[i], + &q_lo[i] * &q_hi[i], + ) + }; + #[cfg(feature = "parallel")] + let (next_p, next_q): (Vec<_>, Vec<_>) = if half >= crate::SERIAL_BELOW { + (0..half).into_par_iter().map(both).unzip() + } else { + (0..half).map(both).unzip() + }; + #[cfg(not(feature = "parallel"))] + let (next_p, next_q): (Vec<_>, Vec<_>) = (0..half).map(both).unzip(); + + Self::new(Mle::new(next_p)?, Mle::new(next_q)?) + } +} + +/// The whole tree, from the input layer down to the single output fraction. +/// +/// `layers[0]` is the output (zero variables); the last entry is the input. +#[derive(Debug)] +pub struct FractionTree { + /// Empty when the tree lives on a device, which holds every layer. + layers: Vec>, + device: Option, + num_layers: usize, + output: (FieldElement, FieldElement), +} + +impl FractionTree { + /// Builds every layer by repeated folding. + /// + /// On a device the layers stay there: the tree is the biggest thing a + /// table's argument holds, and GKR reads every level of it. + pub fn build(input: FractionLayer) -> Result + where + FieldElement: Send + Sync, + { + if let Some(device) = crate::gpu::build_tree(&input.p, &input.q) { + return Self::from_device(device); + } + + let mut layers = vec![input]; + while layers.last().expect("non-empty").num_vars() > 0 { + let next = layers.last().expect("non-empty").fold()?; + layers.push(next); + } + layers.reverse(); + let top = &layers[0]; + let output = (top.p.evals()[0].clone(), top.q.evals()[0].clone()); + let num_layers = layers.len(); + Ok(Self { + layers, + device: None, + num_layers, + output, + }) + } + + /// A tree a device already holds, layers and all. + pub fn from_device(device: crate::gpu::DeviceTree) -> Result { + let num_layers = device.num_layers(); + let output = device.output()?; + Ok(Self { + layers: Vec::new(), + device: Some(device), + num_layers, + output, + }) + } + + /// The output fraction `(p, q)`. The bus balances when `p` is zero. + pub fn output(&self) -> (FieldElement, FieldElement) { + self.output.clone() + } + + pub fn num_layers(&self) -> usize { + self.num_layers + } + + /// The layer, for a tree that kept them here. + pub fn layer(&self, i: usize) -> &FractionLayer { + &self.layers[i] + } + + /// The device holding every layer, when one does. + pub fn device(&self) -> Option<&crate::gpu::DeviceTree> { + self.device.as_ref() + } + + pub fn input_layer(&self) -> &FractionLayer { + self.layers.last().expect("non-empty") + } + + /// The levels near the output, back here — `prefix[i]` is layer `i`. + /// + /// Empty for a tree that already lives here. One download brings the whole + /// prefix: the levels above the deepest of them are its folds, and the + /// deepest is a few kilobytes. + fn host_prefix(&self) -> Vec> + where + FieldElement: Send + Sync, + { + let Some(device) = self.device.as_ref() else { + return Vec::new(); + }; + // Never the input layer: a tree that dropped it writes it again for its + // own sumcheck, and it is the one level too big to walk here. + let deepest = HOST_LAYER_VARS.min(self.num_layers.saturating_sub(2)); + let Some((p, q)) = device.layer_to_host::(deepest) else { + return Vec::new(); + }; + let Ok(layer) = FractionLayer::new(p, q) else { + return Vec::new(); + }; + if layer.num_vars() != deepest { + return Vec::new(); + } + let mut prefix = vec![layer]; + while prefix.last().expect("non-empty").num_vars() > 0 { + let Ok(next) = prefix.last().expect("non-empty").fold() else { + return Vec::new(); + }; + prefix.push(next); + } + prefix.reverse(); + prefix + } +} + +/// The deepest level that comes here whole: the one whose halves are already +/// a cube of [`crate::HOST_CUBE_DIRECT`], so its sumcheck never goes to a device. +const HOST_LAYER_VARS: usize = crate::HOST_CUBE_DIRECT.trailing_zeros() as usize + 1; + +/// The layer relation: `Σ_x eq(r,x)·[p_lo·q_hi + p_hi·q_lo + λ·q_lo·q_hi]`, +/// which equals `p_out(r) + λ·q_out(r)` when the layer really is the fold. +struct LayerRelation { + /// `[eq, p_lo, p_hi, q_lo, q_hi]`. + polys: Vec>, + lambda: FieldElement, + /// The same rule as straight-line code, for the device path. `None` for a + /// relation that is here on purpose — a level the tree handed over, or the + /// tail of one — so the dispatch does not send it back. + program: Option>, +} + +impl LayerRelation { + const EQ: usize = 0; + const P_LO: usize = 1; + const P_HI: usize = 2; + const Q_LO: usize = 3; + const Q_HI: usize = 4; + + fn new( + next: &FractionLayer, + r: &[FieldElement], + lambda: FieldElement, + dispatchable: bool, + ) -> Result { + let half = next.p.len() / 2; + let split = |m: &Mle| -> Result<(Mle, Mle), Error> { + Ok(( + Mle::new(m.evals()[..half].to_vec())?, + Mle::new(m.evals()[half..].to_vec())?, + )) + }; + let (p_lo, p_hi) = split(&next.p)?; + let (q_lo, q_hi) = split(&next.q)?; + Ok(Self { + polys: vec![eq_mle(r)?, p_lo, p_hi, q_lo, q_hi], + program: dispatchable + .then(|| Self::program_for(&lambda)) + .transpose()?, + lambda, + }) + } + + /// The same relation over factors a device already folded: the weight and + /// the four halves as its last round left them. + fn from_factors(polys: Vec>, lambda: FieldElement) -> Result { + if polys.len() != 5 { + return Err(Error::VariableCountMismatch { + expected: 5, + got: polys.len(), + }); + } + Ok(Self { + polys, + program: None, + lambda, + }) + } + + /// `eq·(p_lo·q_hi + p_hi·q_lo + lambda·q_lo·q_hi)`, the same expression + /// [`combine`](SumcheckPolynomial::combine) evaluates. + fn program_for(lambda: &FieldElement) -> Result, Error> { + let mut b = Builder::::new(); + let eq = b.var(Self::EQ); + let p_lo = b.var(Self::P_LO); + let p_hi = b.var(Self::P_HI); + let q_lo = b.var(Self::Q_LO); + let q_hi = b.var(Self::Q_HI); + let first = b.mul(p_lo, q_hi); + let second = b.mul(p_hi, q_lo); + let numerator = b.add(first, second); + let denominator = b.mul(q_lo, q_hi); + let weighted = b.fixed(lambda.clone()); + let scaled = b.mul(weighted, denominator); + let sum = b.add(numerator, scaled); + let root = b.mul(eq, sum); + b.finish(root) + } +} + +/// `eq` times a product of two layer values. +const LAYER_DEGREE: usize = 3; + +impl SumcheckPolynomial for LayerRelation { + fn num_vars(&self) -> usize { + self.polys[Self::EQ].num_vars() + } + + fn degree(&self) -> usize { + LAYER_DEGREE + } + + fn polys(&self) -> &[Mle] { + &self.polys + } + + fn combine(&self, v: &[FieldElement]) -> FieldElement { + let numerator = &v[Self::P_LO] * &v[Self::Q_HI] + &v[Self::P_HI] * &v[Self::Q_LO]; + let denominator = &v[Self::Q_LO] * &v[Self::Q_HI]; + &v[Self::EQ] * (numerator + &self.lambda * denominator) + } + + fn fix_first_variable(&mut self, r: &FieldElement) -> Result<(), Error> { + for p in &mut self.polys { + p.fix_first_variable_in_place(r)?; + } + Ok(()) + } + + fn program(&self) -> Option<&Program> { + self.program.as_ref() + } + + /// The layer relation is four multiplications written out, not a program + /// the host walks — so its rounds are worth taking back much earlier. + fn host_cube(&self) -> usize { + crate::HOST_CUBE_DIRECT + } + + fn accept_folded(&mut self, polys: Vec>) -> Result<(), Error> { + if polys.len() != self.polys.len() { + return Err(Error::VariableCountMismatch { + expected: self.polys.len(), + got: polys.len(), + }); + } + self.polys = polys; + Ok(()) + } +} + +/// One layer's transcript: the sumcheck plus the four values it reduces to. +#[derive( + Clone, + Debug, + PartialEq, + Eq, + serde::Serialize, + serde::Deserialize, + rkyv::Archive, + rkyv::Serialize, + rkyv::Deserialize, +)] +#[serde(bound = "")] +pub struct LayerProof { + pub sumcheck: SumcheckProof, + pub p_lo: FieldElement, + pub p_hi: FieldElement, + pub q_lo: FieldElement, + pub q_hi: FieldElement, +} + +/// A proof for the whole tree, output layer first. +#[derive( + Clone, + Debug, + PartialEq, + Eq, + serde::Serialize, + serde::Deserialize, + rkyv::Archive, + rkyv::Serialize, + rkyv::Deserialize, +)] +#[serde(bound = "")] +pub struct GkrProof { + pub layers: Vec>, +} + +/// What the verifier is left holding about the **input** layer. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct GkrClaim { + pub point: Vec>, + pub p: FieldElement, + pub q: FieldElement, +} + +/// What proving leaves the caller holding. +/// +/// The claim is the same one [`verify`] arrives at: the prover needs it to +/// discharge the input layer, which is where the trace is. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct GkrOutput { + pub proof: GkrProof, + pub claim: GkrClaim, +} + +/// `(1 − c)·lo + c·hi` — the multilinear interpolation that turns the two +/// restricted claims back into one claim on the fuller layer. +fn combine_halves( + lo: &FieldElement, + hi: &FieldElement, + c: &FieldElement, +) -> FieldElement { + lo + c * &(hi - lo) +} + +/// Proves the tree, from the output fraction down to the input layer. +pub fn prove(tree: &FractionTree, transcript: &mut T) -> Result, Error> +where + F: IsField + 'static, + T: IsTranscript, +{ + let mut layers = Vec::with_capacity(tree.num_layers().saturating_sub(1)); + // The output layer has no variables, so the first claim sits at the empty + // point and needs no challenge. + let mut point: Vec> = Vec::new(); + let (mut p_claim, mut q_claim) = tree.output(); + // The levels near the output, fetched once. Their rounds are over cubes a + // core walks in microseconds, and a device pays a launch for each. + let prefix = tree.host_prefix(); + + for i in 0..tree.num_layers() - 1 { + let lambda: FieldElement = transcript.sample_field_element(); + + // What the device ran of this layer, and the relation it left behind. + // A tree the device holds proves its layer where it lies — the halves + // are the sumcheck's factors in place — until the cube reaches the + // crossover, and hands the factors over from there. + let (mut rounds, mut z, mut relation) = match prefix.get(i + 1) { + Some(here) => ( + Vec::new(), + Vec::new(), + LayerRelation::new(here, &point, lambda, false)?, + ), + None => match tree.device() { + Some(device) => { + let program = LayerRelation::program_for(&lambda)?; + let attempt = device.prove_layer( + i + 1, + &point, + &program, + LAYER_DEGREE, + crate::HOST_CUBE_DIRECT, + |sent| { + for value in sent { + transcript.append_field_element(value); + } + transcript.sample_field_element() + }, + ); + let Some(outcome) = attempt else { + return Err(Error::DeviceFailed { + stage: "layer sumcheck", + }); + }; + let (rounds, z, factors) = outcome?; + (rounds, z, LayerRelation::from_factors(factors, lambda)?) + } + None => ( + Vec::new(), + Vec::new(), + LayerRelation::new(tree.layer(i + 1), &point, lambda, true)?, + ), + }, + }; + + // The tail, however much of it is left: all of it for a level that came + // here whole, none for one the device ran out. + let left = relation.num_vars(); + let (tail, tail_z) = sumcheck::prove_rounds(&mut relation, left, transcript)?; + rounds.extend(tail); + z.extend(tail_z); + let sumcheck = SumcheckProof { rounds }; + + // The four restricted values the verifier needs to close the round are + // what the sumcheck's own factors have become: binding every variable + // to `z` is the evaluation at `z`. Evaluating the halves again would be + // a second pass over the layer, and the layers are the biggest thing + // the fraction tree holds. + let bound = |slot: usize| -> Result, Error> { + relation.polys()[slot] + .as_constant() + .cloned() + .ok_or(Error::NoVariablesLeft) + }; + let p_lo = bound(LayerRelation::::P_LO)?; + let p_hi = bound(LayerRelation::::P_HI)?; + let q_lo = bound(LayerRelation::::Q_LO)?; + let q_hi = bound(LayerRelation::::Q_HI)?; + + for v in [&p_lo, &p_hi, &q_lo, &q_hi] { + transcript.append_field_element(v); + } + let c = transcript.sample_field_element(); + + // Next layer's claim lives at (c, z). + p_claim = combine_halves(&p_lo, &p_hi, &c); + q_claim = combine_halves(&q_lo, &q_hi, &c); + point = std::iter::once(c).chain(z).collect(); + + layers.push(LayerProof { + sumcheck, + p_lo, + p_hi, + q_lo, + q_hi, + }); + } + + Ok(GkrOutput { + proof: GkrProof { layers }, + claim: GkrClaim { + point, + p: p_claim, + q: q_claim, + }, + }) +} + +/// Verifies the tree against a claimed output fraction. +pub fn verify( + proof: &GkrProof, + output: (FieldElement, FieldElement), + transcript: &mut T, +) -> Result, Error> +where + F: IsField + 'static, + T: IsTranscript, +{ + let (mut p_claim, mut q_claim) = output; + let mut point: Vec> = Vec::new(); + + for (i, layer) in proof.layers.iter().enumerate() { + let lambda = transcript.sample_field_element(); + let claimed_sum = &p_claim + &lambda * &q_claim; + + let claim = sumcheck::verify(&layer.sumcheck, claimed_sum, point.len(), 3, transcript)?; + + // The sumcheck's residual must be the layer relation at that point. + let eq_at = eq_eval(&point, &claim.point)?; + let numerator = &layer.p_lo * &layer.q_hi + &layer.p_hi * &layer.q_lo; + let denominator = &layer.q_lo * &layer.q_hi; + let expected = eq_at * (numerator + &lambda * denominator); + if expected != claim.expected_evaluation { + return Err(Error::LayerRelationMismatch { layer: i }); + } + + for v in [&layer.p_lo, &layer.p_hi, &layer.q_lo, &layer.q_hi] { + transcript.append_field_element(v); + } + let c = transcript.sample_field_element(); + + p_claim = combine_halves(&layer.p_lo, &layer.p_hi, &c); + q_claim = combine_halves(&layer.q_lo, &layer.q_hi, &c); + point = std::iter::once(c).chain(claim.point).collect(); + } + + Ok(GkrClaim { + point, + p: p_claim, + q: q_claim, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crypto::fiat_shamir::default_transcript::DefaultTranscript; + use math::field::goldilocks::GoldilocksField as F; + + type FE = FieldElement; + + fn transcript() -> DefaultTranscript { + DefaultTranscript::::new(b"gkr-test") + } + + fn mle(vals: &[u64]) -> Mle { + Mle::new(vals.iter().map(|v| FE::from(*v)).collect()).unwrap() + } + + fn layer(p: &[u64], q: &[u64]) -> FractionLayer { + FractionLayer::new(mle(p), mle(q)).unwrap() + } + + /// Σ pᵢ/qᵢ computed directly, for comparison against the tree. + fn direct_sum(p: &[u64], q: &[u64]) -> FE { + p.iter().zip(q).fold(FE::zero(), |acc, (pi, qi)| { + acc + FE::from(*pi) * FE::from(*qi).inv().unwrap() + }) + } + + /// A LogUp-shaped input layer: `mult / (alpha - fingerprint)`, with the + /// sends and receives arranged to cancel. + fn balanced_logup_layer(num_vars: usize) -> FractionLayer { + let size = 1usize << num_vars; + let alpha = FE::from(0x9E37_79B9u64); + let mut p = Vec::with_capacity(size); + let mut q = Vec::with_capacity(size); + for i in 0..size { + // Each fingerprint appears once as a send (+1) and once as a + // receive (−1), so the whole bus balances. + let fingerprint = FE::from((i as u64 / 2) * 7 + 5); + let sign = if i % 2 == 0 { FE::one() } else { -FE::one() }; + p.push(sign); + q.push(alpha - fingerprint); + } + FractionLayer::new(Mle::new(p).unwrap(), Mle::new(q).unwrap()).unwrap() + } + + #[test] + fn folding_adds_the_two_halves() { + let l = layer(&[1, 2, 3, 4], &[5, 6, 7, 8]); + let folded = l.fold().unwrap(); + assert_eq!(folded.num_vars(), 1); + + // Variable 0 is the most significant bit, so index i pairs with i + 2. + let p = [1u64, 2, 3, 4]; + let q = [5u64, 6, 7, 8]; + for (i, (a, b)) in [(0usize, 2usize), (1, 3)].iter().enumerate() { + let expected = FE::from(p[*a]) * FE::from(q[*a]).inv().unwrap() + + FE::from(p[*b]) * FE::from(q[*b]).inv().unwrap(); + let got = folded.p.evals()[i] * folded.q.evals()[i].inv().unwrap(); + assert_eq!(expected, got, "pair ({a}, {b})"); + } + } + + #[test] + fn the_tree_output_is_the_sum_of_every_input_fraction() { + let p = [3u64, 1, 4, 1, 5, 9, 2, 6]; + let q = [2u64, 7, 1, 8, 2, 8, 1, 8]; + let tree = FractionTree::build(layer(&p, &q)).unwrap(); + + let (out_p, out_q) = tree.output(); + assert_eq!(out_p * out_q.inv().unwrap(), direct_sum(&p, &q)); + } + + #[test] + fn layer_count_is_one_per_variable_plus_the_output() { + let tree = FractionTree::build(balanced_logup_layer(4)).unwrap(); + assert_eq!(tree.num_layers(), 5); + assert_eq!(tree.layer(0).num_vars(), 0); + assert_eq!(tree.input_layer().num_vars(), 4); + } + + #[test] + fn a_balanced_bus_has_a_zero_numerator() { + let tree = FractionTree::build(balanced_logup_layer(4)).unwrap(); + let (p, q) = tree.output(); + assert_eq!(p, FE::zero()); + assert_ne!(q, FE::zero(), "denominators must not vanish"); + } + + #[test] + fn an_unbalanced_bus_does_not() { + let mut input = balanced_logup_layer(4); + // Drop one receive: the bus no longer cancels. + let mut p = input.p.evals().to_vec(); + p[3] = FE::zero(); + input = FractionLayer::new(Mle::new(p).unwrap(), input.q).unwrap(); + + let tree = FractionTree::build(input).unwrap(); + assert_ne!(tree.output().0, FE::zero()); + } + + #[test] + fn prove_and_verify_round_trip() { + let tree = FractionTree::build(balanced_logup_layer(5)).unwrap(); + let output = tree.output(); + + let proof = prove(&tree, &mut transcript()).unwrap().proof; + assert_eq!(proof.layers.len(), tree.num_layers() - 1); + + let claim = verify(&proof, output, &mut transcript()).unwrap(); + + // The residual claim must be the input layer at the final point. + let input = tree.input_layer(); + assert_eq!(claim.point.len(), input.num_vars()); + assert_eq!(input.p.evaluate(&claim.point).unwrap(), claim.p); + assert_eq!(input.q.evaluate(&claim.point).unwrap(), claim.q); + } + + #[test] + fn round_trips_on_an_unbalanced_bus_too() { + // GKR proves the sum is whatever it is; deciding that it is zero is the + // caller's check on the output numerator, not part of this protocol. + let mut p = balanced_logup_layer(4).p.evals().to_vec(); + p[3] = FE::from(9); + let input = FractionLayer::new(Mle::new(p).unwrap(), balanced_logup_layer(4).q).unwrap(); + let tree = FractionTree::build(input).unwrap(); + + let proof = prove(&tree, &mut transcript()).unwrap().proof; + let claim = verify(&proof, tree.output(), &mut transcript()).unwrap(); + assert_eq!( + tree.input_layer().p.evaluate(&claim.point).unwrap(), + claim.p + ); + } + + #[test] + fn the_prover_arrives_at_the_claim_the_verifier_does() { + // The prover has to discharge the input-layer claim against the trace, + // so it needs the same claim the verifier ends up holding. + let tree = FractionTree::build(balanced_logup_layer(3)).unwrap(); + let out = prove(&tree, &mut transcript()).unwrap(); + let claim = verify(&out.proof, tree.output(), &mut transcript()).unwrap(); + assert_eq!(out.claim, claim); + + // And it is the input layer's value at that point. + let input = tree.input_layer(); + assert_eq!(claim.p, input.p.evaluate(&claim.point).unwrap()); + assert_eq!(claim.q, input.q.evaluate(&claim.point).unwrap()); + } + + #[test] + fn a_wrong_output_claim_is_rejected() { + let tree = FractionTree::build(balanced_logup_layer(4)).unwrap(); + let (p, q) = tree.output(); + let proof = prove(&tree, &mut transcript()).unwrap().proof; + + assert!(verify(&proof, (p + FE::one(), q), &mut transcript()).is_err()); + } + + #[test] + fn a_tampered_half_value_is_rejected() { + let tree = FractionTree::build(balanced_logup_layer(4)).unwrap(); + let output = tree.output(); + let mut proof = prove(&tree, &mut transcript()).unwrap().proof; + + proof.layers[1].q_lo += FE::one(); + let err = verify(&proof, output, &mut transcript()).unwrap_err(); + assert!(matches!(err, Error::LayerRelationMismatch { layer: 1 })); + } + + #[test] + fn a_tampered_sumcheck_round_is_rejected() { + let tree = FractionTree::build(balanced_logup_layer(4)).unwrap(); + let output = tree.output(); + let mut proof = prove(&tree, &mut transcript()).unwrap().proof; + + proof.layers[2].sumcheck.rounds[0].evaluations[0] += FE::one(); + assert!(verify(&proof, output, &mut transcript()).is_err()); + } + + #[test] + fn a_proof_replayed_under_another_transcript_is_rejected() { + let tree = FractionTree::build(balanced_logup_layer(4)).unwrap(); + let output = tree.output(); + let proof = prove(&tree, &mut transcript()).unwrap().proof; + + verify(&proof, output, &mut transcript()).unwrap(); + let mut other = DefaultTranscript::::new(b"another-statement"); + assert!(verify(&proof, output, &mut other).is_err()); + } + + #[test] + fn a_layer_carries_on_from_its_own_folded_factors() { + // What the crossover relies on: a relation rebuilt from the factors a + // few rounds left behind is the same relation. On a device those + // factors are read back off the layer; here the check is that picking + // them up mid-sumcheck changes nothing the verifier sees. + let tree = FractionTree::build(balanced_logup_layer(6)).unwrap(); + let lambda = FE::from(7); + let layer = tree.layer(4); + + let whole = { + let mut relation = LayerRelation::new(layer, &[FE::from(3)], lambda, true).unwrap(); + let vars = relation.num_vars(); + sumcheck::prove_rounds(&mut relation, vars, &mut transcript()).unwrap() + }; + + let mut relation = LayerRelation::new(layer, &[FE::from(3)], lambda, true).unwrap(); + let mut t = transcript(); + let (mut rounds, mut z) = sumcheck::prove_rounds(&mut relation, 1, &mut t).unwrap(); + let mut carried = LayerRelation::from_factors(relation.polys().to_vec(), lambda).unwrap(); + let left = carried.num_vars(); + let (tail, tail_z) = sumcheck::prove_rounds(&mut carried, left, &mut t).unwrap(); + rounds.extend(tail); + z.extend(tail_z); + + assert_eq!(whole, (rounds, z)); + } + + #[test] + fn a_single_fraction_needs_no_layers() { + let tree = FractionTree::build(layer(&[7], &[3])).unwrap(); + assert_eq!(tree.num_layers(), 1); + + let proof = prove(&tree, &mut transcript()).unwrap().proof; + assert!(proof.layers.is_empty()); + + let claim = verify(&proof, tree.output(), &mut transcript()).unwrap(); + assert!(claim.point.is_empty()); + assert_eq!(claim.p, FE::from(7)); + assert_eq!(claim.q, FE::from(3)); + } +} diff --git a/crypto/multilinear/src/gpu.rs b/crypto/multilinear/src/gpu.rs new file mode 100644 index 000000000..aaad6fbe4 --- /dev/null +++ b/crypto/multilinear/src/gpu.rs @@ -0,0 +1,549 @@ +//! Device dispatch for the multilinear path. +//! +//! Every entry point here returns `None` when the device declines — a field +//! the kernels do not cover, a size below the launch threshold, a kill switch, +//! or any CUDA error — and the caller runs the host path. A dispatch that +//! succeeded is counted, so a bench can tell a GPU number from a CPU one +//! wearing its label. +//! +//! There is no device yet: this is the half that declines, so every entry point +//! returns `None` unconditionally and the counters stay at zero. The half that +//! does the work lands last, and until it does the whole path is host-only — +//! which is how it is meant to be read. + +use core::sync::atomic::{AtomicU64, Ordering}; + +/// Successful device commits of a stacked polynomial. +static COMMIT_CALLS: AtomicU64 = AtomicU64::new(0); +/// Sumchecks whose rounds ran on device. +static SUMCHECK_CALLS: AtomicU64 = AtomicU64::new(0); +/// Rounds within them, so a declined tail shows up. +static SUMCHECK_ROUNDS: AtomicU64 = AtomicU64::new(0); +/// Multilinear evaluations bound on device. +static EVALUATE_CALLS: AtomicU64 = AtomicU64::new(0); +/// Fraction trees built and kept on device. +static TREE_CALLS: AtomicU64 = AtomicU64::new(0); +/// Tables whose factors were uploaded once and reused. +static FACTOR_CALLS: AtomicU64 = AtomicU64::new(0); +/// Openings whose two factors stayed on device across their groups. +static OPEN_CALLS: AtomicU64 = AtomicU64::new(0); + +pub fn commit_calls() -> u64 { + COMMIT_CALLS.load(Ordering::Relaxed) +} + +pub fn sumcheck_calls() -> u64 { + SUMCHECK_CALLS.load(Ordering::Relaxed) +} + +pub fn sumcheck_rounds() -> u64 { + SUMCHECK_ROUNDS.load(Ordering::Relaxed) +} + +pub fn evaluate_calls() -> u64 { + EVALUATE_CALLS.load(Ordering::Relaxed) +} + +pub fn tree_calls() -> u64 { + TREE_CALLS.load(Ordering::Relaxed) +} + +pub fn factor_calls() -> u64 { + FACTOR_CALLS.load(Ordering::Relaxed) +} + +pub fn open_calls() -> u64 { + OPEN_CALLS.load(Ordering::Relaxed) +} + +pub fn reset_call_counters() { + COMMIT_CALLS.store(0, Ordering::Relaxed); + SUMCHECK_CALLS.store(0, Ordering::Relaxed); + SUMCHECK_ROUNDS.store(0, Ordering::Relaxed); + EVALUATE_CALLS.store(0, Ordering::Relaxed); + TREE_CALLS.store(0, Ordering::Relaxed); + FACTOR_CALLS.store(0, Ordering::Relaxed); + OPEN_CALLS.store(0, Ordering::Relaxed); +} + +/// A sumcheck's round proofs, the challenges they drew, and what every slot +/// was bound to — the factors are folded where they lie, so their values at +/// the sumcheck's point are already there when the rounds end. +/// What a resident sumcheck's rounds on device leave: the rounds, the point +/// they drew, and **every** factor — resident and not — as the last fold left +/// it. One value each when the device ran the cube out, a cube when it stopped +/// at the crossover for the caller to finish. +type ResidentRounds = ( + Vec>, + Vec>, + Vec>, +); + +type SumcheckRounds = ( + Vec>, + Vec>, + Vec>, +); + +/// Op tags the sumcheck kernel reads. MUST stay in sync with +/// `crypto/math-cuda/kernels/sumcheck.cu`. +pub mod op { + pub const FIXED: u32 = 0; + pub const VAR: u32 = 1; + pub const ADD: u32 = 2; + pub const SUB: u32 = 3; + pub const MUL: u32 = 4; + pub const NEG: u32 = 5; +} + +/// A program lowered for the device: the nodes (two u64 each, `op | a << 32` +/// then `b | res << 32`), the ext3 constants they read (three u64 each), the +/// slot file's width and the slot the root lands in. +#[derive(Clone, Debug)] +pub struct Lowered { + pub nodes: Vec, + pub consts: Vec, + pub num_slots: usize, + pub root_slot: u32, +} + +/// Slots the round kernel will hold per thread before the dispatch declines. +/// +/// The slot file is `slots * 24 * threads` bytes and the scratch budget is +/// fixed, so a wider program buys fewer threads. This is where that stops +/// being a trade: at 8192 live values a single block of 256 already wants +/// 48 MiB, and a launch of one block is not a launch. +/// +/// It is a cliff, not a dial. The real AIRs peak near a thousand — the widest +/// precompile lowers to 1036 — and a cap below that sends the tables with the +/// *most* work per row to the host, which is where they cost the most. +pub const MAX_SLOTS: usize = 8192; + +pub fn reserve_room(_bytes: u64) -> Option { + None +} + +/// A promise no device made. Never constructed. +#[derive(Debug)] +pub struct DeviceRoom(std::convert::Infallible); + +/// Assigns every step a slot, reusing the slot of a value whose last read has +/// passed. +/// +/// This is what makes the kernel possible at all: the precompile tables compile +/// to tens of thousands of steps, and a slot per step would be a megabyte of +/// scratch per thread. +pub fn lower(program: &crate::program::Program) -> Option +where + E: math::field::traits::IsField + 'static, +{ + use crate::program::Op; + use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField as Ext3; + + // The blob is ext3: the kernel reads three limbs per value, whether or not + // this particular program happens to hold a constant that would say so. + if std::any::TypeId::of::() != std::any::TypeId::of::() { + return None; + } + let steps = program.steps(); + // Last step that reads each value; the root is read by the caller, so it is + // never freed. + let mut last_use: Vec = vec![0; steps.len()]; + for (i, step) in steps.iter().enumerate() { + let mut mark = |operand: u32| last_use[operand as usize] = i; + match *step { + Op::Fixed(_) | Op::Var(_) => {} + Op::Neg(a) => mark(a), + Op::Add(a, b) | Op::Sub(a, b) | Op::Mul(a, b) => { + mark(a); + mark(b); + } + } + } + last_use[program.root() as usize] = usize::MAX; + + let mut slot_of: Vec = vec![u32::MAX; steps.len()]; + let mut free: Vec = Vec::new(); + let mut num_slots = 0usize; + let mut nodes: Vec = Vec::with_capacity(steps.len() * 2); + let mut consts: Vec = Vec::new(); + + for (i, step) in steps.iter().enumerate() { + let (op, a, b, operands) = match *step { + Op::Fixed(ref value) => { + let at = consts.len() / 3; + consts.extend_from_slice(&ext3_raw(value)?); + (op::FIXED, at as u32, 0, [None, None]) + } + Op::Var(slot) => (op::VAR, slot, 0, [None, None]), + Op::Neg(a) => (op::NEG, slot_of[a as usize], 0, [Some(a), None]), + Op::Add(a, b) | Op::Sub(a, b) | Op::Mul(a, b) => { + let tag = match *step { + Op::Add(..) => op::ADD, + Op::Sub(..) => op::SUB, + _ => op::MUL, + }; + ( + tag, + slot_of[a as usize], + slot_of[b as usize], + [Some(a), Some(b)], + ) + } + }; + // Freed before the result is allocated: the kernel loads both operands + // before it stores, so the result may take a slot this step frees. A + // value read twice — `x·x`, which is what a squaring compiles to — + // frees its slot once, or two later values would be handed the same + // one. + if let Some(x) = operands[0].filter(|x| last_use[*x as usize] == i) { + free.push(slot_of[x as usize]); + } + if let Some(y) = + operands[1].filter(|y| last_use[*y as usize] == i && operands[0] != Some(*y)) + { + free.push(slot_of[y as usize]); + } + let res = free.pop().unwrap_or_else(|| { + let slot = num_slots as u32; + num_slots += 1; + slot + }); + if num_slots > MAX_SLOTS { + return None; + } + slot_of[i] = res; + nodes.push(u64::from(op) | (u64::from(a) << 32)); + nodes.push(u64::from(b) | (u64::from(res) << 32)); + } + + Some(Lowered { + nodes, + consts, + num_slots, + root_slot: slot_of[program.root() as usize], + }) +} + +/// An ext3 element's three limbs, or `None` when `E` is not that field. +pub fn ext3_raw(value: &math::field::element::FieldElement) -> Option<[u64; 3]> +where + E: math::field::traits::IsField + 'static, +{ + use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField as Ext3; + if std::any::TypeId::of::() != std::any::TypeId::of::() { + return None; + } + // SAFETY: `E == Ext3`, whose `FieldElement` is a transparent wrapper over + // three Goldilocks limbs, each transparent over its `u64`. + let limbs = unsafe { *(value as *const _ as *const [u64; 3]) }; + Some(limbs) +} + +/// Rebuilds an ext3 element from its three limbs. +pub fn ext3_from_raw(limbs: &[u64]) -> math::field::element::FieldElement +where + E: math::field::traits::IsField + 'static, +{ + use math::field::element::FieldElement; + use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField as Ext3; + use math::field::goldilocks::GoldilocksField as Gl; + + let value = FieldElement::::new([ + FieldElement::::from_raw(limbs[0]), + FieldElement::::from_raw(limbs[1]), + FieldElement::::from_raw(limbs[2]), + ]); + // SAFETY: only called under a TypeId check that `E == Ext3`. + unsafe { core::mem::transmute_copy::, FieldElement>(&value) } +} + +pub(crate) fn prove_sumcheck( + _polys: &[crate::mle::Mle], + _program: &crate::program::Program, + _degree: usize, + _host_cube: usize, + _challenge: impl FnMut( + &[math::field::element::FieldElement], + ) -> math::field::element::FieldElement, +) -> Option, crate::Error>> +where + E: math::field::traits::IsField + 'static, +{ + None +} + +pub(crate) fn prove_sumcheck_resident( + _resident: &DeviceFactors, + _extra: &[crate::mle::Mle], + _program: &crate::program::Program, + _degree: usize, + _challenge: impl FnMut( + &[math::field::element::FieldElement], + ) -> math::field::element::FieldElement, +) -> Option, crate::Error>> +where + E: math::field::traits::IsField + 'static, +{ + None +} + +pub(crate) fn evaluate_many_base( + _columns: &[crate::mle::Mle], + _point: &[math::field::element::FieldElement], + _resident: Option<(&ResidentColumns, usize)>, +) -> Option>> +where + F: math::field::traits::IsField + 'static, + E: math::field::traits::IsField + 'static, +{ + None +} + +pub(crate) fn evaluate_mle( + _evals: &[math::field::element::FieldElement], + _point: &[math::field::element::FieldElement], +) -> Option> +where + C: math::field::traits::IsField + 'static, + E: math::field::traits::IsField + 'static, +{ + None +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::program::Builder; + use math::field::element::FieldElement; + use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField as Ext3; + use math::field::goldilocks::GoldilocksField as Gl; + + type FE = FieldElement; + + /// The kernel's walk, in Rust: the same slot file, the same node encoding. + /// + /// This is what pins the lowering without a device — a slot freed too early + /// or an operand read from the wrong class shows up here as a wrong value, + /// not as a proof that does not verify an hour later. + fn run_lowered(lowered: &Lowered, values: &[FE]) -> FE { + let mut slots = vec![FE::zero(); lowered.num_slots]; + for node in lowered.nodes.chunks_exact(2) { + let op = (node[0] & 0xFFFF_FFFF) as u32; + let a = (node[0] >> 32) as u32 as usize; + let b = (node[1] & 0xFFFF_FFFF) as u32 as usize; + let res = (node[1] >> 32) as u32 as usize; + slots[res] = match op { + op::FIXED => ext3_from_raw::(&lowered.consts[a * 3..a * 3 + 3]), + op::VAR => values[a], + op::ADD => slots[a] + slots[b], + op::SUB => slots[a] - slots[b], + op::MUL => slots[a] * slots[b], + op::NEG => -slots[a], + _ => panic!("unknown op {op}"), + }; + } + slots[lowered.root_slot as usize] + } + + fn values(n: usize) -> Vec { + (0..n as u64) + .map(|i| { + FE::new([ + FieldElement::::from(i * 31 + 7), + FieldElement::::from(i * 17 + 2), + FieldElement::::from(i + 5), + ]) + }) + .collect() + } + + /// Every op, a constant, and a chain long enough that slots have to be + /// recycled. + fn sample_program() -> crate::program::Program { + let mut b = Builder::::new(); + let mut acc = b.var(0); + for slot in 1..6 { + let v = b.var(slot); + let doubled = b.add(v, v); + let scaled = b.mul(doubled, acc); + let shifted = b.sub(scaled, v); + acc = b.neg(shifted); + } + let seven = b.fixed(FE::from(7u64)); + let root = b.add(acc, seven); + b.finish(root).unwrap() + } + + #[test] + fn the_lowered_program_computes_what_the_program_does() { + let program = sample_program(); + let lowered = lower(&program).expect("lowers"); + let v = values(6); + let mut scratch = Vec::new(); + assert_eq!(run_lowered(&lowered, &v), program.eval(&v, &mut scratch)); + } + + /// A value read twice in the step that kills it — `x·x` — must not free its + /// slot twice, or two later values are handed the same one and the second + /// clobbers the first. Squarings are everywhere in a real constraint + /// program, so this is the shape that matters. + #[test] + fn a_value_read_twice_frees_its_slot_once() { + let mut b = Builder::::new(); + let mut acc = b.var(0); + // Each square kills its operand, and the sums below keep enough values + // live that a doubly-freed slot gets reused while it is still needed. + let mut squares = Vec::new(); + for slot in 1..8 { + let v = b.var(slot); + let squared = b.mul(v, v); + let with_acc = b.add(squared, acc); + squares.push(with_acc); + acc = b.mul(with_acc, with_acc); + } + squares.push(acc); + let root = b.sum(&squares); + let program = b.finish(root).unwrap(); + + let lowered = lower(&program).expect("lowers"); + let v = values(8); + let mut scratch = Vec::new(); + assert_eq!(run_lowered(&lowered, &v), program.eval(&v, &mut scratch)); + } + + /// The point of the slot file: a long chain of dead intermediates does not + /// widen it. + #[test] + fn slots_are_reused_once_a_value_is_dead() { + let lowered = lower(&sample_program()).expect("lowers"); + assert!( + lowered.num_slots < lowered.nodes.len() / 2, + "{} slots for {} steps is no reuse at all", + lowered.num_slots, + lowered.nodes.len() / 2 + ); + } + + /// A program wider than the slot file declines rather than asking a device + /// for scratch it cannot have. + #[test] + fn a_program_past_the_slot_ceiling_declines() { + let mut b = Builder::::new(); + // Every value stays live to the end, so the slots cannot be recycled. + let terms: Vec = (0..=MAX_SLOTS).map(|slot| b.var(slot)).collect(); + let root = b.sum(&terms); + let program = b.finish(root).unwrap(); + assert!(lower(&program).is_none()); + } + + #[test] + fn a_field_the_kernel_does_not_cover_declines() { + let mut b = Builder::::new(); + let root = b.var(0); + let program = b.finish(root).unwrap(); + assert!(lower(&program).is_none()); + } +} + +/// A device tree the build declined to make. Never constructed. +#[derive(Debug)] +pub struct DeviceTree(std::convert::Infallible); + +impl DeviceTree { + pub(crate) fn num_layers(&self) -> usize { + match self.0 {} + } + + pub(crate) fn output( + &self, + ) -> Result< + ( + math::field::element::FieldElement, + math::field::element::FieldElement, + ), + crate::Error, + > + where + E: math::field::traits::IsField + 'static, + { + match self.0 {} + } + + pub(crate) fn prove_layer( + &self, + _layer: usize, + _point: &[math::field::element::FieldElement], + _program: &crate::program::Program, + _degree: usize, + _tail: usize, + _challenge: impl FnMut( + &[math::field::element::FieldElement], + ) -> math::field::element::FieldElement, + ) -> Option, crate::Error>> + where + E: math::field::traits::IsField + 'static, + { + match self.0 {} + } + + pub(crate) fn layer_to_host( + &self, + _layer: usize, + ) -> Option<(crate::mle::Mle, crate::mle::Mle)> + where + E: math::field::traits::IsField + 'static, + { + match self.0 {} + } +} + +/// What a layer's rounds on device leave: the rounds themselves, the point +/// they drew, and the five factors as the last fold left them. +/// +/// The factors are one value each when the device ran the layer out, and a +/// cube when it stopped at the crossover for the host to finish — the caller +/// carries on from them either way. +pub type LayerRounds = ( + Vec>, + Vec>, + Vec>, +); + +pub(crate) fn build_tree(_p: &crate::mle::Mle, _q: &crate::mle::Mle) -> Option +where + E: math::field::traits::IsField + 'static, +{ + None +} + +/// One that could not be made. Never constructed. +pub struct ResidentColumns(std::convert::Infallible); + +impl std::fmt::Debug for ResidentColumns { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("ResidentColumns") + } +} + +pub fn upload_columns(_columns: &[&crate::mle::Mle]) -> Option +where + F: math::field::traits::IsField + 'static, +{ + None +} + +/// Factors a build declined to upload. Never constructed. +#[derive(Debug)] +pub struct DeviceFactors(std::convert::Infallible); + +pub fn input_layer_tree( + _factors: std::sync::Arc, + _numerators: Vec>, + _denominators: Vec>, +) -> Option +where + E: math::field::traits::IsField + 'static, +{ + None +} diff --git a/crypto/multilinear/src/lib.rs b/crypto/multilinear/src/lib.rs new file mode 100644 index 000000000..aee786a0a --- /dev/null +++ b/crypto/multilinear/src/lib.rs @@ -0,0 +1,120 @@ +//! Multilinear machinery for a sumcheck-based proof system: extensions over the +//! Boolean hypercube, batched sumcheck, zerocheck and LogUp-GKR. +//! +//! One brick and two uses. [`sumcheck`] reduces a sum over `2^n` points to one +//! evaluation; on top of it [`zerocheck`] says a polynomial vanishes on the +//! whole cube and [`logup`] says a bus balances. [`claim_reduce`] ties the +//! shifted views a table reads back to the columns they read. +//! +//! Nothing here is wired into a prover yet. + +pub mod batch; +pub mod claim_reduce; +pub mod eq; +pub mod gkr; +pub mod gpu; +pub mod logup; +pub mod mle; +pub mod poly; +pub mod program; +pub mod selector; +pub mod sumcheck; +pub mod zerocheck; + +use math::field::{element::FieldElement, traits::IsField}; +use thiserror::Error; + +/// Below this a pass stays on the thread that asked for it: handing a slice to +/// the pool costs tens of microseconds whatever is in it, and a sumcheck's last +/// rounds are over a few hundred values. +#[cfg(feature = "parallel")] +pub(crate) const SERIAL_BELOW: usize = 1 << 12; + +/// The cube below which a sumcheck's rounds belong here, for a rule the host +/// evaluates directly: a device round costs the same whatever the cube — one +/// thread walks the whole rule at each index — and a host round is linear in it. +pub(crate) const HOST_CUBE_DIRECT: usize = 1 << 9; + +/// The same for a rule the host walks through the program interpreter, which +/// costs about ten times a multiplication written out per step while the device +/// round costs the same either way. +pub(crate) const HOST_CUBE_COMPILED: usize = 1 << 5; + +/// `[1, gamma, gamma^2, ..]` — the weights a batching challenge expands into. +pub(crate) fn challenge_powers( + gamma: &FieldElement, + count: usize, +) -> Vec> { + let mut acc = FieldElement::::one(); + (0..count) + .map(|_| { + let current = acc.clone(); + acc = &acc * gamma; + current + }) + .collect() +} + +#[derive(Debug, Error, PartialEq, Eq)] +pub enum Error { + #[error("expected a power-of-two number of evaluations, got {0}")] + NotPowerOfTwo(usize), + #[error("expected {expected} variables, got {got}")] + VariableCountMismatch { expected: usize, got: usize }, + #[error("polynomial has no variables left to fold")] + NoVariablesLeft, + #[error("term references polynomial index {index}, but only {len} are registered")] + UnknownPolynomial { index: usize, len: usize }, + #[error("virtual polynomial has no terms")] + EmptyPolynomial, + #[error("round {round}: expected a degree-{expected} polynomial, got {got} evaluations")] + RoundDegreeMismatch { + round: usize, + expected: usize, + got: usize, + }, + #[error("proof has {got} rounds, expected {expected}")] + RoundCountMismatch { expected: usize, got: usize }, + #[error("final evaluation does not match the oracle: expected {expected}, got {got}")] + FinalEvaluationMismatch { expected: String, got: String }, + #[error("{exemptions} exempted steps exceed the {size}-step trace")] + TooManyExemptions { exemptions: usize, size: usize }, + #[error("layer {layer}: the sumcheck residual does not match the fraction-fold relation")] + LayerRelationMismatch { layer: usize }, + #[error("a column on {column_vars} variables does not fit a {n_stack}-variable stack")] + ColumnTallerThanStack { column_vars: usize, n_stack: usize }, + #[error("no subgroup of order 2^{l_skip} exists (field two-adicity is {two_adicity})")] + SkipDomainUnavailable { l_skip: usize, two_adicity: usize }, + #[error("{coefficients} coefficients do not fit a domain of {domain} points")] + CodewordTooShort { coefficients: usize, domain: usize }, + #[error("query index {index} is outside the {bound} committed leaves")] + QueryOutOfRange { index: usize, bound: usize }, + #[error("expected {expected} query openings, got {got}")] + QueryCountMismatch { expected: usize, got: usize }, + #[error("query {query}: the Merkle opening does not match the commitment")] + OpeningRejected { query: usize }, + #[error("query {query}: the folded block does not match the committed successor")] + FoldInconsistent { query: usize }, + #[error("the folded codeword and the sumcheck disagree on the evaluation")] + EvaluationMismatch, + #[error("eq(z, alpha) vanished, leaving the evaluation unconstrained")] + DegenerateEvaluationPoint, + #[error("the out-of-domain point landed inside the evaluation domain")] + OodPointInDomain, + #[error("no proof-of-work nonce found for {bits} bits")] + GrindingFailed { bits: u8 }, + #[error("a proof-of-work nonce does not carry the required {bits} bits")] + GrindingRejected { bits: u8 }, + #[error("the buses do not balance across the proof")] + BusImbalance, + #[error("the statements rebuilt from the factor values are not what the batch demands")] + BatchMismatch, + #[error("a shifted read is not the shift of the column it claims to shift")] + ShiftedReadMismatch, + #[error("column {column}: its claimed value does not match the commitment")] + ColumnOpeningRejected { column: usize }, + /// A device path that had already drawn a challenge cannot be retried on + /// the host: the transcript has moved. + #[error("the device failed mid-{stage}, after the transcript had moved")] + DeviceFailed { stage: &'static str }, +} diff --git a/crypto/multilinear/src/logup.rs b/crypto/multilinear/src/logup.rs new file mode 100644 index 000000000..9a234c4f5 --- /dev/null +++ b/crypto/multilinear/src/logup.rs @@ -0,0 +1,551 @@ +//! A bus as a fraction tree over the trace, and the rules that settle its +//! input-layer claim against the trace's factors. +//! +//! A bus balances when `Σ ± mult / (z − fingerprint) = 0`. Both parts are +//! **affine** in the trace columns — a multiplicity is a linear combination of +//! flag columns, a fingerprint is limbs weighted by powers of two and bus +//! elements weighted by powers of a challenge — so neither is ever a column of +//! its own. [`FractionTree`](crate::gkr::FractionTree) proves the sum, and what +//! it leaves is one claim about the input layer. +//! +//! That claim reduces to **two rules**, whatever the number of interactions. +//! The input layer is indexed by `(interaction, row)` with the interaction in +//! the high variables, so the claim point splits and +//! `p(hi ‖ lo) = Σ_i eq(hi, i)·p_i(lo)`: the interaction weights are constants +//! the verifier computes, leaving a single affine combination of factors per +//! side. Both rules are degree 2 once the row weight is counted. + +use math::field::{element::FieldElement, traits::IsField}; + +#[cfg(feature = "parallel")] +use rayon::prelude::*; + +use crate::{ + Error, + batch::Rule, + eq::eq_evals, + gkr::FractionLayer, + mle::Mle, + program::{Builder, Program}, +}; + +/// An affine expression over the sumcheck factors: `Σ c_j·f_{s_j} + k`. +#[derive(Clone, Debug)] +pub struct Affine { + terms: Vec<(usize, FieldElement)>, + constant: FieldElement, +} + +impl Affine { + /// `terms` pairs a factor slot with its coefficient. + pub fn new(terms: Vec<(usize, FieldElement)>, constant: FieldElement) -> Self { + Self { terms, constant } + } + + pub fn constant(value: FieldElement) -> Self { + Self::new(Vec::new(), value) + } + + /// A single factor read with coefficient one. + pub fn factor(slot: usize) -> Self { + Self::new(vec![(slot, FieldElement::one())], FieldElement::zero()) + } + + pub fn terms(&self) -> &[(usize, FieldElement)] { + &self.terms + } + + /// The expression's value, given every factor's value at a point. + pub fn evaluate(&self, values: &[FieldElement]) -> FieldElement { + self.terms + .iter() + .fold(self.constant.clone(), |acc, (slot, coefficient)| { + acc + coefficient * &values[*slot] + }) + } + + /// The expression as steps over the factor values, returning the step + /// holding it. + pub fn emit(&self, builder: &mut Builder) -> u32 { + if self.terms.is_empty() { + return builder.fixed(self.constant.clone()); + } + let terms: Vec<(u32, FieldElement)> = self + .terms + .iter() + .map(|(slot, coefficient)| (builder.var(*slot), coefficient.clone())) + .collect(); + let sum = builder.weighted_sum(&terms); + if self.constant == FieldElement::zero() { + return sum; + } + let constant = builder.fixed(self.constant.clone()); + builder.add(sum, constant) + } + + /// The expression's table over the cube. `factors` must not be empty: its + /// first entry sets the height. + pub fn table(&self, factors: &[Mle]) -> Result, Error> + where + FieldElement: Send + Sync, + { + let size = factors.first().ok_or(Error::EmptyPolynomial)?.len(); + let mut evals = vec![self.constant.clone(); size]; + for (slot, coefficient) in &self.terms { + let factor = factors.get(*slot).ok_or(Error::UnknownPolynomial { + index: *slot, + len: factors.len(), + })?; + if factor.len() != size { + return Err(Error::VariableCountMismatch { + expected: size.trailing_zeros() as usize, + got: factor.num_vars(), + }); + } + // One pass per term over the whole column. The bus builds one of + // these per interaction, so it is a hot loop on a real table. + #[cfg(feature = "parallel")] + evals + .par_iter_mut() + .zip(factor.evals().par_iter()) + .for_each(|(slot, value)| *slot += coefficient * value); + #[cfg(not(feature = "parallel"))] + for (slot, value) in evals.iter_mut().zip(factor.evals()) { + *slot += coefficient * value; + } + } + Mle::new(evals) + } +} + +/// One bus interaction on one table. +/// +/// `numerator` is the **signed** multiplicity — a receiver's sign is already in +/// its coefficients — and `denominator` is `z − fingerprint`. +#[derive(Clone, Debug)] +pub struct Interaction { + pub numerator: Affine, + pub denominator: Affine, +} + +impl Interaction { + pub fn new(numerator: Affine, denominator: Affine) -> Self { + Self { + numerator, + denominator, + } + } +} + +/// The number of variables the input layer spans: the rows plus the bits that +/// index the interaction. +pub fn input_layer_vars(interactions: usize, num_row_vars: usize) -> usize { + num_row_vars + interactions.next_power_of_two().trailing_zeros() as usize +} + +/// The fraction tree's input layer, indexed by `(interaction, row)` with the +/// interaction in the high variables. +/// +/// Interactions are padded up to a power of two with `0/1`, which the tree adds +/// without moving the sum. +pub fn input_layer( + interactions: &[Interaction], + factors: &[Mle], +) -> Result, Error> { + if interactions.is_empty() { + return Err(Error::EmptyPolynomial); + } + let size = factors.first().ok_or(Error::EmptyPolynomial)?.len(); + let slots = interactions.len().next_power_of_two(); + + let mut p = Vec::with_capacity(slots * size); + let mut q = Vec::with_capacity(slots * size); + for interaction in interactions { + p.extend(interaction.numerator.table(factors)?.into_evals()); + q.extend(interaction.denominator.table(factors)?.into_evals()); + } + for _ in interactions.len()..slots { + p.extend(std::iter::repeat_n(FieldElement::::zero(), size)); + q.extend(std::iter::repeat_n(FieldElement::::one(), size)); + } + + FractionLayer::new(Mle::new(p)?, Mle::new(q)?) +} + +/// The input layer and the tree above it, built where the factors already are. +/// +/// The layer is `interactions × rows` fractions — the biggest thing a table's +/// argument builds — and every one of its cells is an affine expression over +/// the factors, which is a program the device can run. +/// +/// `None` when the device declines; the caller then builds the layer here. +pub fn resident_tree( + interactions: &[Interaction], + factors: std::sync::Arc, +) -> Option> { + if interactions.is_empty() { + return None; + } + let emit = |side: &Affine| { + let mut builder = Builder::::new(); + let root = side.emit(&mut builder); + builder.finish(root).ok() + }; + let numerators: Vec> = interactions + .iter() + .map(|i| emit(&i.numerator)) + .collect::>()?; + let denominators: Vec> = interactions + .iter() + .map(|i| emit(&i.denominator)) + .collect::>()?; + + let tree = crate::gpu::input_layer_tree(factors, numerators, denominators)?; + crate::gkr::FractionTree::from_device(tree).ok() +} + +/// What the batch needs to settle a bus's input-layer claim. +pub struct BusStatements<'a, E: IsField> { + pub numerator: Rule<'a, E>, + pub denominator: Rule<'a, E>, + /// The row half of the claim point — where the weight table belongs. + pub row_point: Vec>, +} + +/// Turns a GKR input-layer claim into two rules over the trace's factors. +/// +/// `weight` is the factor slot holding `eq(row_point, ·)`, which the caller +/// adds as a public factor. Both sides must build these from the same +/// interactions: they are the bus's structure, not proof data. +pub fn claim_statements<'a, E: IsField + 'static>( + interactions: &'a [Interaction], + claim_point: &[FieldElement], + num_row_vars: usize, + weight: usize, +) -> Result, Error> { + let expected = input_layer_vars(interactions.len(), num_row_vars); + if claim_point.len() != expected { + return Err(Error::VariableCountMismatch { + expected, + got: claim_point.len(), + }); + } + let (interaction_point, row_point) = claim_point.split_at(claim_point.len() - num_row_vars); + + // Weight per interaction; the tail belongs to the 0/1 padding slots, whose + // denominators are one and whose numerators vanish. + let weights = eq_evals(interaction_point); + let padding = weights[interactions.len()..] + .iter() + .fold(FieldElement::::zero(), |acc, w| acc + w); + let live = weights[..interactions.len()].to_vec(); + + // `Σ_i w_i · side_i(f)` (plus the padding, where it belongs), times the row + // weight — the shape both sides of the bus statement take. + let weighted = |sides: Vec<&Affine>, constant: Option>| { + let mut builder = Builder::::new(); + let mut terms: Vec<(u32, FieldElement)> = sides + .into_iter() + .zip(&live) + .map(|(side, w)| (side.emit(&mut builder), w.clone())) + .collect(); + if let Some(constant) = constant { + let step = builder.fixed(constant); + terms.push((step, FieldElement::one())); + } + let sum = builder.weighted_sum(&terms); + let row = builder.var(weight); + let root = builder.mul(row, sum); + builder.finish(root) + }; + + let numerator: Program = + weighted(interactions.iter().map(|i| &i.numerator).collect(), None)?; + let denominator: Program = weighted( + interactions.iter().map(|i| &i.denominator).collect(), + Some(padding), + )?; + let numerator = Rule::compiled(2, numerator); + let denominator = Rule::compiled(2, denominator); + + Ok(BusStatements { + numerator, + denominator, + row_point: row_point.to_vec(), + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crypto::fiat_shamir::default_transcript::DefaultTranscript; + use math::field::goldilocks::GoldilocksField as F; + + use crate::{ + batch, + eq::eq_mle, + gkr::{self, FractionTree}, + }; + + type FE = FieldElement; + + fn transcript() -> DefaultTranscript { + DefaultTranscript::::new(b"logup-test") + } + + /// Factor layout: the value, the two half-flags, and the row weight. + const VALUE: usize = 0; + const LOW: usize = 1; + const HIGH: usize = 2; + const WEIGHT: usize = 3; + + /// The low half sends each value and the high half receives the same one, + /// so the bus balances *because of the values*, not by construction. + fn columns(num_vars: usize) -> Vec> { + let size = 1usize << num_vars; + let half = size / 2; + let value: Vec = (0..size).map(|i| FE::from((i % half) as u64 + 1)).collect(); + let low: Vec = (0..size) + .map(|i| if i < half { FE::one() } else { FE::zero() }) + .collect(); + let high: Vec = low.iter().map(|v| FE::one() - v).collect(); + vec![ + Mle::new(value).unwrap(), + Mle::new(low).unwrap(), + Mle::new(high).unwrap(), + ] + } + + /// `+low / (z − alpha·value)` against `−high / (z − alpha·value)`. + fn balanced(z: FE, alpha: FE, _num_vars: usize) -> Vec> { + vec![ + Interaction::new(Affine::factor(LOW), Affine::new(vec![(VALUE, -alpha)], z)), + Interaction::new( + Affine::new(vec![(HIGH, -FE::one())], FE::zero()), + Affine::new(vec![(VALUE, -alpha)], z), + ), + ] + } + + #[test] + fn an_affine_expression_evaluates_and_tabulates_alike() { + let factors = columns(3); + let expr = Affine::new(vec![(VALUE, FE::from(7)), (LOW, FE::from(11))], FE::from(5)); + let table = expr.table(&factors).unwrap(); + + for i in 0..factors[0].len() { + let values: Vec = factors.iter().map(|f| f.evals()[i]).collect(); + assert_eq!(table.evals()[i], expr.evaluate(&values), "row {i}"); + } + } + + #[test] + fn the_input_layer_lays_interactions_out_row_by_row() { + let num_vars = 3; + let size = 1usize << num_vars; + let factors = columns(num_vars); + let interactions = balanced(FE::from(97), FE::from(31), num_vars); + let layer = input_layer(&interactions, &factors).unwrap(); + + assert_eq!(layer.num_vars(), input_layer_vars(2, num_vars)); + for (i, interaction) in interactions.iter().enumerate() { + let p = interaction.numerator.table(&factors).unwrap(); + let q = interaction.denominator.table(&factors).unwrap(); + for row in 0..size { + // The interaction sits in the high variables. + assert_eq!(layer.p.evals()[i * size + row], p.evals()[row]); + assert_eq!(layer.q.evals()[i * size + row], q.evals()[row]); + } + } + } + + #[test] + fn padding_slots_are_the_zero_fraction() { + let num_vars = 2; + let size = 1usize << num_vars; + let factors = columns(num_vars); + let z = FE::from(97); + let alpha = FE::from(31); + // Three interactions pad up to four slots. + let mut interactions = balanced(z, alpha, num_vars); + interactions.push(Interaction::new( + Affine::constant(FE::zero()), + Affine::constant(z), + )); + + let layer = input_layer(&interactions, &factors).unwrap(); + assert_eq!(layer.num_vars(), num_vars + 2); + for row in 0..size { + assert_eq!(layer.p.evals()[3 * size + row], FE::zero()); + assert_eq!(layer.q.evals()[3 * size + row], FE::one()); + } + } + + #[test] + fn a_balanced_bus_has_a_zero_output_numerator() { + let num_vars = 3; + let factors = columns(num_vars); + let interactions = balanced(FE::from(97), FE::from(31), num_vars); + let tree = FractionTree::build(input_layer(&interactions, &factors).unwrap()).unwrap(); + + let (p, q) = tree.output(); + assert_eq!(p, FE::zero()); + assert_ne!(q, FE::zero()); + } + + #[test] + fn an_unbalanced_bus_does_not() { + let num_vars = 3; + let factors = columns(num_vars); + let z = FE::from(97); + let alpha = FE::from(31); + // Only the send: nothing cancels it. + let interactions = vec![balanced(z, alpha, num_vars).remove(0)]; + let tree = FractionTree::build(input_layer(&interactions, &factors).unwrap()).unwrap(); + + assert_ne!(tree.output().0, FE::zero()); + } + + #[test] + fn a_value_received_that_was_never_sent_unbalances_the_bus() { + // The multisets are what balance, so moving one value apart is enough. + let num_vars = 3; + let mut factors = columns(num_vars); + let mut value = factors[0].evals().to_vec(); + value[5] += FE::one(); + factors[0] = Mle::new(value).unwrap(); + + let interactions = balanced(FE::from(97), FE::from(31), num_vars); + let tree = FractionTree::build(input_layer(&interactions, &factors).unwrap()).unwrap(); + assert_ne!(tree.output().0, FE::zero()); + } + + /// The identity the batch leans on: each rule sums over the row cube to the + /// claim the GKR handed back. + #[test] + fn the_input_claim_is_what_the_rules_sum_to() { + for num_vars in 1..=4usize { + let factors = columns(num_vars); + let interactions = balanced(FE::from(97), FE::from(31), num_vars); + let tree = FractionTree::build(input_layer(&interactions, &factors).unwrap()).unwrap(); + let out = gkr::prove(&tree, &mut transcript()).unwrap(); + + let statements = + claim_statements(&interactions, &out.claim.point, num_vars, WEIGHT).unwrap(); + let mut with_weight = factors.clone(); + with_weight.push(eq_mle(&statements.row_point).unwrap()); + + let sum = |rule: &Rule<'_, F>| { + (0..(1usize << num_vars)).fold(FE::zero(), |acc, i| { + let values: Vec = with_weight.iter().map(|f| f.evals()[i]).collect(); + acc + rule.apply(&values) + }) + }; + assert_eq!(sum(&statements.numerator), out.claim.p, "n={num_vars}"); + assert_eq!(sum(&statements.denominator), out.claim.q, "n={num_vars}"); + } + } + + /// Same identity with the padding in play, where a wrong padding weight + /// would show up in the denominator. + #[test] + fn the_claim_holds_with_padded_interactions() { + let num_vars = 3; + let factors = columns(num_vars); + let z = FE::from(97); + let alpha = FE::from(31); + let mut interactions = balanced(z, alpha, num_vars); + interactions.push(Interaction::new( + Affine::constant(FE::zero()), + Affine::constant(z), + )); + + let tree = FractionTree::build(input_layer(&interactions, &factors).unwrap()).unwrap(); + let out = gkr::prove(&tree, &mut transcript()).unwrap(); + let statements = + claim_statements(&interactions, &out.claim.point, num_vars, WEIGHT).unwrap(); + + let mut with_weight = factors; + with_weight.push(eq_mle(&statements.row_point).unwrap()); + let sum = |rule: &Rule<'_, F>| { + (0..(1usize << num_vars)).fold(FE::zero(), |acc, i| { + let values: Vec = with_weight.iter().map(|f| f.evals()[i]).collect(); + acc + rule.apply(&values) + }) + }; + assert_eq!(sum(&statements.numerator), out.claim.p); + assert_eq!(sum(&statements.denominator), out.claim.q); + } + + /// Two rules, one batched sumcheck, whatever the interaction count. + #[test] + fn the_bus_costs_two_rules_in_one_sumcheck() { + let num_vars = 3; + let factors = columns(num_vars); + let interactions = balanced(FE::from(97), FE::from(31), num_vars); + let tree = FractionTree::build(input_layer(&interactions, &factors).unwrap()).unwrap(); + let output = tree.output(); + + let mut prover = transcript(); + let out = gkr::prove(&tree, &mut prover).unwrap(); + let statements = + claim_statements(&interactions, &out.claim.point, num_vars, WEIGHT).unwrap(); + let mut with_weight = factors.clone(); + with_weight.push(eq_mle(&statements.row_point).unwrap()); + let claims = [out.claim.p, out.claim.q]; + + let (proof, point) = batch::prove( + with_weight.clone(), + vec![statements.numerator, statements.denominator], + &claims, + &mut prover, + ) + .unwrap(); + assert_eq!(proof.rounds.len(), num_vars); + + let mut verifier = transcript(); + let claim = gkr::verify(&out.proof, output, &mut verifier).unwrap(); + let statements = claim_statements(&interactions, &claim.point, num_vars, WEIGHT).unwrap(); + let row_point = statements.row_point.clone(); + + let checked = batch::verify( + &proof, + &[statements.numerator, statements.denominator], + &claims, + |at: &[FE]| { + // Only the trace factors travel; the weight is recomputed. + let mut values: Vec = factors + .iter() + .map(|f| f.evaluate(at)) + .collect::>()?; + values.push(crate::eq::eq_eval(&row_point, at)?); + Ok(values) + }, + num_vars, + &mut verifier, + ) + .unwrap(); + assert_eq!(checked, point); + } + + #[test] + fn a_claim_point_of_the_wrong_arity_is_rejected() { + let interactions = balanced(FE::from(97), FE::from(31), 3); + let result = claim_statements(&interactions, &[FE::one(); 3], 3, WEIGHT); + assert_eq!( + result.err(), + Some(Error::VariableCountMismatch { + expected: 4, + got: 3 + }) + ); + } + + #[test] + fn an_empty_bus_is_rejected() { + assert_eq!( + input_layer::(&[], &columns(3)).unwrap_err(), + Error::EmptyPolynomial + ); + } +} diff --git a/crypto/multilinear/src/mle.rs b/crypto/multilinear/src/mle.rs new file mode 100644 index 000000000..ec2971e37 --- /dev/null +++ b/crypto/multilinear/src/mle.rs @@ -0,0 +1,417 @@ +//! Multilinear extensions, held as their `2^n` hypercube evaluations. +//! +//! Index `i` is read with **variable 0 as the most significant bit**. Every +//! fold in this crate assumes that. + +use math::field::{ + element::FieldElement, + traits::{IsField, IsSubFieldOf}, +}; + +#[cfg(feature = "parallel")] +use rayon::prelude::*; + +use crate::Error; + +/// A multilinear polynomial held by its hypercube evaluations. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct Mle { + evals: Vec>, + num_vars: usize, +} + +impl Mle { + /// Builds an MLE from `2^n` evaluations in hypercube order. + pub fn new(evals: Vec>) -> Result { + let len = evals.len(); + if !len.is_power_of_two() { + return Err(Error::NotPowerOfTwo(len)); + } + Ok(Self { + num_vars: len.trailing_zeros() as usize, + evals, + }) + } + + /// The constant polynomial on zero variables. + pub fn constant(value: FieldElement) -> Self { + Self { + evals: vec![value], + num_vars: 0, + } + } + + pub fn num_vars(&self) -> usize { + self.num_vars + } + + pub fn len(&self) -> usize { + self.evals.len() + } + + pub fn is_empty(&self) -> bool { + self.evals.is_empty() + } + + pub fn evals(&self) -> &[FieldElement] { + &self.evals + } + + pub fn into_evals(self) -> Vec> { + self.evals + } + + /// Fixes variable 0 to `r`, returning a polynomial on `n - 1` variables. + /// + /// `new[j] = (1 - r)·old[j] + r·old[j + 2^(n-1)]`, which is the multilinear + /// interpolation between the two halves of the table. + pub fn fix_first_variable(&self, r: &FieldElement) -> Result { + if self.num_vars == 0 { + return Err(Error::NoVariablesLeft); + } + let half = self.evals.len() / 2; + let evals = (0..half) + .map(|j| { + let lo = &self.evals[j]; + let hi = &self.evals[j + half]; + // lo + r·(hi - lo) — one multiplication instead of two. + lo + r * &(hi - lo) + }) + .collect(); + Ok(Self { + evals, + num_vars: self.num_vars - 1, + }) + } + + /// Fixes variable 0 in place. Same arithmetic as [`Self::fix_first_variable`], + /// but reuses the allocation — sumcheck folds once per round. + pub fn fix_first_variable_in_place(&mut self, r: &FieldElement) -> Result<(), Error> { + if self.num_vars == 0 { + return Err(Error::NoVariablesLeft); + } + let half = self.evals.len() / 2; + let (lo, hi) = self.evals.split_at_mut(half); + // Every index is independent, and the halves are disjoint slices, so the + // split is what lets the rows go out to the pool at all. + let fold = |(a, b): (&mut FieldElement, &FieldElement)| *a = &*a + r * &(b - &*a); + #[cfg(feature = "parallel")] + if half >= crate::SERIAL_BELOW { + lo.par_iter_mut().zip(hi.par_iter()).for_each(fold); + } else { + lo.iter_mut().zip(hi.iter()).for_each(fold); + } + #[cfg(not(feature = "parallel"))] + lo.iter_mut().zip(hi.iter()).for_each(fold); + self.evals.truncate(half); + self.num_vars -= 1; + Ok(()) + } + + /// Fixes the **last** variable to `r`, returning a polynomial on `n - 1` + /// variables. + /// + /// The last variable is the low bit of the index, so this pairs `2j` with + /// `2j + 1`. Codeword folding binds variables from this end, which is why + /// it exists alongside [`Self::fix_first_variable_in_place`]. + pub fn fix_last_variable_in_place(&mut self, r: &FieldElement) -> Result<(), Error> { + if self.num_vars == 0 { + return Err(Error::NoVariablesLeft); + } + let half = self.evals.len() / 2; + for j in 0..half { + let lo = self.evals[2 * j].clone(); + let delta = &self.evals[2 * j + 1] - &lo; + self.evals[j] = lo + r * δ + } + self.evals.truncate(half); + self.num_vars -= 1; + Ok(()) + } + + /// Evaluates the extension at an arbitrary point in `F^n`. + pub fn evaluate(&self, point: &[FieldElement]) -> Result, Error> { + if point.len() != self.num_vars { + return Err(Error::VariableCountMismatch { + expected: self.num_vars, + got: point.len(), + }); + } + Self::evaluate_at(&self.evals, point) + } + + /// The extension of `evals` at `point`, without owning an [`Mle`]. + /// + /// The first fold reads the slice and writes the half-size buffer the rest + /// fold in place, so the table is never copied at full width. A caller + /// holding a slice of a bigger table — a GKR layer's half, say — evaluates + /// it without materializing it at all. + pub fn evaluate_at( + evals: &[FieldElement], + point: &[FieldElement], + ) -> Result, Error> { + if evals.len() != 1usize << point.len() { + return Err(Error::VariableCountMismatch { + expected: point.len(), + got: evals.len().trailing_zeros() as usize, + }); + } + if let Some(value) = crate::gpu::evaluate_mle(evals, point) { + return Ok(value); + } + let Some((first, rest)) = point.split_first() else { + return Ok(evals[0].clone()); + }; + + let half = evals.len() / 2; + let (lo, hi) = evals.split_at(half); + let combine = |(l, h): (&FieldElement, &FieldElement)| l + first * &(h - l); + #[cfg(feature = "parallel")] + let mut current: Vec> = if half >= crate::SERIAL_BELOW { + lo.par_iter().zip(hi.par_iter()).map(combine).collect() + } else { + lo.iter().zip(hi.iter()).map(combine).collect() + }; + #[cfg(not(feature = "parallel"))] + let mut current: Vec> = lo.iter().zip(hi.iter()).map(combine).collect(); + + for r in rest { + let half = current.len() / 2; + let (lo, hi) = current.split_at_mut(half); + let fold = |(a, b): (&mut FieldElement, &FieldElement)| *a = &*a + r * &(b - &*a); + #[cfg(feature = "parallel")] + if half >= crate::SERIAL_BELOW { + lo.par_iter_mut().zip(hi.par_iter()).for_each(fold); + } else { + lo.iter_mut().zip(hi.iter()).for_each(fold); + } + #[cfg(not(feature = "parallel"))] + lo.iter_mut().zip(hi.iter()).for_each(fold); + current.truncate(half); + } + Ok(current.swap_remove(0)) + } + + /// The extension at a point in a **larger** field. + /// + /// A trace column lives in the base field while the challenges do not, so + /// this is how a committed column answers a claim: the first fold lifts, + /// the rest stay up. Lifting the whole table first would instead cost its + /// size times the extension degree. + pub fn evaluate_in(&self, point: &[FieldElement]) -> Result, Error> + where + F: IsSubFieldOf, + E: IsField + 'static, + { + if point.len() != self.num_vars { + return Err(Error::VariableCountMismatch { + expected: self.num_vars, + got: point.len(), + }); + } + if let Some(value) = crate::gpu::evaluate_mle(&self.evals, point) { + return Ok(value); + } + let Some((first, rest)) = point.split_first() else { + return Ok(self.evals[0].clone().to_extension::()); + }; + + let half = self.evals.len() / 2; + let mut current: Vec> = (0..half) + .map(|j| { + let lo = &self.evals[j]; + let hi = &self.evals[j + half]; + // The base element on the left: the only direction the tower + // gives. + lo.clone().to_extension::() + (hi - lo) * first + }) + .collect(); + + for r in rest { + let half = current.len() / 2; + for j in 0..half { + let delta = ¤t[j + half] - ¤t[j]; + current[j] = ¤t[j] + r * δ + } + current.truncate(half); + } + Ok(current.into_iter().next().expect("one value remains")) + } + + /// The single remaining evaluation, once every variable has been fixed. + pub fn as_constant(&self) -> Option<&FieldElement> { + (self.num_vars == 0).then(|| &self.evals[0]) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use math::field::goldilocks::GoldilocksField as F; + + type FE = FieldElement; + + fn mle(vals: &[u64]) -> Mle { + Mle::new(vals.iter().map(|v| FE::from(*v)).collect()).unwrap() + } + + #[test] + fn rejects_non_power_of_two() { + let evals: Vec = (0..3).map(FE::from).collect(); + assert_eq!(Mle::new(evals).unwrap_err(), Error::NotPowerOfTwo(3)); + } + + #[test] + fn num_vars_is_log2_of_the_table() { + assert_eq!(mle(&[1]).num_vars(), 0); + assert_eq!(mle(&[1, 2]).num_vars(), 1); + assert_eq!(mle(&[1, 2, 3, 4]).num_vars(), 2); + assert_eq!(mle(&[0; 256]).num_vars(), 8); + } + + #[test] + fn agrees_with_the_table_on_hypercube_corners() { + // f(x0, x1) with x0 the high bit: [f(00), f(01), f(10), f(11)] + let f = mle(&[7, 11, 13, 17]); + for (i, expected) in [7u64, 11, 13, 17].iter().enumerate() { + let x0 = FE::from(((i >> 1) & 1) as u64); + let x1 = FE::from((i & 1) as u64); + assert_eq!(f.evaluate(&[x0, x1]).unwrap(), FE::from(*expected)); + } + } + + #[test] + fn is_multilinear_in_each_variable() { + // A multilinear polynomial is affine along every axis, so the midpoint + // evaluation is the average of the two endpoints. + let f = mle(&[7, 11, 13, 17]); + let two_inv = FE::from(2).inv().unwrap(); + let x1 = FE::from(5); + + let at_0 = f.evaluate(&[FE::zero(), x1]).unwrap(); + let at_1 = f.evaluate(&[FE::one(), x1]).unwrap(); + let at_mid = f.evaluate(&[two_inv, x1]).unwrap(); + + assert_eq!(at_mid, (at_0 + at_1) * two_inv); + } + + #[test] + fn fixing_a_variable_matches_evaluating_it() { + let f = mle(&[3, 5, 8, 13, 21, 34, 55, 89]); + let r = FE::from(42); + let folded = f.fix_first_variable(&r).unwrap(); + + assert_eq!(folded.num_vars(), 2); + for a in 0..2u64 { + for b in 0..2u64 { + let rest = [FE::from(a), FE::from(b)]; + let via_fold = folded.evaluate(&rest).unwrap(); + let direct = f.evaluate(&[r, FE::from(a), FE::from(b)]).unwrap(); + assert_eq!(via_fold, direct); + } + } + } + + #[test] + fn in_place_fold_matches_the_allocating_one() { + let f = mle(&[3, 5, 8, 13, 21, 34, 55, 89]); + let r = FE::from(9); + let expected = f.fix_first_variable(&r).unwrap(); + + let mut g = f; + g.fix_first_variable_in_place(&r).unwrap(); + assert_eq!(g, expected); + } + + #[test] + fn folding_every_variable_leaves_the_evaluation() { + let f = mle(&[3, 5, 8, 13]); + let point = [FE::from(6), FE::from(7)]; + let expected = f.evaluate(&point).unwrap(); + + let mut g = f; + for r in &point { + g.fix_first_variable_in_place(r).unwrap(); + } + assert_eq!(g.num_vars(), 0); + assert_eq!(g.as_constant().unwrap(), &expected); + } + + #[test] + fn folding_a_constant_is_an_error() { + let mut f = Mle::::constant(FE::from(4)); + assert_eq!( + f.fix_first_variable_in_place(&FE::from(1)).unwrap_err(), + Error::NoVariablesLeft + ); + } + + #[test] + fn evaluate_rejects_a_point_of_the_wrong_arity() { + let f = mle(&[1, 2, 3, 4]); + assert_eq!( + f.evaluate(&[FE::from(1)]).unwrap_err(), + Error::VariableCountMismatch { + expected: 2, + got: 1 + } + ); + } + + #[test] + fn fixing_the_last_variable_matches_evaluating_it() { + let f = mle(&[3, 5, 8, 13, 21, 34, 55, 89]); + let r = FE::from(23); + let mut folded = f.clone(); + folded.fix_last_variable_in_place(&r).unwrap(); + + assert_eq!(folded.num_vars(), 2); + for a in 0..2u64 { + for b in 0..2u64 { + let via_fold = folded.evaluate(&[FE::from(a), FE::from(b)]).unwrap(); + let direct = f.evaluate(&[FE::from(a), FE::from(b), r]).unwrap(); + assert_eq!(via_fold, direct, "a={a}, b={b}"); + } + } + } + + #[test] + fn evaluating_in_a_larger_field_matches_lifting_first() { + use math::field::extensions_goldilocks::Degree3GoldilocksExtensionField as Ext; + type ExtE = FieldElement; + + for num_vars in 0..=4usize { + let f = mle(&(0..(1u64 << num_vars)) + .map(|i| i.wrapping_mul(6364136223846793005) >> 13) + .collect::>()); + let point: Vec = (0..num_vars).map(|i| ExtE::from(101 + i as u64)).collect(); + + let lifted = + Mle::new(f.evals().iter().map(|v| v.to_extension::()).collect()).unwrap(); + + assert_eq!( + f.evaluate_in(&point).unwrap(), + lifted.evaluate(&point).unwrap(), + "num_vars={num_vars}" + ); + } + } + + #[test] + fn evaluating_in_the_same_field_is_evaluating() { + let f = mle(&[3, 5, 8, 13]); + let point = [FE::from(6), FE::from(7)]; + assert_eq!(f.evaluate_in(&point).unwrap(), f.evaluate(&point).unwrap()); + } + + #[test] + fn first_and_last_folds_bind_different_ends() { + let f = mle(&[1, 2, 3, 4]); + let r = FE::from(5); + let mut by_first = f.clone(); + by_first.fix_first_variable_in_place(&r).unwrap(); + let mut by_last = f; + by_last.fix_last_variable_in_place(&r).unwrap(); + assert_ne!(by_first, by_last); + } +} diff --git a/crypto/multilinear/src/poly.rs b/crypto/multilinear/src/poly.rs new file mode 100644 index 000000000..a5ec245da --- /dev/null +++ b/crypto/multilinear/src/poly.rs @@ -0,0 +1,428 @@ +//! What sumcheck needs from a polynomial: its multilinear factors plus a rule +//! for combining their values. Keeps a constraint DAG out of expanded form, +//! which would be exponential in the nesting depth. + +use math::field::{element::FieldElement, traits::IsField}; + +use crate::{Error, mle::Mle, program::Program}; + +/// A polynomial over the hypercube, presented as multilinear factors plus a +/// rule for combining their values. +pub trait SumcheckPolynomial { + /// Variables left to bind. + fn num_vars(&self) -> usize; + + /// Total degree, which bounds each round polynomial's degree. + /// + /// Must be an upper bound: the verifier checks the round polynomial has + /// exactly `degree + 1` evaluations, so understating it rejects honest + /// proofs and overstating it only costs proof size. + fn degree(&self) -> usize; + + /// The multilinear factors, in the order [`combine`](Self::combine) indexes. + fn polys(&self) -> &[Mle]; + + /// The polynomial's value, given each factor's value at the same point. + fn combine(&self, values: &[FieldElement]) -> FieldElement; + + /// The same, with a scratch buffer the caller owns. The sumcheck calls this + /// once per cube index per interpolation node, so an implementation backed + /// by a program has nowhere to put its steps that is not the caller's. + fn combine_in( + &self, + values: &[FieldElement], + scratch: &mut Vec>, + ) -> FieldElement { + let _ = scratch; + self.combine(values) + } + + /// Binds variable 0 to `r` in every factor. + fn fix_first_variable(&mut self, r: &FieldElement) -> Result<(), Error>; + + /// The program this polynomial is, when it can say: the description a + /// device runs in place of [`combine`](Self::combine). + /// + /// An implementation that offers one must also accept its factors back + /// through [`accept_folded`](Self::accept_folded) — the device binds them + /// there, and the sumcheck's contract is that the polynomial comes back + /// folded. + fn program(&self) -> Option<&Program> { + None + } + + /// The cube below which this polynomial's rounds stop going to a device. + /// + /// The default is for a rule the host runs through the program interpreter; + /// one that evaluates its rule directly crosses much later. + fn host_cube(&self) -> usize { + crate::HOST_CUBE_COMPILED + } + + /// Takes factors bound elsewhere, in the order [`polys`](Self::polys) + /// returns them. + fn accept_folded(&mut self, polys: Vec>) -> Result<(), Error> { + let _ = polys; + Err(Error::DeviceFailed { + stage: "write-back", + }) + } + + /// Value at hypercube index `i`. + fn eval_at_index(&self, i: usize) -> FieldElement { + let values: Vec> = + self.polys().iter().map(|p| p.evals()[i].clone()).collect(); + self.combine(&values) + } + + /// Sum over the whole hypercube. Reference implementation — the point of + /// sumcheck is to avoid paying this. + fn sum_over_hypercube(&self) -> FieldElement { + (0..(1usize << self.num_vars())) + .fold(FieldElement::zero(), |acc, i| acc + self.eval_at_index(i)) + } + + /// Value at an arbitrary point, extending every factor multilinearly. + fn evaluate(&self, point: &[FieldElement]) -> Result, Error> { + let values: Vec> = self + .polys() + .iter() + .map(|p| p.evaluate(point)) + .collect::>()?; + Ok(self.combine(&values)) + } +} + +/// Factors plus a closure that combines them. The closure holds no trace data, +/// so the verifier can carry the same one. +pub struct Composed { + polys: Vec>, + combine: C, + degree: usize, + num_vars: usize, + /// The same rule as straight-line code, when the caller can say what it is. + program: Option>, +} + +impl Composed +where + C: Fn(&[FieldElement]) -> FieldElement + 'static, +{ + /// `degree` must upper-bound the closure's total degree in the factors. + pub fn new(polys: Vec>, combine: C, degree: usize) -> Result { + let num_vars = polys.first().map(|p| p.num_vars()).unwrap_or(0); + for p in &polys { + if p.num_vars() != num_vars { + return Err(Error::VariableCountMismatch { + expected: num_vars, + got: p.num_vars(), + }); + } + } + Ok(Self { + polys, + combine, + degree, + num_vars, + program: None, + }) + } + + /// The same, saying which program the closure is. The two must agree: the + /// device runs the program and the host may run either. + pub fn with_program(mut self, program: Program) -> Self { + self.program = Some(program); + self + } +} + +impl Composed { + /// The factors, dropping the rule. + /// + /// A chained WHIR needs them back after each group of rounds: it rebuilds + /// its weight from the folded one, so it cannot keep the polynomial. + pub fn into_polys(self) -> Vec> { + self.polys + } +} + +impl SumcheckPolynomial for Composed +where + C: Fn(&[FieldElement]) -> FieldElement, +{ + fn num_vars(&self) -> usize { + self.num_vars + } + + fn degree(&self) -> usize { + self.degree + } + + fn polys(&self) -> &[Mle] { + &self.polys + } + + fn combine(&self, values: &[FieldElement]) -> FieldElement { + (self.combine)(values) + } + + fn combine_in( + &self, + values: &[FieldElement], + scratch: &mut Vec>, + ) -> FieldElement { + match &self.program { + Some(program) => program.eval(values, scratch), + None => (self.combine)(values), + } + } + + fn fix_first_variable(&mut self, r: &FieldElement) -> Result<(), Error> { + for p in &mut self.polys { + p.fix_first_variable_in_place(r)?; + } + self.num_vars -= 1; + Ok(()) + } + + fn program(&self) -> Option<&Program> { + self.program.as_ref() + } + + fn accept_folded(&mut self, polys: Vec>) -> Result<(), Error> { + if polys.len() != self.polys.len() { + return Err(Error::VariableCountMismatch { + expected: self.polys.len(), + got: polys.len(), + }); + } + self.num_vars = polys.first().map(Mle::num_vars).unwrap_or(0); + self.polys = polys; + Ok(()) + } +} + +/// Multiplies another polynomial by `eq(r, ·)`, appended as one more factor. +/// +/// This is what turns a sumcheck into a zerocheck, and it works for any +/// underlying polynomial rather than only the sum-of-products one. +#[derive(Debug)] +pub struct EqScaled> { + inner: P, + /// `inner`'s factors followed by the `eq` table — the layout `combine` and + /// the sumcheck prover both index. + polys: Vec>, +} + +impl> EqScaled { + /// Wraps `inner` with the `eq(r, ·)` table. + pub fn new(inner: P, eq: Mle) -> Result { + if eq.num_vars() != inner.num_vars() { + return Err(Error::VariableCountMismatch { + expected: inner.num_vars(), + got: eq.num_vars(), + }); + } + let mut polys = inner.polys().to_vec(); + polys.push(eq); + Ok(Self { inner, polys }) + } + + pub fn into_inner(self) -> P { + self.inner + } +} + +impl> SumcheckPolynomial for EqScaled { + fn num_vars(&self) -> usize { + self.inner.num_vars() + } + + fn degree(&self) -> usize { + self.inner.degree() + 1 + } + + fn polys(&self) -> &[Mle] { + &self.polys + } + + fn combine(&self, values: &[FieldElement]) -> FieldElement { + let (inner_values, eq_value) = values.split_at(values.len() - 1); + self.inner.combine(inner_values) * &eq_value[0] + } + + fn fix_first_variable(&mut self, r: &FieldElement) -> Result<(), Error> { + // `inner` keeps its own copies of the factors, so both views must fold. + self.inner.fix_first_variable(r)?; + for p in &mut self.polys { + p.fix_first_variable_in_place(r)?; + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use math::field::goldilocks::GoldilocksField as F; + + use crate::eq::eq_mle; + + type FE = FieldElement; + + fn mle(vals: &[u64]) -> Mle { + Mle::new(vals.iter().map(|v| FE::from(*v)).collect()).unwrap() + } + + /// `2·a·b`. Rebuilt on each call rather than cloned: the rule is a + /// closure, and a closure is not `Clone`. + fn inner() -> Composed FE> { + Composed::new( + vec![mle(&[1, 2, 3, 4]), mle(&[5, 6, 7, 8])], + |v: &[FE]| FE::from(2) * v[0] * v[1], + 2, + ) + .unwrap() + } + + #[test] + fn eq_scaling_adds_one_to_the_degree() { + let r = vec![FE::from(3), FE::from(5)]; + let scaled = EqScaled::new(inner(), eq_mle(&r).unwrap()).unwrap(); + assert_eq!(inner().degree(), 2); + assert_eq!(scaled.degree(), 3); + } + + #[test] + fn eq_scaling_multiplies_pointwise_on_the_cube() { + let r = vec![FE::from(3), FE::from(5)]; + let eq = eq_mle(&r).unwrap(); + let base = inner(); + let scaled = EqScaled::new(inner(), eq.clone()).unwrap(); + + for i in 0..4 { + assert_eq!( + scaled.eval_at_index(i), + base.eval_at_index(i) * eq.evals()[i] + ); + } + } + + #[test] + fn eq_scaled_sum_of_a_multilinear_polynomial_is_its_extension_at_r() { + // Σ_x eq(r,x)·f(x) = f̃(r) — the defining property of eq, and it holds + // only when f is itself multilinear. + let linear = + || Composed::new(vec![mle(&[1, 2, 3, 4])], |v: &[FE]| FE::from(2) * v[0], 1).unwrap(); + let r = vec![FE::from(11), FE::from(13)]; + let scaled = EqScaled::new(linear(), eq_mle(&r).unwrap()).unwrap(); + + assert_eq!(scaled.sum_over_hypercube(), linear().evaluate(&r).unwrap()); + } + + #[test] + fn above_degree_one_the_eq_weighted_sum_is_not_the_pointwise_product() { + // For a product of two columns, `evaluate` gives ã(r)·b̃(r) while the + // eq-weighted sum gives the multilinear extension of the *function* + // a·b. Those are different polynomials, and conflating them is an easy + // way to write a wrong soundness argument. + let r = vec![FE::from(11), FE::from(13)]; + let scaled = EqScaled::new(inner(), eq_mle(&r).unwrap()).unwrap(); + + assert_ne!(scaled.sum_over_hypercube(), inner().evaluate(&r).unwrap()); + } + + #[test] + fn folding_keeps_both_views_in_step() { + let r = vec![FE::from(2), FE::from(4)]; + let mut scaled = EqScaled::new(inner(), eq_mle(&r).unwrap()).unwrap(); + let challenge = FE::from(9); + let rest = FE::from(21); + + let expected = scaled.evaluate(&[challenge, rest]).unwrap(); + scaled.fix_first_variable(&challenge).unwrap(); + + assert_eq!(scaled.num_vars(), 1); + assert_eq!(scaled.evaluate(&[rest]).unwrap(), expected); + } + + #[test] + fn rejects_an_eq_table_of_the_wrong_arity() { + let r = vec![FE::from(3)]; + // `.err()` rather than `unwrap_err()`: a closure has no `Debug`. + let err = EqScaled::new(inner(), eq_mle(&r).unwrap()).err().unwrap(); + assert_eq!( + err, + Error::VariableCountMismatch { + expected: 2, + got: 1 + } + ); + } + + #[test] + fn a_composed_polynomial_applies_its_closure() { + // 2ab + 3c, written as a closure instead of terms. + let polys = vec![ + mle(&[1, 2, 3, 4]), + mle(&[5, 6, 7, 8]), + mle(&[9, 10, 11, 12]), + ]; + let composed = Composed::new( + polys, + |v: &[FE]| FE::from(2) * v[0] * v[1] + FE::from(3) * v[2], + 2, + ) + .unwrap(); + + assert_eq!(composed.degree(), 2); + for i in 0..4 { + let (a, b, c) = ( + FE::from(1 + i as u64), + FE::from(5 + i as u64), + FE::from(9 + i as u64), + ); + assert_eq!( + composed.eval_at_index(i), + FE::from(2) * a * b + FE::from(3) * c + ); + } + } + + #[test] + fn a_composed_polynomial_is_its_rule_over_the_factors_extensions() { + let (a, b) = (mle(&[1, 2, 3, 4]), mle(&[5, 6, 7, 8])); + let composed = Composed::new( + vec![a.clone(), b.clone()], + |v: &[FE]| FE::from(2) * v[0] * v[1], + 2, + ) + .unwrap(); + + // Off the cube, the rule is applied to each factor's own extension — + // not to the extension of the product, which is a different polynomial. + let point = [FE::from(19), FE::from(23)]; + assert_eq!( + composed.evaluate(&point).unwrap(), + FE::from(2) * a.evaluate(&point).unwrap() * b.evaluate(&point).unwrap() + ); + + // On the cube, the sum is the rule applied row by row. + let row_by_row = (0..4).fold(FE::zero(), |acc, i| { + acc + FE::from(2) * a.evals()[i] * b.evals()[i] + }); + assert_eq!(composed.sum_over_hypercube(), row_by_row); + } + + #[test] + fn composed_rejects_factors_of_differing_arity() { + let result = Composed::new(vec![mle(&[1, 2]), mle(&[1, 2, 3, 4])], |v: &[FE]| v[0], 1); + assert!(matches!( + result.err(), + Some(Error::VariableCountMismatch { + expected: 1, + got: 2 + }) + )); + } +} diff --git a/crypto/multilinear/src/program.rs b/crypto/multilinear/src/program.rs new file mode 100644 index 000000000..1e5e2e1e6 --- /dev/null +++ b/crypto/multilinear/src/program.rs @@ -0,0 +1,567 @@ +//! A batch's rule as straight-line code over the factor values. +//! +//! A [`Rule`](crate::batch::Rule) is a closure, which is what makes a statement +//! easy to state and impossible to hand to a device. Every rule this crate +//! proves is nonetheless *data* — a compiled constraint program, an affine +//! expression over factor slots, a fixed formula — so a rule can carry the +//! program it is, and the sumcheck can run that instead: one description, run +//! by the host loop and by the kernel that replaces it. +//! +//! Values are whatever the sumcheck's field is. Two evaluators of the same +//! program may sum in different orders, which the field does not distinguish: +//! Goldilocks compares and serializes canonically, so equal values are equal +//! everywhere the protocol looks. + +use math::field::{element::FieldElement, traits::IsField}; + +use crate::Error; + +/// One step. Operands are indices of earlier steps; `Var` indexes the factor +/// values the sumcheck supplies. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum Op { + Fixed(FieldElement), + Var(u32), + Add(u32, u32), + Sub(u32, u32), + Mul(u32, u32), + Neg(u32), +} + +/// Straight-line code and the step holding its value. +#[derive(Clone, Debug)] +pub struct Program { + steps: Vec>, + root: u32, +} + +impl Program { + pub fn steps(&self) -> &[Op] { + &self.steps + } + + pub fn root(&self) -> u32 { + self.root + } + + /// The program's value, given every factor's value at one point. + /// + /// `scratch` is the caller's, because this runs once per cube index per + /// interpolation node: a buffer of its own per call would be one heap + /// allocation per evaluation. + pub fn eval( + &self, + values: &[FieldElement], + scratch: &mut Vec>, + ) -> FieldElement { + scratch.clear(); + scratch.reserve(self.steps.len()); + for step in &self.steps { + let v = match *step { + Op::Fixed(ref c) => c.clone(), + Op::Var(i) => values[i as usize].clone(), + Op::Add(a, b) => &scratch[a as usize] + &scratch[b as usize], + Op::Sub(a, b) => &scratch[a as usize] - &scratch[b as usize], + Op::Mul(a, b) => &scratch[a as usize] * &scratch[b as usize], + Op::Neg(a) => -&scratch[a as usize], + }; + scratch.push(v); + } + scratch[self.root as usize].clone() + } + + /// The same program with each step emitted once and nothing dead in it. + /// + /// Straight-line code over a field: two steps with the same operator over + /// the same operands compute the same value, adding zero or multiplying by + /// one computes nothing, and a step the root does not reach is not + /// computed at all. A batch's program is spliced together out of pieces + /// that share structure — the same alpha powers, the same column read by + /// several interactions — so what this removes is not a mistake in any one + /// of them, it is the seam between them. + /// + /// The kernel walks this program once per cube index per interpolation + /// node, through a slot file in global memory, so a step removed here is + /// removed from every one of those walks. + pub fn simplify(&self) -> Self + where + FieldElement: PartialEq, + { + use std::collections::HashMap; + + let zero = FieldElement::::zero(); + let one = FieldElement::::one(); + let mut steps: Vec> = Vec::with_capacity(self.steps.len()); + // Where each old step ended up. + let mut moved: Vec = Vec::with_capacity(self.steps.len()); + // Steps that are not constants, keyed by what they compute. + let mut seen: HashMap<(u8, u32, u32), u32> = HashMap::new(); + // Constants, which have to be compared by value. + let mut constants: Vec<(u32, FieldElement)> = Vec::new(); + + let constant_of = |steps: &[Op], at: u32| match &steps[at as usize] { + Op::Fixed(c) => Some(c.clone()), + _ => None, + }; + + for step in &self.steps { + let emit = |steps: &mut Vec>, + seen: &mut HashMap<(u8, u32, u32), u32>, + constants: &mut Vec<(u32, FieldElement)>, + op: Op| + -> u32 { + if let Op::Fixed(value) = &op { + if let Some((at, _)) = constants.iter().find(|(_, c)| c == value) { + return *at; + } + steps.push(op.clone()); + let at = (steps.len() - 1) as u32; + constants.push((at, value.clone())); + return at; + } + // Addition and multiplication do not care which operand came + // first, so the key does not either. + let key = match op { + Op::Var(i) => (0, i, 0), + Op::Add(a, b) => (1, a.min(b), a.max(b)), + Op::Sub(a, b) => (2, a, b), + Op::Mul(a, b) => (3, a.min(b), a.max(b)), + Op::Neg(a) => (4, a, 0), + Op::Fixed(_) => unreachable!("handled above"), + }; + if let Some(at) = seen.get(&key) { + return *at; + } + steps.push(op); + let at = (steps.len() - 1) as u32; + seen.insert(key, at); + at + }; + + let at = match *step { + Op::Fixed(ref c) => { + emit(&mut steps, &mut seen, &mut constants, Op::Fixed(c.clone())) + } + Op::Var(i) => emit(&mut steps, &mut seen, &mut constants, Op::Var(i)), + Op::Neg(a) => { + let a = moved[a as usize]; + match constant_of(&steps, a) { + Some(c) => emit(&mut steps, &mut seen, &mut constants, Op::Fixed(-c)), + None => emit(&mut steps, &mut seen, &mut constants, Op::Neg(a)), + } + } + Op::Add(a, b) => { + let (a, b) = (moved[a as usize], moved[b as usize]); + match (constant_of(&steps, a), constant_of(&steps, b)) { + (Some(x), Some(y)) => { + emit(&mut steps, &mut seen, &mut constants, Op::Fixed(x + y)) + } + (Some(x), None) if x == zero => b, + (None, Some(y)) if y == zero => a, + _ => emit(&mut steps, &mut seen, &mut constants, Op::Add(a, b)), + } + } + Op::Sub(a, b) => { + let (a, b) = (moved[a as usize], moved[b as usize]); + match (constant_of(&steps, a), constant_of(&steps, b)) { + (Some(x), Some(y)) => { + emit(&mut steps, &mut seen, &mut constants, Op::Fixed(x - y)) + } + (None, Some(y)) if y == zero => a, + _ => emit(&mut steps, &mut seen, &mut constants, Op::Sub(a, b)), + } + } + Op::Mul(a, b) => { + let (a, b) = (moved[a as usize], moved[b as usize]); + match (constant_of(&steps, a), constant_of(&steps, b)) { + (Some(x), Some(y)) => { + emit(&mut steps, &mut seen, &mut constants, Op::Fixed(x * y)) + } + (Some(x), _) if x == zero => emit( + &mut steps, + &mut seen, + &mut constants, + Op::Fixed(zero.clone()), + ), + (_, Some(y)) if y == zero => emit( + &mut steps, + &mut seen, + &mut constants, + Op::Fixed(zero.clone()), + ), + (Some(x), None) if x == one => b, + (None, Some(y)) if y == one => a, + _ => emit(&mut steps, &mut seen, &mut constants, Op::Mul(a, b)), + } + } + }; + moved.push(at); + } + + let root = moved[self.root as usize]; + Self { steps, root }.prune() + } + + /// Drops the steps the root does not reach, renumbering the rest. + fn prune(self) -> Self { + let mut live = vec![false; self.steps.len()]; + live[self.root as usize] = true; + for i in (0..self.steps.len()).rev() { + if !live[i] { + continue; + } + match self.steps[i] { + Op::Add(a, b) | Op::Sub(a, b) | Op::Mul(a, b) => { + live[a as usize] = true; + live[b as usize] = true; + } + Op::Neg(a) => live[a as usize] = true, + Op::Fixed(_) | Op::Var(_) => {} + } + } + let mut moved = vec![0u32; self.steps.len()]; + let mut steps = Vec::with_capacity(self.steps.len()); + for (i, step) in self.steps.into_iter().enumerate() { + if !live[i] { + continue; + } + let shifted = match step { + Op::Fixed(c) => Op::Fixed(c), + Op::Var(v) => Op::Var(v), + Op::Add(a, b) => Op::Add(moved[a as usize], moved[b as usize]), + Op::Sub(a, b) => Op::Sub(moved[a as usize], moved[b as usize]), + Op::Mul(a, b) => Op::Mul(moved[a as usize], moved[b as usize]), + Op::Neg(a) => Op::Neg(moved[a as usize]), + }; + steps.push(shifted); + moved[i] = (steps.len() - 1) as u32; + } + let root = moved[self.root as usize]; + Self { steps, root } + } + + /// The highest factor slot the program reads, or `None` when it reads none. + pub fn max_slot(&self) -> Option { + self.steps + .iter() + .filter_map(|step| match step { + Op::Var(i) => Some(*i), + _ => None, + }) + .max() + } +} + +/// Emits steps, handing back the index of each one. +#[derive(Debug, Default)] +pub struct Builder { + steps: Vec>, +} + +impl Builder { + pub fn new() -> Self { + Self { steps: Vec::new() } + } + + fn push(&mut self, op: Op) -> u32 { + self.steps.push(op); + (self.steps.len() - 1) as u32 + } + + pub fn fixed(&mut self, value: FieldElement) -> u32 { + self.push(Op::Fixed(value)) + } + + pub fn var(&mut self, slot: usize) -> u32 { + self.push(Op::Var(slot as u32)) + } + + pub fn add(&mut self, a: u32, b: u32) -> u32 { + self.push(Op::Add(a, b)) + } + + pub fn sub(&mut self, a: u32, b: u32) -> u32 { + self.push(Op::Sub(a, b)) + } + + pub fn mul(&mut self, a: u32, b: u32) -> u32 { + self.push(Op::Mul(a, b)) + } + + pub fn neg(&mut self, a: u32) -> u32 { + self.push(Op::Neg(a)) + } + + /// `Σ terms`, or a zero step when there are none. + pub fn sum(&mut self, terms: &[u32]) -> u32 { + match terms.split_first() { + None => self.fixed(FieldElement::zero()), + Some((&first, rest)) => rest.iter().fold(first, |acc, &t| self.add(acc, t)), + } + } + + /// `Σ coefficient_i · step_i`, folding a coefficient of one away. + pub fn weighted_sum(&mut self, terms: &[(u32, FieldElement)]) -> u32 { + let scaled: Vec = terms + .iter() + .map(|(step, coefficient)| { + if *coefficient == FieldElement::one() { + *step + } else { + let c = self.fixed(coefficient.clone()); + self.mul(c, *step) + } + }) + .collect(); + self.sum(&scaled) + } + + /// Copies `program`'s steps in, renumbering its operands, and returns where + /// its root landed. + pub fn splice(&mut self, program: &Program) -> u32 { + let base = self.steps.len() as u32; + for step in &program.steps { + let shifted = match *step { + Op::Fixed(ref c) => Op::Fixed(c.clone()), + Op::Var(i) => Op::Var(i), + Op::Add(a, b) => Op::Add(a + base, b + base), + Op::Sub(a, b) => Op::Sub(a + base, b + base), + Op::Mul(a, b) => Op::Mul(a + base, b + base), + Op::Neg(a) => Op::Neg(a + base), + }; + self.steps.push(shifted); + } + base + program.root + } + + pub fn len(&self) -> usize { + self.steps.len() + } + + pub fn is_empty(&self) -> bool { + self.steps.is_empty() + } + + /// The program computing `root`, with each step emitted once. + /// + /// A builder is written for whoever emits — a rule at a time, a piece at a + /// time — and the pieces overlap: the same alpha power, the same column, + /// the same difference. [`simplify`](Program::simplify) is run here because + /// this is where a program stops being written and starts being walked, + /// and it is walked once per cube index per interpolation node. + pub fn finish(self, root: u32) -> Result, Error> + where + FieldElement: PartialEq, + { + if root as usize >= self.steps.len() { + return Err(Error::UnknownPolynomial { + index: root as usize, + len: self.steps.len(), + }); + } + Ok(Program { + steps: self.steps, + root, + } + .simplify()) + } + + /// The program as it was emitted, step for step. + pub fn finish_verbatim(self, root: u32) -> Result, Error> { + if root as usize >= self.steps.len() { + return Err(Error::UnknownPolynomial { + index: root as usize, + len: self.steps.len(), + }); + } + Ok(Program { + steps: self.steps, + root, + }) + } +} + +#[cfg(test)] +mod simplify_tests { + use super::*; + use math::field::goldilocks::GoldilocksField as F; + + type FE = FieldElement; + + /// Every simplification has to be invisible to the evaluator: same value at + /// every point, fewer steps to get there. + fn agrees(program: &Program, slots: usize) { + let small = program.simplify(); + assert!( + small.steps().len() <= program.steps().len(), + "simplifying grew the program" + ); + let mut scratch = Vec::new(); + for seed in 0..6u64 { + let values: Vec = (0..slots) + .map(|i| FE::from((seed + 1) * (i as u64 + 3) + 7)) + .collect(); + let before = program.eval(&values, &mut scratch); + let after = small.eval(&values, &mut scratch); + assert_eq!(before, after, "seed {seed}"); + } + } + + #[test] + fn a_repeated_step_is_emitted_once() { + let mut b = Builder::::new(); + let x = b.var(0); + let y = b.var(1); + let first = b.mul(x, y); + // The same product again, and the same one with its operands swapped. + let second = b.mul(x, y); + let third = b.mul(y, x); + let sum = b.add(first, second); + let root = b.add(sum, third); + let program = b.finish_verbatim(root).unwrap(); + + agrees(&program, 2); + // one var, one var, one product, two sums + assert_eq!(program.simplify().steps().len(), 5); + } + + #[test] + fn constants_fold_and_units_disappear() { + let mut b = Builder::::new(); + let x = b.var(0); + let one = b.fixed(FE::one()); + let zero = b.fixed(FE::zero()); + let scaled = b.mul(x, one); + let shifted = b.add(scaled, zero); + let two = b.fixed(FE::from(2)); + let three = b.fixed(FE::from(3)); + let six = b.mul(two, three); + let root = b.add(shifted, six); + let program = b.finish_verbatim(root).unwrap(); + + agrees(&program, 1); + let small = program.simplify(); + // `x`, the constant six, and their sum. + assert_eq!(small.steps().len(), 3); + } + + #[test] + fn what_the_root_does_not_reach_is_dropped() { + let mut b = Builder::::new(); + let x = b.var(0); + let y = b.var(1); + let _dead = b.mul(x, y); + let root = b.add(x, x); + let program = b.finish_verbatim(root).unwrap(); + + agrees(&program, 2); + assert_eq!(program.simplify().steps().len(), 2); + } +} + +/// `Σ_i lambda_i · program_i`, the batch as one program. +pub fn combine( + programs: &[&Program], + lambdas: &[FieldElement], +) -> Result, Error> { + if programs.len() != lambdas.len() { + return Err(Error::VariableCountMismatch { + expected: programs.len(), + got: lambdas.len(), + }); + } + let mut builder = Builder::::new(); + let terms: Vec<(u32, FieldElement)> = programs + .iter() + .zip(lambdas) + .map(|(program, lambda)| (builder.splice(program), lambda.clone())) + .collect(); + let root = builder.weighted_sum(&terms); + builder.finish(root) +} + +#[cfg(test)] +mod tests { + use super::*; + use math::field::goldilocks::GoldilocksField as F; + + type FE = FieldElement; + + fn values(n: usize) -> Vec { + (0..n as u64).map(|i| FE::from(i * 7 + 3)).collect() + } + + #[test] + fn a_program_evaluates_its_root() { + // (f0 + f1) · f2 − 5 + let mut b = Builder::::new(); + let f0 = b.var(0); + let f1 = b.var(1); + let f2 = b.var(2); + let sum = b.add(f0, f1); + let product = b.mul(sum, f2); + let five = b.fixed(FE::from(5u64)); + let root = b.sub(product, five); + let program = b.finish(root).unwrap(); + + let v = values(3); + let mut scratch = Vec::new(); + assert_eq!( + program.eval(&v, &mut scratch), + (v[0] + v[1]) * v[2] - FE::from(5u64) + ); + assert_eq!(program.max_slot(), Some(2)); + } + + #[test] + fn a_spliced_program_keeps_its_meaning() { + let mut b = Builder::::new(); + let f0 = b.var(0); + let f1 = b.var(1); + let root = b.mul(f0, f1); + let inner = b.finish(root).unwrap(); + + let mut outer = Builder::::new(); + // A step before the splice, so the renumbering is not a no-op. + let f2 = outer.var(2); + let spliced = outer.splice(&inner); + let root = outer.add(spliced, f2); + let program = outer.finish(root).unwrap(); + + let v = values(3); + let mut scratch = Vec::new(); + assert_eq!(program.eval(&v, &mut scratch), v[0] * v[1] + v[2]); + } + + #[test] + fn a_combined_program_is_the_weighted_sum() { + let mut b = Builder::::new(); + let root = b.var(0); + let first = b.finish(root).unwrap(); + let mut b = Builder::::new(); + let root = b.var(1); + let second = b.finish(root).unwrap(); + + let lambdas = vec![FE::one(), FE::from(9u64)]; + let program = combine(&[&first, &second], &lambdas).unwrap(); + let v = values(2); + let mut scratch = Vec::new(); + assert_eq!(program.eval(&v, &mut scratch), v[0] + FE::from(9u64) * v[1]); + } + + #[test] + fn a_weighted_sum_of_nothing_is_zero() { + let mut b = Builder::::new(); + let root = b.weighted_sum(&[]); + let program = b.finish(root).unwrap(); + let mut scratch = Vec::new(); + assert_eq!(program.eval(&[], &mut scratch), FE::zero()); + } + + #[test] + fn a_root_past_the_end_is_rejected() { + let mut b = Builder::::new(); + b.var(0); + assert!(b.finish(7).is_err()); + } +} diff --git a/crypto/multilinear/src/selector.rs b/crypto/multilinear/src/selector.rs new file mode 100644 index 000000000..870050935 --- /dev/null +++ b/crypto/multilinear/src/selector.rs @@ -0,0 +1,210 @@ +//! Row selectors: the indicator of `index(x) < 2^n - end_exemptions`. +//! +//! A transition constraint reading the next step cannot hold on the last one. +//! Multiplying by a selector costs exactly one degree. + +use math::field::{element::FieldElement, traits::IsField}; + +use crate::{Error, mle::Mle}; + +/// The indicator of `index(x) < 2^n − end_exemptions`. +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub struct Selector { + pub end_exemptions: usize, +} + +impl Selector { + /// Applies on every step. + pub const ALL: Selector = Selector { end_exemptions: 0 }; + + pub const fn except_last(k: usize) -> Selector { + Selector { end_exemptions: k } + } + + /// True when this selector is the constant one and can be skipped. + pub fn is_trivial(&self) -> bool { + self.end_exemptions == 0 + } + + /// Number of steps the constraint applies to. + pub fn active_steps(&self, num_vars: usize) -> usize { + (1usize << num_vars).saturating_sub(self.end_exemptions) + } + + /// The selector's hypercube table. + pub fn table(&self, num_vars: usize) -> Result, Error> { + let size = 1usize << num_vars; + if self.end_exemptions > size { + return Err(Error::TooManyExemptions { + exemptions: self.end_exemptions, + size, + }); + } + let cutoff = size - self.end_exemptions; + let evals = (0..size) + .map(|i| { + if i < cutoff { + FieldElement::::one() + } else { + FieldElement::::zero() + } + }) + .collect(); + Mle::new(evals) + } + + /// The selector at an arbitrary point, in `O(num_vars)`. + /// + /// `s(x) = 1 − geq(x, cutoff)`, where `geq` is the multilinear indicator of + /// `index(x) >= cutoff`: either `x` matches `cutoff` bit for bit, or they + /// first differ at a position where `cutoff` has a zero and `x` a one. + pub fn evaluate( + &self, + point: &[FieldElement], + ) -> Result, Error> { + let num_vars = point.len(); + let size = 1usize << num_vars; + if self.end_exemptions > size { + return Err(Error::TooManyExemptions { + exemptions: self.end_exemptions, + size, + }); + } + if self.is_trivial() { + return Ok(FieldElement::one()); + } + let cutoff = size - self.end_exemptions; + if cutoff == 0 { + // Every step is exempt: the constraint applies nowhere. + return Ok(FieldElement::zero()); + } + + let one = FieldElement::::one(); + // Running product of "x agrees with cutoff on every earlier bit". + let mut prefix = one.clone(); + let mut geq = FieldElement::::zero(); + + for (i, x_i) in point.iter().enumerate() { + // Variable 0 is the most significant bit. + let bit = (cutoff >> (num_vars - 1 - i)) & 1; + if bit == 0 { + // x exceeds cutoff here: everything above matched, x_i = 1. + geq += &prefix * x_i; + prefix *= &one - x_i; + } else { + prefix *= x_i; + } + } + // The remaining prefix is the "x == cutoff" case. + Ok(one - (geq + prefix)) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use math::field::goldilocks::GoldilocksField as F; + + type FE = FieldElement; + + fn corner(index: usize, num_vars: usize) -> Vec { + (0..num_vars) + .map(|i| FE::from(((index >> (num_vars - 1 - i)) & 1) as u64)) + .collect() + } + + #[test] + fn all_is_the_constant_one() { + let s = Selector::ALL; + assert!(s.is_trivial()); + let table = s.table::(3).unwrap(); + assert!(table.evals().iter().all(|v| *v == FE::one())); + assert_eq!(s.evaluate(&corner(5, 3)).unwrap(), FE::one()); + } + + #[test] + fn except_last_one_masks_only_the_final_step() { + let s = Selector::except_last(1); + let table = s.table::(3).unwrap(); + for i in 0..8 { + let expected = if i < 7 { FE::one() } else { FE::zero() }; + assert_eq!(table.evals()[i], expected, "step {i}"); + } + } + + #[test] + fn except_last_two_masks_the_final_pair() { + let s = Selector::except_last(2); + let table = s.table::(4).unwrap(); + for i in 0..16 { + let expected = if i < 14 { FE::one() } else { FE::zero() }; + assert_eq!(table.evals()[i], expected, "step {i}"); + } + } + + #[test] + fn closed_form_matches_the_table_on_every_corner() { + for num_vars in 1..=5usize { + let size = 1usize << num_vars; + for k in 0..=size { + let s = Selector::except_last(k); + let table = s.table::(num_vars).unwrap(); + for i in 0..size { + assert_eq!( + s.evaluate(&corner(i, num_vars)).unwrap(), + table.evals()[i], + "num_vars={num_vars}, k={k}, step={i}" + ); + } + } + } + } + + #[test] + fn closed_form_is_the_multilinear_extension_off_the_cube() { + // The verifier evaluates the closed form at a random point; it must be + // the same polynomial the prover's table extends to. + for k in [1usize, 2, 3, 5] { + let s = Selector::except_last(k); + let table = s.table::(4).unwrap(); + let point = vec![FE::from(9), FE::from(17), FE::from(3), FE::from(41)]; + assert_eq!( + s.evaluate(&point).unwrap(), + table.evaluate(&point).unwrap(), + "k={k}" + ); + } + } + + #[test] + fn exempting_everything_selects_nothing() { + let s = Selector::except_last(8); + let table = s.table::(3).unwrap(); + assert!(table.evals().iter().all(|v| *v == FE::zero())); + assert_eq!(s.evaluate(&corner(0, 3)).unwrap(), FE::zero()); + assert_eq!(s.evaluate(&[FE::from(7); 3]).unwrap(), FE::zero()); + } + + #[test] + fn active_step_count_matches_the_table() { + for k in 0..=8usize { + let s = Selector::except_last(k); + let table = s.table::(3).unwrap(); + let ones = table.evals().iter().filter(|v| **v == FE::one()).count(); + assert_eq!(ones, s.active_steps(3), "k={k}"); + } + } + + #[test] + fn more_exemptions_than_steps_is_an_error() { + let s = Selector::except_last(9); + assert_eq!( + s.table::(3).unwrap_err(), + Error::TooManyExemptions { + exemptions: 9, + size: 8 + } + ); + assert!(s.evaluate(&corner(0, 3)).is_err()); + } +} diff --git a/crypto/multilinear/src/sumcheck.rs b/crypto/multilinear/src/sumcheck.rs new file mode 100644 index 000000000..c77135649 --- /dev/null +++ b/crypto/multilinear/src/sumcheck.rs @@ -0,0 +1,624 @@ +//! The sumcheck protocol, reducing `Σ_x f(x) = S` to one evaluation of `f`. +//! +//! Round polynomials travel as evaluations at `1, .., d`. `g(0)` is **not** +//! sent: `g(0) + g(1)` is the claim carried into the round, which fixes it. So +//! the prover skips a whole pass over the cube and the proof loses one field +//! element per round — and the rejection that used to happen per round now +//! happens **only** against the final claim, which [`verify`] returns and the +//! caller must discharge. + +use crypto::fiat_shamir::is_transcript::IsTranscript; +use math::field::{element::FieldElement, traits::IsField}; + +#[cfg(feature = "parallel")] +use rayon::prelude::*; + +use crate::{Error, mle::Mle, poly::SumcheckPolynomial}; + +/// One round: the round polynomial as evaluations at `1, .., degree`. +/// +/// `g(0)` is absent by construction — see the module docs. +#[derive( + Clone, + Debug, + PartialEq, + Eq, + serde::Serialize, + serde::Deserialize, + rkyv::Archive, + rkyv::Serialize, + rkyv::Deserialize, +)] +#[serde(bound = "")] +pub struct RoundProof { + pub evaluations: Vec>, +} + +/// A full sumcheck transcript. +#[derive( + Clone, + Debug, + PartialEq, + Eq, + serde::Serialize, + serde::Deserialize, + rkyv::Archive, + rkyv::Serialize, + rkyv::Deserialize, +)] +#[serde(bound = "")] +pub struct SumcheckProof { + pub rounds: Vec>, +} + +/// What the verifier is left holding: `f(point)` must equal `expected_evaluation`. +/// +/// Discharging this is the **only** place a sumcheck rejects, so dropping it +/// silently accepts anything. +#[must_use] +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct SumcheckClaim { + pub point: Vec>, + pub expected_evaluation: FieldElement, +} + +/// Lagrange-interpolates values at `0, 1, .., d` and evaluates at `x`. +/// +/// Round polynomials are small, so the quadratic form is the cheap one. The +/// denominators depend only on how many nodes there are, so they go through one +/// batched inversion instead of one each — and the quadratic numerators are +/// kept rather than switching to the barycentric form, which would divide by +/// `x − x_i` and so need a special case for an `x` that lands on a node. +pub(crate) fn interpolate( + values: &[FieldElement], + x: &FieldElement, +) -> FieldElement { + let n = values.len(); + let node = |j: usize| FieldElement::::from(j as u64); + let others = |i: usize, at: &FieldElement| { + (0..n) + .filter(|&j| j != i) + .fold(FieldElement::::one(), |acc, j| acc * (at - node(j))) + }; + + let mut denominators: Vec> = (0..n).map(|i| others(i, &node(i))).collect(); + // The nodes are distinct, so none of them is zero. + FieldElement::inplace_batch_inverse(&mut denominators).expect("distinct interpolation nodes"); + + values + .iter() + .zip(&denominators) + .enumerate() + .fold(FieldElement::zero(), |acc, (i, (y_i, inv))| { + acc + y_i * others(i, x) * inv + }) +} + +/// Sums `f(r_0..r_{j-1}, t, rest)` over the remaining cube. +/// Sums `f(r_0..r_{j-1}, t, rest)` over the remaining cube. +/// +/// `t` runs over `1..=degree`, or `0..=degree` when `with_zero` — which only +/// the prover's debug self-check asks for. +/// +/// `poly` has already been folded on the earlier variables, so its first +/// variable is the one this round binds. +/// +/// One pass over the cube serves **every** `t`: each factor's `lo` and `hi` are +/// read once per index and the extensions come off `hi - lo`, rather than +/// re-reading the tables once per `t`. On a real trace the factors are hundreds +/// of megabytes, so the reads are the cost, not the arithmetic. +pub(crate) fn round_evaluations( + poly: &P, + degree: usize, + with_zero: bool, +) -> Vec> +where + F: IsField + 'static, + P: SumcheckPolynomial + Sync, + FieldElement: Send + Sync, +{ + let half = 1usize << (poly.num_vars() - 1); + let first = usize::from(!with_zero); + let steps: Vec> = (first..=degree) + .map(|t| FieldElement::::from(t as u64)) + .collect(); + + // A slice of the cube, summed independently. The rounds are the whole cost + // of the prover, and every index is independent, so this is where the cores + // go in. + let slice = |range: std::ops::Range| -> Vec> { + let width = poly.polys().len(); + // Buffers for the whole slice. Collecting a fresh `Vec` per cube index + // would be one heap allocation per index, which on a real trace is + // millions of them per round and dwarfs the arithmetic. + let mut totals = vec![FieldElement::::zero(); steps.len()]; + let mut lo = vec![FieldElement::::zero(); width]; + let mut hi = vec![FieldElement::::zero(); width]; + let mut delta = vec![FieldElement::::zero(); width]; + let mut values = vec![FieldElement::::zero(); width]; + // A compiled rule's steps go here, once for the whole slice. + let mut scratch = Vec::new(); + + for j in range { + for (k, p) in poly.polys().iter().enumerate() { + lo[k] = p.evals()[j].clone(); + hi[k] = p.evals()[j + half].clone(); + delta[k] = &hi[k] - &lo[k]; + } + for (total, t) in totals.iter_mut().zip(&steps) { + // `t = 0` and `t = 1` are the halves as they are, so they skip + // the multiplication entirely. + if t == &FieldElement::::zero() { + values.clone_from(&lo); + } else if t == &FieldElement::::one() { + values.clone_from(&hi); + } else { + for (v, (l, d)) in values.iter_mut().zip(lo.iter().zip(&delta)) { + *v = l + t * d; + } + } + *total += poly.combine_in(&values, &mut scratch); + } + } + totals + }; + + let add = |mut acc: Vec>, part: Vec>| { + for (a, p) in acc.iter_mut().zip(part) { + *a += p; + } + acc + }; + + #[cfg(feature = "parallel")] + { + if half < crate::SERIAL_BELOW { + return slice(0..half); + } + let chunk = half.div_ceil(rayon::current_num_threads().max(1)); + (0..half) + .into_par_iter() + .step_by(chunk) + .map(|start| slice(start..(start + chunk).min(half))) + .reduce(|| vec![FieldElement::::zero(); steps.len()], add) + } + #[cfg(not(feature = "parallel"))] + { + let _ = add; + slice(0..half) + } +} + +/// The round polynomial's evaluations at `1..=degree` for one compiled rule +/// over `factors`, through the host round loop. +/// +/// This is the reference the device round is checked against +/// (`math-cuda/tests/sumcheck.rs`): the comparison has to be against the loop +/// the prover actually runs, not a second copy of the formula. +pub fn round_evaluations_for_program( + factors: &[Mle], + program: &crate::program::Program, + degree: usize, +) -> Result>, Error> +where + F: IsField + 'static, + FieldElement: Send + Sync, +{ + let batched = crate::batch::Batched::new( + factors.to_vec(), + vec![crate::batch::Rule::compiled(degree, program.clone())], + vec![FieldElement::one()], + )?; + Ok(round_evaluations(&batched, degree, false)) +} + +/// Runs the prover, absorbing each round polynomial and drawing each challenge +/// from `transcript`. +/// +/// Returns the proof and the challenge point. The caller is responsible for +/// having absorbed the claimed sum and any statement binding beforehand. +pub fn prove( + mut poly: P, + transcript: &mut T, +) -> Result<(SumcheckProof, Vec>), Error> +where + F: IsField + 'static, + T: IsTranscript, + P: SumcheckPolynomial + Sync, + FieldElement: Send + Sync, +{ + let num_vars = poly.num_vars(); + let (rounds, challenges) = prove_rounds(&mut poly, num_vars, transcript)?; + Ok((SumcheckProof { rounds }, challenges)) +} + +/// A group of round polynomials and the challenges they drew. +pub type RoundGroup = (Vec>, Vec>); + +/// Runs `rounds` rounds, leaving `poly` folded on them. +/// +/// A chained WHIR interleaves: between groups of rounds it folds its codeword +/// and commits the successor, so it cannot run the whole sumcheck in one call. +pub fn prove_rounds( + poly: &mut P, + rounds: usize, + transcript: &mut T, +) -> Result, Error> +where + F: IsField + 'static, + T: IsTranscript, + P: SumcheckPolynomial + Sync, + FieldElement: Send + Sync, +{ + if rounds > poly.num_vars() { + return Err(Error::RoundCountMismatch { + expected: poly.num_vars(), + got: rounds, + }); + } + let degree = poly.degree().max(1); + + // The rounds on device when the polynomial says which program it is and the + // device takes it. Only a whole sumcheck: a group of rounds leaves tables + // the caller needs back, and downloading them between groups costs more + // than the rounds do. + let attempt = if rounds == poly.num_vars() { + match poly.program() { + Some(program) => { + let cube = poly.host_cube(); + crate::gpu::prove_sumcheck(poly.polys(), program, degree, cube, |evaluations| { + for e in evaluations { + transcript.append_field_element(e); + } + transcript.sample_field_element() + }) + } + None => None, + } + } else { + None + }; + let mut proofs = Vec::with_capacity(rounds); + let mut challenges = Vec::with_capacity(rounds); + // What a device ran of this, if it ran any. It stops where the cube stops + // being worth sending and hands the factors back folded, so what is left + // carries on below from exactly where it left off. + if let Some(outcome) = attempt { + let (device_proofs, device_challenges, folded) = outcome?; + poly.accept_folded(folded)?; + proofs = device_proofs; + challenges = device_challenges; + } + let rounds = rounds - proofs.len(); + // The identity the verifier now takes on faith. Checking it costs the pass + // over the cube the protocol exists to skip, so it runs in debug only — + // where it turns a silent prover bug into a local failure. + // + // What that trades, and it is worth knowing which way: **in release a + // prover that produces a wrong round polynomial is not caught here**. It + // surfaces as a proof that does not verify, which for a real trace is + // minutes and fifty tables later and says nothing about where. Reach for a + // debug build when a proof stops verifying and you do not know why. + #[cfg(debug_assertions)] + let mut running: Option> = None; + + for _ in 0..rounds { + let all = round_evaluations(poly, degree, cfg!(debug_assertions)); + let sent = all[all.len() - degree..].to_vec(); + for e in &sent { + transcript.append_field_element(e); + } + let r = transcript.sample_field_element(); + + #[cfg(debug_assertions)] + { + let sum = &all[0] + &all[1]; + if let Some(expected) = &running { + debug_assert_eq!( + &sum, expected, + "sumcheck: g(0) + g(1) is not the claim carried into the round" + ); + } + running = Some(interpolate(&all, &r)); + } + + poly.fix_first_variable(&r)?; + proofs.push(RoundProof { evaluations: sent }); + challenges.push(r); + } + + Ok((proofs, challenges)) +} + +/// Checks every round against the running claim and returns the final claim. +/// +/// Verifying `claimed_sum` requires one more step the caller must perform: +/// obtain `f` at [`SumcheckClaim::point`] and compare it against +/// [`SumcheckClaim::expected_evaluation`]. +pub fn verify( + proof: &SumcheckProof, + claimed_sum: FieldElement, + num_vars: usize, + degree: usize, + transcript: &mut T, +) -> Result, Error> +where + F: IsField, + T: IsTranscript, +{ + if proof.rounds.len() != num_vars { + return Err(Error::RoundCountMismatch { + expected: num_vars, + got: proof.rounds.len(), + }); + } + verify_rounds(&proof.rounds, claimed_sum, degree, transcript) +} + +/// Verifies a group of rounds against a running claim. +/// +/// The returned claim carries this group's challenges and the claim it leaves, +/// which is what the next group starts from. Round indices in errors are +/// relative to the group. +pub fn verify_rounds( + rounds: &[RoundProof], + claimed_sum: FieldElement, + degree: usize, + transcript: &mut T, +) -> Result, Error> +where + F: IsField, + T: IsTranscript, +{ + let degree = degree.max(1); + + let mut current = claimed_sum; + let mut point = Vec::with_capacity(rounds.len()); + + for (round, r_proof) in rounds.iter().enumerate() { + if r_proof.evaluations.len() != degree { + return Err(Error::RoundDegreeMismatch { + round, + expected: degree, + got: r_proof.evaluations.len(), + }); + } + // `g(0)` is recovered rather than checked: `g(0) + g(1)` is the claim + // carried in. So nothing is rejected here, and everything rides on the + // claim this returns. + let mut all = Vec::with_capacity(degree + 1); + all.push(¤t - &r_proof.evaluations[0]); + all.extend(r_proof.evaluations.iter().cloned()); + + for e in &r_proof.evaluations { + transcript.append_field_element(e); + } + let r = transcript.sample_field_element(); + + current = interpolate(&all, &r); + point.push(r); + } + + Ok(SumcheckClaim { + point, + expected_evaluation: current, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crypto::fiat_shamir::default_transcript::DefaultTranscript; + use math::field::goldilocks::GoldilocksField as F; + + use crate::{mle::Mle, poly::Composed}; + + type FE = FieldElement; + + fn transcript() -> DefaultTranscript { + DefaultTranscript::::new(b"sumcheck-test") + } + + fn mle(vals: &[u64]) -> Mle { + Mle::new(vals.iter().map(|v| FE::from(*v)).collect()).unwrap() + } + + /// Pseudo-random table of `2^n` entries, deterministic across runs. + fn pseudo_mle(n: usize, seed: u64) -> Mle { + let vals: Vec = (0..(1u64 << n)) + .map(|i| (i.wrapping_mul(6364136223846793005).wrapping_add(seed)) >> 11) + .collect(); + mle(&vals) + } + + /// `f` itself, as a one-factor composition. Rebuilt rather than cloned: + /// the rule is a closure, and a closure is not `Clone`. + fn single(m: Mle) -> Composed FE> { + Composed::new(vec![m], |v: &[FE]| v[0], 1).unwrap() + } + + #[test] + fn interpolation_reproduces_the_nodes() { + let values: Vec = [3u64, 1, 4, 1].iter().map(|v| FE::from(*v)).collect(); + for (i, v) in values.iter().enumerate() { + assert_eq!(interpolate(&values, &FE::from(i as u64)), *v); + } + } + + #[test] + fn interpolation_of_a_line_is_affine() { + // g(t) = 5 + 2t sampled at 0,1 must give 5 + 2·7 at t = 7. + let values = vec![FE::from(5), FE::from(7)]; + assert_eq!(interpolate(&values, &FE::from(7)), FE::from(19)); + } + + #[test] + fn linear_polynomial_round_trips() { + let f = single(pseudo_mle(4, 1)); + let claimed = f.sum_over_hypercube(); + let num_vars = f.num_vars(); + let degree = f.degree(); + + let (proof, _) = prove(single(pseudo_mle(4, 1)), &mut transcript()).unwrap(); + let claim = verify(&proof, claimed, num_vars, degree, &mut transcript()).unwrap(); + + assert_eq!(f.evaluate(&claim.point).unwrap(), claim.expected_evaluation); + } + + #[test] + fn product_of_three_polynomials_round_trips() { + let product = || { + Composed::new( + vec![pseudo_mle(5, 1), pseudo_mle(5, 2), pseudo_mle(5, 3)], + |v: &[FE]| FE::from(7) * v[0] * v[1] * v[2], + 3, + ) + .unwrap() + }; + let f = product(); + let claimed = f.sum_over_hypercube(); + let (num_vars, degree) = (f.num_vars(), f.degree()); + assert_eq!(degree, 3); + + let (proof, challenges) = prove(product(), &mut transcript()).unwrap(); + let claim = verify(&proof, claimed, num_vars, degree, &mut transcript()).unwrap(); + + // Prover and verifier must derive the same challenges from the transcript. + assert_eq!(challenges, claim.point); + assert_eq!(f.evaluate(&claim.point).unwrap(), claim.expected_evaluation); + } + + #[test] + fn sum_of_terms_of_mixed_degree_round_trips() { + let mixed = || { + Composed::new( + vec![pseudo_mle(4, 10), pseudo_mle(4, 20), pseudo_mle(4, 30)], + |v: &[FE]| FE::from(2) * v[0] * v[1] + FE::from(3) * v[2] + FE::from(5), + 2, + ) + .unwrap() + }; + let f = mixed(); + let claimed = f.sum_over_hypercube(); + let (num_vars, degree) = (f.num_vars(), f.degree()); + + let (proof, _) = prove(mixed(), &mut transcript()).unwrap(); + let claim = verify(&proof, claimed, num_vars, degree, &mut transcript()).unwrap(); + + assert_eq!(f.evaluate(&claim.point).unwrap(), claim.expected_evaluation); + } + + #[test] + fn single_variable_round_trips() { + let f = single(mle(&[3, 11])); + let claimed = f.sum_over_hypercube(); + assert_eq!(claimed, FE::from(14)); + + let (proof, _) = prove(single(mle(&[3, 11])), &mut transcript()).unwrap(); + let claim = verify(&proof, claimed, 1, 1, &mut transcript()).unwrap(); + assert_eq!(f.evaluate(&claim.point).unwrap(), claim.expected_evaluation); + } + + /// `g(0)` is derived from the claim, so a wrong claim is not caught in the + /// round — it is caught by the residual, which stops being the + /// polynomial's value. Every caller discharges that; this is the property + /// they rely on. + #[test] + fn a_wrong_claimed_sum_corrupts_the_residual() { + let f = single(pseudo_mle(3, 1)); + let claimed = f.sum_over_hypercube(); + let (proof, _) = prove(single(pseudo_mle(3, 1)), &mut transcript()).unwrap(); + + let honest = verify(&proof, claimed, 3, 1, &mut transcript()).unwrap(); + assert_eq!( + f.evaluate(&honest.point).unwrap(), + honest.expected_evaluation + ); + + let lied = verify(&proof, claimed + FE::one(), 3, 1, &mut transcript()).unwrap(); + assert_ne!(f.evaluate(&lied.point).unwrap(), lied.expected_evaluation); + } + + #[test] + fn a_tampered_round_polynomial_corrupts_the_residual() { + let pair = || { + Composed::new( + vec![pseudo_mle(4, 1), pseudo_mle(4, 2)], + |v: &[FE]| v[0] * v[1], + 2, + ) + .unwrap() + }; + let f = pair(); + let claimed = f.sum_over_hypercube(); + let (mut proof, _) = prove(pair(), &mut transcript()).unwrap(); + + proof.rounds[2].evaluations[0] += FE::one(); + let claim = verify(&proof, claimed, 4, 2, &mut transcript()).unwrap(); + assert_ne!(f.evaluate(&claim.point).unwrap(), claim.expected_evaluation); + } + + #[test] + fn a_proof_with_the_wrong_round_count_is_rejected() { + let f = single(pseudo_mle(3, 1)); + let claimed = f.sum_over_hypercube(); + let (mut proof, _) = prove(f, &mut transcript()).unwrap(); + proof.rounds.pop(); + + let err = verify(&proof, claimed, 3, 1, &mut transcript()).unwrap_err(); + assert_eq!( + err, + Error::RoundCountMismatch { + expected: 3, + got: 2 + } + ); + } + + #[test] + fn a_round_polynomial_of_the_wrong_degree_is_rejected() { + let f = single(pseudo_mle(3, 1)); + let claimed = f.sum_over_hypercube(); + let (mut proof, _) = prove(f, &mut transcript()).unwrap(); + proof.rounds[0].evaluations.push(FE::from(1)); + + let err = verify(&proof, claimed, 3, 1, &mut transcript()).unwrap_err(); + assert!(matches!(err, Error::RoundDegreeMismatch { round: 0, .. })); + } + + #[test] + fn a_different_transcript_seed_yields_a_different_point() { + // The challenge point is bound to the statement, not just to the + // polynomial: proving the same claim under a different seed must land + // somewhere else. + let (_, a) = prove(single(pseudo_mle(4, 1)), &mut transcript()).unwrap(); + let mut other = DefaultTranscript::::new(b"a-different-statement"); + let (_, b) = prove(single(pseudo_mle(4, 1)), &mut other).unwrap(); + + assert_ne!(a, b); + } + + #[test] + fn a_proof_replayed_under_another_transcript_is_rejected() { + // Fiat-Shamir binding: the verifier redraws challenges, so a proof + // lifted onto a different statement stops matching. + let f = single(pseudo_mle(4, 1)); + let claimed = f.sum_over_hypercube(); + let (proof, _) = prove(single(pseudo_mle(4, 1)), &mut transcript()).unwrap(); + + let honest = verify(&proof, claimed, 4, 1, &mut transcript()).unwrap(); + assert_eq!( + f.evaluate(&honest.point).unwrap(), + honest.expected_evaluation + ); + + // Replayed, the verifier redraws different challenges, so the residual + // stops describing the polynomial. + let mut other = DefaultTranscript::::new(b"a-different-statement"); + let replayed = verify(&proof, claimed, 4, 1, &mut other).unwrap(); + assert_ne!( + f.evaluate(&replayed.point).unwrap(), + replayed.expected_evaluation + ); + } +} diff --git a/crypto/multilinear/src/zerocheck.rs b/crypto/multilinear/src/zerocheck.rs new file mode 100644 index 000000000..c385b90c9 --- /dev/null +++ b/crypto/multilinear/src/zerocheck.rs @@ -0,0 +1,276 @@ +//! `f` vanishes on the whole hypercube, via `Σ_x eq(r, x)·f(x) = 0`. +//! +//! The multilinear replacement for a quotient argument. `eq` adds one degree. +//! The residual claim about `f` is returned, not decided. +//! +//! **This is the reference, not the path.** The prover's zerocheck does not +//! come through here: a table's constraints go in as a compiled +//! [`batch::Rule`](crate::batch::Rule) alongside its two bus claims, so all +//! three share one pass over one factor list — which is the whole point of +//! batching them. What this module is for is saying the argument once, in the +//! shape it has on paper, and being the thing the batched version is checked +//! against. Nothing outside tests calls it. + +use crypto::fiat_shamir::is_transcript::IsTranscript; +use math::field::{element::FieldElement, traits::IsField}; + +use crate::{ + Error, + eq::eq_mle, + poly::{EqScaled, SumcheckPolynomial}, + sumcheck::{self, SumcheckProof}, +}; + +/// A zerocheck proof: the sumcheck transcript for `eq(r, ·)·C`. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ZeroCheckProof { + pub sumcheck: SumcheckProof, +} + +/// What the verifier is left holding. +/// +/// `eq(r, point)` is computable by the verifier alone, so the residual claim is +/// entirely about `C` — [`constraint_evaluation`](Self::constraint_evaluation) +/// is the value the commitment scheme must confirm. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct ZeroCheckClaim { + /// The zerocheck challenge `r`. + pub r: Vec>, + /// The sumcheck challenge point. + pub point: Vec>, + /// Required value of `eq(r, point)·C(point)`. + pub expected_evaluation: FieldElement, + /// `eq(r, point)`, recomputed by the verifier. + pub eq_at_point: FieldElement, +} + +impl ZeroCheckClaim { + /// The value `C(point)` must take, isolated from the `eq` factor. + /// + /// `None` when `eq(r, point)` is zero, which leaves `C(point)` + /// unconstrained by this claim. It happens only if the sumcheck challenges + /// land exactly on a cube corner other than `r`. + pub fn constraint_evaluation(&self) -> Option> { + self.eq_at_point + .inv() + .ok() + .map(|inv| &self.expected_evaluation * inv) + } +} + +/// What proving leaves the caller holding. +pub struct ZeroCheckOutput { + pub proof: ZeroCheckProof, + /// The zerocheck challenge. + pub r: Vec>, + /// The sumcheck challenge point — where the residual claim about the + /// constraint lives, and therefore where the commitment scheme must open. + pub point: Vec>, +} + +/// Proves that `constraint` vanishes on `{0,1}^n`. +/// +/// Draws `r` from the transcript first, so the prover cannot choose the +/// constraint after seeing it. +pub fn prove(constraint: P, transcript: &mut T) -> Result, Error> +where + F: IsField + 'static, + T: IsTranscript, + P: SumcheckPolynomial + Sync, + FieldElement: Send + Sync, +{ + let num_vars = constraint.num_vars(); + let r: Vec> = (0..num_vars) + .map(|_| transcript.sample_field_element()) + .collect(); + + let combined = EqScaled::new(constraint, eq_mle(&r)?)?; + let (sumcheck, point) = sumcheck::prove(combined, transcript)?; + Ok(ZeroCheckOutput { + proof: ZeroCheckProof { sumcheck }, + r, + point, + }) +} + +/// Verifies a zerocheck, returning the residual claim about `C`. +/// +/// `constraint_degree` is the degree of `C`; `eq` adds one on top. +pub fn verify( + proof: &ZeroCheckProof, + num_vars: usize, + constraint_degree: usize, + transcript: &mut T, +) -> Result, Error> +where + F: IsField + 'static, + T: IsTranscript, +{ + let r: Vec> = (0..num_vars) + .map(|_| transcript.sample_field_element()) + .collect(); + + // The claimed sum is zero — that is the whole statement. + let claim = sumcheck::verify( + &proof.sumcheck, + FieldElement::::zero(), + num_vars, + constraint_degree + 1, + transcript, + )?; + + let eq_at_point = crate::eq::eq_eval(&r, &claim.point)?; + Ok(ZeroCheckClaim { + r, + point: claim.point, + expected_evaluation: claim.expected_evaluation, + eq_at_point, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use crypto::fiat_shamir::default_transcript::DefaultTranscript; + use math::field::goldilocks::GoldilocksField as F; + + use crate::{mle::Mle, poly::Composed}; + + type FE = FieldElement; + + fn transcript() -> DefaultTranscript { + DefaultTranscript::::new(b"zerocheck-test") + } + + fn mle(vals: &[u64]) -> Mle { + Mle::new(vals.iter().map(|v| FE::from(*v)).collect()).unwrap() + } + + /// `C = a·b − c`, with `c` set to the product so it vanishes everywhere. + /// This is the shape of a real AIR constraint: a relation among columns. + /// + /// Rebuilt on each call rather than cloned: the rule is a closure, and a + /// closure is not `Clone`. + fn satisfied_constraint(n: usize) -> Composed FE> { + let size = 1usize << n; + let a: Vec = (0..size as u64).map(|i| i * 3 + 1).collect(); + let b: Vec = (0..size as u64).map(|i| i * 5 + 2).collect(); + let c: Vec = a.iter().zip(&b).map(|(x, y)| x * y).collect(); + Composed::new( + vec![mle(&a), mle(&b), mle(&c)], + |v: &[FE]| v[0] * v[1] - v[2], + 2, + ) + .unwrap() + } + + #[test] + fn a_satisfied_constraint_verifies() { + let c = satisfied_constraint(4); + assert_eq!(c.sum_over_hypercube(), FE::zero()); + + let out = prove(satisfied_constraint(4), &mut transcript()).unwrap(); + let (proof, r_prover) = (out.proof, out.r); + let claim = verify(&proof, 4, c.degree(), &mut transcript()).unwrap(); + + assert_eq!(claim.r, r_prover); + // The residual claim must be exactly what C evaluates to there. + assert_eq!( + claim.constraint_evaluation().unwrap(), + c.evaluate(&claim.point).unwrap() + ); + } + + #[test] + fn the_residual_claim_factors_as_eq_times_c() { + let c = satisfied_constraint(3); + let proof = prove(satisfied_constraint(3), &mut transcript()) + .unwrap() + .proof; + let claim = verify(&proof, 3, c.degree(), &mut transcript()).unwrap(); + + let c_at_point = c.evaluate(&claim.point).unwrap(); + assert_eq!(claim.eq_at_point * c_at_point, claim.expected_evaluation); + } + + #[test] + fn a_constraint_violated_in_one_row_is_rejected() { + let size = 8usize; + let a: Vec = (0..size as u64).map(|i| i * 3 + 1).collect(); + let b: Vec = (0..size as u64).map(|i| i * 5 + 2).collect(); + let mut c: Vec = a.iter().zip(&b).map(|(x, y)| x * y).collect(); + c[5] += 1; // one bad row + + let build = || { + Composed::new( + vec![mle(&a), mle(&b), mle(&c)], + |v: &[FE]| v[0] * v[1] - v[2], + 2, + ) + .unwrap() + }; + let broken = build(); + assert_ne!(broken.sum_over_hypercube(), FE::zero()); + + // The prover runs the protocol honestly on a false statement. `g(0)` is + // derived from the claim, so nothing is rejected in the round — the lie + // surfaces in the residual, which stops describing the constraint. + let proof = prove(build(), &mut transcript()).unwrap().proof; + let claim = verify(&proof, 3, broken.degree(), &mut transcript()).unwrap(); + assert_ne!( + claim.constraint_evaluation(), + Some(broken.evaluate(&claim.point).unwrap()), + "a violated constraint produced a consistent claim" + ); + } + + #[test] + fn a_nonzero_polynomial_that_happens_to_sum_to_zero_is_still_rejected() { + // Σ C = 0 but C is not identically zero: exactly the case a plain + // sumcheck-for-zero would miss and eq(r, ·) is there to catch. + let build = || { + Composed::new( + vec![Mle::new(vec![FE::from(7), -FE::from(7)]).unwrap()], + |v: &[FE]| v[0], + 1, + ) + .unwrap() + }; + let c = build(); + assert_eq!(c.sum_over_hypercube(), FE::zero()); + + let proof = prove(build(), &mut transcript()).unwrap().proof; + let result = verify(&proof, 1, c.degree(), &mut transcript()); + + // The residual claim must be inconsistent with the real polynomial: + // that is where a non-vanishing constraint gets caught. + let claim = result.unwrap(); + assert_ne!( + claim.constraint_evaluation(), + Some(c.evaluate(&claim.point).unwrap()), + "a non-vanishing constraint produced a consistent claim" + ); + } + + #[test] + fn degree_accounts_for_the_eq_factor() { + let c = satisfied_constraint(3); + let proof = prove(satisfied_constraint(3), &mut transcript()) + .unwrap() + .proof; + // C has degree 2; with eq the round polynomials are degree 3, so each + // carries three evaluations — `g(1), g(2), g(3)`, with `g(0)` derived. + assert_eq!(c.degree(), 2); + assert_eq!(proof.sumcheck.rounds[0].evaluations.len(), 3); + } + + #[test] + fn verifying_with_the_wrong_degree_is_rejected() { + let c = satisfied_constraint(3); + let proof = prove(satisfied_constraint(3), &mut transcript()) + .unwrap() + .proof; + let err = verify(&proof, 3, c.degree() + 1, &mut transcript()).unwrap_err(); + assert!(matches!(err, Error::RoundDegreeMismatch { round: 0, .. })); + } +}