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, .. })); + } +}