diff --git a/Cargo.lock b/Cargo.lock index d5eb178..601e410 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -382,10 +382,13 @@ checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" name = "common" version = "0.1.0" dependencies = [ + "bytemuck", "crypto-primitives", + "divan", "field", "num-traits", "poly", + "rand 0.10.2", "rayon", "spongefish", ] @@ -754,6 +757,7 @@ dependencies = [ name = "gkr" version = "0.1.0" dependencies = [ + "common", "field", "num-traits", "poly", diff --git a/crates/common/Cargo.toml b/crates/common/Cargo.toml index b3feed8..b22086d 100644 --- a/crates/common/Cargo.toml +++ b/crates/common/Cargo.toml @@ -8,11 +8,23 @@ license.workspace = true [features] default = ["parallel"] parallel = ["dep:rayon"] +# Exposes private kernels to `benches/`; not part of the public API. +bench = [] [dependencies] crypto-primitives = { workspace = true } +bytemuck = {workspace = true} field = { workspace = true, features = ["spongefish"] } spongefish = { workspace = true } poly = { workspace = true } rayon = { workspace = true, optional = true } num-traits = { workspace = true } + +[dev-dependencies] +divan = { workspace = true } +rand = { workspace = true } + +[[bench]] +name = "matrix" +harness = false +required-features = ["bench"] diff --git a/crates/common/benches/matrix.rs b/crates/common/benches/matrix.rs new file mode 100644 index 0000000..7ec9b35 --- /dev/null +++ b/crates/common/benches/matrix.rs @@ -0,0 +1,49 @@ +//! Compare the blocked `bit_transpose` with the bit-at-a-time reference. +//! +//! Run with `cargo bench -p common --features bench --bench matrix`. +//! Both transpose a `128 x DIM2`-bit matrix (`dim2` contiguous), including +//! output allocation. Input generation runs outside the measured loop. + +use common::matrix::bench::{bit_transpose, transpose_reference}; +use divan::counter::BytesCount; +use divan::{Bencher, black_box}; +use rand::rngs::StdRng; +use rand::{RngExt, SeedableRng}; + +const DIM1: usize = 128; +/// `2^20` bits per row: a 16 MiB matrix, well past the last-level cache. +const DIM2: usize = 1 << 20; + +fn main() { + let xs = input(); + assert_eq!( + bit_transpose(&xs, DIM1, DIM2), + transpose_reference::(&xs), + "bit_transpose disagrees with the reference" + ); + divan::main(); +} + +/// Seeded so every run transposes the same matrix. +fn input() -> Vec { + let mut rng = StdRng::seed_from_u64(0); + (0..DIM1 * DIM2 / u128::BITS as usize) + .map(|_| rng.random()) + .collect() +} + +#[divan::bench] +fn blocked(bencher: Bencher) { + let xs = input(); + bencher + .counter(BytesCount::of_slice(&xs)) + .bench(|| bit_transpose(black_box(&xs), DIM1, DIM2)); +} + +#[divan::bench] +fn reference(bencher: Bencher) { + let xs = input(); + bencher + .counter(BytesCount::of_slice(&xs)) + .bench(|| transpose_reference::(black_box(&xs))); +} diff --git a/crates/common/src/fold.rs b/crates/common/src/fold.rs index 5075396..f29d29a 100644 --- a/crates/common/src/fold.rs +++ b/crates/common/src/fold.rs @@ -5,7 +5,7 @@ use poly::DenseMultilinearExtension; #[cfg(feature = "parallel")] use rayon::prelude::*; -use crate::{BitTable, LinearClaim, Shape, table::PACKED_BITS}; +use crate::{BitTable, LinearClaim, Shape}; /// A round whose parts do not describe the shape they belong to. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -30,23 +30,16 @@ pub fn fold_column(table: &BitTable<'_>, exponents: &[u128], column: usize) -> u table .column(column) .iter() - .enumerate() - .map(|(index, element)| { - let base = index * PACKED_BITS; - [(0, element.lo), (64, element.hi)] - .into_iter() - .map(|(half, mut remaining)| { - let mut total = 0u128; - while remaining != 0 { - total += exponents[base + half + remaining.trailing_zeros() as usize]; - // Clears the lowest set bit. - remaining &= remaining - 1; - } - total - }) - .sum::() + .fold((0usize, 0u128), |(base, mut total), &element| { + let mut remaining = element; + while remaining != 0 { + total += exponents[base + remaining.trailing_zeros() as usize]; + // Clears the lowest set bit. + remaining &= remaining - 1; + } + (base + BitTable::BITS, total) }) - .sum() + .1 } /// Every column's fold, in column order. @@ -190,6 +183,7 @@ mod tests { use num_traits::{ConstOne, ConstZero}; use super::*; + const PACKED_BITS: usize = 128; use crate::{BitZParams, Shape}; const Q114: u128 = (1 << 114) - 11; @@ -237,7 +231,7 @@ mod tests { let claim = claim(field_weights, vec![Fq::ONE; shape.columns()]); let packed = witness(&shape, &[(1, 0), (5, 0), (127, 0), (64, 4)]); - let table = BitTable::new(shape, &packed).unwrap(); + let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); assert_eq!( fold_column(&table, &claim.row_exponents(), 0), @@ -257,7 +251,7 @@ mod tests { ); let all: Vec<(usize, usize)> = (0..shape.rows()).map(|row| (row, 0)).collect(); let packed = witness(&shape, &all); - let table = BitTable::new(shape, &packed).unwrap(); + let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); assert_eq!( fold_column(&table, &claim.row_exponents(), 0), @@ -278,7 +272,7 @@ mod tests { .map(|row| (row, 6)) .collect(); let packed = witness(&shape, &bits); - let table = BitTable::new(shape, &packed).unwrap(); + let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); let expected: u128 = (0..shape.rows()) .filter(|&row| table.bit(6, row)) diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs index 1897a00..6ca98e2 100644 --- a/crates/common/src/lib.rs +++ b/crates/common/src/lib.rs @@ -6,6 +6,7 @@ pub mod claim; pub mod fold; +pub mod matrix; pub mod opening; pub mod params; pub mod shape; @@ -16,10 +17,11 @@ pub use claim::{ClaimError, LinearClaim, Root}; pub use fold::{ Fold, FoldError, column_images, fold_column, fold_columns, reconstruct, row_images, }; +pub use matrix::BitMatrix; pub use opening::OpeningQuery; pub use params::{BitZParams, ParamsError, VirtualParams, VirtualParamsError}; pub use shape::{Shape, ShapeError}; -pub use table::{BitTable, TableError, TransposeError, TransposedBitTable}; +pub use table::{BitTable, TableError}; pub use virtual_map::{ TransposedWeights, VirtualMap, VirtualMapError, VirtualStatement, VirtualStatementError, }; diff --git a/crates/common/src/matrix.rs b/crates/common/src/matrix.rs new file mode 100644 index 0000000..2f7f513 --- /dev/null +++ b/crates/common/src/matrix.rs @@ -0,0 +1,439 @@ +//! Bit matrices in memory order, and the transpose between them. + +use std::{borrow::Cow, 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 +/// into a different word size of at least 64 bits. +/// +/// 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`. +pub(crate) type Word = u128; +/// Bits in a [`Word`]. +const BITS: usize = Word::BITS as usize; + +/// Packed bits read as `xs[dim1][dim2]`: bit `(i1, i2)` is bit +/// `(i1 << log2(dim2)) | i2` of the words, least significant bit first. Each +/// `xs[i1]` is contiguous, so that is what reads fast; [`BitMatrix::transpose`] +/// swaps which axis that is. +/// +/// Geometry only: unlike [`crate::BitTable`] it carries no [`crate::Shape`], +/// so no protocol gate applies to it, and an `xs[i1]` shorter than a word +/// shares that word with its neighbours. Only `dim2` is stored; `dim1` is +/// whatever the words hold beyond it, so [`BitMatrix::reshape`] changes one +/// number. +/// +/// The words are owned or borrowed: [`crate::BitTable::as_matrix`] reads the +/// table's words where they lie, and [`BitMatrix::transpose`] always returns +/// an owned copy. +#[derive(Debug, Clone)] +pub struct BitMatrix<'a> { + packed: Cow<'a, [Word]>, + log_dim2: usize, +} + +impl<'a> BitMatrix<'a> { + /// Wraps `packed` as `xs[dim1][2^log_dim2]`. The words must hold a power + /// of two bits, at least `dim2` of them. + pub(crate) fn new(packed: impl Into>, log_dim2: usize) -> Self { + let packed = packed.into(); + let bits = packed.len() * BITS; + assert!( + bits.is_power_of_two() && (1 << log_dim2) <= bits, + "{bits} bits do not split into rows of 2^{log_dim2}" + ); + Self { packed, log_dim2 } + } + + /// Swaps the axes, returning an owned copy: bit `(i2, i1)` of the result + /// is bit `(i1, i2)` of `self` for every `i1` and `i2`. + /// + /// # 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. + /// - **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. + /// - **Involution**: transposing the result returns the original bits. + pub fn transpose(&self) -> BitMatrix<'static> { + let (dim1, dim2) = (self.dim1(), self.dim2()); + // `dim1` is a power of two, so its trailing zeros are its log. + let log_dim1 = dim1.trailing_zeros() as usize; + BitMatrix::new(bit_transpose(&self.packed, dim1, dim2), log_dim1) + } + + /// The outer, strided axis. + pub fn dim1(&self) -> usize { + (self.packed.len() * BITS) >> self.log_dim2 + } + + /// The inner, contiguous axis. + pub fn dim2(&self) -> usize { + 1 << self.log_dim2 + } + + /// The `dim2` bits of `xs[i1]`, in ascending order. + /// + /// # Panics + /// + /// If `dim2` is below `128`: `xs[i1]` then shares a word with its + /// neighbours. + pub fn bits(&self, i1: usize) -> BitsIter<'_> { + assert!( + self.dim2() >= BITS, + "xs[{i1}] of {} bits shares a word", + self.dim2() + ); + let words = self.dim2() / BITS; + BitsIter::new(self.packed[i1 * words..(i1 + 1) * words].iter()) + } + + /// The same words read as `xs[dim1][2^log_dim2]`. No bit moves: bit + /// `(i1 << log_dim2) | i2` stays where it is, only the split between + /// `dim1` and `dim2` changes. `log_dim2` is at most the words' total + /// number of index bits. + pub fn reshape(self, log_dim2: usize) -> Self { + Self::new(self.packed, log_dim2) + } + + /// Unwraps the packed bits, copying them if they were borrowed. + pub fn into_packed(self) -> Vec { + self.packed.into_owned() + } +} + +/// Transposes `xs[dim1][dim2]`, `dim2` contiguous and both in bits, into an +/// owned `[dim2][dim1]`, one `BITS x BITS` 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. +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]; + // 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 + // ~10% slower on a 2^13-column table. + 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]; + } + + transpose_bit_block(&mut block); + + for (o, &word) in block.iter().enumerate() { + out[(g2 * BITS + o) * dim1_words + g1] = word; + } + } + } + + out +} + +/// 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`. +/// +/// 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; + while j > 0 { + let mut k = 0; + while k < BITS { + let t = ((a[k] >> j) ^ a[k | j]) & mask; + a[k | j] ^= t; + a[k] ^= t << j; + k = ((k | j) + 1) & !j; + } + j >>= 1; + mask ^= mask << j; + } +} + +/// Bit-at-a-time transpose of `GS` equal segments of `tt`: output bit `i` +/// is bit `i / GS` of segment `i % GS`. The oracle [`bit_transpose`] is +/// checked (and benchmarked) against. +#[cfg(any(test, feature = "bench"))] +fn transpose_reference(tt: &[Word]) -> Vec { + assert_eq!(tt.len() % GS, 0); + let segment_n = tt.len() / GS; + let mut segments: [_; GS] = + std::array::from_fn(|i| BitsIter::new(tt[i * segment_n..(i + 1) * segment_n].iter())); + + let mut out = Vec::with_capacity(tt.len()); + let mut word: Word = 0; + let mut offset = 0; + 'outer: loop { + for s in &mut segments { + let Some(b) = s.next() else { break 'outer }; + // LSB first, like `StepIter`. + word |= b << offset; + offset += 1; + if offset == BITS { + out.push(word); + word = 0; + offset = 0; + } + } + } + debug_assert_eq!(offset, 0); + + out +} + +/// Benchmark-only access to the private transposes; see `benches/table.rs`. +#[cfg(feature = "bench")] +#[doc(hidden)] +pub mod bench { + pub fn bit_transpose(xs: &[super::Word], dim1: usize, dim2: usize) -> Vec { + super::bit_transpose(xs, dim1, dim2) + } + + pub fn transpose_reference(tt: &[super::Word]) -> Vec { + super::transpose_reference::(tt) + } +} + +pub type BitsIter<'a> = StepIter<'a, 1>; + +/// Iterator over packed bits `S` at a time, least significant first — see +/// [`BitMatrix::bits`]. +pub struct StepIter<'a, const S: usize> { + elements: std::slice::Iter<'a, Word>, + /// Bits not yet emitted; the next one to emit is the LSB. + word: Word, + 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" + ); + const MASK: Word = (1 as Word).unbounded_shl(S as u32).wrapping_sub(1); + fn new(iter: Iter<'a, Word>) -> Self { + let () = Self::CHECK_SIZE; + Self { + elements: iter, + word: 0, + remaining: 0, + } + } +} + +impl Iterator for StepIter<'_, S> { + type Item = Word; + + fn next(&mut self) -> Option { + if self.remaining == 0 { + let element = self.elements.next()?; + self.word = *element; + self.remaining = Word::BITS as u8; + } + let bitset = self.word & Self::MASK; + self.word = self.word.unbounded_shr(S as u32); + self.remaining -= S as u8; + Some(bitset) + } + + fn size_hint(&self) -> (usize, Option) { + let len = self.len(); + (len, Some(len)) + } +} + +impl ExactSizeIterator for StepIter<'_, S> { + fn len(&self) -> usize { + (self.remaining as usize + self.elements.len() * BITS) / S + } +} + +#[cfg(test)] +mod tests { + use super::*; + + /// A word with bits spread over its whole width, distinct for every `i`. + fn pseudo_random_word(i: usize) -> Word { + (i as Word).wrapping_mul(0x2545_F491_4F6C_DD1D_9E37_79B9_7F4A_7C15) ^ 0xA5A5 + } + + /// A pseudo-random `xs[2^log_dim1][2^log_dim2]`. + fn pseudo_random_matrix(log_dim1: usize, log_dim2: usize) -> BitMatrix<'static> { + let packed: Vec = (0..(1 << (log_dim1 + log_dim2)) / BITS) + .map(pseudo_random_word) + .collect(); + BitMatrix::new(packed, log_dim2) + } + + /// The bit `xs[i1][i2]`, read on its own: the oracle the word-level + /// reads and the transpose are checked against. + fn bit(matrix: &BitMatrix, i1: usize, i2: usize) -> bool { + assert!(i1 < matrix.dim1(), "i1 {i1} is outside dim1"); + assert!(i2 < matrix.dim2(), "i2 {i2} is outside dim2"); + + let index = (i1 << matrix.log_dim2) | i2; + (matrix.packed[index / BITS] >> (index % BITS)) & 1 == 1 + } + + #[test] + fn a_bit_lands_where_the_memory_order_says_it_does() { + // xs[1][130] of an xs[2][256]: bit 256 + 130, which is word 3, bit 2. + let mut packed = vec![0 as Word; 4]; + packed[3] = 1 << 2; + let matrix = BitMatrix::new(packed, 8); + assert_eq!((matrix.dim1(), matrix.dim2()), (2, 256)); + + for i1 in 0..matrix.dim1() { + for i2 in 0..matrix.dim2() { + assert_eq!(bit(&matrix, i1, i2), (i1, i2) == (1, 130), "xs[{i1}][{i2}]"); + } + } + } + + /// `dim2` of one word and of two. + #[test] + fn bits_agrees_with_the_bit_accessor() { + for log_dim2 in [7, 8] { + let matrix = pseudo_random_matrix(12 - log_dim2, log_dim2); + + for i1 in 0..matrix.dim1() { + let iter = matrix.bits(i1); + assert_eq!(iter.len(), matrix.dim2()); + let bits: Vec = iter.map(|b| b == 1).collect(); + let expected: Vec = + (0..matrix.dim2()).map(|i2| bit(&matrix, i1, i2)).collect(); + assert_eq!(bits, expected, "log_dim2 {log_dim2}, i1 {i1}"); + } + } + } + + #[test] + fn transpose_bit_block_matches_a_brute_force_reference() { + let mut original = [0 as Word; BITS]; + for (i, word) in original.iter_mut().enumerate() { + *word = pseudo_random_word(i); + } + + let mut transposed = original; + transpose_bit_block(&mut transposed); + + for (i, &transposed_word) in transposed.iter().enumerate() { + for (j, &original_word) in original.iter().enumerate() { + assert_eq!( + (transposed_word >> j) & 1, + (original_word >> i) & 1, + "word {i} bit {j}" + ); + } + } + } + + /// `bit_transpose` of `GS` rows of `words` words each against the + /// bit-at-a-time reference. + fn check_bit_transpose(words: usize) { + let tt: Vec = (0..GS * words).map(pseudo_random_word).collect(); + assert_eq!( + bit_transpose(&tt, GS, words * BITS), + transpose_reference::(&tt), + "dim1 {GS}, {words} words per row" + ); + } + + #[test] + fn bit_transpose_matches_the_reference() { + // One block, then two along each axis and along both, so block + // placement is exercised as well as the block contents. + check_bit_transpose::<128>(1); + check_bit_transpose::<128>(2); + check_bit_transpose::<256>(1); + check_bit_transpose::<256>(2); + } + + /// `xs[512][256]`: two word blocks along each axis, so the block-tiling + /// loop in `bit_transpose` runs more than once both ways. + 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); + + let transposed = matrix.transpose(); + assert_eq!(transposed.dim1(), matrix.dim2()); + assert_eq!(transposed.dim2(), matrix.dim1()); + + for i1 in 0..matrix.dim1() { + for i2 in 0..matrix.dim2() { + assert_eq!( + bit(&transposed, i2, i1), + bit(&matrix, i1, i2), + "xs[{i1}][{i2}] vs its transpose" + ); + } + } + } + + #[test] + fn transposing_twice_recovers_the_original_bits() { + let matrix = pseudo_random_matrix(TRANSPOSE_LOG_DIM1, TRANSPOSE_LOG_DIM2); + + let roundtripped = matrix.transpose().transpose(); + + assert_eq!(roundtripped.dim2(), matrix.dim2()); + assert_eq!(roundtripped.into_packed(), matrix.into_packed()); + } + + /// Every split of the index, from one long `xs[0]` down to `dim2 = 1`. + #[test] + fn reshape_leaves_every_bit_in_place() { + let matrix = pseudo_random_matrix(TRANSPOSE_LOG_DIM1, TRANSPOSE_LOG_DIM2); + let log_bits = TRANSPOSE_LOG_DIM1 + TRANSPOSE_LOG_DIM2; + + for log_dim2 in 0..=log_bits { + let reshaped = matrix.clone().reshape(log_dim2); + assert_eq!(reshaped.dim1(), 1 << (log_bits - log_dim2)); + assert_eq!(reshaped.dim2(), 1 << log_dim2); + + for index in 0..1 << log_bits { + assert_eq!( + bit(&reshaped, index >> log_dim2, index & (reshaped.dim2() - 1)), + bit( + &matrix, + index >> TRANSPOSE_LOG_DIM2, + index & (matrix.dim2() - 1) + ), + "log_dim2 {log_dim2}, bit {index}" + ); + } + } + } + + #[test] + #[should_panic(expected = "do not split")] + fn reshape_rejects_a_dim2_longer_than_the_words() { + let matrix = pseudo_random_matrix(TRANSPOSE_LOG_DIM1, TRANSPOSE_LOG_DIM2); + matrix.reshape(TRANSPOSE_LOG_DIM1 + TRANSPOSE_LOG_DIM2 + 1); + } +} diff --git a/crates/common/src/params.rs b/crates/common/src/params.rs index e8ca41e..51ee4a0 100644 --- a/crates/common/src/params.rs +++ b/crates/common/src/params.rs @@ -8,6 +8,16 @@ use crate::{BitTable, Shape, TableError, VirtualMap}; /// A parameter set one of the pre-claim gates rejects. #[derive(Debug, Clone, Copy, PartialEq, Eq)] 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. + ColumnCountTooNarrow, /// `(k_1 + 1)(Q - 1)` reaches `ord(g)`, so two folds could collide in the /// exponent. FoldBoundExceeded, @@ -35,6 +45,11 @@ impl BitZParams { // `Fq` asserts Q is an odd prime below 2^126 on its own behalf, so // no modulus gate is needed here. + // A limit of the transpose, not of the protocol; see the variant. + if shape.columns() < BitTable::BITS { + return Err(ParamsError::ColumnCountTooNarrow); + } + // `ord(g) > (k_1 + 1)(Q - 1)`, the paper's requisite. A fold is an // integer at most `k_1 (Q - 1)` while the value it is compared against // is at most `Q - 1`, so the two differ by at most the sum; the @@ -64,7 +79,7 @@ impl BitZParams { /// The only way to build a [`BitTable`], so a table can never be shaped by /// anything but a checked parameter set. pub fn table<'a>(&self, packed: &'a [F128]) -> Result, TableError> { - BitTable::new(self.shape, packed) + BitTable::new(self.shape, bytemuck::cast_slice(packed)) } pub fn shape(&self) -> &Shape { @@ -163,7 +178,7 @@ impl VirtualParams { /// Views the committed witness, which is `f` and not the vector the claim /// is about. pub fn table<'a>(&self, packed: &'a [F128]) -> Result, TableError> { - BitTable::new(self.committed, packed) + BitTable::new(self.committed, bytemuck::cast_slice(packed)) } } @@ -305,6 +320,21 @@ mod tests { ); } + #[test] + fn rejects_a_table_too_narrow_to_transpose() { + // Small enough for `2^16` rows, so the column count is the only gate + // that can fail: `Q114` already caps `t` at 13, and with it `s` at 9 + // or more. + const Q100: u128 = (1 << 100) - 15; + let params_q100 = |shape| BitZParams::::new(shape, smallest_generator()); + + assert!(params_q100(Shape::new(15, 7).unwrap()).is_ok()); + assert_eq!( + params_q100(Shape::new(16, 6).unwrap()).err(), + Some(ParamsError::ColumnCountTooNarrow) + ); + } + #[test] fn rejects_a_generator_of_partial_order() { assert_eq!( diff --git a/crates/common/src/table.rs b/crates/common/src/table.rs index 6119698..04edaad 100644 --- a/crates/common/src/table.rs +++ b/crates/common/src/table.rs @@ -1,14 +1,9 @@ //! The committed bit table. -use crate::{Shape, ShapeError}; -use field::F128; - -/// Bits in a packed element, and the shift that divides an index into element -/// and offset. -pub(crate) const PACKED_BITS: usize = 128; -const PACKED_SHIFT: u32 = 7; -/// Bits in each half of a packed element. -const HALF_BITS: usize = 64; +use crate::{ + Shape, + matrix::{BitMatrix, Word}, +}; /// A witness that does not match its shape. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -17,21 +12,6 @@ pub enum TableError { BitCountMismatch, } -/// [`BitTable::transpose`] cannot lay out its result. -#[derive(Debug, Clone, Copy, PartialEq, Eq)] -pub enum TransposeError { - /// The source has fewer than `2^7 = 128` columns. - /// - /// Transposing swaps rows and columns, so the result's row count is the - /// source's column count -- and every `BitTable`'s row count must be a - /// whole number of `128`-bit packed elements (the same rule - /// `ShapeError::RowIndexTooNarrow` enforces when a shape is built). - /// `Shape::new` only enforces that floor on `log_rows`, never on - /// `log_columns`, so a source table can satisfy it and still be too - /// narrow to transpose. - ColumnCountTooNarrow, -} - /// The committed bits `B[c][b]`, read through a shape. /// /// A borrowed view over the prover's packed witness: at `m = 35` a byte per @@ -47,17 +27,19 @@ pub enum TransposeError { #[derive(Debug, Clone, Copy)] pub struct BitTable<'a> { shape: Shape, - packed: &'a [F128], + packed: &'a [Word], } impl<'a> BitTable<'a> { + pub const BITS: usize = Word::BITS as usize; + /// Wraps `packed` in `shape`, least significant bit first inside `lo`. /// /// Crate-private: [`crate::BitZParams::table`] is the only way in, so a /// table is always shaped by a checked parameter set. - pub(crate) fn new(shape: Shape, packed: &'a [F128]) -> Result { - // `m >= 22`, so the bit count is always a whole number of elements. - if packed.len() != (1 << shape.log_bits()) / PACKED_BITS { + pub(crate) fn new(shape: Shape, packed: &'a [Word]) -> Result { + // `m >= 22`, so the bit count is always a whole number of words. + if packed.len() != (1 << shape.log_bits()) / BitTable::BITS { return Err(TableError::BitCountMismatch); } Ok(Self { shape, packed }) @@ -78,15 +60,13 @@ impl<'a> BitTable<'a> { column < self.shape.columns(), "column {column} is outside the table" ); + let index = (column << self.shape.log_rows()) | row; - let element = self.packed[index >> PACKED_SHIFT]; - let offset = index % PACKED_BITS; - let half = if offset < HALF_BITS { - element.lo - } else { - element.hi - }; - (half >> (offset % HALF_BITS)) & 1 == 1 + // /% on compile time constant should be properly optimised away + let element = self.packed[index / BitTable::BITS]; + let offset = index % BitTable::BITS; + + ((element >> offset) & 1) == 1 } /// The `2^(t-7)` elements holding one column, in ascending row order. @@ -95,193 +75,16 @@ impl<'a> BitTable<'a> { /// element boundary and spans whole elements. The fold walks these directly /// rather than calling [`BitTable::bit`] once per row: at `m = 35` that is /// `2^35` calls, each of them a bounds check and two shifts. - pub fn column(&self, column: usize) -> &'a [F128] { - let elements = self.shape.rows() / PACKED_BITS; + pub fn column(&self, column: usize) -> &'a [Word] { + let elements = self.shape.rows() / BitTable::BITS; &self.packed[column * elements..(column + 1) * elements] } - /// The column's bits, in ascending row order. - pub fn column_bits(&self, column: usize) -> ColumnBits<'a> { - ColumnBits { - elements: self.column(column).iter(), - word: 0, - hi_pending: None, - remaining: 0, - } - } - - /// Swaps rows and columns, returning an owned copy: `result.bit(b, c) == - /// self.bit(c, b)` for every row `b` and column `c` of `self`. - /// - /// Built for callers like a leaf construction that must walk the table - /// row by row: the packing is column-major (see the type docs), so a - /// row-major walk over `self` is a scatter, one `bit()` call and cache - /// miss per cell. Transposing once up front turns that into the same - /// sequential access [`BitTable::column_bits`] already gives per column, - /// just over what were originally rows. - /// - /// # Constraints - /// - /// - **Column count**: `self.shape().log_columns()` must be at least `7` - /// (`PACK_BITS`), i.e. at least `128` columns. See - /// [`TransposeError::ColumnCountTooNarrow`] -- this is the only way the - /// operation fails. - /// - **Commitment-size window**: never a separate concern. The total bit - /// count `shape().log_bits() = log_rows + log_columns` is a sum, so - /// swapping its two terms leaves it unchanged; the result automatically - /// sits in `[MIN_LOG_BITS, MAX_LOG_BITS]` whenever `self` did. - /// - **Allocation**: always a full out-of-place copy -- a fresh buffer - /// the same size as `self`'s packed witness (up to 4 GiB at `m = 35`), - /// never a view over the original. - /// - **Involution**: transposing the result returns `self`'s bits - /// exactly. The forward direction is the only place the column-count - /// constraint can fail: once it holds, the result's row count is the - /// source's (old) column count and its column count is the source's - /// (old) row count, and every `BitTable` already has `log_rows >= 7` - /// unconditionally -- so the second transpose is never the one that - /// rejects. - pub fn transpose(&self) -> Result { - let shape = - Shape::new(self.shape.log_columns(), self.shape.log_rows()).map_err(|error| { - debug_assert_eq!( - error, - ShapeError::RowIndexTooNarrow, - "log_bits is preserved by swapping row/column axes, so the \ - commitment-size window cannot be what rejected the swapped shape" - ); - TransposeError::ColumnCountTooNarrow - })?; - - let row_groups = self.shape.rows() / PACKED_BITS; - let column_groups = self.shape.columns() / PACKED_BITS; - debug_assert_eq!(row_groups * column_groups * PACKED_BITS, self.packed.len()); - - let mut packed = F128::zeroed_vec(self.packed.len()); - let mut block = [0u128; PACKED_BITS]; - for row_group in 0..row_groups { - for column_group in 0..column_groups { - for (lane, word) in block.iter_mut().enumerate() { - let column = column_group * PACKED_BITS + lane; - let element = self.packed[column * row_groups + row_group]; - *word = u128::from(element.lo) | (u128::from(element.hi) << HALF_BITS); - } - - transpose_bit_block(&mut block); - - for (lane, &word) in block.iter().enumerate() { - let row = row_group * PACKED_BITS + lane; - packed[row * column_groups + column_group] = - F128::new(word as u64, (word >> HALF_BITS) as u64); - } - } - } - - Ok(TransposedBitTable { shape, packed }) - } -} - -/// Transposes a `PACKED_BITS x PACKED_BITS` bit matrix stored as -/// `PACKED_BITS` words of `PACKED_BITS` 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)` bit-at-a-time approach `BitTable::bit` would need to do the same -/// work. -fn transpose_bit_block(a: &mut [u128; PACKED_BITS]) { - let mut j = PACKED_BITS / 2; - let mut mask: u128 = (1u128 << j) - 1; - while j > 0 { - let mut k = 0; - while k < PACKED_BITS { - let t = ((a[k] >> j) ^ a[k | j]) & mask; - a[k | j] ^= t; - a[k] ^= t << j; - k = ((k | j) + 1) & !j; - } - j >>= 1; - mask ^= mask << j; - } -} - -/// An owned, transposed copy of a [`BitTable`]'s bits -- see -/// [`BitTable::transpose`], which is the only way to build one. -#[derive(Debug)] -pub struct TransposedBitTable { - shape: Shape, - packed: Vec, -} - -impl TransposedBitTable { - /// The transposed shape: rows and columns swapped from the source. - pub fn shape(&self) -> &Shape { - &self.shape - } - - /// Borrows the transposed bits as an ordinary [`BitTable`]. - /// - /// Infallible: [`BitTable::transpose`] already sized this buffer to - /// match `shape`. - pub fn as_table(&self) -> BitTable<'_> { - BitTable::new(self.shape, &self.packed) - .expect("a TransposedBitTable's shape and packed length always agree") - } - - /// Unwraps the packed, transposed bits. - pub fn into_packed(self) -> Vec { - self.packed - } -} - -/// Iterator over one column's bits, in ascending row order — see -/// [`BitTable::column_bits`]. -pub struct ColumnBits<'a> { - elements: std::slice::Iter<'a, F128>, - /// Bits not yet emitted; the next one to emit is the LSB. - word: u64, - /// `hi` of the current element, once `lo` has been fully shifted out of - /// `word` but before it takes `word`'s place. - hi_pending: Option, - /// How many low bits of `word` are still valid, i.e. still unemitted. - remaining: u32, -} - -impl Iterator for ColumnBits<'_> { - type Item = bool; - - fn next(&mut self) -> Option { - match (self.remaining, self.hi_pending) { - (0, Some(hi)) => { - self.word = hi; - self.hi_pending = None; - self.remaining = HALF_BITS as u32; - } - (0, None) => { - let element = self.elements.next()?; - self.word = element.lo; - self.hi_pending = Some(element.hi); - self.remaining = HALF_BITS as u32; - } - _ => {} - } - let bit = self.word & 1 == 1; - self.word >>= 1; - self.remaining -= 1; - Some(bit) - } - - fn size_hint(&self) -> (usize, Option) { - let len = self.len(); - (len, Some(len)) - } -} - -impl ExactSizeIterator for ColumnBits<'_> { - fn len(&self) -> usize { - self.remaining as usize - + self.hi_pending.map_or(0, |_| HALF_BITS) - + self.elements.len() * PACKED_BITS + /// The table as a [`BitMatrix`] with `dim1` the columns and `dim2` the + /// rows, borrowing the packed words: bit `(c, b)` of the result is + /// `self.bit(c, b)`. No bit moves. + pub fn as_matrix(&self) -> BitMatrix<'a> { + BitMatrix::new(self.packed, self.shape.log_rows()) } } @@ -289,6 +92,12 @@ impl ExactSizeIterator for ColumnBits<'_> { mod tests { use num_traits::ConstZero; + // We want to use compare against + use field::F128; + + /// Bits in a packed element. + const PACKED_BITS: usize = 128; + use super::*; /// `m = 22`: 128 rows per column, 32768 columns. @@ -317,7 +126,7 @@ mod tests { let shape = small_shape(); let set = [(0, 0), (2, 0), (3, 0), (65, 0), (1, 9)]; let packed = with_bits(&shape, &set); - let table = BitTable::new(shape, &packed).unwrap(); + let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); for column in [0, 9, 10] { for row in 0..shape.rows() { @@ -338,7 +147,7 @@ mod tests { let shape = small_shape(); let set = [(0, 0), (5, 0), (127, 0), (64, 3), (2, 9)]; let packed = with_bits(&shape, &set); - let table = BitTable::new(shape, &packed).unwrap(); + let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); assert_eq!(packed.len(), 1 << shape.log_packed_len()); @@ -361,7 +170,7 @@ mod tests { fn rejects_a_witness_of_the_wrong_length() { let shape = small_shape(); assert_eq!( - BitTable::new(shape, &[F128::ZERO; 8]).err(), + BitTable::new(shape, bytemuck::cast_slice(&[F128::ZERO; 8])).err(), Some(TableError::BitCountMismatch) ); } @@ -370,27 +179,11 @@ mod tests { fn a_column_is_element_aligned_and_spans_whole_elements() { let shape = small_shape(); let packed = with_bits(&shape, &[(0, 3), (2, 3), (64, 3)]); - let table = BitTable::new(shape, &packed).unwrap(); + let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); - assert_eq!(table.column(3).len(), shape.rows() / PACKED_BITS); - assert_eq!(table.column(3), [F128::new(0b101, 1)]); - assert!(table.column(4).iter().all(|element| *element == F128::ZERO)); - } - - #[test] - fn column_bits_agrees_with_the_bit_accessor() { - let shape = small_shape(); - let rows: Vec<(usize, usize)> = (0..shape.rows()).step_by(7).map(|row| (row, 2)).collect(); - let packed = with_bits(&shape, &rows); - let table = BitTable::new(shape, &packed).unwrap(); - - let iter = table.column_bits(2); - assert_eq!(iter.len(), shape.rows()); - let bits: Vec = iter.collect(); - assert_eq!(bits.len(), shape.rows()); - for (row, bit) in bits.into_iter().enumerate() { - assert_eq!(bit, table.bit(2, row), "row {row}"); - } + assert_eq!(table.column(3).len(), shape.rows() / BitTable::BITS); + assert_eq!(table.column(3), [0b101 | (1 << 64)]); + assert!(table.column(4).iter().all(|&element| element == 0)); } #[test] @@ -398,128 +191,68 @@ mod tests { let shape = small_shape(); let rows: Vec<(usize, usize)> = (0..shape.rows()).step_by(7).map(|row| (row, 2)).collect(); let packed = with_bits(&shape, &rows); - let table = BitTable::new(shape, &packed).unwrap(); + let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); - for (index, element) in table.column(2).iter().enumerate() { - let bits = element.to_u128(); - for offset in 0..PACKED_BITS { - let row = index * PACKED_BITS + offset; - assert_eq!(table.bit(2, row), (bits >> offset) & 1 == 1, "row {row}"); + for (index, &element) in table.column(2).iter().enumerate() { + for offset in 0..BitTable::BITS { + let row = index * BitTable::BITS + offset; + assert_eq!(table.bit(2, row), (element >> offset) & 1 == 1, "row {row}"); } } } - /// A cheap, deterministic stand-in for randomness: no `rand` dependency, - /// but different `seed`s still produce unrelated-looking bit patterns. - fn pseudo_random_bit(seed: u64) -> bool { - (seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) >> 63) & 1 == 1 - } - + /// The matrix has the table's columns along `dim1`, each one contiguous. #[test] - fn transpose_bit_block_matches_a_brute_force_reference() { - let mut original = [0u128; PACKED_BITS]; - for (i, word) in original.iter_mut().enumerate() { - let hi = (i as u64).wrapping_mul(0xD1B5_4A32_D192_ED03) ^ 0xA5A5; - let lo = (i as u64).wrapping_mul(0x2545_F491_4F6C_DD1D); - *word = u128::from(lo) | (u128::from(hi) << 64); - } - - let mut transposed = original; - transpose_bit_block(&mut transposed); - - for (i, &transposed_word) in transposed.iter().enumerate() { - for (j, &original_word) in original.iter().enumerate() { - assert_eq!( - (transposed_word >> j) & 1, - (original_word >> i) & 1, - "word {i} bit {j}" - ); - } - } - } - - /// Two row groups (`log_rows = 8`), two column groups (`log_columns = - /// 15`), so the block-tiling loop in `transpose` runs more than once on - /// both axes. - fn transpose_shape() -> Shape { - Shape::new(8, 15).unwrap() - } - - fn pseudo_random_table(shape: &Shape) -> Vec { - let mut packed = vec![F128::ZERO; (1 << shape.log_bits()) / PACKED_BITS]; - for column in 0..shape.columns() { - for row in 0..shape.rows() { - let seed = (column as u64).wrapping_mul(2654435761) ^ (row as u64); - if pseudo_random_bit(seed) { - let index = (column << shape.log_rows()) | row; - let element = &mut packed[index >> PACKED_SHIFT]; - let offset = index % PACKED_BITS; - if offset < HALF_BITS { - element.lo |= 1u64 << offset; - } else { - element.hi |= 1u64 << (offset - HALF_BITS); - } - } - } + fn the_matrix_has_the_columns_contiguous() { + let shape = small_shape(); + let set = [(0, 0), (127, 0), (64, 3), (5, 9), (1, shape.columns() - 1)]; + let packed = with_bits(&shape, &set); + let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); + + let matrix = table.as_matrix(); + assert_eq!(matrix.dim1(), shape.columns()); + assert_eq!(matrix.dim2(), shape.rows()); + + for column in [0, 3, 9, shape.columns() - 1] { + let bits: Vec = matrix.bits(column).map(|b| b == 1).collect(); + let expected: Vec = (0..shape.rows()) + .map(|row| table.bit(column, row)) + .collect(); + assert_eq!(bits, expected, "column {column}"); } - packed } + /// The transpose has the table's rows along `dim1`, each one contiguous. #[test] - fn transpose_swaps_row_and_column_bits() { - let shape = transpose_shape(); - let packed = pseudo_random_table(&shape); - let table = BitTable::new(shape, &packed).unwrap(); - - let transposed = table.transpose().unwrap(); - assert_eq!(transposed.shape().log_rows(), shape.log_columns()); - assert_eq!(transposed.shape().log_columns(), shape.log_rows()); - - let transposed = transposed.as_table(); - - // Full columns at both ends of each axis, including a block - // boundary (128 rows per group here), rather than every column -- - // this shape alone already has 2^15 of them. - for column in [0, 1, shape.columns() / 2, shape.columns() - 1] { - for row in [0, 1, PACKED_BITS - 1, PACKED_BITS, shape.rows() - 1] { - assert_eq!( - transposed.bit(row, column), - table.bit(column, row), - "B[{column}][{row}] vs its transpose" - ); - } + fn the_transpose_has_the_rows_contiguous() { + let shape = small_shape(); + let set = [(0, 0), (127, 0), (64, 3), (5, 9), (1, shape.columns() - 1)]; + let packed = with_bits(&shape, &set); + let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); + + let transposed = table.as_matrix().transpose(); + assert_eq!(transposed.dim1(), shape.rows()); + assert_eq!(transposed.dim2(), shape.columns()); + + for row in [0, 1, 5, 64, 127] { + let bits: Vec = transposed.bits(row).map(|b| b == 1).collect(); + let expected: Vec = (0..shape.columns()) + .map(|column| table.bit(column, row)) + .collect(); + assert_eq!(bits, expected, "row {row}"); } } #[test] - fn transposing_twice_recovers_the_original_bits() { - let shape = transpose_shape(); - let packed = pseudo_random_table(&shape); - let table = BitTable::new(shape, &packed).unwrap(); - - let roundtripped = table - .transpose() - .unwrap() - .as_table() - .transpose() - .unwrap() - .into_packed(); - - assert_eq!(roundtripped, packed); - } - - #[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 packed element -- even - // though this shape is perfectly admissible for `BitTable` itself. + // transposed row would need to fill one word -- 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, &packed).unwrap(); + let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); - assert_eq!( - table.transpose().err(), - Some(TransposeError::ColumnCountTooNarrow) - ); + table.as_matrix().transpose(); } } diff --git a/crates/field/Cargo.toml b/crates/field/Cargo.toml index 13807f9..7aa7aa9 100644 --- a/crates/field/Cargo.toml +++ b/crates/field/Cargo.toml @@ -6,7 +6,7 @@ rust-version.workspace = true license.workspace = true [dependencies] -bytemuck = { workspace = true, features = ["extern_crate_alloc"] } +bytemuck = { workspace = true, features = ["extern_crate_alloc", "derive"] } num-traits = { workspace = true } crypto-primitives = { workspace = true } crypto-primitives-proc-macros = { workspace = true } diff --git a/crates/field/src/gf128.rs b/crates/field/src/gf128.rs index c5e44e0..2f15c1b 100644 --- a/crates/field/src/gf128.rs +++ b/crates/field/src/gf128.rs @@ -12,6 +12,7 @@ use std::fmt::{Display, Formatter, Result as FmtResult}; use std::iter::{Product, Sum}; use std::ops::{Add, AddAssign, Div, DivAssign, Mul, MulAssign, Neg, Sub, SubAssign}; +use bytemuck::NoUninit; use crypto_primitives::{WithAssociatedInteger, WithExtensionDegree}; use crypto_primitives_proc_macros::InfallibleCheckedOp; use num_traits::{ @@ -59,6 +60,7 @@ pub const REDUCTION: u64 = 0x87; #[derive(Clone, Copy, Debug, Default, PartialEq, Eq, Hash, InfallibleCheckedOp)] #[infallible_checked_unary_op((CheckedNeg, neg))] #[infallible_checked_binary_op((CheckedAdd, add), (CheckedSub, sub), (CheckedMul, mul))] +#[derive(NoUninit)] #[repr(C, align(16))] pub struct F128 { pub lo: u64, diff --git a/crates/gkr/Cargo.toml b/crates/gkr/Cargo.toml index 7781751..26913ec 100644 --- a/crates/gkr/Cargo.toml +++ b/crates/gkr/Cargo.toml @@ -6,12 +6,14 @@ rust-version.workspace = true license.workspace = true [dependencies] -field = { path = "../field", features = ["spongefish"] } -transcript = {path ="../transcript"} +field = { workspace= true } +transcript = { workspace = true} tracing = { workspace = true } -poly = { path = "../poly" } +poly = { workspace = true } num-traits = { workspace = true } +common = {workspace=true} rayon = { workspace = true } + [dev-dependencies] proptest = "1.11.0" diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs index 1f390e6..a00ebfe 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs @@ -1,5 +1,6 @@ use std::collections::VecDeque; +use common::BitTable; use field::{F128, Wide256}; use num_traits::{ConstOne, ConstZero}; use rayon::prelude::*; @@ -342,8 +343,25 @@ pub struct GrandProductCircuit { } impl GrandProductCircuit { - // Fails if the leafs are 0. - pub fn new(mut leafs: Vec) -> Self { + pub fn new(row_images: &[F128], bittable: &BitTable) -> Self { + // 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 }; + } + } + + 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); } @@ -428,7 +446,7 @@ 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::new(leaves).batched_eval(1); + let (top, mut witnesses) = GrandProductCircuit::from_leafs(leaves).batched_eval(1); prop_assert_eq!(product(&top), expected); @@ -450,7 +468,7 @@ mod tests { let mut padded = leaves.clone(); padded.resize(padded_len, Field::ONE); - let circuit = GrandProductCircuit::new(leaves); + let circuit = GrandProductCircuit::from_leafs(leaves); let (top, _witnesses) = circuit.batched_eval(groups); prop_assert_eq!(top.len(), groups); @@ -511,7 +529,7 @@ mod tests { [Field::ZERO, Field::ONE], [Field::ONE, Field::ZERO], ] { - let circuit = GrandProductCircuit::new(leaves.clone()); + let circuit = GrandProductCircuit::from_leafs(leaves.clone()); let (output, witnesses) = circuit.batched_eval(4); let claim = mle(output, &point); let instance = (leaves.clone(), point.to_vec()); @@ -557,7 +575,7 @@ mod tests { } pub fn prove(input: Vec, log_groups: usize) -> (Vec, transcript::Proof) { - let circuit = GrandProductCircuit::new(input); + let circuit = GrandProductCircuit::from_leafs(input); // Fake hashing let n = circuit.leafs.len() as u128; @@ -581,7 +599,7 @@ mod tests { } pub fn verify(input: Vec, output: Vec, proof: transcript::Proof) -> bool { - let circuit = GrandProductCircuit::new(input); + let circuit = GrandProductCircuit::from_leafs(input); // Fake hashing let instance = (&output, circuit.leafs.len() as u128); let mut verifier = transcript::build_verifier("gkr", &instance, &proof); diff --git a/crates/prover/src/reduce.rs b/crates/prover/src/reduce.rs index 4364cf2..17e03b7 100644 --- a/crates/prover/src/reduce.rs +++ b/crates/prover/src/reduce.rs @@ -5,7 +5,7 @@ //! `sum(row, column) u1[row] * u2[column] * table.bit(column, row)`. //! The caller must discharge this claim through the commitment opening. -use common::{BitTable, ClaimError, Fold, LinearClaim, OpeningQuery, TransposeError}; +use common::{BitTable, ClaimError, Fold, LinearClaim, OpeningQuery}; use field::F128; use gkr::{GrandProductCircuit, gpgkr_prove}; use num_traits::ConstOne; @@ -15,39 +15,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 { - let columns = table.shape().columns(); - let dim = columns * table.shape().rows(); - let mut leafs = F128::zeroed_vec(dim); - - // TODO optimisation: Handle the leafs and the two layers above it lazily. - // Columns occupy the low index bits, so each product tree reduces one column. - match table.transpose() { - Ok(transposed) => { - // Transpose wide tables so each row can be read sequentially. - let transposed = transposed.as_table(); - for (b, &row_image) in fold.row_images.iter().enumerate() { - let leafs = &mut leafs[b * columns..(b + 1) * columns]; - for (leaf, bit) in leafs.iter_mut().zip(transposed.column_bits(b)) { - *leaf = if bit { row_image } else { F128::ONE }; - } - } - } - Err(TransposeError::ColumnCountTooNarrow) => { - // Narrow tables cannot form packed columns after transposition. - for (b, &row_image) in fold.row_images.iter().enumerate() { - let leafs = &mut leafs[b * columns..(b + 1) * columns]; - for (c, leaf) in leafs.iter_mut().enumerate() { - *leaf = if table.bit(c, b) { - row_image - } else { - F128::ONE - }; - } - } - } - } - - GrandProductCircuit::new(leafs) + GrandProductCircuit::new(&fold.row_images, table) } /// Reduces the grand-product circuit to a factored claim on the committed bits. @@ -76,7 +44,7 @@ pub fn gkr_reduce( .row_images .iter() .zip(poly::eq_table(&alfa_b)) - .map(|(a, b)| (*a - F128::ONE) * b) // Does the later step benefit from wide mul? + .map(|(a, b)| (*a - F128::ONE) * b) .collect(); let u2 = eq_table(&alfa_c); @@ -152,8 +120,9 @@ mod order_check_ai_test { fn factored_claim_matches_for_narrow_tables() { const Q100: u128 = (1 << 100) - 15; - // Cover one column and both sides of the 128-column transpose boundary. - for log_columns in [0, 6, 7] { + // The narrowest tables the transpose takes: one block of 128 columns, + // then two. + for log_columns in [7, 8] { let shape = Shape::new(22 - log_columns, log_columns).unwrap(); let params = BitZParams::::new(shape, smallest_generator()).unwrap(); let packed = packed_witness(&shape, |column, row| { diff --git a/crates/verifier/src/reduce.rs b/crates/verifier/src/reduce.rs index eec8a59..f684d5b 100644 --- a/crates/verifier/src/reduce.rs +++ b/crates/verifier/src/reduce.rs @@ -56,7 +56,7 @@ mod round_trip_ai_test { use common::{BitZParams, Fold, Shape}; use field::gf128::smallest_generator; use gkr::{GrandProductCircuit, gpgkr_prove}; - use num_traits::{ConstOne, ConstZero, identities::Zero}; + use num_traits::{ConstOne, ConstZero}; use super::*; @@ -100,18 +100,7 @@ mod round_trip_ai_test { .map(|i| F128::from((i as u128 + 3) * 0xABCDEF12345u128 + 1)) .collect(); - // Build affine leaves directly so the verifier test needs no prover dependency. - let mut leafs = vec![F128::zero(); shape.columns() * shape.rows()]; - for b in 0..shape.rows() { - for c in 0..shape.columns() { - leafs[b * shape.columns() + c] = if table.bit(c, b) { - row_images[b] - } else { - F128::ONE - }; - } - } - let circuit = GrandProductCircuit::new(leafs); + let circuit = GrandProductCircuit::new(&row_images, &table); let (top_layer, witnesses) = circuit.batched_eval(shape.columns()); // Fold::new evaluates top_layer at zeta to obtain GKR's initial claim.