From 48fb7c88dd8aa4cf4a0625b1aaaeba9e9aa4365b Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Thu, 8 Oct 2026 15:54:30 +0800 Subject: [PATCH 01/13] gkr: start looking up the bottom three layers in 16-entry tables (WIP, does not build) --- crates/common/src/matrix.rs | 16 +++++- crates/gkr/src/lib.rs | 103 ++++++++++++++++++++++++------------ 2 files changed, 84 insertions(+), 35 deletions(-) diff --git a/crates/common/src/matrix.rs b/crates/common/src/matrix.rs index 2f7f513..7d8b9aa 100644 --- a/crates/common/src/matrix.rs +++ b/crates/common/src/matrix.rs @@ -1,6 +1,6 @@ //! Bit matrices in memory order, and the transpose between them. -use std::{borrow::Cow, slice::Iter}; +use std::{borrow::Cow, ops::Deref, slice::Iter}; /// The word `packed` is sliced into: one `F128` element. [`crate::BitTable::BITS`] /// is derived from this, so changing it is the only step needed to repack @@ -75,6 +75,10 @@ impl<'a> BitMatrix<'a> { 1 << self.log_dim2 } + pub fn log_dim2(&self) -> usize { + self.log_dim2 + } + /// The `dim2` bits of `xs[i1]`, in ascending order. /// /// # Panics @@ -233,7 +237,7 @@ impl<'a, const S: usize> StepIter<'a, S> { "step size must be nonzero, at most Word::BITS, and divide Word::BITS evenly" ); const MASK: Word = (1 as Word).unbounded_shl(S as u32).wrapping_sub(1); - fn new(iter: Iter<'a, Word>) -> Self { + pub fn new(iter: Iter<'a, Word>) -> Self { let () = Self::CHECK_SIZE; Self { elements: iter, @@ -270,6 +274,14 @@ impl ExactSizeIterator for StepIter<'_, S> { } } +impl<'a> Deref for BitMatrix<'a> { + type Target = [Word]; + + fn deref(&self) -> &Self::Target { + &self.packed + } +} + #[cfg(test)] mod tests { use super::*; diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs index a00ebfe..a7649ba 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs @@ -1,6 +1,6 @@ use std::collections::VecDeque; -use common::BitTable; +use common::{BitMatrix, BitTable, matrix::StepIter}; use field::{F128, Wide256}; use num_traits::{ConstOne, ConstZero}; use rayon::prelude::*; @@ -338,44 +338,76 @@ fn verify_layer(vs: &mut VerifierState, mut claim: Field, point: Point) -> Optio } // TODO: Optimise for circuits that are padded. -pub struct GrandProductCircuit { - leafs: Vec, +pub struct GrandProductCircuit<'a> { + matrix: BitMatrix<'static>, + images: &'a [F128], } -impl GrandProductCircuit { - pub fn new(row_images: &[F128], bittable: &BitTable) -> Self { +impl<'a> GrandProductCircuit<'a> { + pub fn new(row_images: &'a [F128], bittable: &BitTable) -> GrandProductCircuit<'a> { // Columns occupy the low index bits, so each product tree reduces one // column. Transposed so each row can be read sequentially. - let matrix = bittable.as_matrix().transpose(); - let dim2 = matrix.dim2(); - let dim = matrix.dim1() * dim2; - let mut leafs = F128::zeroed_vec(dim); - - for (b, &row_image) in row_images.iter().enumerate() { - let leafs = &mut leafs[b * dim2..(b + 1) * dim2]; - for (leaf, bit) in leafs.iter_mut().zip(matrix.bits(b)) { - *leaf = if bit == 1 { row_image } else { F128::ONE }; - } + let matrix = bittable.as_matrix(); + // c b = c q b'. Since we combine within subtrees and want to start from d-3 that is 4 bits + let new_dim2 = matrix.log_dim2() - 2; // should this be -2 or -4 if we want the lookup to work? + let matrix = matrix.reshape(new_dim2).transpose(); + + Self { + matrix: matrix, + images: row_images, } - - Self::from_leafs(leafs) - } - - fn from_leafs(mut leafs: Vec) -> Self { - if !leafs.is_empty() { - leafs.resize(leafs.len().next_power_of_two(), Field::ONE); - } - GrandProductCircuit { leafs } } // Returns the final evaluation and the witnesses of the intermediate layers // 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) { + // Lower + /// `lut[s]` is the product of the inputs whose bit is set in `s`, `a` + /// on bit 0 through `d` on bit 3 (`0b0000 -> 1`, `0b0001 -> a`, + /// `0b0011 -> a*b`, ...), so a 4-bit lookup string indexes its + /// multilinear value directly. + /// + /// Built by doubling: entries `2^i..2^(i+1)` are entries `0..2^i` + /// times the `i`-th input. The first entry of each new half is the + /// input itself, so it is copied rather than multiplied by one, + /// leaving 11 multiplications -- one per entry with two or more + /// factors, which is the minimum. + fn create_lut(c: [Field; 4]) -> [Field; 16] { + let mut lut = [Field::ONE; 16]; + for (i, x) in c.into_iter().enumerate() { + let half = 1 << i; + lut[half] = x; + for k in 1..half { + lut[half + k] = lut[k] * x; + } + } + lut + } + + let n = self.images.len(); + // deal with the case that it isn't divisible by 4. What does that mean? Can we just put an assert. A certain invariant is passed down + let quart = n / 4; + // Each LUT takes one image from each quarter of the rows. + let (a, rest) = self.images.split_at(quart); + let (b, rest) = rest.split_at(quart); + let (c, d) = rest.split_at(quart); + let lower_lut = a + .iter() + .zip(b) + .zip(c) + .zip(d) + .map(|(((&a, &b), &c), &d)| create_lut([a, b, c, d])) + .collect(); + + // Upper // +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); + // derive based on matrix length and subtract the last three layers + let mut witnesses = Vec::with_capacity((self.matrix.len() + 1).ilog2() as usize); - let mut prev_eval = self.leafs; + // split the matrix in the middle for left and right and a table is shared among c. + // TODO This will start with a what was created from the last LUT layer + let mut prev_eval = StepIter::<4>::new(self.matrix.iter()).map(|bs| lower_lut[]); // Stop when there is one output per group while prev_eval.len() > groups { @@ -386,25 +418,30 @@ impl GrandProductCircuit { prev_eval = eval; } - (prev_eval, LayerWitnesses(witnesses)) + ( + prev_eval, + LayerWitnesses { + lower_lut: lower_lut, + upper_nodes: witnesses, + }, + ) } } // TODO: This can be a single vec that is allocated at the start. -pub struct LayerWitnesses(Vec>); - -impl LayerWitnesses { - pub fn pop(&mut self) -> Option> { - self.0.pop() - } +pub struct LayerWitnesses { + lower_lut: Vec<[Field; 16]>, + upper_nodes: Vec>, } +// Can I still use into iterator or should something else be needed? impl IntoIterator for LayerWitnesses { type Item = Vec; type IntoIter = std::iter::Rev>>; fn into_iter(self) -> Self::IntoIter { self.0.into_iter().rev() + // TODO: do we add the LUT handling here? Good for an intermediate step such that the sumcheck part keeps working. They will be a lookup from } } From afb9dad4b9c601b43f97f2aecf1dff1d492e75d7 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Fri, 9 Oct 2026 12:30:12 +0800 Subject: [PATCH 02/13] common: let bit_transpose drop to u64..u8 blocks for an axis shorter than a word --- crates/common/src/matrix.rs | 144 ++++++++++++++++++++++++++---------- crates/common/src/params.rs | 13 ++-- crates/common/src/table.rs | 6 +- 3 files changed, 115 insertions(+), 48 deletions(-) diff --git a/crates/common/src/matrix.rs b/crates/common/src/matrix.rs index 7d8b9aa..4f428a9 100644 --- a/crates/common/src/matrix.rs +++ b/crates/common/src/matrix.rs @@ -1,6 +1,10 @@ //! Bit matrices in memory order, and the transpose between them. -use std::{borrow::Cow, ops::Deref, slice::Iter}; +use std::{ + borrow::Cow, + ops::{BitAnd, BitXor, BitXorAssign, Deref, Not, Shl, Shr}, + slice::Iter, +}; /// The word `packed` is sliced into: one `F128` element. [`crate::BitTable::BITS`] /// is derived from this, so changing it is the only step needed to repack @@ -8,7 +12,8 @@ use std::{borrow::Cow, ops::Deref, slice::Iter}; /// /// 128 rather than 64 bits because `bit_transpose` is bound by its strided /// loads and stores, and each one then moves 16 bytes instead of 8: the same -/// code with `u64` measured 20-55% slower on `BitMatrix::transpose`. +/// code with `u64` measured 20-55% slower on `BitMatrix::transpose`. It +/// drops to a narrower word only for an axis shorter than this. pub(crate) type Word = u128; /// Bits in a [`Word`]. const BITS: usize = Word::BITS as usize; @@ -51,9 +56,9 @@ impl<'a> BitMatrix<'a> { /// /// # Constraints /// - /// - **Both axes**: at least `128`, since `bit_transpose` moves whole - /// `128 x 128` blocks. Its doc comment says where a version for a - /// shorter `dim1` lives. + /// - **Both axes**: at least `8`, since `bit_transpose` moves square + /// blocks of the widest word both axes fill, `u128` down to `u8`. Its + /// doc comment says where a version for a shorter `dim1` lives. /// - **Allocation**: always a full out-of-place copy -- a fresh owned /// buffer the same size as the packed words (up to 4 GiB at `m = 35`), /// whether `self` owned or borrowed them. @@ -110,27 +115,51 @@ impl<'a> BitMatrix<'a> { } /// Transposes `xs[dim1][dim2]`, `dim2` contiguous and both in bits, into an -/// owned `[dim2][dim1]`, one `BITS x BITS` block at a time. +/// owned `[dim2][dim1]`, one square block at a time. /// -/// Both axes must be whole words. Commit -/// 336e8097225aca4f8146f35594740b0e5054d839 has a version that also takes -/// any power-of-two `dim1` below a word (`bit_transpose_blocks` and -/// `rotate_bit_index` in this file): it gathers `BITS / dim1` words of each -/// row into one block and interleaves them by rotating the bit index. Run in -/// reverse, the same steps would take `dim2` below a word, if that is ever -/// needed. Bringing it back also lifts [`crate::ParamsError::ColumnCountTooNarrow`], -/// the gate this limit puts on tables. +/// The block is the widest word both axes are whole multiples of: a [`Word`] +/// whenever they can be, otherwise the words are read in place as `u64` down +/// to `u8` and the same kernel runs at that width. An axis of 1, 2 or 4 bits +/// fills no word. For a short `dim1`, commit +/// 336e8097225aca4f8146f35594740b0e5054d839 has a version that reaches those +/// too (`bit_transpose_blocks` and `rotate_bit_index` in this file): it +/// gathers `BITS / dim1` words of each row into one block and interleaves +/// them by rotating the bit index. fn bit_transpose(xs: &[Word], dim1: usize, dim2: usize) -> Vec { assert_eq!(xs.len() * BITS, dim1 * dim2); - assert!( - dim1.is_multiple_of(BITS) && dim2.is_multiple_of(BITS), - "{dim1} x {dim2} bits is not whole words along both axes" - ); - let dim1_words = dim1 / BITS; - let dim2_words = dim2 / BITS; let mut out = bytemuck::zeroed_vec(xs.len()); - let mut block = [0; BITS]; + // The lowest set bit of `dim1 | dim2` is the largest power of two both + // are multiples of; or-ing in `BITS` caps it at a `Word`. + match 1usize << (dim1 | dim2 | BITS).trailing_zeros() { + 8 => transpose_blocks::(xs, &mut out, dim1, dim2), + 16 => transpose_blocks::(xs, &mut out, dim1, dim2), + 32 => transpose_blocks::(xs, &mut out, dim1, dim2), + 64 => transpose_blocks::(xs, &mut out, dim1, dim2), + BITS => transpose_blocks::(xs, &mut out, dim1, dim2), + _ => panic!("{dim1} x {dim2} bits is not whole words along both axes"), + } + + out +} + +/// [`bit_transpose`] in `N x N` blocks of the `N`-bit word `W`. +fn transpose_blocks( + xs: &[Word], + out: &mut [Word], + dim1: usize, + dim2: usize, +) { + const { assert!(N == 8 * size_of::()) }; + // Read in place: on a little-endian target, bit `i` of the words is bit + // `i` of the narrower words too. + const { assert!(cfg!(target_endian = "little")) }; + let xs: &[W] = bytemuck::cast_slice(xs); + let out: &mut [W] = bytemuck::cast_slice_mut(out); + let dim1_words = dim1 / N; + let dim2_words = dim2 / N; + + let mut block = [W::zeroed(); N]; // g1 innermost, so consecutive blocks write neighbouring words of the // same output rows while those lines are still cached, and it is the // input that is revisited after a sweep; the other way round measured @@ -138,34 +167,51 @@ fn bit_transpose(xs: &[Word], dim1: usize, dim2: usize) -> Vec { for g2 in 0..dim2_words { for g1 in 0..dim1_words { for (a, word) in block.iter_mut().enumerate() { - *word = xs[(g1 * BITS + a) * dim2_words + g2]; + *word = xs[(g1 * N + a) * dim2_words + g2]; } transpose_bit_block(&mut block); for (o, &word) in block.iter().enumerate() { - out[(g2 * BITS + o) * dim1_words + g1] = word; + out[(g2 * N + o) * dim1_words + g1] = word; } } } +} - out +/// An unsigned word [`transpose_bit_block`] can work in. +trait BlockWord: + bytemuck::Pod + + BitAnd + + BitXor + + BitXorAssign + + Not + + Shl + + Shr +{ } -/// Transposes a `BITS x BITS` bit matrix stored as `BITS` words of `BITS` -/// bits: after the call, bit `j` of word `i` is what bit `i` of word `j` held -/// before it, for every `i, j`. +impl BlockWord for u8 {} +impl BlockWord for u16 {} +impl BlockWord for u32 {} +impl BlockWord for u64 {} +impl BlockWord for u128 {} + +/// Transposes an `N x N` bit matrix stored as `N` words of `N` bits: after +/// the call, bit `j` of word `i` is what bit `i` of word `j` held before it, +/// for every `i, j`. /// /// The recursive-doubling bit-matrix transpose (Hacker's Delight, 2nd ed., -/// §7-3), generalized from its usual 64x64 form to the 128-bit lane width -/// `F128` packs bits into. `O(n log n)` word operations rather than the -/// `O(n^2)` of moving one bit at a time. -fn transpose_bit_block(a: &mut [Word; BITS]) { - let mut j = BITS / 2; - let mut mask: Word = (1 << j) - 1; +/// §7-3), generalized from its usual 64x64 form to any word width, `u128` +/// being the lane width `F128` packs bits into. `O(n log n)` word operations +/// rather than the `O(n^2)` of moving one bit at a time. +fn transpose_bit_block(a: &mut [W; N]) { + let mut j = N / 2; + // The low half of the word. + let mut mask = !W::zeroed() >> j; while j > 0 { let mut k = 0; - while k < BITS { + while k < N { let t = ((a[k] >> j) ^ a[k | j]) & mask; a[k | j] ^= t; a[k] ^= t << j; @@ -388,9 +434,10 @@ mod tests { const TRANSPOSE_LOG_DIM1: usize = 9; const TRANSPOSE_LOG_DIM2: usize = 8; - #[test] - fn transpose_swaps_the_axes() { - let matrix = pseudo_random_matrix(TRANSPOSE_LOG_DIM1, TRANSPOSE_LOG_DIM2); + /// The transpose of a pseudo-random `xs[2^log_dim1][2^log_dim2]` against + /// the bit accessor. + fn check_transpose(log_dim1: usize, log_dim2: usize) { + let matrix = pseudo_random_matrix(log_dim1, log_dim2); let transposed = matrix.transpose(); assert_eq!(transposed.dim1(), matrix.dim2()); @@ -401,12 +448,33 @@ mod tests { assert_eq!( bit(&transposed, i2, i1), bit(&matrix, i1, i2), - "xs[{i1}][{i2}] vs its transpose" + "xs[{i1}][{i2}] vs its transpose, {log_dim1} x {log_dim2}" ); } } } + #[test] + fn transpose_swaps_the_axes() { + check_transpose(TRANSPOSE_LOG_DIM1, TRANSPOSE_LOG_DIM2); + } + + /// An axis of 8, 16, 32 and 64 bits, so each narrower block word runs, + /// with the short axis on either side. + #[test] + fn transpose_swaps_the_axes_of_a_short_axis() { + for log_short in 3..7 { + check_transpose(log_short, 9); + check_transpose(9, log_short); + } + } + + #[test] + #[should_panic(expected = "whole words")] + fn transpose_rejects_an_axis_below_a_byte() { + pseudo_random_matrix(2, 10).transpose(); + } + #[test] fn transposing_twice_recovers_the_original_bits() { let matrix = pseudo_random_matrix(TRANSPOSE_LOG_DIM1, TRANSPOSE_LOG_DIM2); diff --git a/crates/common/src/params.rs b/crates/common/src/params.rs index 51ee4a0..26285ac 100644 --- a/crates/common/src/params.rs +++ b/crates/common/src/params.rs @@ -10,13 +10,12 @@ use crate::{BitTable, Shape, TableError, VirtualMap}; pub enum ParamsError { /// Fewer than `128` columns (`s < 7`). /// - /// Not a protocol bound: it flows up from the prover's transpose. The - /// reduction reads the GKR leaves row by row through the - /// [`crate::BitMatrix::transpose`] of [`BitTable::as_matrix`], where each - /// transposed row holds one bit per column, and `bit_transpose` in - /// `crates/common/src/matrix.rs` only moves whole `128 x 128` blocks. Teaching it rows shorter than a word - /// lifts this gate; commit 336e8097225aca4f8146f35594740b0e5054d839 has - /// that version. + /// Not a protocol bound: it flowed up from the prover's transpose, when + /// `bit_transpose` in `crates/common/src/matrix.rs` moved only whole + /// `128 x 128` blocks. It now drops to narrower words, and the GKR + /// relayout transposes `4 * columns` by `rows / 4` bits, so the prover + /// itself only needs `s >= 1`. This gate is stricter than that until it + /// is revisited. ColumnCountTooNarrow, /// `(k_1 + 1)(Q - 1)` reaches `ord(g)`, so two folds could collide in the /// exponent. diff --git a/crates/common/src/table.rs b/crates/common/src/table.rs index 04edaad..3196f6c 100644 --- a/crates/common/src/table.rs +++ b/crates/common/src/table.rs @@ -246,9 +246,9 @@ mod tests { #[test] #[should_panic(expected = "whole words")] fn transpose_rejects_a_table_with_too_few_columns() { - // `log_columns = 0`: a single column, well under the 128 a - // transposed row would need to fill one word -- even though this - // shape is perfectly admissible for `BitTable` itself. + // `log_columns = 0`: a single column, under the 8 a transposed row + // would need to fill even a byte -- even though this shape is + // perfectly admissible for `BitTable` itself. let shape = Shape::new(22, 0).unwrap(); let packed = vec![F128::ZERO; (1 << shape.log_bits()) / PACKED_BITS]; let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); From 7da7355fccd3000ade72303278b98e6c6593d366 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Fri, 9 Oct 2026 12:53:01 +0800 Subject: [PATCH 03/13] gkr: build layer 3 up from 16-entry tables and nibbles, rebuilding layers 0 to 2 for the sumcheck --- crates/gkr/src/lib.rs | 252 ++++++++++++++++++++++++++++-------- crates/prover/src/reduce.rs | 2 +- 2 files changed, 197 insertions(+), 57 deletions(-) diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs index a7649ba..65add32 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs @@ -18,8 +18,8 @@ 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, + // All the intermediate witnesses + the input layer, from the top down. Doesn't contain the output layer + witnesses: impl IntoIterator>, ) -> (Vec, Field) { // Edge cases // - empty witnesses -> single constant circuit -> one verifier message that permutes the proof state, but a single constant can't have an MLE @@ -29,9 +29,9 @@ pub fn gpgkr_prove( let mut point = VecDeque::from(point); let mut claim = Field::ZERO; - let mut storage = SuffixTable::alloc_storage(log_bits); + let mut suffix_storage = SuffixTable::alloc_storage(log_bits); for wnext in witnesses.into_iter() { - (point, claim) = prove_layer(ps, &mut storage, point, wnext); + (point, claim) = prove_layer(ps, &mut suffix_storage, point, wnext); } let mut point = Vec::from(point); @@ -345,21 +345,35 @@ pub struct GrandProductCircuit<'a> { impl<'a> GrandProductCircuit<'a> { pub fn new(row_images: &'a [F128], bittable: &BitTable) -> GrandProductCircuit<'a> { - // Columns occupy the low index bits, so each product tree reduces one - // column. Transposed so each row can be read sequentially. + assert_eq!( + row_images.len(), + bittable.shape().rows(), + "one image per row of the table" + ); + // Leaf `(b << s) | c` is `row_images[b]` where `bit(c, b)` is set and + // one elsewhere: columns occupy the low index bits, so each product + // tree reduces one column, folding the row bits from the top. + // + // The bottom three layers are looked up in 16-entry tables, one per + // four rows, so the top two row bits are split off and moved beside + // the columns. After the transpose, row group `b' < rows/4` is one + // contiguous run, and in it the bits of rows `b'`, `b' + rows/4`, + // `b' + 2 rows/4` and `b' + 3 rows/4` of one column sit next to each + // other as a nibble: the four leaves under node `(b' << s) | c` of + // layer 2, columns in ascending order. let matrix = bittable.as_matrix(); - // c b = c q b'. Since we combine within subtrees and want to start from d-3 that is 4 bits - let new_dim2 = matrix.log_dim2() - 2; // should this be -2 or -4 if we want the lookup to work? - let matrix = matrix.reshape(new_dim2).transpose(); + let log_dim2 = matrix.log_dim2() - 2; + let matrix = matrix.reshape(log_dim2).transpose(); Self { - matrix: matrix, + matrix, images: row_images, } } - // Returns the final evaluation and the witnesses of the intermediate layers - // TODO: replace with leaf lookups and add multithreading + /// Returns the output layer, one value per group, and the witnesses of + /// the layers below it. + // TODO: add multithreading #[tracing::instrument(name = "Evaluate grand-product circuit", level = "debug", skip_all)] pub fn batched_eval(self, groups: usize) -> (Vec, LayerWitnesses) { // Lower @@ -385,14 +399,13 @@ impl<'a> GrandProductCircuit<'a> { lut } - let n = self.images.len(); - // deal with the case that it isn't divisible by 4. What does that mean? Can we just put an assert. A certain invariant is passed down - let quart = n / 4; + // `new` checked there is one image per row, so four per row group. + let quart = self.matrix.dim1(); // Each LUT takes one image from each quarter of the rows. let (a, rest) = self.images.split_at(quart); let (b, rest) = rest.split_at(quart); let (c, d) = rest.split_at(quart); - let lower_lut = a + let lower_lut: Vec<[Field; 16]> = a .iter() .zip(b) .zip(c) @@ -400,54 +413,127 @@ impl<'a> GrandProductCircuit<'a> { .map(|(((&a, &b), &c), &d)| create_lut([a, b, c, d])) .collect(); + // Layers 0 to 2 are all in the tables, so the first layer built in + // full is layer 3. Its node `k` multiplies layer-2 nodes `k` and + // `k + n/2`, each a whole nibble looked up, so the matrix and the + // tables split in the middle into the left and the right child. + let (left, right) = self.matrix.split_at(self.matrix.len() / 2); + let (left_luts, right_luts) = lower_lut.split_at(quart / 2); + let layer3: Vec = lut_nodes(left, left_luts) + .zip(lut_nodes(right, right_luts)) + .map(|((l, l_bs), (r, r_bs))| l[l_bs] * r[r_bs]) + .collect(); + assert!( + groups <= layer3.len(), + "{groups} groups is more than the {} nodes of layer 3", + layer3.len() + ); + // Upper - // +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 - // derive based on matrix length and subtract the last three layers - let mut witnesses = Vec::with_capacity((self.matrix.len() + 1).ilog2() as usize); - - // split the matrix in the middle for left and right and a table is shared among c. - // TODO This will start with a what was created from the last LUT layer - let mut prev_eval = StepIter::<4>::new(self.matrix.iter()).map(|bs| lower_lut[]); - - // Stop when there is one output per group - while prev_eval.len() > groups { - let mid = prev_eval.len() / 2; - let (l, r) = prev_eval.split_at(mid); - let eval: Vec = l.iter().zip(r).map(|(&a, &b)| a * b).collect(); - witnesses.push(prev_eval); - prev_eval = eval; - } + let (output, upper_nodes) = fold_layers(layer3, groups); ( - prev_eval, + output, LayerWitnesses { - lower_lut: lower_lut, - upper_nodes: witnesses, + matrix: self.matrix, + lower_lut, + upper_nodes, }, ) } } +/// Each layer-2 node of a run of row groups as the table of its row group and +/// the nibble of its four leaves, in node order: node `k = (b' << s) | c` +/// reads nibble `k`, and all `2^s` columns of row group `b'` share its table. +/// `nibbles` and `luts` cover the same row groups. +fn lut_nodes<'a>( + nibbles: &'a [u128], + luts: &'a [[Field; 16]], +) -> impl ExactSizeIterator { + let nibbles = StepIter::<4>::new(nibbles.iter()); + // Every table serves the same run of nibbles, one per column. + // TODO check assembly that this isn't reloaded constantly. Or maybe it is unavoidable but might still save a shift instruction. Also might generate too much indirection. + let log_columns = (nibbles.len() / luts.len()).ilog2(); + nibbles + .enumerate() + .map(move |(k, bs)| (&luts[k >> log_columns], bs as usize)) +} + +/// Multiplies `layer` pairwise, left half against right half, until one +/// value per group remains. Returns that output layer and every layer before +/// it, `layer` first. +fn fold_layers(mut layer: Vec, groups: usize) -> (Vec, Vec>) { + let rounds = (layer.len() / groups).checked_ilog2().unwrap_or(0); + let mut witnesses = Vec::with_capacity(rounds as usize); + + // Stop when there is one output per group + while layer.len() > groups { + let (l, r) = layer.split_at(layer.len() / 2); + let next = l.iter().zip(r).map(|(&a, &b)| a * b).collect(); + witnesses.push(layer); + layer = next; + } + + (layer, witnesses) +} + +/// The layers below the output, yielded from the top down. The bottom three +/// are kept as the tables they are looked up in, the rest in full. // TODO: This can be a single vec that is allocated at the start. pub struct LayerWitnesses { + /// The nibbles of the layer-2 nodes; see [`GrandProductCircuit::new`]. + matrix: BitMatrix<'static>, + /// One table per row group, shared by all of its columns. lower_lut: Vec<[Field; 16]>, + /// Layer 3 and up, layer 3 first. upper_nodes: Vec>, } -// Can I still use into iterator or should something else be needed? +impl LayerWitnesses { + /// The leaves and layers 1 and 2, looked up in full. A node is the + /// product of the nibble bits its subtree covers: leaf quarter `q` is bit + /// `q` alone, each half of layer 1 pairs the quarters two apart, and + /// layer 2 takes all four. + fn lower_layers(&self) -> [Vec; 3] { + let nodes = lut_nodes(&self.matrix, &self.lower_lut); + let n = nodes.len(); + let mut leaves = Field::zeroed_vec(4 * n); + let mut layer1 = Field::zeroed_vec(2 * n); + let mut layer2 = Field::zeroed_vec(n); + + for (k, (lut, bs)) in nodes.enumerate() { + for q in 0..4 { + leaves[q * n + k] = lut[bs & (1 << q)]; + } + layer1[k] = lut[bs & 0b0101]; + layer1[n + k] = lut[bs & 0b1010]; + layer2[k] = lut[bs]; + } + + [leaves, layer1, layer2] + } +} + +// The sumcheck still reads every layer in full, so for now the bottom three +// are rebuilt from the tables here; layer 3 and up are already in +// `upper_nodes`. impl IntoIterator for LayerWitnesses { type Item = Vec; type IntoIter = std::iter::Rev>>; fn into_iter(self) -> Self::IntoIter { - self.0.into_iter().rev() - // TODO: do we add the LUT handling here? Good for an intermediate step such that the sumcheck part keeps working. They will be a lookup from + let mut layers = Vec::from(self.lower_layers()); + layers.extend(self.upper_nodes); + layers.into_iter().rev() } } #[cfg(test)] mod tests { use super::*; + use common::{BitZParams, Shape}; + use field::gf128::smallest_generator; use proptest::prelude::*; fn field() -> impl Strategy { @@ -460,6 +546,16 @@ mod tests { xs.iter().fold(Field::ONE, |acc, &x| acc * x) } + /// `leaves` padded with ones to a power of two, so they have an MLE. The + /// sumcheck tests below take any leaves, not only those a bit table + /// gives, so they build the layers with `fold_layers` directly. + fn pad(mut leaves: Vec) -> Vec { + if !leaves.is_empty() { + leaves.resize(leaves.len().next_power_of_two(), Field::ONE); + } + leaves + } + /// Proves a grand-product circuit for `leaves` batched into `2^log_groups` /// groups, then checks the verifier accepts the proof against the /// prover's own claimed output layer. Returns that output layer so @@ -474,7 +570,7 @@ mod tests { } proptest! { - // batched_eval, unbatched (groups = 1): pairwise multiplication is + // fold_layers, unbatched (groups = 1): pairwise multiplication is // associative, so folding-by-multiply any witness layer (including // the leaves, padded with the multiplicative identity 1) must equal // folding the original input. Checks every intermediate layer, not @@ -483,18 +579,18 @@ mod tests { fn eval_preserves_product_across_layers(leaves in prop::collection::vec(field(), 0..12)) { let expected = product(&leaves); - let (top, mut witnesses) = GrandProductCircuit::from_leafs(leaves).batched_eval(1); + let (top, witnesses) = fold_layers(pad(leaves), 1); prop_assert_eq!(product(&top), expected); - while let Some(layer) = witnesses.pop() { + for layer in witnesses { prop_assert_eq!(product(&layer), expected); } } #[test] - fn batched_eval_top_layer_is_interleaved_group_products( + fn fold_layers_top_layer_is_interleaved_group_products( leaves in prop::collection::vec(field(), 1..32), groups_log2 in 0usize..4, ) { @@ -505,8 +601,7 @@ mod tests { let mut padded = leaves.clone(); padded.resize(padded_len, Field::ONE); - let circuit = GrandProductCircuit::from_leafs(leaves); - let (top, _witnesses) = circuit.batched_eval(groups); + let (top, _witnesses) = fold_layers(pad(leaves), groups); prop_assert_eq!(top.len(), groups); for (j, &value) in top.iter().enumerate() { @@ -566,8 +661,7 @@ mod tests { [Field::ZERO, Field::ONE], [Field::ONE, Field::ZERO], ] { - let circuit = GrandProductCircuit::from_leafs(leaves.clone()); - let (output, witnesses) = circuit.batched_eval(4); + let (output, witnesses) = fold_layers(leaves.clone(), 4); let claim = mle(output, &point); let instance = (leaves.clone(), point.to_vec()); let mut prover = transcript::build_prover("gkr-zero", &instance); @@ -575,7 +669,7 @@ mod tests { &mut prover, leaves.len().ilog2() as usize, &point, - witnesses, + witnesses.into_iter().rev(), ); assert_eq!(terminal.1, mle(leaves.clone(), &terminal.0)); let proof = prover.finish(); @@ -592,6 +686,52 @@ mod tests { } } + /// Small enough that `BitZParams` accepts every shape below. + const Q: u128 = (1 << 100) - 15; + + /// The bottom three layers looked up in the tables, and every layer + /// above, match the circuit built one leaf at a time from the same table. + #[test] + fn lut_layers_match_the_leaf_construction() { + // At t = 7 and 8 the transposed row groups are shorter than a `u128`. + for (log_rows, log_columns) in [(7, 15), (8, 14), (9, 13), (15, 7)] { + let shape = Shape::new(log_rows, log_columns).unwrap(); + let params = BitZParams::::new(shape, smallest_generator()).unwrap(); + let packed: Vec = (0..(1u128 << shape.log_bits()) / 128) + .map(|i| { + F128::from(i.wrapping_mul(0x2545_F491_4F6C_DD1D_9E37_79B9_7F4A_7C15) ^ 0xA5A5) + }) + .collect(); + let table = params.table(&packed).unwrap(); + let row_images: Vec = (0..shape.rows() as u128) + .map(|b| F128::from((b + 1) * 0x9E37_79B9_7F4A_7C15 + 7)) + .collect(); + + // Leaf `(b << s) | c`, as `GrandProductCircuit::new` describes it. + let mut leaves = vec![Field::ONE; 1 << shape.log_bits()]; + for (b, &image) in row_images.iter().enumerate() { + for c in 0..shape.columns() { + if table.bit(c, b) { + leaves[(b << log_columns) | c] = image; + } + } + } + let (expected_output, expected_layers) = fold_layers(leaves, shape.columns()); + + let (output, witnesses) = + GrandProductCircuit::new(&row_images, &table).batched_eval(shape.columns()); + // Witnesses come top-down; the reference is bottom-up. + let mut layers: Vec> = witnesses.into_iter().collect(); + layers.reverse(); + + assert!(output == expected_output, "output, t = {log_rows}"); + assert_eq!(layers.len(), expected_layers.len(), "t = {log_rows}"); + for (i, (layer, expected)) in layers.iter().zip(&expected_layers).enumerate() { + assert!(layer == expected, "layer {i}, t = {log_rows}"); + } + } + } + #[test] fn gpgkr_verify_rejects_every_truncated_prefix() { let leaves: Vec = (1u128..=8).map(Field::from).collect(); @@ -612,13 +752,13 @@ mod tests { } pub fn prove(input: Vec, log_groups: usize) -> (Vec, transcript::Proof) { - let circuit = GrandProductCircuit::from_leafs(input); + let leaves = pad(input); // Fake hashing - let n = circuit.leafs.len() as u128; + let n = leaves.len() as u128; - let log_bits = circuit.leafs.len().max(1).ilog2() as usize; + let log_bits = leaves.len().max(1).ilog2() as usize; let groups = 1usize << log_groups; - let (last_value, witnesses) = circuit.batched_eval(groups); + let (last_value, witnesses) = fold_layers(leaves, groups); // TODO instance is a bit loose and should be replaced by PCS let instance = (last_value.clone(), n); @@ -631,14 +771,14 @@ 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, log_bits, &point, witnesses); + gpgkr_prove(&mut prover, log_bits, &point, witnesses.into_iter().rev()); (last_value, prover.finish()) } pub fn verify(input: Vec, output: Vec, proof: transcript::Proof) -> bool { - let circuit = GrandProductCircuit::from_leafs(input); + let leaves = pad(input); // Fake hashing - let instance = (&output, circuit.leafs.len() as u128); + let instance = (&output, leaves.len() as u128); let mut verifier = transcript::build_verifier("gkr", &instance, &proof); let log_groups = output.len().max(1).ilog2(); @@ -647,12 +787,12 @@ mod tests { .collect(); let claim = mle(output, &point); let log_groups = point.len() as u32; - let log_leafs = circuit.leafs.len().max(1).ilog2(); + let log_leafs = leaves.len().max(1).ilog2(); let rounds = log_leafs.saturating_sub(log_groups); match gpgkr_verify(&mut verifier, claim, &point, rounds) { Some((point, claim)) => { - let leaf_check = mle(circuit.leafs, &point); + let leaf_check = mle(leaves, &point); let ok = leaf_check == claim; diff --git a/crates/prover/src/reduce.rs b/crates/prover/src/reduce.rs index 17e03b7..c6c69ad 100644 --- a/crates/prover/src/reduce.rs +++ b/crates/prover/src/reduce.rs @@ -14,7 +14,7 @@ use transcript::ProverState; #[inline(never)] #[tracing::instrument(name = "Build grand-product circuit", level = "debug", skip_all)] -fn init_circuit(table: &BitTable, fold: &Fold) -> GrandProductCircuit { +fn init_circuit<'a>(table: &BitTable, fold: &'a Fold) -> GrandProductCircuit<'a> { GrandProductCircuit::new(&fold.row_images, table) } From af95cf5aa2e83dea0da5ac196332f1e8c8d1868c Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Fri, 9 Oct 2026 13:09:45 +0800 Subject: [PATCH 04/13] common: read StepIter's words as u64, as a u128 made the GKR's layer 3 ~25% slower --- crates/common/src/matrix.rs | 27 +++++++++++++++++---------- 1 file changed, 17 insertions(+), 10 deletions(-) diff --git a/crates/common/src/matrix.rs b/crates/common/src/matrix.rs index 4f428a9..57fb68d 100644 --- a/crates/common/src/matrix.rs +++ b/crates/common/src/matrix.rs @@ -239,7 +239,7 @@ fn transpose_reference(tt: &[Word]) -> Vec { for s in &mut segments { let Some(b) = s.next() else { break 'outer }; // LSB first, like `StepIter`. - word |= b << offset; + word |= Word::from(b) << offset; offset += 1; if offset == BITS { out.push(word); @@ -270,23 +270,30 @@ pub type BitsIter<'a> = StepIter<'a, 1>; /// Iterator over packed bits `S` at a time, least significant first — see /// [`BitMatrix::bits`]. +/// +/// Reads the words as `u64`, since no step needs more. A `u128` word was +/// kept on the stack between steps and made the GKR's layer 3 ~25% slower. pub struct StepIter<'a, const S: usize> { - elements: std::slice::Iter<'a, Word>, + elements: std::slice::Iter<'a, u64>, /// Bits not yet emitted; the next one to emit is the LSB. - word: Word, + word: u64, remaining: u8, } impl<'a, const S: usize> StepIter<'a, S> { const CHECK_SIZE: () = assert!( - S != 0 && S <= Word::BITS as usize && (Word::BITS as usize).is_multiple_of(S), - "step size must be nonzero, at most Word::BITS, and divide Word::BITS evenly" + S != 0 && S <= u64::BITS as usize && (u64::BITS as usize).is_multiple_of(S), + "step size must be nonzero, at most 64, and divide 64 evenly" ); - const MASK: Word = (1 as Word).unbounded_shl(S as u32).wrapping_sub(1); + const MASK: u64 = 1u64.unbounded_shl(S as u32).wrapping_sub(1); pub fn new(iter: Iter<'a, Word>) -> Self { let () = Self::CHECK_SIZE; + // On a little-endian target, bit `i` of the words is bit `i` of the + // `u64`s too. + const { assert!(cfg!(target_endian = "little")) }; + let elements: &[u64] = bytemuck::cast_slice(iter.as_slice()); Self { - elements: iter, + elements: elements.iter(), word: 0, remaining: 0, } @@ -294,13 +301,13 @@ impl<'a, const S: usize> StepIter<'a, S> { } impl Iterator for StepIter<'_, S> { - type Item = Word; + type Item = u64; fn next(&mut self) -> Option { if self.remaining == 0 { let element = self.elements.next()?; self.word = *element; - self.remaining = Word::BITS as u8; + self.remaining = u64::BITS as u8; } let bitset = self.word & Self::MASK; self.word = self.word.unbounded_shr(S as u32); @@ -316,7 +323,7 @@ impl Iterator for StepIter<'_, S> { impl ExactSizeIterator for StepIter<'_, S> { fn len(&self) -> usize { - (self.remaining as usize + self.elements.len() * BITS) / S + (self.remaining as usize + self.elements.len() * u64::BITS as usize) / S } } From a10bfca8071f75e3648dafd92ff4b0d8bbfc3888 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Fri, 9 Oct 2026 14:08:13 +0800 Subject: [PATCH 05/13] gkr: build layer 3 a word of nibbles at a time per pair of tables, drop batched_eval's groups argument --- Cargo.lock | 1 + crates/gkr/Cargo.toml | 1 + crates/gkr/src/lib.rs | 56 +++++++++++++++++++++++++---------- crates/prover/src/reduce.rs | 2 +- crates/verifier/src/reduce.rs | 2 +- 5 files changed, 44 insertions(+), 18 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 601e410..0f412ba 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -757,6 +757,7 @@ dependencies = [ name = "gkr" version = "0.1.0" dependencies = [ + "bytemuck", "common", "field", "num-traits", diff --git a/crates/gkr/Cargo.toml b/crates/gkr/Cargo.toml index 26913ec..7943048 100644 --- a/crates/gkr/Cargo.toml +++ b/crates/gkr/Cargo.toml @@ -6,6 +6,7 @@ rust-version.workspace = true license.workspace = true [dependencies] +bytemuck = { workspace = true } field = { workspace= true } transcript = { workspace = true} tracing = { workspace = true } diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs index 65add32..e2ef71d 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs @@ -371,11 +371,11 @@ impl<'a> GrandProductCircuit<'a> { } } - /// Returns the output layer, one value per group, and the witnesses of - /// the layers below it. + /// Returns the output layer, the product of each column of the table, and + /// the witnesses of the layers below it. // TODO: 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) -> (Vec, LayerWitnesses) { // Lower /// `lut[s]` is the product of the inputs whose bit is set in `s`, `a` /// on bit 0 through `d` on bit 3 (`0b0000 -> 1`, `0b0001 -> a`, @@ -405,6 +405,7 @@ impl<'a> GrandProductCircuit<'a> { let (a, rest) = self.images.split_at(quart); let (b, rest) = rest.split_at(quart); let (c, d) = rest.split_at(quart); + // LUT for the bottom three layers (0,1,2) let lower_lut: Vec<[Field; 16]> = a .iter() .zip(b) @@ -417,20 +418,44 @@ impl<'a> GrandProductCircuit<'a> { // full is layer 3. Its node `k` multiplies layer-2 nodes `k` and // `k + n/2`, each a whole nibble looked up, so the matrix and the // tables split in the middle into the left and the right child. - let (left, right) = self.matrix.split_at(self.matrix.len() / 2); + // Read as `u64`s: on a little-endian target, bit `i` of the words is + // bit `i` of the `u64`s too. + const { assert!(cfg!(target_endian = "little")) }; + let bitstring: &[u64] = bytemuck::cast_slice(&self.matrix[..]); + let (left, right) = bitstring.split_at(bitstring.len() / 2); let (left_luts, right_luts) = lower_lut.split_at(quart / 2); - let layer3: Vec = lut_nodes(left, left_luts) - .zip(lut_nodes(right, right_luts)) - .map(|((l, l_bs), (r, r_bs))| l[l_bs] * r[r_bs]) - .collect(); - assert!( - groups <= layer3.len(), - "{groups} groups is more than the {} nodes of layer 3", - layer3.len() - ); + // Each table serves one run of whole words, so a pair of tables is + // picked once and its nibbles are read a word at a time. + // A table's run is one `xs[i1]` of the matrix, a nibble per node. + const NIBBLES_IN_WORD: usize = u64::BITS as usize / 4; + let lut_words = left.len() / left_luts.len(); // This introduces a bound. StepIterator can prevent that but is more expensive than doing it this way + + assert!(lut_words > 0, "each table needs at least a word of nibbles"); + let runs = left + .chunks_exact(lut_words) + .zip(right.chunks_exact(lut_words)) + .zip(left_luts.iter().zip(right_luts)); + + // Access nible + let nibble = |word: u64, i: usize| (word >> (4 * i)) as usize & 15; + + // One node per nibble of each left table. + let mut layer3 = Vec::with_capacity(left_luts.len() * lut_words * NIBBLES_IN_WORD); + for ((l_words, r_words), (l_lut, r_lut)) in runs { + for (&l, &r) in l_words.iter().zip(r_words) { + // `extend` from a `Range` knows its length, so there is no + // capacity check or length update per node. + layer3.extend( + (0..NIBBLES_IN_WORD).map(|i| l_lut[nibble(l, i)] * r_lut[nibble(r, i)]), + ); + } + } // Upper - let (output, upper_nodes) = fold_layers(layer3, groups); + // A table serves one node of every product tree, so there is a tree + // per nibble of a table. + let trees = lut_words * NIBBLES_IN_WORD; + let (output, upper_nodes) = fold_layers(layer3, trees); ( output, @@ -718,8 +743,7 @@ mod tests { } let (expected_output, expected_layers) = fold_layers(leaves, shape.columns()); - let (output, witnesses) = - GrandProductCircuit::new(&row_images, &table).batched_eval(shape.columns()); + let (output, witnesses) = GrandProductCircuit::new(&row_images, &table).batched_eval(); // Witnesses come top-down; the reference is bottom-up. let mut layers: Vec> = witnesses.into_iter().collect(); layers.reverse(); diff --git a/crates/prover/src/reduce.rs b/crates/prover/src/reduce.rs index c6c69ad..f79a134 100644 --- a/crates/prover/src/reduce.rs +++ b/crates/prover/src/reduce.rs @@ -26,7 +26,7 @@ pub fn gkr_reduce( table: &BitTable, ) -> Result { let circuit = init_circuit(table, fold); - let (_last_value, witnesses) = circuit.batched_eval(table.shape().columns()); + let (_last_value, witnesses) = circuit.batched_eval(); let (mut point, claim) = gpgkr_prove(transcript, table.shape().log_bits(), &fold.zeta, witnesses); diff --git a/crates/verifier/src/reduce.rs b/crates/verifier/src/reduce.rs index f684d5b..032191e 100644 --- a/crates/verifier/src/reduce.rs +++ b/crates/verifier/src/reduce.rs @@ -101,7 +101,7 @@ mod round_trip_ai_test { .collect(); let circuit = GrandProductCircuit::new(&row_images, &table); - let (top_layer, witnesses) = circuit.batched_eval(shape.columns()); + let (top_layer, witnesses) = circuit.batched_eval(); // Fold::new evaluates top_layer at zeta to obtain GKR's initial claim. let folds = vec![0u128; shape.columns()]; From 387c9e78a0b2e893240b4e7435ae067446421257 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Mon, 19 Oct 2026 14:18:41 +0800 Subject: [PATCH 06/13] gkr: make points little-endian, rename SuffixTable to EqTables, split prove_layer and batched_eval into helpers --- crates/gkr/src/lib.rs | 581 +++++++++++++++++++++--------------------- 1 file changed, 293 insertions(+), 288 deletions(-) diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs index e2ef71d..2449159 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs @@ -1,5 +1,3 @@ -use std::collections::VecDeque; - use common::{BitMatrix, BitTable, matrix::StepIter}; use field::{F128, Wide256}; use num_traits::{ConstOne, ConstZero}; @@ -8,10 +6,11 @@ use transcript::{ProverState, VerifierState}; pub type Field = F128; -type Point = VecDeque; - /// Proves the layer-by-layer sumcheck reduction from a claim at `point` /// (an evaluation point on the output layer) down to a claim on the leaves. +/// +/// Points are little-endian, as in `poly`'s tables: `point[i]` is the +/// coordinate of index bit `i`. Every layer adds one coordinate, at the top. // TODO #[must_use], requires changing the test suite #[tracing::instrument(name = "Prove GKR", skip_all)] pub fn gpgkr_prove( @@ -24,170 +23,119 @@ pub fn gpgkr_prove( // Edge cases // - empty witnesses -> single constant circuit -> one verifier message that permutes the proof state, but a single constant can't have an MLE let mut point = point.to_owned(); - point.reverse(); - - let mut point = VecDeque::from(point); - let mut claim = Field::ZERO; - let mut suffix_storage = SuffixTable::alloc_storage(log_bits); - for wnext in witnesses.into_iter() { - (point, claim) = prove_layer(ps, &mut suffix_storage, point, wnext); + let mut eq_storage = EqTables::alloc_storage(log_bits); + for children in witnesses { + (point, claim) = prove_layer(ps, &mut eq_storage, &point, children); } - - let mut point = Vec::from(point); - point.reverse(); (point, claim) } +/// Proves `c = Σ_x eq(z, x) · L(x) · R(x)` for the halves `L`, `R` of +/// `children`, one round per coordinate of `point`, top first. Round `i` +/// sends the sum `sum_i` at one endpoint and its leading coefficient, where +/// `s_i(X) = eq(z_i, X) · sum_i(X)` and `sum_i` carries +/// `eq_bound = Π_{j > i} eq(z_j, r_j)` from the rounds already bound. +/// Returns the point `(r, t)` on `children` and the claim +/// `W(r, t) = L(r) + t · (R(r) - L(r))`. fn prove_layer( ps: &mut ProverState, storage: &mut [Field], - mut point: Point, - mut wnext: Vec, -) -> (Point, Field) { - let suffix_table = SuffixTable::new(storage, &point); - let mut factor = Field::ONE; - - let mid = wnext.len() / 2; - // Tree is encoded in LSB order - let (mut mle_l, mut mle_r) = wnext.split_at_mut(mid); - - let rounds = point.len(); - let mut next_point = VecDeque::with_capacity(rounds + 1); - - // `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) { - 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); - - // `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) = if send_one { - reduce_sumcheck_round::(lo_l, lo_r, hi_l, hi_r, eq) - } else { - reduce_sumcheck_round::(lo_l, lo_r, hi_l, hi_r, eq) - }; - - ps.prover_message(&[factor * sum_endpoint.reduce(), factor * sum_inf.reduce()]); - - let r = ps.verifier_message(); - next_point.push_back(r); - - lo_l.par_iter_mut() - .zip(lo_r.par_iter_mut()) - .zip(hi_l.par_iter()) - .zip(hi_r.par_iter()) - .with_min_len(PARALLEL_MIN_LANES) - .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. - let (d_l, d_r) = (l_hi - *l_lo, r_hi - *r_lo); - - // narrow multiply as the next round will multiply with this result - *l_lo += r * d_l; - *r_lo += r * d_r; - }); - mle_l = &mut mle_l[..h]; - mle_r = &mut mle_r[..h]; - - 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) + point: &[Field], + mut children: Vec, +) -> (Vec, Field) { + let n = point.len(); + let eq_tables = EqTables::new(storage, point); + // The top index bit of `children` picks the half: `L`, then `R`. + let mid = children.len() / 2; + let (mut left, mut right) = children.split_at_mut(mid); + let mut next_point = vec![Field::ZERO; n + 1]; + let mut eq_bound = Field::ONE; + + for (i, &z) in point.iter().enumerate().rev() { + // `x_i = 0` in the low half of each, `x_i = 1` in the high half. + let half = 1 << i; + let (left0, left1) = left.split_at_mut(half); + let (right0, right1) = right.split_at_mut(half); + + // Branch once per round: `round_sums` takes the endpoint as a const + // generic, so the hot fold has no per-element check. + let eq = eq_tables.below(i); + let (sum_end, sum_inf) = if sends_one(z) { + round_sums::(left0, left1, right0, right1, eq) } else { - (l_lo, r_lo) + round_sums::(left0, left1, right0, right1, eq) }; - 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()]); + ps.prover_message(&[eq_bound * sum_end.reduce(), eq_bound * 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]; + bind(r, left0, left1, right0, right1); + left = &mut left[..half]; + right = &mut right[..half]; - factor *= eq_factor(r, z); + next_point[i] = r; + eq_bound *= eq_factor(r, z); } - ps.prover_message(&[mle_l[0], mle_r[0]]); - let r = ps.verifier_message(); - next_point.push_front(r); - let claim = mle_l[0] + r * (mle_r[0] - mle_l[0]); - (next_point, claim) + let (left_r, right_r) = (left[0], right[0]); + ps.prover_message(&[left_r, right_r]); + let t = ps.verifier_message(); + next_point[n] = t; + (next_point, interpolate(left_r, right_r, t)) } -fn reduce_sumcheck_round( - lo_l: &[Field], - lo_r: &[Field], - hi_l: &[Field], - hi_r: &[Field], +/// `Σ eq · L(e) · R(e)` and `Σ eq · ΔL · ΔR` over the free variables: the +/// round's sum at the endpoint `e` (`1` if `AT_ONE`) and its leading +/// coefficient, without `eq_bound` and left unreduced. +fn round_sums( + left0: &[Field], + left1: &[Field], + right0: &[Field], + right1: &[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)) - .zip(hi_r.par_chunks_exact(2)) + let at = |v0: Field, v1: Field| if AT_ONE { v1 } else { v0 }; + + // The last round has a single summand, and `eq` there is the empty + // product `1`: no multiplication by it, and too few for two lanes. + if let [_] = eq { + let (l0, l1, r0, r1) = (left0[0], left1[0], right0[0], right1[0]); + return ( + Wide256::mul(at(l0, l1), at(r0, r1)), + Wide256::mul(l1 - l0, r1 - r0), + ); + } + + // Two lanes per step: independent multiplication chains. Chunking also + // reduces variance in benchmarking. + left0 + .par_chunks_exact(2) + .zip(left1.par_chunks_exact(2)) + .zip(right0.par_chunks_exact(2)) + .zip(right1.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) + |(mut end, mut inf), ((((l0, l1), r0), r1), e)| { + // `e · L(e) · R(e)` into `end` and `e · ΔL · ΔR` into `inf`, + // with `e` joining the left factor through a reduced multiply + // so that the right one can stay wide. + // + // Steps, products, then sums, each for both lanes: this order + // puts every load of the step ahead of the first multiply. + // Lane by lane is 1.1-1.4% slower single-threaded. + let (d_l0, d_r0) = (l1[0] - l0[0], r1[0] - r0[0]); + let (d_l1, d_r1) = (l1[1] - l0[1], r1[1] - r0[1]); + let (a0, b0) = (e[0] * at(l0[0], l1[0]), e[0] * d_l0); + let (a1, b1) = (e[1] * at(l0[1], l1[1]), e[1] * d_l1); + end += Wide256::mul(a0, at(r0[0], r1[0])); + inf += Wide256::mul(b0, d_r0); + end += Wide256::mul(a1, at(r0[1], r1[1])); + inf += Wide256::mul(b1, d_r1); + (end, inf) }, ) .reduce( @@ -196,6 +144,52 @@ fn reduce_sumcheck_round( ) } +/// Binds the variable that splits each pair of halves to `r`, writing the +/// result over the low halves. +fn bind(r: Field, left0: &mut [Field], left1: &[Field], right0: &mut [Field], right1: &[Field]) { + left0 + .par_iter_mut() + .zip(right0.par_iter_mut()) + .zip(left1.par_iter()) + .zip(right1.par_iter()) + .with_min_len(PARALLEL_MIN_LANES) + .for_each(|(((l0, r0), &l1), &r1)| { + // Recalculating the step rather than writing it in the high half + // in `round_sums` to save on writes. No benchmarking has been done + // to check the difference. + let (d_l, d_r) = (l1 - *l0, r1 - *r0); + + // `interpolate`, spelled out: as `*r0 = interpolate(*r0, r1, r)` + // the compiler loads `*r0` as a vector and moves it back to + // general registers for the subtraction, one extra instruction + // per pair on the way into the multiply (~3.7% single-threaded). + *l0 += r * d_l; + *r0 += r * d_r; + }); +} + +/// The line through `at0` (at `0`) and `at1` (at `1`), at `r`: a multilinear +/// function of one variable. +fn interpolate(at0: Field, at1: Field, r: Field) -> Field { + // Narrow multiply, as the next round will multiply with this result. + at0 + r * (at1 - at0) +} + +/// The quadratic through `at0` (at `0`) and `at1` (at `1`) with leading +/// coefficient `lead`, at `r`: `at0 + r·(at1 - at0) + r·(r - 1)·lead`, with +/// `r` factored out. +fn quadratic_at(at0: Field, at1: Field, lead: Field, r: Field) -> Field { + at0 + r * ((at1 - at0) + (r - Field::ONE) * lead) +} + +/// Whether round `i` sends `sum_i(1)` rather than `sum_i(0)`. The verifier +/// recovers the other one from `c = (1 - z_i) · sum_i(0) + z_i · sum_i(1)`, +/// which divides by `z_i` to get `sum_i(1)`; at `z_i = 0` it gives +/// `sum_i(0) = c` directly instead. +fn sends_one(z: Field) -> bool { + z == Field::ZERO +} + /// `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 @@ -204,23 +198,24 @@ fn eq_factor(r: Field, z: Field) -> Field { Field::ONE + r + z } -// TODO: Split suffix table -struct SuffixTable<'a> { +/// `eq(z_{ { storage: &'a [Field], } -impl<'a> SuffixTable<'a> { - fn layer(&self, i: usize) -> &[Field] { +impl<'a> EqTables<'a> { + /// `eq(z_{ &[Field] { let start = (1 << i) - 1; let end = (1 << (i + 1)) - 1; &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> { + fn new(storage: &'a mut [Field], point: &[Field]) -> EqTables<'a> { // 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); @@ -228,14 +223,12 @@ impl<'a> 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 - // range start in-bounds there instead of panicking). - let c = point.range(point.len().min(1)..); - - // Suffix table is in the reverse order of the point - for &z in c.rev() { + // Round `i` reads only the coordinates below `i`, so the top one + // never enters a table. + let below_top = &point[..point.len().saturating_sub(1)]; + + // Each coordinate doubles the table, as its new top index bit. + for &z in below_top { let size = prev.len() << 1; let (entry, next) = remaining.split_at_mut(size); let (low, hi) = entry.split_at_mut(entry.len() >> 1); @@ -249,7 +242,7 @@ impl<'a> SuffixTable<'a> { prev = entry; remaining = next; } - SuffixTable { storage } + EqTables { storage } } fn alloc_storage(log_bits: usize) -> Vec { @@ -268,7 +261,7 @@ pub fn gpgkr_verify( vs: &mut VerifierState, mut claim: Field, point: &[F128], - rounds: u32, + layers: u32, ) -> Option<(Vec, Field)> { // Edge cases around input lenghts, 0 meaning empty // | circuit | last value | @@ -281,60 +274,50 @@ pub fn gpgkr_verify( // -> false // let mut point = point.to_owned(); - point.reverse(); - let mut point = VecDeque::from(point); - - for _i in 0..rounds { - (point, claim) = verify_layer(vs, claim, point)? + for _ in 0..layers { + (point, claim) = verify_layer(vs, claim, &point)? } - - let mut point = Vec::from(point); - point.reverse(); - Some((point, claim)) } -/// Point's orientation is the reverse of gpgkr_verify -fn verify_layer(vs: &mut VerifierState, mut claim: Field, point: Point) -> Option<(Point, Field)> { - let mut prefix = Field::ONE; - - let mut next_point: Point = VecDeque::new(); - - for z in point { - let [sum_endpoint, suminf]: [Field; 2] = vs.prover_message().ok()?; - // claim = (1 - z) * sum0 + z * sum1. At z = 0, sum0 is known - // and the prover supplies sum1; otherwise recover sum1 by division. - let (sum0, sum1) = if z == Field::ZERO { - (claim, sum_endpoint) +/// Per round, recovers the other endpoint from +/// `c = (1 - z_i) · sum_i(0) + z_i · sum_i(1)` and moves the claim to +/// `s_i(r_i)`. +/// Then checks `c = eq(z, r) · L(r) · R(r)` and continues at +/// `W(r, t) = L(r) + t · (R(r) - L(r))`. +fn verify_layer( + vs: &mut VerifierState, + mut claim: Field, + point: &[Field], +) -> Option<(Vec, Field)> { + let n = point.len(); + let mut next_point = vec![Field::ZERO; n + 1]; + let mut eq_bound = Field::ONE; + + for (i, &z) in point.iter().enumerate().rev() { + let [sum_end, sum_inf]: [Field; 2] = vs.prover_message().ok()?; + // `claim = (1 - z) * sum0 + z * sum1`. + let (sum0, sum1) = if sends_one(z) { + (claim, sum_end) } else { - let eqjsum0 = (Field::ONE - z) * sum_endpoint; - (sum_endpoint, (claim - eqjsum0) / z) + (sum_end, (claim - (Field::ONE - z) * sum_end) / z) }; let r = vs.verifier_message(); - next_point.push_back(r); - let factor = eq_factor(r, z); - - // `claim = factor * (sum0 + r*(sum1 - sum0) + r*(r - 1)*suminf)`. - let bracket = (sum1 - sum0) + (r - Field::ONE) * suminf; - claim = factor * (sum0 + r * bracket); - - prefix *= factor; + next_point[i] = r; + let eq_r = eq_factor(r, z); + claim = eq_r * quadratic_at(sum0, sum1, sum_inf, r); + eq_bound *= eq_r; } - let elem_lr: [Field; 2] = vs.prover_message().ok()?; - // Check if line polynomial hits same spot as sumcheck check - if (prefix * elem_lr[0] * elem_lr[1]) != claim { - None - } else { - let r = vs.verifier_message(); - next_point.push_front(r); - - // Reduce both claims to a single claim - claim = elem_lr[0] + r * (elem_lr[1] - elem_lr[0]); - - Some((next_point, claim)) + // Every coordinate is bound now: `eq_bound = eq(z, r)`. + let [left_r, right_r]: [Field; 2] = vs.prover_message().ok()?; + if eq_bound * left_r * right_r != claim { + return None; } + let t = vs.verifier_message(); + next_point[n] = t; + Some((next_point, interpolate(left_r, right_r, t))) } // TODO: Optimise for circuits that are padded. @@ -373,88 +356,15 @@ impl<'a> GrandProductCircuit<'a> { /// Returns the output layer, the product of each column of the table, and /// the witnesses of the layers below it. - // TODO: add multithreading + /// + /// Layers 0 to 2 stay as lookup tables, layer 3 is the first one built in + /// full, and the layers above it fold as usual. #[tracing::instrument(name = "Evaluate grand-product circuit", level = "debug", skip_all)] pub fn batched_eval(self) -> (Vec, LayerWitnesses) { - // Lower - /// `lut[s]` is the product of the inputs whose bit is set in `s`, `a` - /// on bit 0 through `d` on bit 3 (`0b0000 -> 1`, `0b0001 -> a`, - /// `0b0011 -> a*b`, ...), so a 4-bit lookup string indexes its - /// multilinear value directly. - /// - /// Built by doubling: entries `2^i..2^(i+1)` are entries `0..2^i` - /// times the `i`-th input. The first entry of each new half is the - /// input itself, so it is copied rather than multiplied by one, - /// leaving 11 multiplications -- one per entry with two or more - /// factors, which is the minimum. - fn create_lut(c: [Field; 4]) -> [Field; 16] { - let mut lut = [Field::ONE; 16]; - for (i, x) in c.into_iter().enumerate() { - let half = 1 << i; - lut[half] = x; - for k in 1..half { - lut[half + k] = lut[k] * x; - } - } - lut - } - - // `new` checked there is one image per row, so four per row group. - let quart = self.matrix.dim1(); - // Each LUT takes one image from each quarter of the rows. - let (a, rest) = self.images.split_at(quart); - let (b, rest) = rest.split_at(quart); - let (c, d) = rest.split_at(quart); - // LUT for the bottom three layers (0,1,2) - let lower_lut: Vec<[Field; 16]> = a - .iter() - .zip(b) - .zip(c) - .zip(d) - .map(|(((&a, &b), &c), &d)| create_lut([a, b, c, d])) - .collect(); - - // Layers 0 to 2 are all in the tables, so the first layer built in - // full is layer 3. Its node `k` multiplies layer-2 nodes `k` and - // `k + n/2`, each a whole nibble looked up, so the matrix and the - // tables split in the middle into the left and the right child. - // Read as `u64`s: on a little-endian target, bit `i` of the words is - // bit `i` of the `u64`s too. - const { assert!(cfg!(target_endian = "little")) }; - let bitstring: &[u64] = bytemuck::cast_slice(&self.matrix[..]); - let (left, right) = bitstring.split_at(bitstring.len() / 2); - let (left_luts, right_luts) = lower_lut.split_at(quart / 2); - - // Each table serves one run of whole words, so a pair of tables is - // picked once and its nibbles are read a word at a time. - // A table's run is one `xs[i1]` of the matrix, a nibble per node. - const NIBBLES_IN_WORD: usize = u64::BITS as usize / 4; - let lut_words = left.len() / left_luts.len(); // This introduces a bound. StepIterator can prevent that but is more expensive than doing it this way - - assert!(lut_words > 0, "each table needs at least a word of nibbles"); - let runs = left - .chunks_exact(lut_words) - .zip(right.chunks_exact(lut_words)) - .zip(left_luts.iter().zip(right_luts)); - - // Access nible - let nibble = |word: u64, i: usize| (word >> (4 * i)) as usize & 15; - - // One node per nibble of each left table. - let mut layer3 = Vec::with_capacity(left_luts.len() * lut_words * NIBBLES_IN_WORD); - for ((l_words, r_words), (l_lut, r_lut)) in runs { - for (&l, &r) in l_words.iter().zip(r_words) { - // `extend` from a `Range` knows its length, so there is no - // capacity check or length update per node. - layer3.extend( - (0..NIBBLES_IN_WORD).map(|i| l_lut[nibble(l, i)] * r_lut[nibble(r, i)]), - ); - } - } - // Upper - // A table serves one node of every product tree, so there is a tree - // per nibble of a table. - let trees = lut_words * NIBBLES_IN_WORD; + let lower_lut = lower_luts(self.images); + let layer3 = layer3(&self.matrix, &lower_lut); + // A table's run holds one nibble per product tree. + let trees = self.matrix.dim2() / NIBBLE_BITS; let (output, upper_nodes) = fold_layers(layer3, trees); ( @@ -468,6 +378,87 @@ impl<'a> GrandProductCircuit<'a> { } } +const NIBBLE_BITS: usize = 4; +const NIBBLES_IN_WORD: usize = u64::BITS as usize / NIBBLE_BITS; + +/// Layers 0 to 2 as tables, one per four images: `lut[s]` is the layer-2 node +/// whose four leaves are the bits of `s`. +/// +/// Each table takes one image from each quarter of `images`. +fn lower_luts(images: &[Field]) -> Vec<[Field; 16]> { + /// `lut[s]` is the product of the inputs whose bit is set in `s`, `a` + /// on bit 0 through `d` on bit 3 (`0b0000 -> 1`, `0b0001 -> a`, + /// `0b0011 -> a*b`, ...), so a 4-bit lookup string indexes its + /// multilinear value directly. + /// + /// Built by doubling: entries `2^i..2^(i+1)` are entries `0..2^i` + /// times the `i`-th input. The first entry of each new half is the + /// input itself, so it is copied rather than multiplied by one, + /// leaving 11 multiplications -- one per entry with two or more + /// factors, which is the minimum. + fn create_lut(c: [Field; 4]) -> [Field; 16] { + let mut lut = [Field::ONE; 16]; + for (i, x) in c.into_iter().enumerate() { + let half = 1 << i; + lut[half] = x; + for k in 1..half { + lut[half + k] = lut[k] * x; + } + } + lut + } + + let quart = images.len() / 4; + let (a, rest) = images.split_at(quart); + let (b, rest) = rest.split_at(quart); + let (c, d) = rest.split_at(quart); + a.iter() + .zip(b) + .zip(c) + .zip(d) + .map(|(((&a, &b), &c), &d)| create_lut([a, b, c, d])) + .collect() +} + +/// Layer 3, the first layer built in full: node `k` multiplies layer-2 nodes +/// `k` and `k + n/2`, each a table entry looked up by its nibble. +/// +/// The matrix and the tables split in the middle into the left and the right +/// child. +fn layer3(matrix: &BitMatrix, lower_lut: &[[Field; 16]]) -> Vec { + // Read as `u64`s: on a little-endian target, bit `i` of the words is + // bit `i` of the `u64`s too. + const { assert!(cfg!(target_endian = "little")) }; + let bitstring: &[u64] = bytemuck::cast_slice(&matrix[..]); + let (left, right) = bitstring.split_at(bitstring.len() / 2); + let (left_luts, right_luts) = lower_lut.split_at(lower_lut.len() / 2); + + // Each table serves one run of whole words, so a pair of tables is + // picked once and its nibbles are read a word at a time. + // A table's run is one `xs[i1]` of the matrix, a nibble per node. + let lut_words = left.len() / left_luts.len(); // This introduces a bound. StepIterator can prevent that but is more expensive than doing it this way + + assert!(lut_words > 0, "each table needs at least a word of nibbles"); + let runs = left + .chunks_exact(lut_words) + .zip(right.chunks_exact(lut_words)) + .zip(left_luts.iter().zip(right_luts)); + + // The layer-2 nodes whose four leaves are the nibbles of `word`. + fn layer2(lut: &[Field; 16], word: u64) -> impl Iterator + '_ { + (0..NIBBLES_IN_WORD).map(move |i| lut[(word >> (NIBBLE_BITS * i)) as usize & 15]) + } + + // One node per nibble of each left table. + let mut layer3 = Vec::with_capacity(left_luts.len() * lut_words * NIBBLES_IN_WORD); + for ((l_words, r_words), (l_lut, r_lut)) in runs { + for (&l, &r) in l_words.iter().zip(r_words) { + fold_into(&mut layer3, layer2(l_lut, l), layer2(r_lut, r)); + } + } + layer3 +} + /// Each layer-2 node of a run of row groups as the table of its row group and /// the nibble of its four leaves, in node order: node `k = (b' << s) | c` /// reads nibble `k`, and all `2^s` columns of row group `b'` share its table. @@ -495,7 +486,8 @@ fn fold_layers(mut layer: Vec, groups: usize) -> (Vec, Vec groups { let (l, r) = layer.split_at(layer.len() / 2); - let next = l.iter().zip(r).map(|(&a, &b)| a * b).collect(); + let mut next = Vec::with_capacity(l.len()); + fold_into(&mut next, l.iter().copied(), r.iter().copied()); witnesses.push(layer); layer = next; } @@ -503,6 +495,18 @@ fn fold_layers(mut layer: Vec, groups: usize) -> (Vec, Vec, + left: impl Iterator, + right: impl Iterator, +) { + // `extend` from zipped ranges or slices knows its length, so there is no + // capacity check or length update per node. + next.extend(left.zip(right).map(|(a, b)| a * b)); +} + /// The layers below the output, yielded from the top down. The bottom three /// are kept as the tables they are looked up in, the rest in full. // TODO: This can be a single vec that is allocated at the start. @@ -518,8 +522,8 @@ pub struct LayerWitnesses { impl LayerWitnesses { /// The leaves and layers 1 and 2, looked up in full. A node is the /// product of the nibble bits its subtree covers: leaf quarter `q` is bit - /// `q` alone, each half of layer 1 pairs the quarters two apart, and - /// layer 2 takes all four. + /// `q` alone, layer-1 half `h` pairs bits `h` and `h + 2`, and layer 2 + /// takes all four. fn lower_layers(&self) -> [Vec; 3] { let nodes = lut_nodes(&self.matrix, &self.lower_lut); let n = nodes.len(); @@ -531,8 +535,9 @@ impl LayerWitnesses { for q in 0..4 { leaves[q * n + k] = lut[bs & (1 << q)]; } - layer1[k] = lut[bs & 0b0101]; - layer1[n + k] = lut[bs & 0b1010]; + for h in 0..2 { + layer1[h * n + k] = lut[bs & (0b0101 << h)]; + } layer2[k] = lut[bs]; } From 37ce8fda8c1f0a0a3f101f4f236a3e0db29c9e69 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Mon, 19 Oct 2026 14:30:32 +0800 Subject: [PATCH 07/13] gkr: document why EqTables keeps every table rather than halving the largest --- crates/gkr/src/lib.rs | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs index 2449159..ac63a20 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs @@ -199,6 +199,15 @@ fn eq_factor(r: Field, z: Field) -> Field { } /// `eq(z_{ { storage: &'a [Field], } From b95bc5942228ea8da544f2a776d23577e2de3d99 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Tue, 20 Oct 2026 13:01:07 +0800 Subject: [PATCH 08/13] gkr: prove_rounds for the dense rounds, tests on bit tables gpgkr_prove takes the LayerWitnesses. The round loop moves out of prove_layer into prove_rounds, so a table-backed round 1 can hand its bound halves to the dense rounds through round_one: Option. The gkr tests build bit tables at (7, 15), the smallest shape the gates accept, and check the prover against the leaves built one bit at a time. --- crates/gkr/src/lib.rs | 335 +++++++++++++++++++++++++----------------- 1 file changed, 201 insertions(+), 134 deletions(-) diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs index ac63a20..bb314e3 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs @@ -18,14 +18,19 @@ pub fn gpgkr_prove( log_bits: usize, point: &[F128], // All the intermediate witnesses + the input layer, from the top down. Doesn't contain the output layer - witnesses: impl IntoIterator>, + witnesses: LayerWitnesses, ) -> (Vec, Field) { // Edge cases // - empty witnesses -> single constant circuit -> one verifier message that permutes the proof state, but a single constant can't have an MLE let mut point = point.to_owned(); let mut claim = Field::ZERO; let mut eq_storage = EqTables::alloc_storage(log_bits); - for children in witnesses { + let lower = witnesses.lower_layers(); + let upper = witnesses.upper_nodes; + for children in upper.into_iter().rev() { + (point, claim) = prove_layer(ps, &mut eq_storage, &point, children); + } + for children in lower.into_iter().rev() { (point, claim) = prove_layer(ps, &mut eq_storage, &point, children); } (point, claim) @@ -44,15 +49,39 @@ fn prove_layer( point: &[Field], mut children: Vec, ) -> (Vec, Field) { - let n = point.len(); let eq_tables = EqTables::new(storage, point); // The top index bit of `children` picks the half: `L`, then `R`. let mid = children.len() / 2; - let (mut left, mut right) = children.split_at_mut(mid); + let (left, right) = children.split_at_mut(mid); + prove_rounds(ps, &eq_tables, point, None, left, right) +} + +/// The dense rounds of a layer, then its last message. A table-backed layer +/// has already proved round 1, which binds the top coordinate of `point`, and +/// passes that round's challenge as `round_one`; its `left` and `right` are +/// then the bound halves. A dense layer passes `None` and proves every round +/// here. +fn prove_rounds( + ps: &mut ProverState, + eq_tables: &EqTables, + point: &[Field], + round_one: Option, + mut left: &mut [Field], + mut right: &mut [Field], +) -> (Vec, Field) { + let n = point.len(); let mut next_point = vec![Field::ZERO; n + 1]; - let mut eq_bound = Field::ONE; + // The coordinates whose rounds run here, and `eq` of those already bound. + let (remaining, mut eq_bound) = match round_one { + Some(r) => { + let top = n - 1; + next_point[top] = r; + (&point[..top], eq_factor(r, point[top])) + } + None => (point, Field::ONE), + }; - for (i, &z) in point.iter().enumerate().rev() { + for (i, &z) in remaining.iter().enumerate().rev() { // `x_i = 0` in the low half of each, `x_i = 1` in the high half. let half = 1 << i; let (left0, left1) = left.split_at_mut(half); @@ -586,8 +615,8 @@ mod tests { } /// `leaves` padded with ones to a power of two, so they have an MLE. The - /// sumcheck tests below take any leaves, not only those a bit table - /// gives, so they build the layers with `fold_layers` directly. + /// `fold_layers` tests below take any leaves, not only those a bit table + /// gives. fn pad(mut leaves: Vec) -> Vec { if !leaves.is_empty() { leaves.resize(leaves.len().next_power_of_two(), Field::ONE); @@ -595,19 +624,6 @@ mod tests { leaves } - /// Proves a grand-product circuit for `leaves` batched into `2^log_groups` - /// groups, then checks the verifier accepts the proof against the - /// prover's own claimed output layer. Returns that output layer so - /// callers can assert further properties of it. - fn prove_and_verify( - leaves: Vec, - log_groups: usize, - ) -> Result, TestCaseError> { - let (last_value, proof) = prove(leaves.clone(), log_groups); - prop_assert!(verify(leaves, last_value.clone(), proof)); - Ok(last_value) - } - proptest! { // fold_layers, unbatched (groups = 1): pairwise multiplication is // associative, so folding-by-multiply any witness layer (including @@ -652,110 +668,172 @@ mod tests { prop_assert_eq!(value, expected); } } + } - // Full protocol, unbatched: proves and verifies through `prove`/ - // `verify`, then checks the output against an independently computed - // product - // Starts at 1 since Dense MLE can't handle zero - #[test] - fn gpgkr_round_trip(leaves in prop::collection::vec(field(), 1..100)) { - let expected = product(&leaves); + /// Small enough that `BitZParams` accepts every shape below. + const Q: u128 = (1 << 100) - 15; + + /// The smallest shape the gates accept, `(t, s)` with `t + s = 22`. + const SMALLEST: (usize, usize) = (7, 15); + + /// How the bits of a test table are set. + #[derive(Clone, Copy, Debug)] + enum Bits { + /// Every leaf is one. + Clear, + /// Every leaf is its row's image. + Set, + /// Bits drawn from the seed. + Drawn, + } + + fn bits() -> impl Strategy { + prop_oneof![Just(Bits::Clear), Just(Bits::Set), Just(Bits::Drawn)] + } + + /// A splitmix64 stream: test data that is fixed by its seed. + fn stream(seed: u64) -> impl Iterator { + let mut x = seed; + std::iter::repeat_with(move || { + x = x.wrapping_add(0x9E37_79B9_7F4A_7C15); + let z = (x ^ (x >> 30)).wrapping_mul(0xBF58_476D_1CE4_E5B9); + let z = (z ^ (z >> 27)).wrapping_mul(0x94D0_49BB_1331_11EB); + z ^ (z >> 31) + }) + } + + fn element(words: &mut impl Iterator) -> F128 { + let mut word = || words.next().unwrap(); + F128::from([word(), word()]) + } + + /// The packed bits of a `(t, s)` table and one image per row, from `seed`. + fn instance((t, s): (usize, usize), bits: Bits, seed: u64) -> (Shape, Vec, Vec) { + let shape = Shape::new(t, s).unwrap(); + let mut words = stream(seed); + let packed = (0..(1usize << shape.log_bits()) / 128) + .map(|_| match bits { + Bits::Clear => F128::ZERO, + Bits::Set => F128::from(u128::MAX), + Bits::Drawn => element(&mut words), + }) + .collect(); + let row_images = (0..shape.rows()).map(|_| element(&mut words)).collect(); + (shape, packed, row_images) + } - let last_value = prove_and_verify(leaves, 0)?; - prop_assert_eq!(last_value, vec![expected]); + fn table(shape: Shape, packed: &[F128]) -> BitTable<'_> { + BitZParams::::new(shape, smallest_generator()) + .unwrap() + .table(packed) + .unwrap() + } + + /// Leaf `(b << s) | c` is `row_images[b]` where `bit(c, b)` is set and one + /// elsewhere, as `GrandProductCircuit::new` describes, built one leaf at a + /// time: the reference the circuit is checked against. + fn leaves(table: &BitTable, row_images: &[Field]) -> Vec { + let shape = table.shape(); + let mut leaves = vec![Field::ONE; 1 << shape.log_bits()]; + for (b, &image) in row_images.iter().enumerate() { + for c in 0..shape.columns() { + if table.bit(c, b) { + leaves[(b << shape.log_columns()) | c] = image; + } + } } + leaves + } + + proptest! { + // Every case proves a whole circuit at the smallest shape. + #![proptest_config(ProptestConfig::with_cases(4))] - // Full protocol, batched (`groups > 1`) + // Full protocol: the verifier accepts, and the output is the product + // of each column's leaves. #[test] - fn gpgkr_round_trip_batched( - leaves in prop::collection::vec(field(), 4..33), - groups_log2 in 1usize..4, - ) { - let padded_len = leaves.len().next_power_of_two(); - let groups_log2 = groups_log2.min(padded_len.ilog2() as usize - 1).max(1); + fn gpgkr_round_trip(bits in bits(), seed in any::()) { + let (shape, packed, row_images) = instance(SMALLEST, bits, seed); + let table = table(shape, &packed); + let leaves = leaves(&table, &row_images); - prove_and_verify(leaves, groups_log2)?; + let (output, proof) = prove(&table, &row_images); + prop_assert!(verify(&leaves, &output, &proof)); + + let columns = shape.columns(); + let expected: Vec = (0..columns) + .map(|c| leaves[c..].iter().step_by(columns).product()) + .collect(); + prop_assert_eq!(output, expected); } // negative/soundness test #[test] - fn gpgkr_verify_rejects_tampered_output(leaves in prop::collection::vec(field(), 2..33)) { - let (last_value, proof) = prove(leaves.clone(), 0); + fn gpgkr_verify_rejects_tampered_output(bits in bits(), seed in any::()) { + let (shape, packed, row_images) = instance(SMALLEST, bits, seed); + let table = table(shape, &packed); + let (output, proof) = prove(&table, &row_images); - let mut tampered = last_value; + let mut tampered = output; tampered[0] += Field::ONE; - prop_assert!(!verify(leaves, tampered, proof)); + prop_assert!(!verify(&leaves(&table, &row_images), &tampered, &proof)); } } + // Zero and one coordinates in the output point take the `sends_one` + // branch in the top layer's rounds. #[test] fn gpgkr_round_trip_with_zero_coordinates() { - let leaves: Vec = (1u128..=16).map(Field::from).collect(); + let (shape, packed, row_images) = instance(SMALLEST, Bits::Drawn, 1); + let table = table(shape, &packed); + let leaves = leaves(&table, &row_images); + let s = shape.log_columns(); + let layers = shape.log_rows() as u32; + + // `value` at coordinate `i`, and a nonzero value at every other. let nonzero = Field::from(5u128); + let with = |i: usize, value: Field| -> Vec { + (0..s) + .map(|j| if j == i { value } else { nonzero }) + .collect() + }; for point in [ - [Field::ZERO, Field::ZERO], - [Field::ZERO, nonzero], - [nonzero, Field::ZERO], - [Field::ZERO, Field::ONE], - [Field::ONE, Field::ZERO], + vec![Field::ZERO; s], + with(s - 1, Field::ZERO), + with(0, Field::ZERO), + with(s - 1, Field::ONE), + with(0, Field::ONE), ] { - let (output, witnesses) = fold_layers(leaves.clone(), 4); + let (output, witnesses) = GrandProductCircuit::new(&row_images, &table).batched_eval(); 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.into_iter().rev(), - ); + let mut prover = transcript::build_prover("gkr-zero", &point); + let terminal = gpgkr_prove(&mut prover, shape.log_bits(), &point, witnesses); assert_eq!(terminal.1, mle(leaves.clone(), &terminal.0)); let proof = prover.finish(); - let mut verifier = transcript::build_verifier("gkr-zero", &instance, &proof); + let mut verifier = transcript::build_verifier("gkr-zero", &point, &proof); assert_eq!( - gpgkr_verify(&mut verifier, claim, &point, 2), + gpgkr_verify(&mut verifier, claim, &point, layers), Some(terminal) ); verifier.check_eof().unwrap(); - let mut verifier = transcript::build_verifier("gkr-zero", &instance, &proof); - assert!(gpgkr_verify(&mut verifier, claim + Field::ONE, &point, 2).is_none()); + let mut verifier = transcript::build_verifier("gkr-zero", &point, &proof); + assert!(gpgkr_verify(&mut verifier, claim + Field::ONE, &point, layers).is_none()); } } - /// Small enough that `BitZParams` accepts every shape below. - const Q: u128 = (1 << 100) - 15; - /// The bottom three layers looked up in the tables, and every layer /// above, match the circuit built one leaf at a time from the same table. #[test] fn lut_layers_match_the_leaf_construction() { // At t = 7 and 8 the transposed row groups are shorter than a `u128`. for (log_rows, log_columns) in [(7, 15), (8, 14), (9, 13), (15, 7)] { - let shape = Shape::new(log_rows, log_columns).unwrap(); - let params = BitZParams::::new(shape, smallest_generator()).unwrap(); - let packed: Vec = (0..(1u128 << shape.log_bits()) / 128) - .map(|i| { - F128::from(i.wrapping_mul(0x2545_F491_4F6C_DD1D_9E37_79B9_7F4A_7C15) ^ 0xA5A5) - }) - .collect(); - let table = params.table(&packed).unwrap(); - let row_images: Vec = (0..shape.rows() as u128) - .map(|b| F128::from((b + 1) * 0x9E37_79B9_7F4A_7C15 + 7)) - .collect(); - - // Leaf `(b << s) | c`, as `GrandProductCircuit::new` describes it. - let mut leaves = vec![Field::ONE; 1 << shape.log_bits()]; - for (b, &image) in row_images.iter().enumerate() { - for c in 0..shape.columns() { - if table.bit(c, b) { - leaves[(b << log_columns) | c] = image; - } - } - } - let (expected_output, expected_layers) = fold_layers(leaves, shape.columns()); + let (shape, packed, row_images) = instance((log_rows, log_columns), Bits::Drawn, 3); + let table = table(shape, &packed); + let (expected_output, expected_layers) = + fold_layers(leaves(&table, &row_images), shape.columns()); let (output, witnesses) = GrandProductCircuit::new(&row_images, &table).batched_eval(); // Witnesses come top-down; the reference is bottom-up. @@ -772,69 +850,58 @@ mod tests { #[test] fn gpgkr_verify_rejects_every_truncated_prefix() { - let leaves: Vec = (1u128..=8).map(Field::from).collect(); - for log_groups in [0, 1] { - let (output, proof) = prove(leaves.clone(), log_groups); - assert!(verify(leaves.clone(), output.clone(), proof.clone())); - - // Cover partial field elements and missing messages in every layer. - for len in 0..proof.narg_string.len() { - let mut truncated = proof.clone(); - truncated.narg_string.truncate(len); - assert!( - !verify(leaves.clone(), output.clone(), truncated), - "accepted {len} bytes with log_groups = {log_groups}", - ); - } + let (shape, packed, row_images) = instance(SMALLEST, Bits::Drawn, 2); + let table = table(shape, &packed); + let leaves = leaves(&table, &row_images); + let (output, proof) = prove(&table, &row_images); + assert!(verify(&leaves, &output, &proof)); + + // Cover partial field elements and missing messages in every layer. + for len in 0..proof.narg_string.len() { + let mut truncated = proof.clone(); + truncated.narg_string.truncate(len); + assert!( + !verify(&leaves, &output, &truncated), + "accepted {len} bytes" + ); } } - pub fn prove(input: Vec, log_groups: usize) -> (Vec, transcript::Proof) { - let leaves = pad(input); - // Fake hashing - let n = leaves.len() as u128; - - let log_bits = leaves.len().max(1).ilog2() as usize; - let groups = 1usize << log_groups; - let (last_value, witnesses) = fold_layers(leaves, groups); + /// Proves the circuit of `table` at a point the verifier draws, as + /// `gkr_reduce` does. Returns the output layer, one product per column. + fn prove(table: &BitTable, row_images: &[Field]) -> (Vec, transcript::Proof) { + let (output, witnesses) = GrandProductCircuit::new(row_images, table).batched_eval(); + let log_bits = table.shape().log_bits(); // TODO instance is a bit loose and should be replaced by PCS - let instance = (last_value.clone(), n); - + let instance = (output.clone(), 1u128 << log_bits); let mut prover = transcript::build_prover("gkr", &instance); - // Extension point chosen by the verifier to evaluate the output layer. - // Mirrors `verify`'s `output.len().max(1).ilog2()` exactly, since both - // sides must draw the same number of challenges here. - let log_groups = last_value.len().max(1).ilog2(); - let point: Vec = (0..log_groups).map(|_| prover.verifier_message()).collect(); + // Mirrors `verify`, since both sides must draw the same number of + // challenges here. + let point: Vec = (0..output.len().ilog2()) + .map(|_| prover.verifier_message()) + .collect(); - gpgkr_prove(&mut prover, log_bits, &point, witnesses.into_iter().rev()); - (last_value, prover.finish()) + gpgkr_prove(&mut prover, log_bits, &point, witnesses); + (output, prover.finish()) } - pub fn verify(input: Vec, output: Vec, proof: transcript::Proof) -> bool { - let leaves = pad(input); - // Fake hashing - let instance = (&output, leaves.len() as u128); - let mut verifier = transcript::build_verifier("gkr", &instance, &proof); + /// Verifies `proof` for `output`, then checks the final claim against the + /// MLE of `leaves`. + fn verify(leaves: &[Field], output: &[Field], proof: &transcript::Proof) -> bool { + let instance = (output.to_vec(), leaves.len() as u128); + let mut verifier = transcript::build_verifier("gkr", &instance, proof); - let log_groups = output.len().max(1).ilog2(); - let point: Vec = (0..log_groups) + let point: Vec = (0..output.len().ilog2()) .map(|_| verifier.verifier_message()) .collect(); - let claim = mle(output, &point); - let log_groups = point.len() as u32; - let log_leafs = leaves.len().max(1).ilog2(); - let rounds = log_leafs.saturating_sub(log_groups); + let claim = mle(output.to_vec(), &point); + let layers = leaves.len().ilog2() - output.len().ilog2(); - match gpgkr_verify(&mut verifier, claim, &point, rounds) { + match gpgkr_verify(&mut verifier, claim, &point, layers) { Some((point, claim)) => { - let leaf_check = mle(leaves, &point); - - let ok = leaf_check == claim; - - ok && verifier.check_eof().is_ok() + mle(leaves.to_vec(), &point) == claim && verifier.check_eof().is_ok() } None => false, } From d4c0ecd76d1b477736ccb745e0204da7eb44af7a Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Wed, 21 Oct 2026 12:51:29 +0800 Subject: [PATCH 09/13] gkr: LUT-backed round 1 for layers 2, 1 and the leaves MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Layers 2, 1 and the leaves are no longer materialised. Round 1 of each reads the bit matrix and the 16-entry lookup tables through LutLayers and writes only the children it has bound, [left′ | right′], into one allocation sized for the leaves'. prove_rounds takes over from round 2. The leaves and layer 1 sum round 1's eq weights in buckets, with eq split into eq_hi · eq_lo; layer 2 works node by node. The sent sums are named after their points: sum_b at the boolean point b, sum_inf at ∞. gkr bench at (10, 15): 116.2 -> 71.5 ms on all threads, 215.7 -> 166-168 ms on one. --- crates/gkr/src/lib.rs | 517 +++++++++++++++++++++++++++++++++++++----- 1 file changed, 455 insertions(+), 62 deletions(-) diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs index bb314e3..b6ad44d 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs @@ -2,6 +2,7 @@ use common::{BitMatrix, BitTable, matrix::StepIter}; use field::{F128, Wide256}; use num_traits::{ConstOne, ConstZero}; use rayon::prelude::*; +use std::ops::Range; use transcript::{ProverState, VerifierState}; pub type Field = F128; @@ -23,40 +24,50 @@ pub fn gpgkr_prove( // Edge cases // - empty witnesses -> single constant circuit -> one verifier message that permutes the proof state, but a single constant can't have an MLE let mut point = point.to_owned(); - let mut claim = Field::ZERO; - let mut eq_storage = EqTables::alloc_storage(log_bits); - let lower = witnesses.lower_layers(); - let upper = witnesses.upper_nodes; - for children in upper.into_iter().rev() { - (point, claim) = prove_layer(ps, &mut eq_storage, &point, children); + let mut eq_memory = EqTables::alloc_storage(log_bits); + let LayerWitnesses { + matrix, + lower_lut, + upper_nodes, + } = witnesses; + for children in upper_nodes.into_iter().rev() { + (point, _) = prove_layer(ps, &mut eq_memory, &point, children); } - for children in lower.into_iter().rev() { - (point, claim) = prove_layer(ps, &mut eq_storage, &point, children); - } - (point, claim) + + // Layers 2, 1 and 0 stay in their lookup tables. Round 1 of each writes the + // children it has bound, `[left′ | right′]`, into one allocation sized + // for the leaves'. + let lut_layers = LutLayers::new(matrix, lower_lut); + // number of nodes at the leaf level is lut_layer.nodes() * 4. We only need half of that since we materialise after round 1. + let mut bound_children = Field::zeroed_vec(2 * lut_layers.nodes()); + let (point, _) = + prove_lut_layer::(ps, &mut eq_memory, &point, &lut_layers, &mut bound_children); + let (point, _) = + prove_lut_layer::(ps, &mut eq_memory, &point, &lut_layers, &mut bound_children); + prove_lut_layer::(ps, &mut eq_memory, &point, &lut_layers, &mut bound_children) } /// Proves `c = Σ_x eq(z, x) · L(x) · R(x)` for the halves `L`, `R` of /// `children`, one round per coordinate of `point`, top first. Round `i` -/// sends the sum `sum_i` at one endpoint and its leading coefficient, where -/// `s_i(X) = eq(z_i, X) · sum_i(X)` and `sum_i` carries -/// `eq_bound = Π_{j > i} eq(z_j, r_j)` from the rounds already bound. -/// Returns the point `(r, t)` on `children` and the claim +/// sends `sum_i(b)` and its leading coefficient `sum_i(∞)`, where `b` is the +/// boolean point that [`sends_one`] picks, `s_i(X) = eq(z_i, X) · sum_i(X)`, +/// and `sum_i` carries `eq_bound = Π_{j > i} eq(z_j, r_j)` from the rounds +/// already bound. Returns the point `(r, t)` on `children` and the claim /// `W(r, t) = L(r) + t · (R(r) - L(r))`. fn prove_layer( ps: &mut ProverState, - storage: &mut [Field], + eq_memory: &mut [Field], point: &[Field], mut children: Vec, ) -> (Vec, Field) { - let eq_tables = EqTables::new(storage, point); + let eq_tables = EqTables::new(eq_memory, point); // The top index bit of `children` picks the half: `L`, then `R`. let mid = children.len() / 2; let (left, right) = children.split_at_mut(mid); prove_rounds(ps, &eq_tables, point, None, left, right) } -/// The dense rounds of a layer, then its last message. A table-backed layer +/// The dense rounds of a layer, then its last message. A LUT-backed layer /// has already proved round 1, which binds the top coordinate of `point`, and /// passes that round's challenge as `round_one`; its `left` and `right` are /// then the bound halves. A dense layer passes `None` and proves every round @@ -87,15 +98,15 @@ fn prove_rounds( let (left0, left1) = left.split_at_mut(half); let (right0, right1) = right.split_at_mut(half); - // Branch once per round: `round_sums` takes the endpoint as a const - // generic, so the hot fold has no per-element check. + // Branch once per round: `round_sums` takes `b` as a const generic, + // so the hot fold has no per-element check. let eq = eq_tables.below(i); - let (sum_end, sum_inf) = if sends_one(z) { + let (sum_b, sum_inf) = if sends_one(z) { round_sums::(left0, left1, right0, right1, eq) } else { round_sums::(left0, left1, right0, right1, eq) }; - ps.prover_message(&[eq_bound * sum_end.reduce(), eq_bound * sum_inf.reduce()]); + ps.prover_message(&[eq_bound * sum_b.reduce(), eq_bound * sum_inf.reduce()]); let r = ps.verifier_message(); bind(r, left0, left1, right0, right1); @@ -113,8 +124,8 @@ fn prove_rounds( (next_point, interpolate(left_r, right_r, t)) } -/// `Σ eq · L(e) · R(e)` and `Σ eq · ΔL · ΔR` over the free variables: the -/// round's sum at the endpoint `e` (`1` if `AT_ONE`) and its leading +/// `Σ eq · L(b) · R(b)` and `Σ eq · ΔL · ΔR` over the free variables: the +/// round's sum at the boolean point `b` (`1` if `AT_ONE`) and its leading /// coefficient, without `eq_bound` and left unreduced. fn round_sums( left0: &[Field], @@ -148,10 +159,10 @@ fn round_sums( .with_min_len(PARALLEL_MIN_LANES / 2) .fold( || (Wide256::zero(), Wide256::zero()), - |(mut end, mut inf), ((((l0, l1), r0), r1), e)| { - // `e · L(e) · R(e)` into `end` and `e · ΔL · ΔR` into `inf`, - // with `e` joining the left factor through a reduced multiply - // so that the right one can stay wide. + |(mut sum_b, mut sum_inf), ((((l0, l1), r0), r1), e)| { + // `e · L(b) · R(b)` into `sum_b` and `e · ΔL · ΔR` into + // `sum_inf`, with `e` joining the left factor through a reduced + // multiply so that the right one can stay wide. // // Steps, products, then sums, each for both lanes: this order // puts every load of the step ahead of the first multiply. @@ -160,16 +171,16 @@ fn round_sums( let (d_l1, d_r1) = (l1[1] - l0[1], r1[1] - r0[1]); let (a0, b0) = (e[0] * at(l0[0], l1[0]), e[0] * d_l0); let (a1, b1) = (e[1] * at(l0[1], l1[1]), e[1] * d_l1); - end += Wide256::mul(a0, at(r0[0], r1[0])); - inf += Wide256::mul(b0, d_r0); - end += Wide256::mul(a1, at(r0[1], r1[1])); - inf += Wide256::mul(b1, d_r1); - (end, inf) + sum_b += Wide256::mul(a0, at(r0[0], r1[0])); + sum_inf += Wide256::mul(b0, d_r0); + sum_b += Wide256::mul(a1, at(r0[1], r1[1])); + sum_inf += Wide256::mul(b1, d_r1); + (sum_b, sum_inf) }, ) .reduce( || (Wide256::zero(), Wide256::zero()), - |(a0, ainf), (b0, binf)| (a0 + b0, ainf + binf), + |(x_b, x_inf), (y_b, y_inf)| (x_b + y_b, x_inf + y_inf), ) } @@ -197,6 +208,327 @@ fn bind(r: Field, left0: &mut [Field], left1: &[Field], right0: &mut [Field], ri }); } +/// The LUT-backed layers, by their index in the tree. Each names the +/// sumcheck whose children it is. +const LEAVES: usize = 0; +const LAYER1: usize = 1; +const LAYER2: usize = 2; + +/// Where round 1 of a LUT-backed sumcheck reads its operands `l0, l1, r0, +/// r1`, by the layer of its children: operand `o` at position `k` is nibble +/// `k + STREAM[o] · stride`, masked by `MASK[o]`, in that nibble's lookup +/// table. [`LutLayers::pairs`] says why. +const STREAM: [[usize; 4]; 3] = [[0, 0, 0, 0], [0, 1, 0, 1], [0, 1, 2, 3]]; +const MASK: [[usize; 4]; 3] = [ + [0b0001, 0b0010, 0b0100, 0b1000], + [0b0101, 0b0101, 0b1010, 0b1010], + [0b1111; 4], +]; + +/// Layers 0 to 2 in their lookup tables, as round 1 of their sumchecks reads +/// them: layer-2 node `k = (g << s) | c` is nibble `k` of the matrix, looked +/// up in the lookup table of node position `g`, which all `2^s` trees share. +struct LutLayers { + /// The nibbles of the layer-2 nodes; see [`GrandProductCircuit::new`]. + matrix: BitMatrix<'static>, + luts: Vec<[Field; 16]>, + /// `s`: a lookup table's run of nibbles holds one per tree. + log_trees: usize, +} + +impl LutLayers { + fn new(matrix: BitMatrix<'static>, luts: Vec<[Field; 16]>) -> Self { + let log_trees = (matrix.dim2() / NIBBLE_BITS).ilog2() as usize; + Self { + matrix, + luts, + log_trees, + } + } + + /// The nibbles, a word at a time. + fn words(&self) -> &[u64] { + // Read as `u64`s: on a little-endian target, bit `i` of the words is + // bit `i` of the `u64`s too. + const { assert!(cfg!(target_endian = "little")) }; + bytemuck::cast_slice(&self.matrix[..]) + } + + /// The layer-2 nodes, one per nibble. + fn nodes(&self) -> usize { + self.words().len() * NIBBLES_IN_WORD + } + + /// The positions of round 1 handled as one piece: a lookup table's run, + /// or a part of it, so that a piece never spans two lookup tables. + fn piece(&self) -> usize { + (1 << self.log_trees).min(PARALLEL_MIN_LANES) + } + + /// Round 1's positions `0..half` in pieces, as a parallel iterator + /// taking at least `PARALLEL_MIN_LANES` positions per task, whatever the + /// number of trees. + fn pieces(&self, half: usize) -> impl IndexedParallelIterator> + use<> { + let piece = self.piece(); + (0..half / piece) + .into_par_iter() + .with_min_len(PARALLEL_MIN_LANES / piece) + .map(move |p| p * piece..(p + 1) * piece) + } + + /// The pairs that round 1 binds at `positions`, which lie in one lookup + /// table's run and start and end on a word: per position, the nibbles of + /// its operands `l0, l1, r0, r1`, unmasked, and each operand's lookup + /// table. + /// + /// Round 1 binds the top bit of the parent's index, so each node pairs + /// with its partner `stride` nodes later, where `stride` is the number of + /// round-1 positions: + /// - Layer 2: a nibble is one node. The left pair is nibbles `k` and + /// `k + stride`, and the right pair the two half the layer on, + /// `k + 2·stride` and `k + 3·stride`. + /// - Layer 1: a nibble holds a left node, bits 0 and 2, and a right node, + /// bits 1 and 3. Both of their partners sit in nibble `k + stride`. + /// - Leaves: a nibble holds both pairs, bits 0 and 1 for the left child + /// and bits 2 and 3 for the right, so every operand reads nibble `k`. + /// + /// The words each operand reads are zipped, and every tuple of words + /// gives its 16 positions. + fn pairs( + &self, + stride: usize, + positions: Range, + ) -> ([&[Field; 16]; 4], impl Iterator + '_) { + let (words, len) = (self.words(), positions.len() / NIBBLES_IN_WORD); + let at = STREAM[LAYER].map(|q| positions.start + q * stride); + let luts = at.map(|k| &self.luts[k >> self.log_trees]); + let [l0, l1, r0, r1] = at.map(|k| &words[k / NIBBLES_IN_WORD..][..len]); + let pairs = l0 + .iter() + .zip(l1) + .zip(r0) + .zip(r1) + .flat_map(|(((&l0, &l1), &r0), &r1)| { + (0..NIBBLES_IN_WORD).map(move |i| { + [l0, l1, r0, r1].map(|word| (word >> (NIBBLE_BITS * i)) as usize & 0xF) + }) + }); + (luts, pairs) + } +} + +/// The operands `[l0, l1, r0, r1]` of one position of round 1 of layer +/// `LAYER`, from their lookup tables and nibbles. +fn operands(luts: &[&[Field; 16]; 4], nibbles: [usize; 4]) -> [Field; 4] { + std::array::from_fn(|o| luts[o][nibbles[o] & MASK[LAYER][o]]) +} + +/// Proves the layer whose children are layer `LAYER` of the lookup tables, +/// like [`prove_layer`]. Round 1 reads its operands from the lookup tables +/// and writes only the children it has bound, `[left′ | right′]`, into +/// `bound_children`; the rounds after it are dense. +fn prove_lut_layer( + ps: &mut ProverState, + eq_memory: &mut [Field], + point: &[Field], + lut_layers: &LutLayers, + bound_children: &mut [Field], +) -> (Vec, Field) { + let eq_tables = EqTables::new(eq_memory, point); + let top = point.len() - 1; + + // Nothing is bound before round 1, so its sums go out as they are. + let (sum_b, sum_inf) = if sends_one(point[top]) { + round_one_sums::(lut_layers, &eq_tables, point) + } else { + round_one_sums::(lut_layers, &eq_tables, point) + }; + ps.prover_message(&[sum_b.reduce(), sum_inf.reduce()]); + + let r = ps.verifier_message(); + let half = 1 << top; + let (left, right) = bound_children[..2 * half].split_at_mut(half); + lut_bind::(r, lut_layers, left, right); + prove_rounds(ps, &eq_tables, point, Some(r), left, right) +} + +/// Round 1's sum at the boolean point `b` (`1` if `AT_ONE`) and its leading +/// coefficient, unreduced, as [`round_sums`] gives them. The leaves and +/// layer 1 sum their eq weights in buckets. Layer 2 pairs whole nibbles of +/// two lookup tables, 256 keys per product, which costs more to finish than +/// a piece saves at small `s`, so it works node by node. +fn round_one_sums( + lut_layers: &LutLayers, + eq_tables: &EqTables, + point: &[Field], +) -> (Wide256, Wide256) { + match LAYER { + LEAVES => bucket_sums(lut_layers, eq_tables, point, leaf_buckets::), + LAYER1 => bucket_sums(lut_layers, eq_tables, point, layer1_buckets::), + _ => layer2_sums::(lut_layers, eq_tables.below(point.len() - 1)), + } +} + +/// Round 1's sums a piece at a time, with `buckets` summing each piece's eq +/// weights by the nibble bits that `L·R` reads and multiplying each bucket +/// once. The weight splits as `eq(z_{, &[Field]) -> (Wide256, Wide256) + Sync, +) -> (Wide256, Wide256) { + let top = point.len() - 1; + let (half, s) = (1 << top, lut_layers.log_trees); + let eq_lo = eq_tables.below(s); + let eq_hi = poly::eq_table(&point[s..top]); + lut_layers + .pieces(half) + .map(|positions| { + let h = eq_hi[positions.start >> s]; + let eq_lo = &eq_lo[positions.start % (1 << s)..][..positions.len()]; + let (sum_b, sum_inf) = buckets(lut_layers, half, positions, eq_lo); + ( + Wide256::mul(h, sum_b.reduce()), + Wide256::mul(h, sum_inf.reduce()), + ) + }) + .reduce( + || (Wide256::zero(), Wide256::zero()), + |(x_b, x_inf), (y_b, y_inf)| (x_b + y_b, x_inf + y_inf), + ) +} + +/// The piece of round 1 for the leaves at `positions`, its weights in `eq`. +/// Every operand is one bit of the same nibble, so both sums are functions of +/// the nibble: each position adds its weight into its nibble's bucket, and +/// each of the 16 buckets is multiplied once. +fn leaf_buckets( + lut_layers: &LutLayers, + stride: usize, + positions: Range, + eq: &[Field], +) -> (Wide256, Wide256) { + let ([lut, ..], pairs) = lut_layers.pairs::(stride, positions); + let mut bucket = [Field::ZERO; 16]; + for ([a, ..], &e) in pairs.zip(eq) { + bucket[a] += e; + } + // `L(b)·R(b)` is the node of bits 0 and 2, or 1 and 3, and `L(∞)·R(∞)` is + // `(l1 − l0)(r1 − r0)`. + let b_mask = if AT_ONE { 0b1010 } else { 0b0101 }; + let at_b = std::array::from_fn(|a| lut[a & b_mask]); + let at_inf = std::array::from_fn(|a| { + (lut[a & 0b0010] - lut[a & 0b0001]) * (lut[a & 0b1000] - lut[a & 0b0100]) + }); + (dot(&bucket, &at_b), dot(&bucket, &at_inf)) +} + +/// The piece of round 1 for layer 1 at `positions`, its weights in `eq`: +/// nibble `n0` and the nibble `n1` half the layer later. In characteristic 2, +/// `(l1 − l0)(r1 − r0)` is `l0·r0 + l1·r1 + l1·r0 + l0·r1`, and each product +/// reads two disjoint 2-bit masks: 16 buckets each, keyed by the bits it +/// reads. +fn layer1_buckets( + lut_layers: &LutLayers, + stride: usize, + positions: Range, + eq: &[Field], +) -> (Wide256, Wide256) { + // `l0` and `r0` read nibble `n0` in its lookup table, `l1` and `r1` + // nibble `n1`. + let ([lut0, lut1, ..], pairs) = lut_layers.pairs::(stride, positions); + let mut bucket = [[Field::ZERO; 16]; 4]; + // A position's keys in the buckets of `l0·r0`, `l1·r1`, `l1·r0` and + // `l0·r1`: bits 0 and 2 from the nibble of the `l` factor, bits 1 and 3 + // from that of the `r` factor. + for ([n0, n1, ..], &e) in pairs.zip(eq) { + bucket[0][n0] += e; + bucket[1][n1] += e; + bucket[2][(n1 & 0b0101) | (n0 & 0b1010)] += e; + bucket[3][(n0 & 0b0101) | (n1 & 0b1010)] += e; + } + // `l0·r0`, `l1·r1`, `l1·r0` and `l0·r1` at every key. + let product: [[Field; 16]; 4] = [ + *lut0, + *lut1, + std::array::from_fn(|x| lut1[x & 0b0101] * lut0[x & 0b1010]), + std::array::from_fn(|x| lut0[x & 0b0101] * lut1[x & 0b1010]), + ]; + // Product `b` is `l_b·r_b`. + let b = usize::from(AT_ONE); + let sum_inf = (0..4).fold(Wide256::zero(), |acc, j| acc + dot(&bucket[j], &product[j])); + (dot(&bucket[b], &product[b]), sum_inf) +} + +/// `Σ bucket · value`, unreduced. +fn dot(bucket: &[Field; 16], value: &[Field; 16]) -> Wide256 { + bucket + .iter() + .zip(value) + .fold(Wide256::zero(), |acc, (&b, &v)| acc + Wide256::mul(b, v)) +} + +/// Round 1's sums node by node, as [`round_sums`] forms them, with the +/// operands read from the lookup tables. +fn layer2_sums( + lut_layers: &LutLayers, + eq: &[Field], +) -> (Wide256, Wide256) { + let at = |v0: Field, v1: Field| if AT_ONE { v1 } else { v0 }; + let half = eq.len(); + lut_layers + .pieces(half) + .map(|positions| { + let (luts, pairs) = lut_layers.pairs::(half, positions.clone()); + pairs.zip(&eq[positions]).fold( + (Wide256::zero(), Wide256::zero()), + |(sum_b, sum_inf), (n, &e)| { + let [l0, l1, r0, r1] = operands::(&luts, n); + ( + sum_b + Wide256::mul(e * at(l0, l1), at(r0, r1)), + sum_inf + Wide256::mul(e * (l1 - l0), r1 - r0), + ) + }, + ) + }) + .reduce( + || (Wide256::zero(), Wide256::zero()), + |(x_b, x_inf), (y_b, y_inf)| (x_b + y_b, x_inf + y_inf), + ) +} + +/// Round 1's bind, `left′[k] = l0 + r·(l1 − l0)` and +/// `right′[k] = r0 + r·(r1 − r0)`, with the operands read from the lookup +/// tables. +/// +/// What position `k` combines depends on the layer; see [`LutLayers::pairs`]: +/// - Layer 2: four nibbles, one node each, the left pair at `k` and +/// `k + stride` and the right pair half the layer on. +/// - Layer 1: two nibbles, `k` and `k + stride`, each holding a left and a +/// right node, so one scan binds both without jumping ahead. +/// - Leaves: one nibble, which holds both pairs. +fn lut_bind( + r: Field, + lut_layers: &LutLayers, + left: &mut [Field], + right: &mut [Field], +) { + let (half, piece) = (left.len(), lut_layers.piece()); + left.par_chunks_mut(piece) + .zip(right.par_chunks_mut(piece)) + .zip(lut_layers.pieces(half)) + .for_each(|((left, right), positions)| { + let (luts, pairs) = lut_layers.pairs::(half, positions); + for ((l, rt), n) in left.iter_mut().zip(right).zip(pairs) { + let [l0, l1, r0, r1] = operands::(&luts, n); + (*l, *rt) = (interpolate(l0, l1, r), interpolate(r0, r1, r)); + } + }); +} + /// The line through `at0` (at `0`) and `at1` (at `1`), at `r`: a multilinear /// function of one variable. fn interpolate(at0: Field, at1: Field, r: Field) -> Field { @@ -205,10 +537,10 @@ fn interpolate(at0: Field, at1: Field, r: Field) -> Field { } /// The quadratic through `at0` (at `0`) and `at1` (at `1`) with leading -/// coefficient `lead`, at `r`: `at0 + r·(at1 - at0) + r·(r - 1)·lead`, with +/// coefficient `at_inf`, at `r`: `at0 + r·(at1 - at0) + r·(r - 1)·at_inf`, with /// `r` factored out. -fn quadratic_at(at0: Field, at1: Field, lead: Field, r: Field) -> Field { - at0 + r * ((at1 - at0) + (r - Field::ONE) * lead) +fn quadratic_at(at0: Field, at1: Field, at_inf: Field, r: Field) -> Field { + at0 + r * ((at1 - at0) + (r - Field::ONE) * at_inf) } /// Whether round `i` sends `sum_i(1)` rather than `sum_i(0)`. The verifier @@ -318,7 +650,7 @@ pub fn gpgkr_verify( Some((point, claim)) } -/// Per round, recovers the other endpoint from +/// Per round, recovers the sum at the other boolean point from /// `c = (1 - z_i) · sum_i(0) + z_i · sum_i(1)` and moves the claim to /// `s_i(r_i)`. /// Then checks `c = eq(z, r) · L(r) · R(r)` and continues at @@ -333,12 +665,12 @@ fn verify_layer( let mut eq_bound = Field::ONE; for (i, &z) in point.iter().enumerate().rev() { - let [sum_end, sum_inf]: [Field; 2] = vs.prover_message().ok()?; + let [sum_b, sum_inf]: [Field; 2] = vs.prover_message().ok()?; // `claim = (1 - z) * sum0 + z * sum1`. let (sum0, sum1) = if sends_one(z) { - (claim, sum_end) + (claim, sum_b) } else { - (sum_end, (claim - (Field::ONE - z) * sum_end) / z) + (sum_b, (claim - (Field::ONE - z) * sum_b) / z) }; let r = vs.verifier_message(); @@ -375,8 +707,8 @@ impl<'a> GrandProductCircuit<'a> { // one elsewhere: columns occupy the low index bits, so each product // tree reduces one column, folding the row bits from the top. // - // The bottom three layers are looked up in 16-entry tables, one per - // four rows, so the top two row bits are split off and moved beside + // The bottom three layers are looked up in 16-entry lookup tables, one + // per four rows, so the top two row bits are split off and moved beside // the columns. After the transpose, row group `b' < rows/4` is one // contiguous run, and in it the bits of rows `b'`, `b' + rows/4`, // `b' + 2 rows/4` and `b' + 3 rows/4` of one column sit next to each @@ -401,7 +733,7 @@ impl<'a> GrandProductCircuit<'a> { pub fn batched_eval(self) -> (Vec, LayerWitnesses) { let lower_lut = lower_luts(self.images); let layer3 = layer3(&self.matrix, &lower_lut); - // A table's run holds one nibble per product tree. + // A lookup table's run holds one nibble per product tree. let trees = self.matrix.dim2() / NIBBLE_BITS; let (output, upper_nodes) = fold_layers(layer3, trees); @@ -419,10 +751,10 @@ impl<'a> GrandProductCircuit<'a> { const NIBBLE_BITS: usize = 4; const NIBBLES_IN_WORD: usize = u64::BITS as usize / NIBBLE_BITS; -/// Layers 0 to 2 as tables, one per four images: `lut[s]` is the layer-2 node -/// whose four leaves are the bits of `s`. +/// Layers 0 to 2 as lookup tables, one per four images: `lut[s]` is the +/// layer-2 node whose four leaves are the bits of `s`. /// -/// Each table takes one image from each quarter of `images`. +/// Each lookup table takes one image from each quarter of `images`. fn lower_luts(images: &[Field]) -> Vec<[Field; 16]> { /// `lut[s]` is the product of the inputs whose bit is set in `s`, `a` /// on bit 0 through `d` on bit 3 (`0b0000 -> 1`, `0b0001 -> a`, @@ -459,10 +791,10 @@ fn lower_luts(images: &[Field]) -> Vec<[Field; 16]> { } /// Layer 3, the first layer built in full: node `k` multiplies layer-2 nodes -/// `k` and `k + n/2`, each a table entry looked up by its nibble. +/// `k` and `k + n/2`, each a lookup-table entry looked up by its nibble. /// -/// The matrix and the tables split in the middle into the left and the right -/// child. +/// The matrix and the lookup tables split in the middle into the left and +/// the right child. fn layer3(matrix: &BitMatrix, lower_lut: &[[Field; 16]]) -> Vec { // Read as `u64`s: on a little-endian target, bit `i` of the words is // bit `i` of the `u64`s too. @@ -471,12 +803,15 @@ fn layer3(matrix: &BitMatrix, lower_lut: &[[Field; 16]]) -> Vec { let (left, right) = bitstring.split_at(bitstring.len() / 2); let (left_luts, right_luts) = lower_lut.split_at(lower_lut.len() / 2); - // Each table serves one run of whole words, so a pair of tables is + // Each lookup table serves one run of whole words, so a pair of them is // picked once and its nibbles are read a word at a time. - // A table's run is one `xs[i1]` of the matrix, a nibble per node. + // A lookup table's run is one `xs[i1]` of the matrix, a nibble per node. let lut_words = left.len() / left_luts.len(); // This introduces a bound. StepIterator can prevent that but is more expensive than doing it this way - assert!(lut_words > 0, "each table needs at least a word of nibbles"); + assert!( + lut_words > 0, + "each lookup table needs at least a word of nibbles" + ); let runs = left .chunks_exact(lut_words) .zip(right.chunks_exact(lut_words)) @@ -487,7 +822,7 @@ fn layer3(matrix: &BitMatrix, lower_lut: &[[Field; 16]]) -> Vec { (0..NIBBLES_IN_WORD).map(move |i| lut[(word >> (NIBBLE_BITS * i)) as usize & 15]) } - // One node per nibble of each left table. + // One node per nibble of each left lookup table. let mut layer3 = Vec::with_capacity(left_luts.len() * lut_words * NIBBLES_IN_WORD); for ((l_words, r_words), (l_lut, r_lut)) in runs { for (&l, &r) in l_words.iter().zip(r_words) { @@ -497,16 +832,17 @@ fn layer3(matrix: &BitMatrix, lower_lut: &[[Field; 16]]) -> Vec { layer3 } -/// Each layer-2 node of a run of row groups as the table of its row group and -/// the nibble of its four leaves, in node order: node `k = (b' << s) | c` -/// reads nibble `k`, and all `2^s` columns of row group `b'` share its table. -/// `nibbles` and `luts` cover the same row groups. +/// Each layer-2 node of a run of row groups as the lookup table of its row +/// group and the nibble of its four leaves, in node order: node +/// `k = (b' << s) | c` reads nibble `k`, and all `2^s` columns of row group +/// `b'` share its lookup table. `nibbles` and `luts` cover the same row +/// groups. fn lut_nodes<'a>( nibbles: &'a [u128], luts: &'a [[Field; 16]], ) -> impl ExactSizeIterator { let nibbles = StepIter::<4>::new(nibbles.iter()); - // Every table serves the same run of nibbles, one per column. + // Every lookup table serves the same run of nibbles, one per column. // TODO check assembly that this isn't reloaded constantly. Or maybe it is unavoidable but might still save a shift instruction. Also might generate too much indirection. let log_columns = (nibbles.len() / luts.len()).ilog2(); nibbles @@ -546,12 +882,12 @@ fn fold_into( } /// The layers below the output, yielded from the top down. The bottom three -/// are kept as the tables they are looked up in, the rest in full. +/// are kept as the lookup tables they are looked up in, the rest in full. // TODO: This can be a single vec that is allocated at the start. pub struct LayerWitnesses { /// The nibbles of the layer-2 nodes; see [`GrandProductCircuit::new`]. matrix: BitMatrix<'static>, - /// One table per row group, shared by all of its columns. + /// One lookup table per row group, shared by all of its columns. lower_lut: Vec<[Field; 16]>, /// Layer 3 and up, layer 3 first. upper_nodes: Vec>, @@ -584,7 +920,7 @@ impl LayerWitnesses { } // The sumcheck still reads every layer in full, so for now the bottom three -// are rebuilt from the tables here; layer 3 and up are already in +// are rebuilt from the lookup tables here; layer 3 and up are already in // `upper_nodes`. impl IntoIterator for LayerWitnesses { type Item = Vec; @@ -824,7 +1160,64 @@ mod tests { } } - /// The bottom three layers looked up in the tables, and every layer + /// Each LUT-backed layer proved on its own, from a point whose top + /// coordinate is zero (round 1's `sends_one` branch) or not: the verifier + /// accepts it, and the final claim is the MLE of the children. + #[test] + fn lut_layers_prove_from_any_top_coordinate() { + let (shape, packed, row_images) = instance(SMALLEST, Bits::Drawn, 4); + let table = table(shape, &packed); + // The leaves and every layer above, built one leaf at a time. + let (_, layers) = fold_layers(leaves(&table, &row_images), shape.columns()); + let (_, witnesses) = GrandProductCircuit::new(&row_images, &table).batched_eval(); + let lut_layers = LutLayers::new(witnesses.matrix, witnesses.lower_lut); + let mut bound_children = Field::zeroed_vec(2 * lut_layers.nodes()); + let mut eq_memory = EqTables::alloc_storage(shape.log_bits()); + + let mut draws = stream(5); + for layer in [LEAVES, LAYER1, LAYER2] { + type Prove = fn( + &mut ProverState, + &mut [Field], + &[Field], + &LutLayers, + &mut [Field], + ) -> (Vec, Field); + let prove: Prove = match layer { + LEAVES => prove_lut_layer::, + LAYER1 => prove_lut_layer::, + _ => prove_lut_layer::, + }; + let (children, parent) = (&layers[layer], &layers[layer + 1]); + for top in [Field::ZERO, Field::from(7u128)] { + let mut point: Vec = (0..parent.len().ilog2()) + .map(|_| element(&mut draws)) + .collect(); + *point.last_mut().unwrap() = top; + + let mut prover = transcript::build_prover("lut-layer", &point); + let terminal = prove( + &mut prover, + &mut eq_memory, + &point, + &lut_layers, + &mut bound_children, + ); + let proof = prover.finish(); + + let mut verifier = transcript::build_verifier("lut-layer", &point, &proof); + let claim = mle(parent.clone(), &point); + assert_eq!( + verify_layer(&mut verifier, claim, &point), + Some(terminal.clone()), + "layer {layer}, top coordinate {top}" + ); + assert_eq!(terminal.1, mle(children.clone(), &terminal.0)); + } + } + } + + /// The bottom three layers looked up in the lookup tables, and every layer /// above, match the circuit built one leaf at a time from the same table. #[test] fn lut_layers_match_the_leaf_construction() { From 4e2c34cbff119cfdecedd3eaf90927b9af892b9d Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Wed, 21 Oct 2026 13:03:58 +0800 Subject: [PATCH 10/13] gkr: pairs yields per word, its users loop over the 16 positions Flattened in pairs and zipped per position, every round-1 loop over the lookup tables compiled to a per-position state machine: a nibble counter, a branch and a variable shift per position, and for layer 2 the iterator state on the stack. pairs now yields one item per word, and each user zips the words with 16-position chunks of its own data and loops inside: the leaf sums unroll, the others become counted loops. gkr bench at (10, 15), one thread: 124.6e9 -> 116.8e9 instructions, 168 -> 159.5 ms fastest; all threads: 71 -> 70.3 ms fastest. --- crates/gkr/src/lib.rs | 62 +++++++++++++++++++++++++++---------------- 1 file changed, 39 insertions(+), 23 deletions(-) diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs index b6ad44d..08332f3 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs @@ -293,12 +293,18 @@ impl LutLayers { /// and bits 2 and 3 for the right, so every operand reads nibble `k`. /// /// The words each operand reads are zipped, and every tuple of words - /// gives its 16 positions. + /// gives its 16 positions, one item per word. Callers zip the words with + /// 16-position chunks of their own data and loop over the positions + /// inside: flattened here and zipped per position, the loop becomes a + /// per-position state machine that LLVM does not unroll. fn pairs( &self, stride: usize, positions: Range, - ) -> ([&[Field; 16]; 4], impl Iterator + '_) { + ) -> ( + [&[Field; 16]; 4], + impl Iterator> + '_, + ) { let (words, len) = (self.words(), positions.len() / NIBBLES_IN_WORD); let at = STREAM[LAYER].map(|q| positions.start + q * stride); let luts = at.map(|k| &self.luts[k >> self.log_trees]); @@ -308,7 +314,7 @@ impl LutLayers { .zip(l1) .zip(r0) .zip(r1) - .flat_map(|(((&l0, &l1), &r0), &r1)| { + .map(|(((&l0, &l1), &r0), &r1)| { (0..NIBBLES_IN_WORD).map(move |i| { [l0, l1, r0, r1].map(|word| (word >> (NIBBLE_BITS * i)) as usize & 0xF) }) @@ -413,8 +419,10 @@ fn leaf_buckets( ) -> (Wide256, Wide256) { let ([lut, ..], pairs) = lut_layers.pairs::(stride, positions); let mut bucket = [Field::ZERO; 16]; - for ([a, ..], &e) in pairs.zip(eq) { - bucket[a] += e; + for (word, eq) in pairs.zip(eq.chunks_exact(NIBBLES_IN_WORD)) { + for ([a, ..], &e) in word.zip(eq) { + bucket[a] += e; + } } // `L(b)·R(b)` is the node of bits 0 and 2, or 1 and 3, and `L(∞)·R(∞)` is // `(l1 − l0)(r1 − r0)`. @@ -444,11 +452,13 @@ fn layer1_buckets( // A position's keys in the buckets of `l0·r0`, `l1·r1`, `l1·r0` and // `l0·r1`: bits 0 and 2 from the nibble of the `l` factor, bits 1 and 3 // from that of the `r` factor. - for ([n0, n1, ..], &e) in pairs.zip(eq) { - bucket[0][n0] += e; - bucket[1][n1] += e; - bucket[2][(n1 & 0b0101) | (n0 & 0b1010)] += e; - bucket[3][(n0 & 0b0101) | (n1 & 0b1010)] += e; + for (word, eq) in pairs.zip(eq.chunks_exact(NIBBLES_IN_WORD)) { + for ([n0, n1, ..], &e) in word.zip(eq) { + bucket[0][n0] += e; + bucket[1][n1] += e; + bucket[2][(n1 & 0b0101) | (n0 & 0b1010)] += e; + bucket[3][(n0 & 0b0101) | (n1 & 0b1010)] += e; + } } // `l0·r0`, `l1·r1`, `l1·r0` and `l0·r1` at every key. let product: [[Field; 16]; 4] = [ @@ -483,16 +493,18 @@ fn layer2_sums( .pieces(half) .map(|positions| { let (luts, pairs) = lut_layers.pairs::(half, positions.clone()); - pairs.zip(&eq[positions]).fold( - (Wide256::zero(), Wide256::zero()), - |(sum_b, sum_inf), (n, &e)| { - let [l0, l1, r0, r1] = operands::(&luts, n); - ( - sum_b + Wide256::mul(e * at(l0, l1), at(r0, r1)), - sum_inf + Wide256::mul(e * (l1 - l0), r1 - r0), - ) - }, - ) + let eq = eq[positions].chunks_exact(NIBBLES_IN_WORD); + pairs + .zip(eq) + .fold((Wide256::zero(), Wide256::zero()), |sums, (word, eq)| { + word.zip(eq).fold(sums, |(sum_b, sum_inf), (n, &e)| { + let [l0, l1, r0, r1] = operands::(&luts, n); + ( + sum_b + Wide256::mul(e * at(l0, l1), at(r0, r1)), + sum_inf + Wide256::mul(e * (l1 - l0), r1 - r0), + ) + }) + }) }) .reduce( || (Wide256::zero(), Wide256::zero()), @@ -522,9 +534,13 @@ fn lut_bind( .zip(lut_layers.pieces(half)) .for_each(|((left, right), positions)| { let (luts, pairs) = lut_layers.pairs::(half, positions); - for ((l, rt), n) in left.iter_mut().zip(right).zip(pairs) { - let [l0, l1, r0, r1] = operands::(&luts, n); - (*l, *rt) = (interpolate(l0, l1, r), interpolate(r0, r1, r)); + let left = left.chunks_exact_mut(NIBBLES_IN_WORD); + let right = right.chunks_exact_mut(NIBBLES_IN_WORD); + for ((left, right), word) in left.zip(right).zip(pairs) { + for ((l, rt), n) in left.iter_mut().zip(right).zip(word) { + let [l0, l1, r0, r1] = operands::(&luts, n); + (*l, *rt) = (interpolate(l0, l1, r), interpolate(r0, r1, r)); + } } }); } From e2691332363792097c8d7a8f79cc331a42dba3be Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Wed, 21 Oct 2026 15:12:06 +0800 Subject: [PATCH 11/13] gkr: build layer 3 and the layers above in parallel LayerWitnesses holds the LutLayers that batched_eval builds, and layer3 reads the lookup tables through LutLayers::pieces and pairs, as round 1 of layer 1 does. lower_luts, layer3 and fold_layers run on rayon with the same PARALLEL_MIN_LANES threshold as the sumcheck. LutLayers::new now asserts that each lookup table's run holds at least a word of nibbles, which layer3 used to check. fold_into is gone. gkr bench at (10, 15), all threads: 70.35 -> 66.1-66.6 ms fastest; one thread unchanged (159.4 ms, 116.8e9 instructions). --- crates/gkr/src/lib.rs | 109 ++++++++++++++++-------------------------- 1 file changed, 41 insertions(+), 68 deletions(-) diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs index 08332f3..5ee74a7 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs @@ -26,8 +26,7 @@ pub fn gpgkr_prove( let mut point = point.to_owned(); let mut eq_memory = EqTables::alloc_storage(log_bits); let LayerWitnesses { - matrix, - lower_lut, + lut_layers, upper_nodes, } = witnesses; for children in upper_nodes.into_iter().rev() { @@ -37,7 +36,6 @@ pub fn gpgkr_prove( // Layers 2, 1 and 0 stay in their lookup tables. Round 1 of each writes the // children it has bound, `[left′ | right′]`, into one allocation sized // for the leaves'. - let lut_layers = LutLayers::new(matrix, lower_lut); // number of nodes at the leaf level is lut_layer.nodes() * 4. We only need half of that since we materialise after round 1. let mut bound_children = Field::zeroed_vec(2 * lut_layers.nodes()); let (point, _) = @@ -239,6 +237,11 @@ struct LutLayers { impl LutLayers { fn new(matrix: BitMatrix<'static>, luts: Vec<[Field; 16]>) -> Self { let log_trees = (matrix.dim2() / NIBBLE_BITS).ilog2() as usize; + // `pairs` reads whole words of one lookup table's run. + assert!( + 1 << log_trees >= NIBBLES_IN_WORD, + "each lookup table needs at least a word of nibbles" + ); Self { matrix, luts, @@ -747,17 +750,14 @@ impl<'a> GrandProductCircuit<'a> { /// full, and the layers above it fold as usual. #[tracing::instrument(name = "Evaluate grand-product circuit", level = "debug", skip_all)] pub fn batched_eval(self) -> (Vec, LayerWitnesses) { - let lower_lut = lower_luts(self.images); - let layer3 = layer3(&self.matrix, &lower_lut); - // A lookup table's run holds one nibble per product tree. - let trees = self.matrix.dim2() / NIBBLE_BITS; - let (output, upper_nodes) = fold_layers(layer3, trees); + let lut_layers = LutLayers::new(self.matrix, lower_luts(self.images)); + let layer3 = layer3(&lut_layers); + let (output, upper_nodes) = fold_layers(layer3, 1 << lut_layers.log_trees); ( output, LayerWitnesses { - matrix: self.matrix, - lower_lut, + lut_layers, upper_nodes, }, ) @@ -798,53 +798,34 @@ fn lower_luts(images: &[Field]) -> Vec<[Field; 16]> { let (a, rest) = images.split_at(quart); let (b, rest) = rest.split_at(quart); let (c, d) = rest.split_at(quart); - a.iter() + a.par_iter() .zip(b) .zip(c) .zip(d) + // A lookup table writes 16 entries. + .with_min_len(PARALLEL_MIN_LANES / 16) .map(|(((&a, &b), &c), &d)| create_lut([a, b, c, d])) .collect() } /// Layer 3, the first layer built in full: node `k` multiplies layer-2 nodes /// `k` and `k + n/2`, each a lookup-table entry looked up by its nibble. -/// -/// The matrix and the lookup tables split in the middle into the left and -/// the right child. -fn layer3(matrix: &BitMatrix, lower_lut: &[[Field; 16]]) -> Vec { - // Read as `u64`s: on a little-endian target, bit `i` of the words is - // bit `i` of the `u64`s too. - const { assert!(cfg!(target_endian = "little")) }; - let bitstring: &[u64] = bytemuck::cast_slice(&matrix[..]); - let (left, right) = bitstring.split_at(bitstring.len() / 2); - let (left_luts, right_luts) = lower_lut.split_at(lower_lut.len() / 2); - - // Each lookup table serves one run of whole words, so a pair of them is - // picked once and its nibbles are read a word at a time. - // A lookup table's run is one `xs[i1]` of the matrix, a nibble per node. - let lut_words = left.len() / left_luts.len(); // This introduces a bound. StepIterator can prevent that but is more expensive than doing it this way - - assert!( - lut_words > 0, - "each lookup table needs at least a word of nibbles" - ); - let runs = left - .chunks_exact(lut_words) - .zip(right.chunks_exact(lut_words)) - .zip(left_luts.iter().zip(right_luts)); - - // The layer-2 nodes whose four leaves are the nibbles of `word`. - fn layer2(lut: &[Field; 16], word: u64) -> impl Iterator + '_ { - (0..NIBBLES_IN_WORD).map(move |i| lut[(word >> (NIBBLE_BITS * i)) as usize & 15]) - } - - // One node per nibble of each left lookup table. - let mut layer3 = Vec::with_capacity(left_luts.len() * lut_words * NIBBLES_IN_WORD); - for ((l_words, r_words), (l_lut, r_lut)) in runs { - for (&l, &r) in l_words.iter().zip(r_words) { - fold_into(&mut layer3, layer2(l_lut, l), layer2(r_lut, r)); - } - } +/// Those are the two nibbles that round 1 of layer 1 reads at position `k`, +/// here read whole. +fn layer3(lut_layers: &LutLayers) -> Vec { + let half = lut_layers.nodes() / 2; + let mut layer3 = Field::zeroed_vec(half); + layer3 + .par_chunks_mut(lut_layers.piece()) + .zip(lut_layers.pieces(half)) + .for_each(|(nodes, positions)| { + let ([lut0, lut1, ..], pairs) = lut_layers.pairs::(half, positions); + for (nodes, word) in nodes.chunks_exact_mut(NIBBLES_IN_WORD).zip(pairs) { + for (node, [n0, n1, ..]) in nodes.iter_mut().zip(word) { + *node = lut0[n0] * lut1[n1]; + } + } + }); layer3 } @@ -875,9 +856,15 @@ fn fold_layers(mut layer: Vec, groups: usize) -> (Vec, Vec groups { + // Node `k` multiplies node `k` of the left half and node `k` of the + // right half below it. let (l, r) = layer.split_at(layer.len() / 2); - let mut next = Vec::with_capacity(l.len()); - fold_into(&mut next, l.iter().copied(), r.iter().copied()); + let next = l + .par_iter() + .zip(r) + .with_min_len(PARALLEL_MIN_LANES) + .map(|(&a, &b)| a * b) + .collect(); witnesses.push(layer); layer = next; } @@ -885,26 +872,12 @@ fn fold_layers(mut layer: Vec, groups: usize) -> (Vec, Vec, - left: impl Iterator, - right: impl Iterator, -) { - // `extend` from zipped ranges or slices knows its length, so there is no - // capacity check or length update per node. - next.extend(left.zip(right).map(|(a, b)| a * b)); -} - /// The layers below the output, yielded from the top down. The bottom three /// are kept as the lookup tables they are looked up in, the rest in full. // TODO: This can be a single vec that is allocated at the start. pub struct LayerWitnesses { - /// The nibbles of the layer-2 nodes; see [`GrandProductCircuit::new`]. - matrix: BitMatrix<'static>, - /// One lookup table per row group, shared by all of its columns. - lower_lut: Vec<[Field; 16]>, + /// Layers 0 to 2, in their lookup tables. + lut_layers: LutLayers, /// Layer 3 and up, layer 3 first. upper_nodes: Vec>, } @@ -915,7 +888,7 @@ impl LayerWitnesses { /// `q` alone, layer-1 half `h` pairs bits `h` and `h + 2`, and layer 2 /// takes all four. fn lower_layers(&self) -> [Vec; 3] { - let nodes = lut_nodes(&self.matrix, &self.lower_lut); + let nodes = lut_nodes(&self.lut_layers.matrix, &self.lut_layers.luts); let n = nodes.len(); let mut leaves = Field::zeroed_vec(4 * n); let mut layer1 = Field::zeroed_vec(2 * n); @@ -1186,7 +1159,7 @@ mod tests { // The leaves and every layer above, built one leaf at a time. let (_, layers) = fold_layers(leaves(&table, &row_images), shape.columns()); let (_, witnesses) = GrandProductCircuit::new(&row_images, &table).batched_eval(); - let lut_layers = LutLayers::new(witnesses.matrix, witnesses.lower_lut); + let lut_layers = witnesses.lut_layers; let mut bound_children = Field::zeroed_vec(2 * lut_layers.nodes()); let mut eq_memory = EqTables::alloc_storage(shape.log_bits()); From 94e02539d4a417e4a8596be5c9cbc04da00e39c5 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Thu, 22 Oct 2026 11:25:32 +0800 Subject: [PATCH 12/13] gkr: drop the full lower layers, docs in tree terms Nothing reads layers 0 to 2 in full any more, so lower_layers, IntoIterator for LayerWitnesses and lut_nodes are gone. lut_layers_match_the_leaf_construction compares the output and layer 3 and up; lut_layers_prove_from_any_top_coordinate covers layers 0 to 2 through the sumchecks that read them. The docs of LayerWitnesses, gpgkr_prove's witnesses, EqTables::new and alloc_storage say what they hold now, and the layout comment in GrandProductCircuit::new and the test comments speak of trees, node positions and nibbles instead of rows and columns. gkr bench at (10, 15) unchanged: 116.8e9 instructions on one thread. --- crates/gkr/src/lib.rs | 141 ++++++++++++++---------------------------- 1 file changed, 47 insertions(+), 94 deletions(-) diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs index 5ee74a7..635aa98 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs @@ -1,4 +1,4 @@ -use common::{BitMatrix, BitTable, matrix::StepIter}; +use common::{BitMatrix, BitTable}; use field::{F128, Wide256}; use num_traits::{ConstOne, ConstZero}; use rayon::prelude::*; @@ -18,11 +18,10 @@ pub fn gpgkr_prove( ps: &mut ProverState, log_bits: usize, point: &[F128], - // All the intermediate witnesses + the input layer, from the top down. Doesn't contain the output layer + // Every layer below the output: layer 3 and up in full, layers 0 to 2 in + // their lookup tables. witnesses: LayerWitnesses, ) -> (Vec, Field) { - // Edge cases - // - empty witnesses -> single constant circuit -> one verifier message that permutes the proof state, but a single constant can't have an MLE let mut point = point.to_owned(); let mut eq_memory = EqTables::alloc_storage(log_bits); let LayerWitnesses { @@ -587,14 +586,14 @@ fn eq_factor(r: Field, z: Field) -> Field { /// only the largest and halving it between rounds (`low[k] += high[k]`) /// would save that half, but it measured 7.5-12% slower, as it adds a pass /// over memory. Splitting the table makes the trade moot. -// TODO: Split the table. +// TODO: Split the table. The buckets already do a kind of splitting struct EqTables<'a> { storage: &'a [Field], } impl<'a> EqTables<'a> { /// `eq(z_{ &[Field] { let start = (1 << i) - 1; @@ -602,8 +601,9 @@ impl<'a> EqTables<'a> { &self.storage[start..end] } - /// Allocates all directly as it is as much space as a double buffer approach would take. - #[inline(never)] + /// Builds every table of `point` in `storage`, a buffer from + /// [`EqTables::alloc_storage`]; nothing is allocated here. + #[inline(never)] // Inlining or not doesn't improve performance, but it does hide it in profiling therefore it has been annotated with inline never. fn new(storage: &'a mut [Field], point: &[Field]) -> EqTables<'a> { // Establishes the whole table's space budget up front so LLVM can // prove the `low`/`hi` writes in the loop below are in bounds @@ -634,6 +634,9 @@ impl<'a> EqTables<'a> { EqTables { storage } } + /// Room for every table of a point of up to `log_bits - 1` coordinates, + /// the leaves' sumcheck's, allocated once for all layers. All the tables + /// in one buffer take as much space as a double buffer would. fn alloc_storage(log_bits: usize) -> Vec { Field::zeroed_vec(1 << (log_bits.saturating_sub(1)).max(1)) } @@ -709,7 +712,6 @@ fn verify_layer( Some((next_point, interpolate(left_r, right_r, t))) } -// TODO: Optimise for circuits that are padded. pub struct GrandProductCircuit<'a> { matrix: BitMatrix<'static>, images: &'a [F128], @@ -723,16 +725,19 @@ impl<'a> GrandProductCircuit<'a> { "one image per row of the table" ); // Leaf `(b << s) | c` is `row_images[b]` where `bit(c, b)` is set and - // one elsewhere: columns occupy the low index bits, so each product - // tree reduces one column, folding the row bits from the top. + // one elsewhere: the low `s` index bits pick one of the `2^s` trees + // and the rest are the leaf's position in it. Each layer folds the + // top position bit, so tree `c` multiplies the images `bit(c, ·)` + // selects. // - // The bottom three layers are looked up in 16-entry lookup tables, one - // per four rows, so the top two row bits are split off and moved beside - // the columns. After the transpose, row group `b' < rows/4` is one - // contiguous run, and in it the bits of rows `b'`, `b' + rows/4`, - // `b' + 2 rows/4` and `b' + 3 rows/4` of one column sit next to each - // other as a nibble: the four leaves under node `(b' << s) | c` of - // layer 2, columns in ascending order. + // Layers 0 to 2 are looked up in 16-entry lookup tables, one per four + // images, and the four leaves under a layer-2 node are a quarter of + // the positions apart, so the top two position bits are split off and + // moved below the tree bits. After the transpose, layer-2 node + // position `g` is one contiguous run, a nibble per tree in ascending + // order, and nibble `c` of it holds `bit(c, g + q·n)` at bit `q`, `n` + // being a quarter of the images: the four leaves under layer-2 node + // `(g << s) | c`. let matrix = bittable.as_matrix(); let log_dim2 = matrix.log_dim2() - 2; let matrix = matrix.reshape(log_dim2).transpose(); @@ -767,14 +772,14 @@ impl<'a> GrandProductCircuit<'a> { const NIBBLE_BITS: usize = 4; const NIBBLES_IN_WORD: usize = u64::BITS as usize / NIBBLE_BITS; -/// Layers 0 to 2 as lookup tables, one per four images: `lut[s]` is the -/// layer-2 node whose four leaves are the bits of `s`. +/// Layers 0 to 2 as lookup tables, one per four images: `lut[nibble]` is the +/// layer-2 node whose four leaves are the bits of `nibble`. /// /// Each lookup table takes one image from each quarter of `images`. fn lower_luts(images: &[Field]) -> Vec<[Field; 16]> { - /// `lut[s]` is the product of the inputs whose bit is set in `s`, `a` - /// on bit 0 through `d` on bit 3 (`0b0000 -> 1`, `0b0001 -> a`, - /// `0b0011 -> a*b`, ...), so a 4-bit lookup string indexes its + /// `lut[nibble]` is the product of the inputs whose bit is set in + /// `nibble`, `a` on bit 0 through `d` on bit 3 (`0b0000 -> 1`, + /// `0b0001 -> a`, `0b0011 -> a*b`, ...), so a nibble indexes its /// multilinear value directly. /// /// Built by doubling: entries `2^i..2^(i+1)` are entries `0..2^i` @@ -829,24 +834,6 @@ fn layer3(lut_layers: &LutLayers) -> Vec { layer3 } -/// Each layer-2 node of a run of row groups as the lookup table of its row -/// group and the nibble of its four leaves, in node order: node -/// `k = (b' << s) | c` reads nibble `k`, and all `2^s` columns of row group -/// `b'` share its lookup table. `nibbles` and `luts` cover the same row -/// groups. -fn lut_nodes<'a>( - nibbles: &'a [u128], - luts: &'a [[Field; 16]], -) -> impl ExactSizeIterator { - let nibbles = StepIter::<4>::new(nibbles.iter()); - // Every lookup table serves the same run of nibbles, one per column. - // TODO check assembly that this isn't reloaded constantly. Or maybe it is unavoidable but might still save a shift instruction. Also might generate too much indirection. - let log_columns = (nibbles.len() / luts.len()).ilog2(); - nibbles - .enumerate() - .map(move |(k, bs)| (&luts[k >> log_columns], bs as usize)) -} - /// Multiplies `layer` pairwise, left half against right half, until one /// value per group remains. Returns that output layer and every layer before /// it, `layer` first. @@ -872,9 +859,9 @@ fn fold_layers(mut layer: Vec, groups: usize) -> (Vec, Vec>, } -impl LayerWitnesses { - /// The leaves and layers 1 and 2, looked up in full. A node is the - /// product of the nibble bits its subtree covers: leaf quarter `q` is bit - /// `q` alone, layer-1 half `h` pairs bits `h` and `h + 2`, and layer 2 - /// takes all four. - fn lower_layers(&self) -> [Vec; 3] { - let nodes = lut_nodes(&self.lut_layers.matrix, &self.lut_layers.luts); - let n = nodes.len(); - let mut leaves = Field::zeroed_vec(4 * n); - let mut layer1 = Field::zeroed_vec(2 * n); - let mut layer2 = Field::zeroed_vec(n); - - for (k, (lut, bs)) in nodes.enumerate() { - for q in 0..4 { - leaves[q * n + k] = lut[bs & (1 << q)]; - } - for h in 0..2 { - layer1[h * n + k] = lut[bs & (0b0101 << h)]; - } - layer2[k] = lut[bs]; - } - - [leaves, layer1, layer2] - } -} - -// The sumcheck still reads every layer in full, so for now the bottom three -// are rebuilt from the lookup tables here; layer 3 and up are already in -// `upper_nodes`. -impl IntoIterator for LayerWitnesses { - type Item = Vec; - type IntoIter = std::iter::Rev>>; - - fn into_iter(self) -> Self::IntoIter { - let mut layers = Vec::from(self.lower_layers()); - layers.extend(self.upper_nodes); - layers.into_iter().rev() - } -} - #[cfg(test)] mod tests { use super::*; @@ -1075,7 +1022,7 @@ mod tests { #![proptest_config(ProptestConfig::with_cases(4))] // Full protocol: the verifier accepts, and the output is the product - // of each column's leaves. + // of each tree's leaves. #[test] fn gpgkr_round_trip(bits in bits(), seed in any::()) { let (shape, packed, row_images) = instance(SMALLEST, bits, seed); @@ -1206,11 +1153,18 @@ mod tests { } } - /// The bottom three layers looked up in the lookup tables, and every layer - /// above, match the circuit built one leaf at a time from the same table. + /// The output and every layer from layer 3 up match the circuit built one + /// leaf at a time from the same table. Layers 0 to 2 stay in their lookup + /// tables; `lut_layers_prove_from_any_top_coordinate` checks them through + /// the sumchecks that read them. #[test] fn lut_layers_match_the_leaf_construction() { - // At t = 7 and 8 the transposed row groups are shorter than a `u128`. + // At t = 7 and 8 there are 32 and 64 lookup tables, fewer than the bits + // of a `u128`, so the transpose in `GrandProductCircuit::new` moves + // `u32` and `u64` blocks; from t = 9 on it moves `u128` ones. At s = 7 + // a lookup table's run is 128 nibbles, 8 words, below + // `PARALLEL_MIN_LANES`, so a piece of `layer3` is a whole run rather + // than part of one. for (log_rows, log_columns) in [(7, 15), (8, 14), (9, 13), (15, 7)] { let (shape, packed, row_images) = instance((log_rows, log_columns), Bits::Drawn, 3); let table = table(shape, &packed); @@ -1218,14 +1172,13 @@ mod tests { fold_layers(leaves(&table, &row_images), shape.columns()); let (output, witnesses) = GrandProductCircuit::new(&row_images, &table).batched_eval(); - // Witnesses come top-down; the reference is bottom-up. - let mut layers: Vec> = witnesses.into_iter().collect(); - layers.reverse(); + // Both start at layer 3: the reference after its first three. + let (layers, expected_layers) = (&witnesses.upper_nodes, &expected_layers[3..]); assert!(output == expected_output, "output, t = {log_rows}"); assert_eq!(layers.len(), expected_layers.len(), "t = {log_rows}"); - for (i, (layer, expected)) in layers.iter().zip(&expected_layers).enumerate() { - assert!(layer == expected, "layer {i}, t = {log_rows}"); + for (i, (layer, expected)) in layers.iter().zip(expected_layers).enumerate() { + assert!(layer == expected, "layer {}, t = {log_rows}", i + 3); } } } @@ -1250,7 +1203,7 @@ mod tests { } /// Proves the circuit of `table` at a point the verifier draws, as - /// `gkr_reduce` does. Returns the output layer, one product per column. + /// `gkr_reduce` does. Returns the output layer, one product per tree. fn prove(table: &BitTable, row_images: &[Field]) -> (Vec, transcript::Proof) { let (output, witnesses) = GrandProductCircuit::new(row_images, table).batched_eval(); let log_bits = table.shape().log_bits(); From 02604b063236d85bdbfe4567dca4a9e0eb85f822 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Thu, 22 Oct 2026 11:48:43 +0800 Subject: [PATCH 13/13] gkr: refresh single_threaded.patch for round_sums and bind --- single_threaded.patch | 187 ++++++++++++++++++++++++++++++++---------- 1 file changed, 143 insertions(+), 44 deletions(-) diff --git a/single_threaded.patch b/single_threaded.patch index 012e19f..1c51ab4 100644 --- a/single_threaded.patch +++ b/single_threaded.patch @@ -1,60 +1,159 @@ diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs -index 89a1bf0..bb364ff 100644 +index e553149..626eaa5 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs -@@ -87,11 +87,11 @@ fn prove_layer( - let r = ps.verifier_message(); - next_point.push_back(r); +@@ -1,7 +1,6 @@ + use common::{BitMatrix, BitTable}; + use field::{F128, Wide256}; + use num_traits::{ConstOne, ConstZero}; +-use rayon::prelude::*; + use std::ops::Range; + use transcript::{ProverState, VerifierState}; -- lo_l.par_iter_mut() -- .zip(lo_r.par_iter_mut()) -- .zip(hi_l.par_iter()) -- .zip(hi_r.par_iter()) -- .with_min_len(PARALLEL_MIN_LANES) -+ lo_l.iter_mut() -+ .zip(lo_r.iter_mut()) -+ .zip(hi_l.iter()) -+ .zip(hi_r.iter()) -+ // .with_min_len(PARALLEL_MIN_LANES) - .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)) +@@ -148,16 +147,13 @@ fn round_sums( + // Two lanes per step: independent multiplication chains. Chunking also + // reduces variance in benchmarking. + left0 +- .par_chunks_exact(2) +- .zip(left1.par_chunks_exact(2)) +- .zip(right0.par_chunks_exact(2)) +- .zip(right1.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. +- // 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) ++ .chunks_exact(2) ++ .zip(left1.chunks_exact(2)) ++ .zip(right0.chunks_exact(2)) ++ .zip(right1.chunks_exact(2)) ++ .zip(eq.chunks_exact(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) + |(mut sum_b, mut sum_inf), ((((l0, l1), r0), r1), e)| { + // `e · L(b) · R(b)` into `sum_b` and `e · ΔL · ΔR` into + // `sum_inf`, with `e` joining the left factor through a reduced +@@ -177,21 +173,16 @@ fn round_sums( + (sum_b, sum_inf) }, ) - .reduce( - || (Wide256::zero(), Wide256::zero()), -- |(a0, ainf), (b0, binf)| (a0 + b0, ainf + binf), +- |(x_b, x_inf), (y_b, y_inf)| (x_b + y_b, x_inf + y_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 + /// Binds the variable that splits each pair of halves to `r`, writing the + /// result over the low halves. + fn bind(r: Field, left0: &mut [Field], left1: &[Field], right0: &mut [Field], right1: &[Field]) { + left0 +- .par_iter_mut() +- .zip(right0.par_iter_mut()) +- .zip(left1.par_iter()) +- .zip(right1.par_iter()) +- .with_min_len(PARALLEL_MIN_LANES) ++ .iter_mut() ++ .zip(right0.iter_mut()) ++ .zip(left1.iter()) ++ .zip(right1.iter()) + .for_each(|(((l0, r0), &l1), &r1)| { + // Recalculating the step rather than writing it in the high half + // in `round_sums` to save on writes. No benchmarking has been done +@@ -269,15 +260,10 @@ impl LutLayers { + (1 << self.log_trees).min(PARALLEL_MIN_LANES) + } + +- /// Round 1's positions `0..half` in pieces, as a parallel iterator +- /// taking at least `PARALLEL_MIN_LANES` positions per task, whatever the +- /// number of trees. +- fn pieces(&self, half: usize) -> impl IndexedParallelIterator> + use<> { ++ /// Round 1's positions `0..half` in pieces. ++ fn pieces(&self, half: usize) -> impl Iterator> + use<> { + let piece = self.piece(); +- (0..half / piece) +- .into_par_iter() +- .with_min_len(PARALLEL_MIN_LANES / piece) +- .map(move |p| p * piece..(p + 1) * piece) ++ (0..half / piece).map(move |p| p * piece..(p + 1) * piece) + } + + /// The pairs that round 1 binds at `positions`, which lie in one lookup +@@ -388,7 +374,7 @@ fn bucket_sums( + lut_layers: &LutLayers, + eq_tables: &EqTables, + point: &[Field], +- buckets: impl Fn(&LutLayers, usize, Range, &[Field]) -> (Wide256, Wide256) + Sync, ++ buckets: impl Fn(&LutLayers, usize, Range, &[Field]) -> (Wide256, Wide256), + ) -> (Wide256, Wide256) { + let top = point.len() - 1; + let (half, s) = (1 << top, lut_layers.log_trees); +@@ -405,8 +391,8 @@ fn bucket_sums( + Wide256::mul(h, sum_inf.reduce()), + ) + }) +- .reduce( +- || (Wide256::zero(), Wide256::zero()), ++ .fold( ++ (Wide256::zero(), Wide256::zero()), + |(x_b, x_inf), (y_b, y_inf)| (x_b + y_b, x_inf + y_inf), + ) + } +@@ -510,8 +496,8 @@ fn layer2_sums( + }) + }) + }) +- .reduce( +- || (Wide256::zero(), Wide256::zero()), ++ .fold( ++ (Wide256::zero(), Wide256::zero()), + |(x_b, x_inf), (y_b, y_inf)| (x_b + y_b, x_inf + y_inf), + ) + } +@@ -533,8 +519,8 @@ fn lut_bind( + right: &mut [Field], + ) { + let (half, piece) = (left.len(), lut_layers.piece()); +- left.par_chunks_mut(piece) +- .zip(right.par_chunks_mut(piece)) ++ left.chunks_mut(piece) ++ .zip(right.chunks_mut(piece)) + .zip(lut_layers.pieces(half)) + .for_each(|((left, right), positions)| { + let (luts, pairs) = lut_layers.pairs::(half, positions); +@@ -805,12 +791,10 @@ fn lower_luts(images: &[Field]) -> Vec<[Field; 16]> { + let (a, rest) = images.split_at(quart); + let (b, rest) = rest.split_at(quart); + let (c, d) = rest.split_at(quart); +- a.par_iter() ++ a.iter() + .zip(b) + .zip(c) + .zip(d) +- // A lookup table writes 16 entries. +- .with_min_len(PARALLEL_MIN_LANES / 16) + .map(|(((&a, &b), &c), &d)| create_lut([a, b, c, d])) + .collect() + } +@@ -823,7 +807,7 @@ fn layer3(lut_layers: &LutLayers) -> Vec { + let half = lut_layers.nodes() / 2; + let mut layer3 = Field::zeroed_vec(half); + layer3 +- .par_chunks_mut(lut_layers.piece()) ++ .chunks_mut(lut_layers.piece()) + .zip(lut_layers.pieces(half)) + .for_each(|(nodes, positions)| { + let ([lut0, lut1, ..], pairs) = lut_layers.pairs::(half, positions); +@@ -848,12 +832,7 @@ fn fold_layers(mut layer: Vec, groups: usize) -> (Vec, Vec