From eb71289651aa4f2384d9414de9ff9488496f0765 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Mon, 14 Sep 2026 14:00:15 +0800 Subject: [PATCH 01/10] prover: move single_threaded.patch out of the crate to the repo root --- crates/prover/single_threaded.patch => single_threaded.patch | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename crates/prover/single_threaded.patch => single_threaded.patch (100%) diff --git a/crates/prover/single_threaded.patch b/single_threaded.patch similarity index 100% rename from crates/prover/single_threaded.patch rename to single_threaded.patch From 9264e24b6fab9172809c65af2fcf6a9d952a9ac9 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Mon, 14 Sep 2026 14:15:37 +0800 Subject: [PATCH 02/10] gkr: flatten SuffixTable into one preallocated buffer, dropping per-layer allocs --- crates/gkr/src/lib.rs | 59 ++++++++++++++++++++++------------- crates/prover/src/reduce.rs | 3 +- crates/verifier/src/reduce.rs | 3 +- 3 files changed, 42 insertions(+), 23 deletions(-) diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs index 4a7f9617..7cb1adf5 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs @@ -15,6 +15,7 @@ type Point = VecDeque; #[tracing::instrument(name = "Prove GKR", skip_all)] pub fn gpgkr_prove( ps: &mut ProverState, + log_bits: usize, point: &[F128], // All the intermediate witnesses + the input layer. Doesn't contain the output layer witnesses: LayerWitnesses, @@ -27,8 +28,9 @@ pub fn gpgkr_prove( let mut point = VecDeque::from(point); let mut claim = Field::ZERO; + let mut storage = Field::zeroed_vec(1 << log_bits); for wnext in witnesses.into_iter() { - (point, claim) = prove_layer(ps, point, wnext); + (point, claim) = prove_layer(ps, &mut storage, point, wnext); } let mut point = Vec::from(point); @@ -36,8 +38,13 @@ pub fn gpgkr_prove( (point, claim) } -fn prove_layer(ps: &mut ProverState, point: Point, mut wnext: Vec) -> (Point, Field) { - let mut suffix_table = SuffixTable::new(&point); +fn prove_layer( + ps: &mut ProverState, + storage: &mut [Field], + point: Point, + mut wnext: Vec, +) -> (Point, Field) { + let mut suffix_table = SuffixTable::new(storage, &point); let mut factor = Field::ONE; let mid = wnext.len() / 2; @@ -48,11 +55,11 @@ fn prove_layer(ps: &mut ProverState, point: Point, mut wnext: Vec) -> (Po // Can go up to ~21 allocations assuming input of 2^35 and 6:4 split let mut next_point = VecDeque::with_capacity(point.len() + 1); - for z in point { + for (i, z) in (0..point.len()).rev().zip(point) { // TODO: special-case eq.len() == 1 (final round) to skip the `eq[i] *` // multiplications below entirely. // TODO: unwrap will be dealt with in upcoming approach to SuffixTable - let eq = suffix_table.pop().unwrap(); + let eq = suffix_table.layer(i); let h = mle_l.len() / 2; debug_assert_eq!(eq.len(), h); @@ -145,14 +152,25 @@ fn mul3_wide(a: Field, b: Field, c: Field) -> Wide256 { // TODO: SuffixTable becomes a wrapper around a preallocated vector that is large enough for all rounds. // SuffixTable can be 'created' each round / destroyed to ensure proper truncation of the underlying vector // TODO: Split suffix table -struct SuffixTable(Vec>); +struct SuffixTable<'a> { + storage: &'a mut [Field], + offset: usize, +} -impl SuffixTable { - /// Allocates all directly as it is as much space as a double buffer approach would take. - fn new(point: &Point) -> SuffixTable { - let mut table = Vec::with_capacity(point.len().max(1)); - let mut prev = Vec::from([Field::ONE]); +impl<'a> SuffixTable<'a> { + fn layer(&mut self, i: usize) -> &mut [Field] { + let start = (1 << i) - 1 - self.offset; + let end = (1 << (i + 1)) - 1 - self.offset; + &mut self.storage[start..end] + } +} +impl<'a> SuffixTable<'a> { + /// Allocates all directly as it is as much space as a double buffer approach would take. + #[inline(never)] + fn new(storage: &'a mut [Field], point: &Point) -> SuffixTable<'a> { + storage[0] = F128::ONE; + let (mut prev, mut remaining) = storage.split_at_mut(1); // The selector is the first entry of the point and we need to skip // that -- except when `point` is itself empty, in which case there // is no selector and `c` must stay empty too (`min(1)` keeps the @@ -162,7 +180,7 @@ impl SuffixTable { // Suffix table is in the reverse order of the point for &z in c.rev() { let size = prev.len() << 1; - let mut entry = Field::zeroed_vec(size); + let (entry, next) = remaining.split_at_mut(size); let (low, hi) = entry.split_at_mut(size >> 1); for (i, &e) in prev.iter().enumerate() { @@ -171,15 +189,13 @@ impl SuffixTable { (low[i], hi[i]) = (e - tmp, tmp) } - table.push(prev); prev = entry; + remaining = next; + } + SuffixTable { + storage: storage, + offset: 0, } - table.push(prev); - SuffixTable(table) - } - - fn pop(&mut self) -> Option> { - self.0.pop() } } @@ -450,7 +466,7 @@ mod tests { let claim = mle(output, &point); let instance = (leaves.clone(), point.to_vec()); let mut prover = transcript::build_prover("gkr-zero", &instance); - let terminal = gpgkr_prove(&mut prover, &point, witnesses); + let terminal = gpgkr_prove(&mut prover, leaves.len().ilog2() as usize, &point, witnesses); assert_eq!(terminal.1, mle(leaves.clone(), &terminal.0)); let proof = prover.finish(); @@ -486,6 +502,7 @@ mod tests { } pub fn prove(input: Vec, log_groups: usize) -> (Vec, transcript::Proof) { + let log_bits = input.len().ilog2() as usize; let circuit = GrandProductCircuit::new(input); let groups = 1usize << log_groups; let (last_value, witnesses) = circuit.batched_eval(groups); @@ -501,7 +518,7 @@ mod tests { let log_groups = last_value.len().max(1).ilog2(); let point: Vec = (0..log_groups).map(|_| prover.verifier_message()).collect(); - gpgkr_prove(&mut prover, &point, witnesses); + gpgkr_prove(&mut prover, log_bits, &point, witnesses); (last_value, prover.finish()) } diff --git a/crates/prover/src/reduce.rs b/crates/prover/src/reduce.rs index 985c9f22..4364cf26 100644 --- a/crates/prover/src/reduce.rs +++ b/crates/prover/src/reduce.rs @@ -60,7 +60,8 @@ pub fn gkr_reduce( let circuit = init_circuit(table, fold); let (_last_value, witnesses) = circuit.batched_eval(table.shape().columns()); - let (mut point, claim) = gpgkr_prove(transcript, &fold.zeta, witnesses); + let (mut point, claim) = + gpgkr_prove(transcript, table.shape().log_bits(), &fold.zeta, witnesses); // The multilinear extension of the constant-one table is one at every point. let inner_product_claim = claim - F128::ONE; diff --git a/crates/verifier/src/reduce.rs b/crates/verifier/src/reduce.rs index d0e75776..eec8a599 100644 --- a/crates/verifier/src/reduce.rs +++ b/crates/verifier/src/reduce.rs @@ -119,7 +119,8 @@ mod round_trip_ai_test { let fold = Fold::new(&shape, folds, top_layer, row_images.clone(), zeta.clone()).unwrap(); let mut prover = transcript::build_prover("verifier-round-trip", &F128::ZERO); - let (mut point, claim) = gpgkr_prove(&mut prover, &zeta, witnesses); + let (mut point, claim) = + gpgkr_prove(&mut prover, table.shape().log_bits(), &zeta, witnesses); let proof = prover.finish(); // Derive the expected factors from the prover's terminal point. From d546f6d0fd6c888b55ccd2693388dd464389d769 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Tue, 22 Sep 2026 19:36:37 +0800 Subject: [PATCH 03/10] gkr: avoid cloning leafs in batched_eval by consuming self --- crates/gkr/src/lib.rs | 20 +++++++++++++++----- 1 file changed, 15 insertions(+), 5 deletions(-) diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs index 7cb1adf5..de2dc47f 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs @@ -299,12 +299,12 @@ impl GrandProductCircuit { // Can't consume the input as the circuit is necessary for the initialisation of fiat shamir // TODO: replace with leaf lookups and add multithreading #[tracing::instrument(name = "Evaluate grand-product circuit", level = "debug", skip_all)] - pub fn batched_eval(&self, groups: usize) -> (Vec, LayerWitnesses) { + pub fn batched_eval(self, groups: usize) -> (Vec, LayerWitnesses) { // +1 to deal with the possible case that the leafs are empty. Given that otherwise the constructor padded it to a power of two, and ilog rounds it down, it becomes a noop let mut witnesses = Vec::with_capacity((self.leafs.len() + 1).ilog2() as usize); // TODO expensive clone going to get replaced by leaf lookups - let mut prev_eval = self.leafs.clone(); + let mut prev_eval = self.leafs; // Stop when there is one output per group while prev_eval.len() > groups { @@ -466,7 +466,12 @@ mod tests { let claim = mle(output, &point); let instance = (leaves.clone(), point.to_vec()); let mut prover = transcript::build_prover("gkr-zero", &instance); - let terminal = gpgkr_prove(&mut prover, leaves.len().ilog2() as usize, &point, witnesses); + let terminal = gpgkr_prove( + &mut prover, + leaves.len().ilog2() as usize, + &point, + witnesses, + ); assert_eq!(terminal.1, mle(leaves.clone(), &terminal.0)); let proof = prover.finish(); @@ -504,11 +509,15 @@ mod tests { pub fn prove(input: Vec, log_groups: usize) -> (Vec, transcript::Proof) { let log_bits = input.len().ilog2() as usize; let circuit = GrandProductCircuit::new(input); + // Fake hashing + let n = circuit.leafs.len() as u128; + + let log_bits = circuit.leafs.len().max(1).ilog2() as usize; let groups = 1usize << log_groups; let (last_value, witnesses) = circuit.batched_eval(groups); // TODO instance is a bit loose and should be replaced by PCS - let instance = (last_value.clone(), circuit.leafs); + let instance = (last_value.clone(), n); let mut prover = transcript::build_prover("gkr", &instance); @@ -524,7 +533,8 @@ mod tests { pub fn verify(input: Vec, output: Vec, proof: transcript::Proof) -> bool { let circuit = GrandProductCircuit::new(input); - let instance = (&output, &circuit.leafs); + // Fake hashing + let instance = (&output, circuit.leafs.len() as u128); let mut verifier = transcript::build_verifier("gkr", &instance, &proof); let log_groups = output.len().max(1).ilog2(); From dc9f073665710d78a32e0cbe9bbb2d8ce2677aab Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Mon, 14 Sep 2026 16:02:51 +0800 Subject: [PATCH 04/10] gkr: inline eq_factor as 1 + r + z, collect eval via iterator instead of a push loop --- crates/gkr/src/lib.rs | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs index de2dc47f..8edff2d3 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs @@ -137,7 +137,7 @@ fn prove_layer( /// this is just `1 + r + z` -- no multiplication at all, and so nothing for /// widemul to help with. fn eq_factor(r: Field, z: Field) -> Field { - poly::eq::eq_eval(&[r], &[z]) + Field::ONE + r + z } /// `a * b * c`: two multiplications in a row. The first is reduced -- it has @@ -308,12 +308,9 @@ impl GrandProductCircuit { // Stop when there is one output per group while prev_eval.len() > groups { - let mut eval = Vec::with_capacity(prev_eval.len() >> 1); let mid = prev_eval.len() / 2; let (l, r) = prev_eval.split_at(mid); - for (&a, &b) in l.iter().zip(r) { - eval.push(a * b) - } + let eval: Vec = l.iter().zip(r).map(|(&a, &b)| a * b).collect(); witnesses.push(prev_eval); prev_eval = eval; } From cced627b342cbec6161a89d86f7cdf45f7084774 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Mon, 14 Sep 2026 16:10:55 +0800 Subject: [PATCH 05/10] gkr: use Field's fused multiply-reduce in mul3_wide, one fewer PMULL --- crates/gkr/src/lib.rs | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs index 8edff2d3..29a7bc08 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs @@ -145,8 +145,12 @@ fn eq_factor(r: Field, z: Field) -> Field { /// multiply -- but the second is left unreduced, so callers can batch its /// reduction with the rest of a running wide sum instead of paying for it on /// every term. +/// +/// The first step uses `Field`'s own fused multiply-reduce (`a * b`, 6 PMULL +/// on aarch64) rather than `Wide256::mul(a, b).reduce()` (4 PMULL to widen + +/// 3 more to reduce = 7): same result, one fewer PMULL. fn mul3_wide(a: Field, b: Field, c: Field) -> Wide256 { - Wide256::mul(Wide256::mul(a, b).reduce(), c) + Wide256::mul(a * b, c) } // TODO: SuffixTable becomes a wrapper around a preallocated vector that is large enough for all rounds. From 0a5031ac00c6772b974aa1b9a9248ee172e332cc Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Mon, 14 Sep 2026 20:49:09 +0800 Subject: [PATCH 06/10] gkr: replace rayon fold with 2-wide interleaved PMULL chains, peel the h==1 last round --- crates/gkr/src/lib.rs | 160 ++++++++++++++++++++++++++++-------------- 1 file changed, 109 insertions(+), 51 deletions(-) diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs index 29a7bc08..f1902eaf 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs @@ -41,7 +41,7 @@ pub fn gpgkr_prove( fn prove_layer( ps: &mut ProverState, storage: &mut [Field], - point: Point, + mut point: Point, mut wnext: Vec, ) -> (Point, Field) { let mut suffix_table = SuffixTable::new(storage, &point); @@ -51,54 +51,37 @@ fn prove_layer( // Tree is encoded in LSB order let (mut mle_l, mut mle_r) = wnext.split_at_mut(mid); + let rounds = point.len(); // TODO: use a double buffer or override approach? Now there is a point allocation each layer // Can go up to ~21 allocations assuming input of 2^35 and 6:4 split - let mut next_point = VecDeque::with_capacity(point.len() + 1); + let mut next_point = VecDeque::with_capacity(rounds + 1); - for (i, z) in (0..point.len()).rev().zip(point) { - // TODO: special-case eq.len() == 1 (final round) to skip the `eq[i] *` - // multiplications below entirely. + // `SuffixTable::layer(i)` has length `2^i`, so `h` is even on every round + // except the last (`i == 0`). Peel that round out instead of checking + // for it on every iteration. + let last_z = point.pop_back(); + + for (i, z) in (1..rounds).rev().zip(point) { // TODO: unwrap will be dealt with in upcoming approach to SuffixTable let eq = suffix_table.layer(i); let h = mle_l.len() / 2; debug_assert_eq!(eq.len(), h); + debug_assert_eq!(h % 2, 0); let (lo_l, hi_l) = mle_l.split_at_mut(h); let (lo_r, hi_r) = mle_r.split_at_mut(h); - // At z = 0 the incoming claim determines the value at zero, so send - // the value at one. Otherwise send the value at zero as usual. + // `send_one` is constant for the whole round: at z = 0 the incoming + // claim determines the value at zero, so this round sends the value + // at one instead. Branch on it once here -- SEND_ONE below is a + // const generic, so each instantiation gets only its own branch, + // not a per-element check inside the hot fold. let send_one = z == Field::ZERO; - let (sum_endpoint, sum_inf) = lo_l - .par_iter_mut() - .zip(lo_r.par_iter_mut()) - .zip(hi_l.par_iter()) - .zip(hi_r.par_iter()) - .zip(eq.par_iter()) - .with_min_len(PARALLEL_MIN_LANES) - .fold( - || (Wide256::zero(), Wide256::zero()), - |(mut sum_endpoint, mut sum_inf), ((((l_lo, r_lo), &l_hi), &r_hi), &e)| { - let (d_l, d_r) = (l_hi - *l_lo, r_hi - *r_lo); - let (l_endpoint, r_endpoint) = if send_one { - (l_hi, r_hi) - } else { - (*l_lo, *r_lo) - }; - - // The endpoint product and `e * (l_hi-l0) * (r_hi-r0)`: each is - // two multiplications in a row, deferred into the - // running wide sums by `mul3_wide` - sum_endpoint += mul3_wide(e, l_endpoint, r_endpoint); - sum_inf += mul3_wide(e, d_l, d_r); - - (sum_endpoint, sum_inf) - }, - ) - .reduce( - || (Wide256::zero(), Wide256::zero()), - |(a0, ainf), (b0, binf)| (a0 + b0, ainf + binf), - ); + let (sum_endpoint, sum_inf) = if send_one { + reduce_round::(lo_l, lo_r, hi_l, hi_r, eq) + } else { + reduce_round::(lo_l, lo_r, hi_l, hi_r, eq) + }; ps.prover_message(&[factor * sum_endpoint.reduce(), factor * sum_inf.reduce()]); @@ -125,6 +108,39 @@ fn prove_layer( factor *= eq_factor(r, z); } + if let Some(z) = last_z { + // i == 0, h == 1: `SuffixTable::new` always seeds the base layer + // (`eq[0]`) with `Field::ONE`, so the `eq[i] *` multiplication every + // other round needs is the identity here -- skipped rather than + // spent multiplying by one. + debug_assert_eq!(mle_l.len(), 2); + let (l_lo, l_hi) = (mle_l[0], mle_l[1]); + let (r_lo, r_hi) = (mle_r[0], mle_r[1]); + let (d_l, d_r) = (l_hi - l_lo, r_hi - r_lo); + + // Same `send_one` swap as the main loop above: at z = 0 the incoming + // claim determines the value at zero, so send the value at one. + let (l_endpoint, r_endpoint) = if z == Field::ZERO { + (l_hi, r_hi) + } else { + (l_lo, r_lo) + }; + let sum_endpoint = Wide256::mul(l_endpoint, r_endpoint); + let sum_inf = Wide256::mul(d_l, d_r); + + ps.prover_message(&[factor * sum_endpoint.reduce(), factor * sum_inf.reduce()]); + + let r = ps.verifier_message(); + next_point.push_back(r); + + mle_l[0] = l_lo + r * d_l; + mle_r[0] = r_lo + r * d_r; + mle_l = &mut mle_l[..1]; + mle_r = &mut mle_r[..1]; + + factor *= eq_factor(r, z); + } + ps.prover_message(&[mle_l[0], mle_r[0]]); let r = ps.verifier_message(); next_point.push_front(r); @@ -132,6 +148,62 @@ fn prove_layer( (next_point, claim) } +/// Sums the round's endpoint and inf products, two elements at a time so +/// both elements' first-stage (fused) multiplies are issued before either +/// second-stage widening multiply -- two independent PMULL chains in flight +/// instead of one, for the CPU (or LLVM's scheduler) to overlap. `SEND_ONE` +/// selects which endpoint this round sends -- `l_lo`/`r_lo` (value at zero) +/// normally, or `l_hi`/`r_hi` (value at one) when `z == 0`; see +/// `prove_layer`. +fn reduce_round( + lo_l: &[Field], + lo_r: &[Field], + hi_l: &[Field], + hi_r: &[Field], + eq: &[Field], +) -> (Wide256, Wide256) { + lo_l.par_chunks_exact(2) + .zip(lo_r.par_chunks_exact(2)) + .zip(hi_l.par_chunks_exact(2)) + .zip(hi_r.par_chunks_exact(2)) + .zip(eq.par_chunks_exact(2)) + // Chunked by 2, so the element-count threshold below is halved in + // terms of chunks. + .with_min_len(PARALLEL_MIN_LANES / 2) + .fold( + || (Wide256::zero(), Wide256::zero()), + |(mut sum_endpoint, mut sum_inf), ((((l_lo, r_lo), l_hi), r_hi), e)| { + let (d_l0, d_r0) = (l_hi[0] - l_lo[0], r_hi[0] - r_lo[0]); + let (d_l1, d_r1) = (l_hi[1] - l_lo[1], r_hi[1] - r_lo[1]); + + let (l_e0, r_e0) = if SEND_ONE { + (l_hi[0], r_hi[0]) + } else { + (l_lo[0], r_lo[0]) + }; + let (l_e1, r_e1) = if SEND_ONE { + (l_hi[1], r_hi[1]) + } else { + (l_lo[1], r_lo[1]) + }; + + let (a0, b0) = (e[0] * l_e0, e[0] * d_l0); + let (a1, b1) = (e[1] * l_e1, e[1] * d_l1); + + sum_endpoint += Wide256::mul(a0, r_e0); + sum_inf += Wide256::mul(b0, d_r0); + sum_endpoint += Wide256::mul(a1, r_e1); + sum_inf += Wide256::mul(b1, d_r1); + + (sum_endpoint, sum_inf) + }, + ) + .reduce( + || (Wide256::zero(), Wide256::zero()), + |(a0, ainf), (b0, binf)| (a0 + b0, ainf + binf), + ) +} + /// `eq(r, z) = r*z + (1 - r)*(1 - z)`. Expanding gives /// `1 + r + z + 2*r*z`, and in characteristic 2 `2*r*z = r*z + r*z = 0`, so /// this is just `1 + r + z` -- no multiplication at all, and so nothing for @@ -140,19 +212,6 @@ fn eq_factor(r: Field, z: Field) -> Field { Field::ONE + r + z } -/// `a * b * c`: two multiplications in a row. The first is reduced -- it has -/// to come back down to a field element to feed the second carryless -/// multiply -- but the second is left unreduced, so callers can batch its -/// reduction with the rest of a running wide sum instead of paying for it on -/// every term. -/// -/// The first step uses `Field`'s own fused multiply-reduce (`a * b`, 6 PMULL -/// on aarch64) rather than `Wide256::mul(a, b).reduce()` (4 PMULL to widen + -/// 3 more to reduce = 7): same result, one fewer PMULL. -fn mul3_wide(a: Field, b: Field, c: Field) -> Wide256 { - Wide256::mul(a * b, c) -} - // TODO: SuffixTable becomes a wrapper around a preallocated vector that is large enough for all rounds. // SuffixTable can be 'created' each round / destroyed to ensure proper truncation of the underlying vector // TODO: Split suffix table @@ -508,7 +567,6 @@ mod tests { } pub fn prove(input: Vec, log_groups: usize) -> (Vec, transcript::Proof) { - let log_bits = input.len().ilog2() as usize; let circuit = GrandProductCircuit::new(input); // Fake hashing let n = circuit.leafs.len() as u128; From 0b619d3c96bd8556d263b291f902d8824485d1a3 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Wed, 16 Sep 2026 22:13:48 +0800 Subject: [PATCH 07/10] update single threaded patch --- single_threaded.patch | 79 ++++++++++++++++++++++--------------------- 1 file changed, 41 insertions(+), 38 deletions(-) diff --git a/single_threaded.patch b/single_threaded.patch index 0f6cf85f..012e19f1 100644 --- a/single_threaded.patch +++ b/single_threaded.patch @@ -1,44 +1,8 @@ diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs +index 89a1bf0..bb364ff 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs -@@ -62,14 +62,14 @@ - // the value at one. Otherwise send the value at zero as usual. - let send_one = z == Field::ZERO; - let (sum_endpoint, sum_inf) = lo_l -- .par_iter_mut() -- .zip(lo_r.par_iter_mut()) -- .zip(hi_l.par_iter()) -- .zip(hi_r.par_iter()) -- .zip(eq.par_iter()) -- .with_min_len(PARALLEL_MIN_LANES) -+ .iter_mut() -+ .zip(lo_r.iter_mut()) -+ .zip(hi_l.iter()) -+ .zip(hi_r.iter()) -+ .zip(eq.iter()) -+ // .with_min_len(PARALLEL_MIN_LANES) - .fold( -- || (Wide256::zero(), Wide256::zero()), -+ (Wide256::zero(), Wide256::zero()), - |(mut sum_endpoint, mut sum_inf), ((((l_lo, r_lo), &l_hi), &r_hi), &e)| { - let (d_l, d_r) = (l_hi - *l_lo, r_hi - *r_lo); - let (l_endpoint, r_endpoint) = if send_one { -@@ -86,22 +86,22 @@ - - (sum_endpoint, sum_inf) - }, -- ) -- .reduce( -- || (Wide256::zero(), Wide256::zero()), -- |(a0, ainf), (b0, binf)| (a0 + b0, ainf + binf), - ); -+ // .reduce( -+ // || (Wide256::zero(), Wide256::zero()), -+ // |(a0, ainf), (b0, binf)| (a0 + b0, ainf + binf), -+ // ); - - ps.prover_message(&[factor * sum_endpoint.reduce(), factor * sum_inf.reduce()]); - +@@ -87,11 +87,11 @@ fn prove_layer( let r = ps.verifier_message(); next_point.push_back(r); @@ -55,3 +19,42 @@ diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs .for_each(|(((l_lo, r_lo), &l_hi), &r_hi)| { // Recalculating delta rather than writing it in the top half of the array in the previous loop to save on writes. // No benchmarking has been done to check the difference. +@@ -161,16 +161,16 @@ fn reduce_round( + hi_r: &[Field], + eq: &[Field], + ) -> (Wide256, Wide256) { +- lo_l.par_chunks_exact(2) +- .zip(lo_r.par_chunks_exact(2)) +- .zip(hi_l.par_chunks_exact(2)) +- .zip(hi_r.par_chunks_exact(2)) +- .zip(eq.par_chunks_exact(2)) ++ lo_l.chunks_exact(2) ++ .zip(lo_r.chunks_exact(2)) ++ .zip(hi_l.chunks_exact(2)) ++ .zip(hi_r.chunks_exact(2)) ++ .zip(eq.chunks_exact(2)) + // Chunked by 2, so the element-count threshold below is halved in + // terms of chunks. +- .with_min_len(PARALLEL_MIN_LANES / 2) ++ // .with_min_len(PARALLEL_MIN_LANES / 2) + .fold( +- || (Wide256::zero(), Wide256::zero()), ++ (Wide256::zero(), Wide256::zero()), + |(mut sum_endpoint, mut sum_inf), ((((l_lo, r_lo), l_hi), r_hi), e)| { + let (d_l0, d_r0) = (l_hi[0] - l_lo[0], r_hi[0] - r_lo[0]); + let (d_l1, d_r1) = (l_hi[1] - l_lo[1], r_hi[1] - r_lo[1]); +@@ -197,10 +197,10 @@ fn reduce_round( + (sum_endpoint, sum_inf) + }, + ) +- .reduce( +- || (Wide256::zero(), Wide256::zero()), +- |(a0, ainf), (b0, binf)| (a0 + b0, ainf + binf), +- ) ++ // .reduce( ++ // || (Wide256::zero(), Wide256::zero()), ++ // |(a0, ainf), (b0, binf)| (a0 + b0, ainf + binf), ++ // ) + } + + /// `eq(r, z) = r*z + (1 - r)*(1 - z)`. Expanding gives From 368597697a9b769ba4f49e6a39ca47faeaa4dda6 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Tue, 22 Sep 2026 16:58:33 +0800 Subject: [PATCH 08/10] gkr: make SuffixTable immutable and bounds-check free via zip loop, add alloc_storage helper --- crates/gkr/src/lib.rs | 45 +++++++++++++++++++++++++------------------ 1 file changed, 26 insertions(+), 19 deletions(-) diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs index f1902eaf..0c7be252 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs @@ -28,7 +28,7 @@ pub fn gpgkr_prove( let mut point = VecDeque::from(point); let mut claim = Field::ZERO; - let mut storage = Field::zeroed_vec(1 << log_bits); + let mut storage = SuffixTable::alloc_storage(log_bits); for wnext in witnesses.into_iter() { (point, claim) = prove_layer(ps, &mut storage, point, wnext); } @@ -44,7 +44,7 @@ fn prove_layer( mut point: Point, mut wnext: Vec, ) -> (Point, Field) { - let mut suffix_table = SuffixTable::new(storage, &point); + let suffix_table = SuffixTable::new(storage, &point); let mut factor = Field::ONE; let mid = wnext.len() / 2; @@ -52,8 +52,6 @@ fn prove_layer( let (mut mle_l, mut mle_r) = wnext.split_at_mut(mid); let rounds = point.len(); - // TODO: use a double buffer or override approach? Now there is a point allocation each layer - // Can go up to ~21 allocations assuming input of 2^35 and 6:4 split let mut next_point = VecDeque::with_capacity(rounds + 1); // `SuffixTable::layer(i)` has length `2^i`, so `h` is even on every round @@ -62,7 +60,6 @@ fn prove_layer( let last_z = point.pop_back(); for (i, z) in (1..rounds).rev().zip(point) { - // TODO: unwrap will be dealt with in upcoming approach to SuffixTable let eq = suffix_table.layer(i); let h = mle_l.len() / 2; debug_assert_eq!(eq.len(), h); @@ -162,6 +159,7 @@ fn reduce_round( hi_r: &[Field], eq: &[Field], ) -> (Wide256, Wide256) { + // chunking reduces variance in benchmarking lo_l.par_chunks_exact(2) .zip(lo_r.par_chunks_exact(2)) .zip(hi_l.par_chunks_exact(2)) @@ -216,15 +214,14 @@ fn eq_factor(r: Field, z: Field) -> Field { // SuffixTable can be 'created' each round / destroyed to ensure proper truncation of the underlying vector // TODO: Split suffix table struct SuffixTable<'a> { - storage: &'a mut [Field], - offset: usize, + storage: &'a [Field], } impl<'a> SuffixTable<'a> { - fn layer(&mut self, i: usize) -> &mut [Field] { - let start = (1 << i) - 1 - self.offset; - let end = (1 << (i + 1)) - 1 - self.offset; - &mut self.storage[start..end] + fn layer(&self, i: usize) -> &[Field] { + let start = (1 << i) - 1; + let end = (1 << (i + 1)) - 1; + &self.storage[start..end] } } @@ -232,6 +229,15 @@ impl<'a> SuffixTable<'a> { /// Allocates all directly as it is as much space as a double buffer approach would take. #[inline(never)] fn new(storage: &'a mut [Field], point: &Point) -> SuffixTable<'a> { + // The three asserts below (here and in the loop) are load-bearing + // for codegen, not just documentation: without them, LLVM can't + // prove `storage` is big enough for the splits and per-element + // `low`/`hi` writes below, so every one of those carries its own + // bounds check -- including inside the hot per-element loop. With + // them, none of it does; verified via disassembly. + let needed = (1usize << point.len()).saturating_sub(1).max(1); + assert!(storage.len() >= needed); + storage[0] = F128::ONE; let (mut prev, mut remaining) = storage.split_at_mut(1); // The selector is the first entry of the point and we need to skip @@ -243,22 +249,25 @@ impl<'a> SuffixTable<'a> { // Suffix table is in the reverse order of the point for &z in c.rev() { let size = prev.len() << 1; + assert!(remaining.len() >= size); let (entry, next) = remaining.split_at_mut(size); + assert!(entry.len() >= size >> 1); let (low, hi) = entry.split_at_mut(size >> 1); - for (i, &e) in prev.iter().enumerate() { + for ((l, h), &e) in low.iter_mut().zip(hi.iter_mut()).zip(prev.iter()) { let tmp = z * e; // (1-z)*e, z*e - (low[i], hi[i]) = (e - tmp, tmp) + (*l, *h) = (e - tmp, tmp) } prev = entry; remaining = next; } - SuffixTable { - storage: storage, - offset: 0, - } + SuffixTable { storage } + } + + fn alloc_storage(log_bits: usize) -> Vec { + Field::zeroed_vec(1 << (log_bits.saturating_sub(1)).max(1)) } } @@ -359,14 +368,12 @@ impl GrandProductCircuit { } // Returns the final evaluation and the witnesses of the intermediate layers - // Can't consume the input as the circuit is necessary for the initialisation of fiat shamir // TODO: replace with leaf lookups and add multithreading #[tracing::instrument(name = "Evaluate grand-product circuit", level = "debug", skip_all)] pub fn batched_eval(self, groups: usize) -> (Vec, LayerWitnesses) { // +1 to deal with the possible case that the leafs are empty. Given that otherwise the constructor padded it to a power of two, and ilog rounds it down, it becomes a noop let mut witnesses = Vec::with_capacity((self.leafs.len() + 1).ilog2() as usize); - // TODO expensive clone going to get replaced by leaf lookups let mut prev_eval = self.leafs; // Stop when there is one output per group From 68906d3c816d95a12441acea361549c9ee5333a2 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Tue, 22 Sep 2026 19:06:34 +0800 Subject: [PATCH 09/10] gkr: drop now-redundant per-round bounds asserts in SuffixTable::new --- crates/gkr/src/lib.rs | 16 +++------------- 1 file changed, 3 insertions(+), 13 deletions(-) diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs index 0c7be252..ad53ca17 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs @@ -210,8 +210,6 @@ fn eq_factor(r: Field, z: Field) -> Field { Field::ONE + r + z } -// TODO: SuffixTable becomes a wrapper around a preallocated vector that is large enough for all rounds. -// SuffixTable can be 'created' each round / destroyed to ensure proper truncation of the underlying vector // TODO: Split suffix table struct SuffixTable<'a> { storage: &'a [Field], @@ -229,12 +227,8 @@ impl<'a> SuffixTable<'a> { /// Allocates all directly as it is as much space as a double buffer approach would take. #[inline(never)] fn new(storage: &'a mut [Field], point: &Point) -> SuffixTable<'a> { - // The three asserts below (here and in the loop) are load-bearing - // for codegen, not just documentation: without them, LLVM can't - // prove `storage` is big enough for the splits and per-element - // `low`/`hi` writes below, so every one of those carries its own - // bounds check -- including inside the hot per-element loop. With - // them, none of it does; verified via disassembly. + // Establishes the whole table's space budget up front so LLVM can + // prove the `low`/`hi` writes in the loop below are in bounds let needed = (1usize << point.len()).saturating_sub(1).max(1); assert!(storage.len() >= needed); @@ -249,10 +243,8 @@ impl<'a> SuffixTable<'a> { // Suffix table is in the reverse order of the point for &z in c.rev() { let size = prev.len() << 1; - assert!(remaining.len() >= size); let (entry, next) = remaining.split_at_mut(size); - assert!(entry.len() >= size >> 1); - let (low, hi) = entry.split_at_mut(size >> 1); + let (low, hi) = entry.split_at_mut(entry.len() >> 1); for ((l, h), &e) in low.iter_mut().zip(hi.iter_mut()).zip(prev.iter()) { let tmp = z * e; @@ -351,8 +343,6 @@ fn verify_layer(vs: &mut VerifierState, mut claim: Field, point: Point) -> Optio } } -//TODO circuit and circuit eval can't have their innards directly available as that would break power of 2 requirements for the rest. -// A circuit is defined by its leaf value only because it is a balanced tree // TODO: Optimise for circuits that are padded. pub struct GrandProductCircuit { leafs: Vec, From 168acd1f19f13fc1331bf5c9728f4b8a820fd574 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Tue, 22 Sep 2026 19:40:13 +0800 Subject: [PATCH 10/10] fixup! gkr: drop now-redundant per-round bounds asserts in SuffixTable::new --- crates/gkr/src/lib.rs | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs index ad53ca17..1f390e6a 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs @@ -75,9 +75,9 @@ fn prove_layer( // not a per-element check inside the hot fold. let send_one = z == Field::ZERO; let (sum_endpoint, sum_inf) = if send_one { - reduce_round::(lo_l, lo_r, hi_l, hi_r, eq) + reduce_sumcheck_round::(lo_l, lo_r, hi_l, hi_r, eq) } else { - reduce_round::(lo_l, lo_r, hi_l, hi_r, eq) + reduce_sumcheck_round::(lo_l, lo_r, hi_l, hi_r, eq) }; ps.prover_message(&[factor * sum_endpoint.reduce(), factor * sum_inf.reduce()]); @@ -145,14 +145,7 @@ fn prove_layer( (next_point, claim) } -/// Sums the round's endpoint and inf products, two elements at a time so -/// both elements' first-stage (fused) multiplies are issued before either -/// second-stage widening multiply -- two independent PMULL chains in flight -/// instead of one, for the CPU (or LLVM's scheduler) to overlap. `SEND_ONE` -/// selects which endpoint this round sends -- `l_lo`/`r_lo` (value at zero) -/// normally, or `l_hi`/`r_hi` (value at one) when `z == 0`; see -/// `prove_layer`. -fn reduce_round( +fn reduce_sumcheck_round( lo_l: &[Field], lo_r: &[Field], hi_l: &[Field],