From f869b88233a5f580a9135bc82ba615b7bddfb04b Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Wed, 23 Sep 2026 16:43:09 +0800 Subject: [PATCH 01/22] common: back BitTable with a machine-sized Word instead of F128 --- Cargo.lock | 2 + crates/common/Cargo.toml | 1 + crates/common/src/fold.rs | 28 +++++------ crates/common/src/table.rs | 92 ++++++++++++++++++------------------- crates/field/Cargo.toml | 2 +- crates/field/src/gf128.rs | 2 + crates/gkr/Cargo.toml | 8 ++-- crates/gkr/src/lib.rs | 1 + crates/prover/src/reduce.rs | 2 +- 9 files changed, 68 insertions(+), 70 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 082feb8d..a8f44b35 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -385,6 +385,7 @@ checksum = "0c9ea0ac24bc397ab3c98583a3c9ba74fa56b09a4449bbe172b9b1ddb016027a" name = "common" version = "0.1.0" dependencies = [ + "bytemuck", "crypto-primitives", "field", "num-traits", @@ -750,6 +751,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 b3feed87..c63b48e8 100644 --- a/crates/common/Cargo.toml +++ b/crates/common/Cargo.toml @@ -11,6 +11,7 @@ parallel = ["dep:rayon"] [dependencies] crypto-primitives = { workspace = true } +bytemuck = {workspace = true} field = { workspace = true, features = ["spongefish"] } spongefish = { workspace = true } poly = { workspace = true } diff --git a/crates/common/src/fold.rs b/crates/common/src/fold.rs index 50753969..c2760e5e 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::*; + use crate::table::PACKED_BITS; use crate::{BitZParams, Shape}; const Q114: u128 = (1 << 114) - 11; diff --git a/crates/common/src/table.rs b/crates/common/src/table.rs index c2e63051..30c8306f 100644 --- a/crates/common/src/table.rs +++ b/crates/common/src/table.rs @@ -9,6 +9,10 @@ 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; +/// The machine word `packed` is sliced into. [`BitTable::BITS`] is derived +/// from this, so changing it is the only step needed to repack into a +/// different word size. +type Word = u64; /// A witness that does not match its shape. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -47,10 +51,11 @@ 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 @@ -60,7 +65,10 @@ impl<'a> BitTable<'a> { if packed.len() != (1 << shape.log_bits()) / PACKED_BITS { return Err(TableError::BitCountMismatch); } - Ok(Self { shape, packed }) + Ok(Self { + shape, + packed: bytemuck::cast_slice(packed), + }) } pub fn shape(&self) -> &Shape { @@ -78,15 +86,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 optimise 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,8 +101,8 @@ 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] } @@ -105,7 +111,6 @@ impl<'a> BitTable<'a> { ColumnBits { elements: self.column(column).iter(), word: 0, - hi_pending: None, remaining: 0, } } @@ -154,16 +159,23 @@ impl<'a> BitTable<'a> { 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()); + // Two `u64` words (lo, hi) per 128-row group in `self.packed`, now + // that it holds words rather than `F128` elements. + let words_per_column = row_groups * 2; + debug_assert_eq!( + row_groups * column_groups * PACKED_BITS, + self.packed.len() / 2 + ); - let mut packed = F128::zeroed_vec(self.packed.len()); + let mut packed = F128::zeroed_vec(row_groups * column_groups * PACKED_BITS); 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); + let element_lo = self.packed[column * words_per_column + row_group * 2]; + let element_hi = self.packed[column * words_per_column + row_group * 2 + 1]; + *word = u128::from(element_lo) | (u128::from(element_hi) << HALF_BITS); } transpose_bit_block(&mut block); @@ -237,33 +249,20 @@ impl TransposedBitTable { /// Iterator over one column's bits, in ascending row order — see /// [`BitTable::column_bits`]. pub struct ColumnBits<'a> { - elements: std::slice::Iter<'a, F128>, + elements: std::slice::Iter<'a, Word>, /// 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, + word: Word, + remaining: u8, } 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; - } - _ => {} + if self.remaining == 0 { + let element = self.elements.next()?; + self.word = *element; + self.remaining = BitTable::BITS as u8; } let bit = self.word & 1 == 1; self.word >>= 1; @@ -279,9 +278,7 @@ impl Iterator for ColumnBits<'_> { 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 + self.remaining as usize + self.elements.len() * (BitTable::BITS as usize) } } @@ -372,9 +369,9 @@ mod tests { let packed = with_bits(&shape, &[(0, 3), (2, 3), (64, 3)]); let table = BitTable::new(shape, &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)); + assert_eq!(table.column(3).len(), shape.rows() / BitTable::BITS); + assert_eq!(table.column(3), [0b101u64, 1u64]); + assert!(table.column(4).iter().all(|&element| element == 0)); } #[test] @@ -400,11 +397,10 @@ mod tests { let packed = with_bits(&shape, &rows); let table = BitTable::new(shape, &packed).unwrap(); - for (index, element) in table.column(2).iter().enumerate() { - let bits = u128::from(element.lo) | (u128::from(element.hi) << 64); - 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}"); } } } diff --git a/crates/field/Cargo.toml b/crates/field/Cargo.toml index 13807f98..7aa7aa9e 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 53c3bce2..88058649 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 7781751f..26913ece 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 1f390e6a..818e76fb 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::*; diff --git a/crates/prover/src/reduce.rs b/crates/prover/src/reduce.rs index 4364cf26..fdd74911 100644 --- a/crates/prover/src/reduce.rs +++ b/crates/prover/src/reduce.rs @@ -76,7 +76,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); From e6a2e4cd15a19371a4081dae987aba49bf725aed Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Wed, 23 Sep 2026 19:47:08 +0800 Subject: [PATCH 02/22] common: rename ColumnBits to BitsIter --- crates/common/src/table.rs | 65 ++++++++++++++++++++++++++++++++------ 1 file changed, 55 insertions(+), 10 deletions(-) diff --git a/crates/common/src/table.rs b/crates/common/src/table.rs index 30c8306f..c5a946c4 100644 --- a/crates/common/src/table.rs +++ b/crates/common/src/table.rs @@ -1,5 +1,7 @@ //! The committed bit table. +use std::{array, slice::Iter}; + use crate::{Shape, ShapeError}; use field::F128; @@ -88,7 +90,7 @@ impl<'a> BitTable<'a> { ); let index = (column << self.shape.log_rows()) | row; - // /% on compile time constant should be properly optimise away + // /% on compile time constant should be properly optimised away let element = self.packed[index / BitTable::BITS]; let offset = index % BitTable::BITS; @@ -107,12 +109,8 @@ impl<'a> BitTable<'a> { } /// 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, - remaining: 0, - } + pub fn column_bits(&self, column: usize) -> BitsIter<'a> { + BitsIter::new(self.column(column).iter()) } /// Swaps rows and columns, returning an owned copy: `result.bit(b, c) == @@ -248,14 +246,61 @@ impl TransposedBitTable { /// Iterator over one column's bits, in ascending row order — see /// [`BitTable::column_bits`]. -pub struct ColumnBits<'a> { +pub struct BitsIter<'a> { elements: std::slice::Iter<'a, Word>, /// Bits not yet emitted; the next one to emit is the LSB. word: Word, remaining: u8, } -impl Iterator for ColumnBits<'_> { +// Find conditions +// at most group of size 64 -> would be a transpose again? +// make grouping a const size that is given to unroll +fn group_pack(tt: &Vec) -> Vec { + let tt: &[Word] = bytemuck::cast_slice(tt); + let n = tt.len(); + let segment_n = n / GS; + + let mut remaining = BitTable::BITS; + let mut word = 0; + let mut segments: [_; GS] = + array::from_fn(|i| BitsIter::new((&tt[i * segment_n..(i + 1) * segment_n]).iter())); + + // Possible without shifting in just selecting the index might be faster + // Same goes for the bits iterator. Directly selecting position might be faster. + // Bit select library? + let mut out = Vec::new(); + 'outer: loop { + for s in &mut segments { + match s.next() { + Some(b) => { + if remaining == 0 { + out.push(word); + word = 0; + remaining = BitTable::BITS + }; + word = (word << 1) | (b as Word); + remaining -= 1; + } + None => break 'outer, + } + } + } + + out +} + +impl<'a> BitsIter<'a> { + fn new(iter: Iter<'a, Word>) -> Self { + Self { + elements: iter, + word: 0, + remaining: 0, + } + } +} + +impl Iterator for BitsIter<'_> { type Item = bool; fn next(&mut self) -> Option { @@ -276,7 +321,7 @@ impl Iterator for ColumnBits<'_> { } } -impl ExactSizeIterator for ColumnBits<'_> { +impl ExactSizeIterator for BitsIter<'_> { fn len(&self) -> usize { self.remaining as usize + self.elements.len() * (BitTable::BITS as usize) } From dd89f3f6bb85d3a3a8837ccbe4d2cd993729fb89 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Wed, 23 Sep 2026 20:32:16 +0800 Subject: [PATCH 03/22] common: generalise BitsIter into StepIter yielding S-bit chunks --- crates/common/src/table.rs | 49 ++++++++++++++++++++++--------------- crates/prover/src/reduce.rs | 2 +- 2 files changed, 30 insertions(+), 21 deletions(-) diff --git a/crates/common/src/table.rs b/crates/common/src/table.rs index c5a946c4..daea7338 100644 --- a/crates/common/src/table.rs +++ b/crates/common/src/table.rs @@ -244,15 +244,6 @@ impl TransposedBitTable { } } -/// Iterator over one column's bits, in ascending row order — see -/// [`BitTable::column_bits`]. -pub struct BitsIter<'a> { - elements: std::slice::Iter<'a, Word>, - /// Bits not yet emitted; the next one to emit is the LSB. - word: Word, - remaining: u8, -} - // Find conditions // at most group of size 64 -> would be a transpose again? // make grouping a const size that is given to unroll @@ -290,8 +281,26 @@ fn group_pack(tt: &Vec) -> Vec { out } -impl<'a> BitsIter<'a> { +pub type BitsIter<'a> = StepIter<'a, 1>; + +/// Iterator over one column's bits, in ascending row order — see +/// [`BitTable::column_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, +} + +// TODO +impl<'a, const S: usize> StepIter<'a, S> { + const CHECK_SIZE: () = assert!( + S != 0 && S <= Word::BITS as usize && Word::BITS as usize % S == 0, + "step size must be nonzero, at most 64, and divide 64 evenly" + ); + const MASK: u64 = 1u64.unbounded_shl(S as u32).wrapping_sub(1); fn new(iter: Iter<'a, Word>) -> Self { + let () = Self::CHECK_SIZE; Self { elements: iter, word: 0, @@ -300,19 +309,19 @@ impl<'a> BitsIter<'a> { } } -impl Iterator for BitsIter<'_> { - type Item = bool; +impl<'a, const S: usize> Iterator for StepIter<'_, S> { + type Item = u64; - fn next(&mut self) -> Option { + fn next(&mut self) -> Option { if self.remaining == 0 { let element = self.elements.next()?; self.word = *element; - self.remaining = BitTable::BITS as u8; + self.remaining = Word::BITS as u8; } - let bit = self.word & 1 == 1; - self.word >>= 1; - self.remaining -= 1; - Some(bit) + 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) { @@ -321,9 +330,9 @@ impl Iterator for BitsIter<'_> { } } -impl ExactSizeIterator for BitsIter<'_> { +impl<'a, const S: usize> ExactSizeIterator for StepIter<'_, S> { fn len(&self) -> usize { - self.remaining as usize + self.elements.len() * (BitTable::BITS as usize) + (self.remaining as usize + self.elements.len() * (BitTable::BITS as usize)) / S } } diff --git a/crates/prover/src/reduce.rs b/crates/prover/src/reduce.rs index fdd74911..bf879336 100644 --- a/crates/prover/src/reduce.rs +++ b/crates/prover/src/reduce.rs @@ -28,7 +28,7 @@ fn init_circuit(table: &BitTable, fold: &Fold) -> GrandProductCircuit { 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 }; + *leaf = if bit == 1 { row_image } else { F128::ONE }; } } } From 448fe41e304271d608cefc37e4d669504ef97c04 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Mon, 28 Sep 2026 10:04:53 +0800 Subject: [PATCH 04/22] common: transpose BitTable via word-level 64x64 bit block transpose --- crates/common/src/fold.rs | 6 +- crates/common/src/params.rs | 4 +- crates/common/src/table.rs | 144 ++++++++++++++++++------------------ 3 files changed, 78 insertions(+), 76 deletions(-) diff --git a/crates/common/src/fold.rs b/crates/common/src/fold.rs index c2760e5e..4d703b9d 100644 --- a/crates/common/src/fold.rs +++ b/crates/common/src/fold.rs @@ -231,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), @@ -251,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), @@ -272,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/params.rs b/crates/common/src/params.rs index e8ca41e1..fab4cf5d 100644 --- a/crates/common/src/params.rs +++ b/crates/common/src/params.rs @@ -64,7 +64,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 +163,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)) } } diff --git a/crates/common/src/table.rs b/crates/common/src/table.rs index daea7338..0a0f79d9 100644 --- a/crates/common/src/table.rs +++ b/crates/common/src/table.rs @@ -62,15 +62,12 @@ impl<'a> BitTable<'a> { /// /// 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: bytemuck::cast_slice(packed), - }) + Ok(Self { shape, packed }) } pub fn shape(&self) -> &Shape { @@ -155,62 +152,67 @@ impl<'a> BitTable<'a> { TransposeError::ColumnCountTooNarrow })?; - let row_groups = self.shape.rows() / PACKED_BITS; - let column_groups = self.shape.columns() / PACKED_BITS; - // Two `u64` words (lo, hi) per 128-row group in `self.packed`, now - // that it holds words rather than `F128` elements. - let words_per_column = row_groups * 2; - debug_assert_eq!( - row_groups * column_groups * PACKED_BITS, - self.packed.len() / 2 - ); + let packed = bit_transpose(self.packed, self.shape.columns(), self.shape.rows()); - let mut packed = F128::zeroed_vec(row_groups * column_groups * PACKED_BITS); - 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_lo = self.packed[column * words_per_column + row_group * 2]; - let element_hi = self.packed[column * words_per_column + row_group * 2 + 1]; - *word = u128::from(element_lo) | (u128::from(element_hi) << HALF_BITS); - } + Ok(TransposedBitTable { shape, packed }) + } +} - transpose_bit_block(&mut block); +// Bits packed into machine words +// out of place variant. +// dim1 and dim2 should be given in bits +// dim2 is the axis over which the data is adjacent, think xs[dim1][dim2]. +fn bit_transpose(xs: &[Word], dim1: usize, dim2: usize) -> Vec { + let mut out = bytemuck::zeroed_vec(xs.len()); + + // Bit transpose approach is dividing the matrix up in blocks of WORD::BITS x WORD::BITS. Perform transpose on the block and write the block to the transposed position. + + // TODO generalise to different bit sizes + // generalisation doesn't handle sub byte sizes. + let dim1_groups = dim1 / (Word::BITS as usize); // chunk size in dim1 (same as number of words in line) + let dim2_groups = dim2 / (Word::BITS as usize); // chunk size in dim2 + for d1 in 0..dim1_groups { + for d2 in 0..dim2_groups { + let mut block = [0; Word::BITS as usize]; + + for (lane, word) in block.iter_mut().enumerate() { + *word = xs[(d1 * Word::BITS as usize + lane) * dim2_groups + d2] + } - 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); - } + bit_block_transpose(&mut block); + + let dim1t = d2; + let dim2t = d1; + let dim2t_groups = dim1_groups; + + for (lane, word) in block.iter().enumerate() { + out[(dim1t * Word::BITS as usize + lane) * dim2t_groups + dim2t] = *word; } } - - Ok(TransposedBitTable { shape, packed }) } + + out } -/// 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; +fn bit_block_transpose(xs: &mut [Word; Word::BITS as usize]) { + let mut j = (Word::BITS / 2) as usize; + let mut mask = (1 << j) - 1; // Lower half high 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; + let mut k: usize = 0; + while k < Word::BITS as usize { + let t = xs[k]; + let b = xs[k | j]; + // Alternative is a delta swap, however this is more readable. + xs[k] = (t & mask) | ((b << j) & !mask); + xs[k | j] = ((t >> j) & mask) | (b & !mask); + // (k|j) count the upper part. +1 advances the upper part. !j converts it into the lower index again except for when it hits the next round. In that case the | j was k = ((k | j) + 1) & !j; } j >>= 1; + // Generate the alternating bit pattern + // 00001111 + // 00110011 + // 01010101 mask ^= mask << j; } } @@ -220,7 +222,7 @@ fn transpose_bit_block(a: &mut [u128; PACKED_BITS]) { #[derive(Debug)] pub struct TransposedBitTable { shape: Shape, - packed: Vec, + packed: Vec, } impl TransposedBitTable { @@ -239,7 +241,7 @@ impl TransposedBitTable { } /// Unwraps the packed, transposed bits. - pub fn into_packed(self) -> Vec { + pub fn into_packed(self) -> Vec { self.packed } } @@ -247,6 +249,8 @@ impl TransposedBitTable { // Find conditions // at most group of size 64 -> would be a transpose again? // make grouping a const size that is given to unroll +// Incorrect order +// doesn't have the check. fn group_pack(tt: &Vec) -> Vec { let tt: &[Word] = bytemuck::cast_slice(tt); let n = tt.len(); @@ -260,6 +264,7 @@ fn group_pack(tt: &Vec) -> Vec { // Possible without shifting in just selecting the index might be faster // Same goes for the bits iterator. Directly selecting position might be faster. // Bit select library? + // Looks like a zero vec for the outer is also an option. let mut out = Vec::new(); 'outer: loop { for s in &mut segments { @@ -292,7 +297,6 @@ pub struct StepIter<'a, const S: usize> { remaining: u8, } -// TODO impl<'a, const S: usize> StepIter<'a, S> { const CHECK_SIZE: () = assert!( S != 0 && S <= Word::BITS as usize && Word::BITS as usize % S == 0, @@ -368,7 +372,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() { @@ -389,7 +393,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()); @@ -412,7 +416,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) ); } @@ -421,7 +425,7 @@ 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() / BitTable::BITS); assert_eq!(table.column(3), [0b101u64, 1u64]); @@ -433,11 +437,11 @@ 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(); let iter = table.column_bits(2); assert_eq!(iter.len(), shape.rows()); - let bits: Vec = iter.collect(); + let bits: Vec = iter.map(|b| b == 1).collect(); assert_eq!(bits.len(), shape.rows()); for (row, bit) in bits.into_iter().enumerate() { assert_eq!(bit, table.bit(2, row), "row {row}"); @@ -449,7 +453,7 @@ 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() { for offset in 0..BitTable::BITS { @@ -466,16 +470,14 @@ mod tests { } #[test] - fn transpose_bit_block_matches_a_brute_force_reference() { - let mut original = [0u128; PACKED_BITS]; + fn bit_block_transpose_matches_a_brute_force_reference() { + let mut original = [0 as Word; BitTable::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); + *word = (i as Word).wrapping_mul(0x2545_F491_4F6C_DD1D) ^ 0xA5A5; } let mut transposed = original; - transpose_bit_block(&mut transposed); + bit_block_transpose(&mut transposed); for (i, &transposed_word) in transposed.iter().enumerate() { for (j, &original_word) in original.iter().enumerate() { @@ -519,7 +521,7 @@ mod tests { 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 table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); let transposed = table.transpose().unwrap(); assert_eq!(transposed.shape().log_rows(), shape.log_columns()); @@ -545,7 +547,7 @@ mod tests { 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 table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); let roundtripped = table .transpose() @@ -555,7 +557,7 @@ mod tests { .unwrap() .into_packed(); - assert_eq!(roundtripped, packed); + assert_eq!(roundtripped, bytemuck::cast_slice::(&packed)); } #[test] @@ -565,7 +567,7 @@ mod tests { // 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(), From 0abcaf0dc16bc22dcd9190ab0f8f3cdc66d75b2a Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Mon, 28 Sep 2026 10:43:46 +0800 Subject: [PATCH 05/22] common: add a test-only transpose_reference and drop F128 packing constants --- crates/common/src/fold.rs | 2 +- crates/common/src/table.rs | 109 ++++++++++++++++++++----------------- 2 files changed, 59 insertions(+), 52 deletions(-) diff --git a/crates/common/src/fold.rs b/crates/common/src/fold.rs index 4d703b9d..f29d29ab 100644 --- a/crates/common/src/fold.rs +++ b/crates/common/src/fold.rs @@ -183,7 +183,7 @@ mod tests { use num_traits::{ConstOne, ConstZero}; use super::*; - use crate::table::PACKED_BITS; + const PACKED_BITS: usize = 128; use crate::{BitZParams, Shape}; const Q114: u128 = (1 << 114) - 11; diff --git a/crates/common/src/table.rs b/crates/common/src/table.rs index 0a0f79d9..493408ae 100644 --- a/crates/common/src/table.rs +++ b/crates/common/src/table.rs @@ -1,16 +1,8 @@ //! The committed bit table. -use std::{array, slice::Iter}; +use std::slice::Iter; 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; /// The machine word `packed` is sliced into. [`BitTable::BITS`] is derived /// from this, so changing it is the only step needed to repack into a /// different word size. @@ -163,12 +155,11 @@ impl<'a> BitTable<'a> { // dim1 and dim2 should be given in bits // dim2 is the axis over which the data is adjacent, think xs[dim1][dim2]. fn bit_transpose(xs: &[Word], dim1: usize, dim2: usize) -> Vec { + assert_eq!(xs.len() * Word::BITS as usize, dim1 * dim2); let mut out = bytemuck::zeroed_vec(xs.len()); // Bit transpose approach is dividing the matrix up in blocks of WORD::BITS x WORD::BITS. Perform transpose on the block and write the block to the transposed position. - // TODO generalise to different bit sizes - // generalisation doesn't handle sub byte sizes. let dim1_groups = dim1 / (Word::BITS as usize); // chunk size in dim1 (same as number of words in line) let dim2_groups = dim2 / (Word::BITS as usize); // chunk size in dim2 for d1 in 0..dim1_groups { @@ -246,46 +237,6 @@ impl TransposedBitTable { } } -// Find conditions -// at most group of size 64 -> would be a transpose again? -// make grouping a const size that is given to unroll -// Incorrect order -// doesn't have the check. -fn group_pack(tt: &Vec) -> Vec { - let tt: &[Word] = bytemuck::cast_slice(tt); - let n = tt.len(); - let segment_n = n / GS; - - let mut remaining = BitTable::BITS; - let mut word = 0; - let mut segments: [_; GS] = - array::from_fn(|i| BitsIter::new((&tt[i * segment_n..(i + 1) * segment_n]).iter())); - - // Possible without shifting in just selecting the index might be faster - // Same goes for the bits iterator. Directly selecting position might be faster. - // Bit select library? - // Looks like a zero vec for the outer is also an option. - let mut out = Vec::new(); - 'outer: loop { - for s in &mut segments { - match s.next() { - Some(b) => { - if remaining == 0 { - out.push(word); - word = 0; - remaining = BitTable::BITS - }; - word = (word << 1) | (b as Word); - remaining -= 1; - } - None => break 'outer, - } - } - } - - out -} - pub type BitsIter<'a> = StepIter<'a, 1>; /// Iterator over one column's bits, in ascending row order — see @@ -342,8 +293,20 @@ impl<'a, const S: usize> ExactSizeIterator for StepIter<'_, S> { #[cfg(test)] mod tests { + use std::array; + use num_traits::ConstZero; + // We want to use compare against + use field::F128; + + /// Bits in a packed element, and the shift that divides an index into element + /// and offset. + const PACKED_BITS: usize = 128; + const PACKED_SHIFT: u32 = 7; + /// Bits in each half of a packed element. + const HALF_BITS: usize = 64; + use super::*; /// `m = 22`: 128 rows per column, 32768 columns. @@ -490,6 +453,23 @@ mod tests { } } + #[test] + fn transpose_reference_matches_bit_transpose() { + // 128 segments of 4 words each: two block groups on both axes. + const GS: usize = 128; + let segment_n = 4; + let tt: Vec = (0..GS * segment_n) + .map(|i| (i as Word).wrapping_mul(0x2545_F491_4F6C_DD1D) ^ 0xA5A5) + .collect(); + + let reference = transpose_reference::(&tt); + assert_eq!(reference.len(), tt.len()); + assert_eq!( + reference, + bit_transpose(&tt, GS, segment_n * BitTable::BITS) + ); + } + /// 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. @@ -574,4 +554,31 @@ mod tests { Some(TransposeError::ColumnCountTooNarrow) ); } + + fn transpose_reference(tt: &[Word]) -> Vec { + assert_eq!(tt.len() % GS, 0); + let segment_n = tt.len() / GS; + let mut segments: [_; GS] = + 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 `bit()` and `StepIter`. + word |= b << offset; + offset += 1; + if offset == BitTable::BITS { + out.push(word); + word = 0; + offset = 0; + } + } + } + debug_assert_eq!(offset, 0); + + out + } } From d626b3d2d97b6b799746d521e8b3dcd9b63f952c Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Mon, 28 Sep 2026 10:57:43 +0800 Subject: [PATCH 06/22] common: let bit_block_transpose handle dim1-row tiles and test d = 4 --- crates/common/src/table.rs | 40 ++++++++++++++++++++++++++++++++++---- 1 file changed, 36 insertions(+), 4 deletions(-) diff --git a/crates/common/src/table.rs b/crates/common/src/table.rs index 493408ae..0ed1fa8d 100644 --- a/crates/common/src/table.rs +++ b/crates/common/src/table.rs @@ -170,7 +170,7 @@ fn bit_transpose(xs: &[Word], dim1: usize, dim2: usize) -> Vec { *word = xs[(d1 * Word::BITS as usize + lane) * dim2_groups + d2] } - bit_block_transpose(&mut block); + bit_block_transpose(Word::BITS as usize, &mut block); let dim1t = d2; let dim2t = d1; @@ -185,9 +185,15 @@ fn bit_transpose(xs: &[Word], dim1: usize, dim2: usize) -> Vec { out } -fn bit_block_transpose(xs: &mut [Word; Word::BITS as usize]) { +fn bit_block_transpose(dim1: usize, xs: &mut [Word; Word::BITS as usize]) { let mut j = (Word::BITS / 2) as usize; let mut mask = (1 << j) - 1; // Lower half high + + while j > dim1 / 2 { + j >>= 1; + mask ^= mask << j; + } + while j > 0 { let mut k: usize = 0; while k < Word::BITS as usize { @@ -200,7 +206,7 @@ fn bit_block_transpose(xs: &mut [Word; Word::BITS as usize]) { k = ((k | j) + 1) & !j; } j >>= 1; - // Generate the alternating bit pattern + // Generate alternating bit pattern // 00001111 // 00110011 // 01010101 @@ -440,7 +446,7 @@ mod tests { } let mut transposed = original; - bit_block_transpose(&mut transposed); + bit_block_transpose(Word::BITS as usize, &mut transposed); for (i, &transposed_word) in transposed.iter().enumerate() { for (j, &original_word) in original.iter().enumerate() { @@ -453,6 +459,32 @@ mod tests { } } + #[test] + fn bit_block_transpose_4_matches_reference() { + // With `dim1 = 4` every 4x4 tile -- words `4g..4g+4`, bits + // `4n..4n+4` -- is transposed on its own. + const D: usize = 4; + let mut original = [0 as Word; BitTable::BITS]; + for (i, word) in original.iter_mut().enumerate() { + *word = (i as Word).wrapping_mul(0x2545_F491_4F6C_DD1D) ^ 0xA5A5; + } + + let mut transposed = original; + bit_block_transpose(D, &mut transposed); + + for (i, &transposed_word) in transposed.iter().enumerate() { + for j in 0..BitTable::BITS { + let source_word = original[i - i % D + j % D]; + let source_bit = j - j % D + i % D; + assert_eq!( + (transposed_word >> j) & 1, + (source_word >> source_bit) & 1, + "word {i} bit {j}" + ); + } + } + } + #[test] fn transpose_reference_matches_bit_transpose() { // 128 segments of 4 words each: two block groups on both axes. From e5543a043c57f76fc579a1ea48db702d938ad29c Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Mon, 28 Sep 2026 17:29:49 +0800 Subject: [PATCH 07/22] common: support any power-of-two dim1 in bit_transpose --- Cargo.lock | 1 + crates/common/Cargo.toml | 10 ++ crates/common/src/table.rs | 247 +++++++++++++++++++++++++------------ 3 files changed, 180 insertions(+), 78 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index a8f44b35..c645d1a1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -387,6 +387,7 @@ version = "0.1.0" dependencies = [ "bytemuck", "crypto-primitives", + "divan", "field", "num-traits", "poly", diff --git a/crates/common/Cargo.toml b/crates/common/Cargo.toml index c63b48e8..c6bf33a6 100644 --- a/crates/common/Cargo.toml +++ b/crates/common/Cargo.toml @@ -8,6 +8,8 @@ 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 } @@ -17,3 +19,11 @@ spongefish = { workspace = true } poly = { workspace = true } rayon = { workspace = true, optional = true } num-traits = { workspace = true } + +[dev-dependencies] +divan = { workspace = true } + +[[bench]] +name = "table" +harness = false +required-features = ["bench"] diff --git a/crates/common/src/table.rs b/crates/common/src/table.rs index 0ed1fa8d..6b6ca915 100644 --- a/crates/common/src/table.rs +++ b/crates/common/src/table.rs @@ -1,5 +1,6 @@ //! The committed bit table. +use std::ops::Range; use std::slice::Iter; use crate::{Shape, ShapeError}; @@ -154,30 +155,93 @@ impl<'a> BitTable<'a> { // out of place variant. // dim1 and dim2 should be given in bits // dim2 is the axis over which the data is adjacent, think xs[dim1][dim2]. +// +// A block is always BITS words, built from `d = min(dim1, BITS)` rows of +// `r = BITS / d` consecutive words each: word i = c * d + a is row a, word c. +// Writing a block's index bits as (word | bit), the input is +// ([c | a] | [p_hi | p_lo]) with a and p_hi log2(d) bits wide, and the +// output needs ([c | p_hi] | [p_lo | a]): the bit index rotates by log2(d). +// A butterfly stage only swaps a word index bit with the bit index bit in the +// same place, and a word reorder only permutes the word index, so p_lo has to +// go down into the word index, across, and back up: +// +// stages s < max(log d, log r) p_lo (and a) swap into the word index +// reorder words rotate the word index left by log2(d) +// stages s >= log d p_lo lands above a in the bit index +// indexed write swap two word index fields into [c | p_hi] +// +// The reorder is never done: the words stay where the gather put them, the +// second pass pairs slots through the rotation (see `bit_block_transpose`'s +// `rot`), and the indexed write applies it. That is 6 stages for d >= 8, +// where the first pass is just the d x d tile transpose. For d = BITS the rotation and +// the write are the identity and the second pass runs no stages, which leaves +// the plain blocked transpose. fn bit_transpose(xs: &[Word], dim1: usize, dim2: usize) -> Vec { - assert_eq!(xs.len() * Word::BITS as usize, dim1 * dim2); - let mut out = bytemuck::zeroed_vec(xs.len()); - - // Bit transpose approach is dividing the matrix up in blocks of WORD::BITS x WORD::BITS. Perform transpose on the block and write the block to the transposed position. + assert_eq!(xs.len() * BitTable::BITS, dim1 * dim2); + assert!(dim1.is_power_of_two(), "dim1 {dim1} is not a power of two"); + + // One body, compiled once per block height. With `log2(d)` a runtime + // value the stage loops cannot be unrolled with constant masks and + // distances, which cost ~30% at d = 64. + match dim1.min(BitTable::BITS).trailing_zeros() { + 0 => bit_transpose_blocks::<0>(xs, dim1, dim2), + 1 => bit_transpose_blocks::<1>(xs, dim1, dim2), + 2 => bit_transpose_blocks::<2>(xs, dim1, dim2), + 3 => bit_transpose_blocks::<3>(xs, dim1, dim2), + 4 => bit_transpose_blocks::<4>(xs, dim1, dim2), + 5 => bit_transpose_blocks::<5>(xs, dim1, dim2), + _ => bit_transpose_blocks::<6>(xs, dim1, dim2), + } +} - let dim1_groups = dim1 / (Word::BITS as usize); // chunk size in dim1 (same as number of words in line) - let dim2_groups = dim2 / (Word::BITS as usize); // chunk size in dim2 - for d1 in 0..dim1_groups { - for d2 in 0..dim2_groups { - let mut block = [0; Word::BITS as usize]; +/// [`bit_transpose`] for blocks of `d = 2^LOG_D` rows. +fn bit_transpose_blocks(xs: &[Word], dim1: usize, dim2: usize) -> Vec { + const BITS: usize = Word::BITS as usize; + const LOG_BITS: u32 = Word::BITS.trailing_zeros(); + + let d = 1 << LOG_D; // rows per block + let log_r = LOG_BITS - LOG_D; + let r = 1 << log_r; // words per row per block + let row_words = dim2 / BITS; + // Output words between consecutive words of a block; 1 when d < BITS. + let out_stride = dim1 / d; + assert!( + dim2.is_multiple_of(BITS) && row_words.is_multiple_of(r), + "dim2 {dim2} does not fill whole {d}-row blocks" + ); - for (lane, word) in block.iter_mut().enumerate() { - *word = xs[(d1 * Word::BITS as usize + lane) * dim2_groups + d2] - } + // The indexed write swaps the word index fields at `lo` and `hi`, each + // `width` bits wide. `width = 0` when d = BITS. + let width = LOG_D.min(log_r); + let (lo, hi) = (LOG_D - width, LOG_BITS - width); + let field = (1 << width) - 1; - bit_block_transpose(Word::BITS as usize, &mut block); + let mut out = bytemuck::zeroed_vec(xs.len()); - let dim1t = d2; - let dim2t = d1; - let dim2t_groups = dim1_groups; + for g1 in 0..dim1 / d { + for g2 in 0..row_words / r { + // Word i = [c | a] is row a, word c. Two loops rather than one + // over i: splitting i back into a and c measured 4-12% slower + // for d < BITS. + let mut block = [0; BITS]; + for c in 0..r { + for a in 0..d { + block[c * d + a] = xs[(g1 * d + a) * row_words + g2 * r + c]; + } + } - for (lane, word) in block.iter().enumerate() { - out[(dim1t * Word::BITS as usize + lane) * dim2t_groups + dim2t] = *word; + bit_block_transpose(0..LOG_D.max(log_r), 0, &mut block); + // Word i belongs at rotl(i, LOG_D) for this pass; pair it there. + bit_block_transpose(LOG_D..LOG_BITS, log_r, &mut block); + + // Word i is word rotl(i, LOG_D) of the second pass. After the + // swap, o is output row g2 * BITS + o, word g1 in [dim2][dim1]; + // g1 is only nonzero when d = BITS. + for (i, &word) in block.iter().enumerate() { + let i = ((i << LOG_D) | (i >> log_r)) & (BITS - 1); + let t = ((i >> lo) ^ (i >> hi)) & field; + let o = i ^ (t << lo) ^ (t << hi); + out[(g2 * BITS + o) * out_stride + g1] = word; } } } @@ -185,25 +249,35 @@ fn bit_transpose(xs: &[Word], dim1: usize, dim2: usize) -> Vec { out } -fn bit_block_transpose(dim1: usize, xs: &mut [Word; Word::BITS as usize]) { +/// Runs the butterfly stages `s` in `stages` on a BITS x BITS block. Stage +/// `s` swaps word index bit `s` with bit index bit `s`, so `0..6` is the full +/// transpose and `0..log2(d)` transposes every `d x d` tile in place. +/// +/// `rot` is for a block whose words are stored at `rotl(i, rot)` over the +/// 6-bit word index rather than at `i`: stage `s` then pairs slots +/// `2^((s + rot) % 6)` apart, while the bits it moves are unchanged. +#[inline(always)] +fn bit_block_transpose(stages: Range, rot: u32, xs: &mut [Word; Word::BITS as usize]) { let mut j = (Word::BITS / 2) as usize; let mut mask = (1 << j) - 1; // Lower half high - while j > dim1 / 2 { + while j >= 1 << stages.end { j >>= 1; mask ^= mask << j; } - while j > 0 { + while j >= 1 << stages.start { + // Distance between the paired words. + let jw = 1 << ((j.trailing_zeros() + rot) % Word::BITS.trailing_zeros()); let mut k: usize = 0; while k < Word::BITS as usize { let t = xs[k]; - let b = xs[k | j]; + let b = xs[k | jw]; // Alternative is a delta swap, however this is more readable. xs[k] = (t & mask) | ((b << j) & !mask); - xs[k | j] = ((t >> j) & mask) | (b & !mask); - // (k|j) count the upper part. +1 advances the upper part. !j converts it into the lower index again except for when it hits the next round. In that case the | j was - k = ((k | j) + 1) & !j; + xs[k | jw] = ((t >> j) & mask) | (b & !mask); + // (k|jw) count the upper part. +1 advances the upper part. !jw converts it into the lower index again except for when it hits the next round. In that case the | jw was + k = ((k | jw) + 1) & !jw; } j >>= 1; // Generate alternating bit pattern @@ -214,6 +288,50 @@ fn bit_block_transpose(dim1: usize, xs: &mut [Word; Word::BITS as usize]) { } } +/// 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 `bit()` and `StepIter`. + word |= b << offset; + offset += 1; + if offset == BitTable::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: &[u64], dim1: usize, dim2: usize) -> Vec { + super::bit_transpose(xs, dim1, dim2) + } + + pub fn transpose_reference(tt: &[u64]) -> Vec { + super::transpose_reference::(tt) + } +} + /// An owned, transposed copy of a [`BitTable`]'s bits -- see /// [`BitTable::transpose`], which is the only way to build one. #[derive(Debug)] @@ -299,8 +417,6 @@ impl<'a, const S: usize> ExactSizeIterator for StepIter<'_, S> { #[cfg(test)] mod tests { - use std::array; - use num_traits::ConstZero; // We want to use compare against @@ -446,7 +562,7 @@ mod tests { } let mut transposed = original; - bit_block_transpose(Word::BITS as usize, &mut transposed); + bit_block_transpose(0..6, 0, &mut transposed); for (i, &transposed_word) in transposed.iter().enumerate() { for (j, &original_word) in original.iter().enumerate() { @@ -459,30 +575,32 @@ mod tests { } } - #[test] - fn bit_block_transpose_4_matches_reference() { - // With `dim1 = 4` every 4x4 tile -- words `4g..4g+4`, bits - // `4n..4n+4` -- is transposed on its own. - const D: usize = 4; - let mut original = [0 as Word; BitTable::BITS]; - for (i, word) in original.iter_mut().enumerate() { - *word = (i as Word).wrapping_mul(0x2545_F491_4F6C_DD1D) ^ 0xA5A5; - } - - let mut transposed = original; - bit_block_transpose(D, &mut transposed); + /// `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(|i| (i as Word).wrapping_mul(0x2545_F491_4F6C_DD1D) ^ 0xA5A5) + .collect(); + assert_eq!( + bit_transpose(&tt, GS, words * BitTable::BITS), + transpose_reference::(&tt), + "dim1 {GS}, {words} words per row" + ); + } - for (i, &transposed_word) in transposed.iter().enumerate() { - for j in 0..BitTable::BITS { - let source_word = original[i - i % D + j % D]; - let source_bit = j - j % D + i % D; - assert_eq!( - (transposed_word >> j) & 1, - (source_word >> source_bit) & 1, - "word {i} bit {j}" - ); - } - } + #[test] + fn bit_transpose_matches_reference_for_every_row_count() { + // Two blocks along dim2 at every row count, so block placement is + // exercised as well as the block contents. + check_bit_transpose::<1>(128); + check_bit_transpose::<2>(64); + check_bit_transpose::<4>(32); + check_bit_transpose::<8>(16); + check_bit_transpose::<16>(8); + check_bit_transpose::<32>(4); + check_bit_transpose::<64>(2); + check_bit_transpose::<128>(2); + check_bit_transpose::<256>(2); } #[test] @@ -586,31 +704,4 @@ mod tests { Some(TransposeError::ColumnCountTooNarrow) ); } - - fn transpose_reference(tt: &[Word]) -> Vec { - assert_eq!(tt.len() % GS, 0); - let segment_n = tt.len() / GS; - let mut segments: [_; GS] = - 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 `bit()` and `StepIter`. - word |= b << offset; - offset += 1; - if offset == BitTable::BITS { - out.push(word); - word = 0; - offset = 0; - } - } - } - debug_assert_eq!(offset, 0); - - out - } } From b8418a69146f95744381f8e01e6a78a780784133 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Mon, 28 Sep 2026 17:46:36 +0800 Subject: [PATCH 08/22] common: reduce bit_transpose_blocks to a single butterfly sweep --- crates/common/src/table.rs | 99 ++++++++++++++++---------------------- 1 file changed, 41 insertions(+), 58 deletions(-) diff --git a/crates/common/src/table.rs b/crates/common/src/table.rs index 6b6ca915..54d7ec97 100644 --- a/crates/common/src/table.rs +++ b/crates/common/src/table.rs @@ -1,6 +1,5 @@ //! The committed bit table. -use std::ops::Range; use std::slice::Iter; use crate::{Shape, ShapeError}; @@ -161,21 +160,15 @@ impl<'a> BitTable<'a> { // Writing a block's index bits as (word | bit), the input is // ([c | a] | [p_hi | p_lo]) with a and p_hi log2(d) bits wide, and the // output needs ([c | p_hi] | [p_lo | a]): the bit index rotates by log2(d). -// A butterfly stage only swaps a word index bit with the bit index bit in the -// same place, and a word reorder only permutes the word index, so p_lo has to -// go down into the word index, across, and back up: // -// stages s < max(log d, log r) p_lo (and a) swap into the word index -// reorder words rotate the word index left by log2(d) -// stages s >= log d p_lo lands above a in the bit index -// indexed write swap two word index fields into [c | p_hi] -// -// The reorder is never done: the words stay where the gather put them, the -// second pass pairs slots through the rotation (see `bit_block_transpose`'s -// `rot`), and the indexed write applies it. That is 6 stages for d >= 8, -// where the first pass is just the d x d tile transpose. For d = BITS the rotation and -// the write are the identity and the second pass runs no stages, which leaves -// the plain blocked transpose. +// A butterfly stage swaps one word index bit with one bit index bit. Stages +// s < log d swap a with the low bits of p: every d x d tile is transposed and +// each word holds correct d-bit pieces. The later stages move whole pieces +// with the same word index bits, s % log d: stage s puts the bit that stage +// s - log d took out of the bit index back in, log d places higher. That is +// one sweep of 6 stages for any d >= 2, and it has to run upwards. The word +// index ends as [c | p_hi] with p_hi's bits rotated by 6 % log d, which the +// indexed write undoes. For d = BITS this is the plain blocked transpose. fn bit_transpose(xs: &[Word], dim1: usize, dim2: usize) -> Vec { assert_eq!(xs.len() * BitTable::BITS, dim1 * dim2); assert!(dim1.is_power_of_two(), "dim1 {dim1} is not a power of two"); @@ -184,7 +177,8 @@ fn bit_transpose(xs: &[Word], dim1: usize, dim2: usize) -> Vec { // value the stage loops cannot be unrolled with constant masks and // distances, which cost ~30% at d = 64. match dim1.min(BitTable::BITS).trailing_zeros() { - 0 => bit_transpose_blocks::<0>(xs, dim1, dim2), + // A single row is its own transpose. + 0 => xs.to_vec(), 1 => bit_transpose_blocks::<1>(xs, dim1, dim2), 2 => bit_transpose_blocks::<2>(xs, dim1, dim2), 3 => bit_transpose_blocks::<3>(xs, dim1, dim2), @@ -194,14 +188,13 @@ fn bit_transpose(xs: &[Word], dim1: usize, dim2: usize) -> Vec { } } -/// [`bit_transpose`] for blocks of `d = 2^LOG_D` rows. +/// [`bit_transpose`] for blocks of `d = 2^LOG_D` rows, `LOG_D >= 1`. fn bit_transpose_blocks(xs: &[Word], dim1: usize, dim2: usize) -> Vec { const BITS: usize = Word::BITS as usize; const LOG_BITS: u32 = Word::BITS.trailing_zeros(); let d = 1 << LOG_D; // rows per block - let log_r = LOG_BITS - LOG_D; - let r = 1 << log_r; // words per row per block + let r = BITS / d; // words per row per block let row_words = dim2 / BITS; // Output words between consecutive words of a block; 1 when d < BITS. let out_stride = dim1 / d; @@ -210,11 +203,9 @@ fn bit_transpose_blocks(xs: &[Word], dim1: usize, dim2: usize) "dim2 {dim2} does not fill whole {d}-row blocks" ); - // The indexed write swaps the word index fields at `lo` and `hi`, each - // `width` bits wide. `width = 0` when d = BITS. - let width = LOG_D.min(log_r); - let (lo, hi) = (LOG_D - width, LOG_BITS - width); - let field = (1 << width) - 1; + // The sweep leaves p_hi in the low LOG_D word index bits rotated left + // by `rot`; the indexed write rotates them back. + let rot = LOG_BITS % LOG_D; let mut out = bytemuck::zeroed_vec(xs.len()); @@ -230,17 +221,13 @@ fn bit_transpose_blocks(xs: &[Word], dim1: usize, dim2: usize) } } - bit_block_transpose(0..LOG_D.max(log_r), 0, &mut block); - // Word i belongs at rotl(i, LOG_D) for this pass; pair it there. - bit_block_transpose(LOG_D..LOG_BITS, log_r, &mut block); + bit_block_transpose(LOG_D, &mut block); - // Word i is word rotl(i, LOG_D) of the second pass. After the - // swap, o is output row g2 * BITS + o, word g1 in [dim2][dim1]; - // g1 is only nonzero when d = BITS. + // o is output row g2 * BITS + o, word g1 in [dim2][dim1]; g1 is + // only nonzero when d = BITS. for (i, &word) in block.iter().enumerate() { - let i = ((i << LOG_D) | (i >> log_r)) & (BITS - 1); - let t = ((i >> lo) ^ (i >> hi)) & field; - let o = i ^ (t << lo) ^ (t << hi); + let low = i & (d - 1); + let o = (i & !(d - 1)) | ((low >> rot) | (low << (LOG_D - rot))) & (d - 1); out[(g2 * BITS + o) * out_stride + g1] = word; } } @@ -249,26 +236,28 @@ fn bit_transpose_blocks(xs: &[Word], dim1: usize, dim2: usize) out } -/// Runs the butterfly stages `s` in `stages` on a BITS x BITS block. Stage -/// `s` swaps word index bit `s` with bit index bit `s`, so `0..6` is the full -/// transpose and `0..log2(d)` transposes every `d x d` tile in place. -/// -/// `rot` is for a block whose words are stored at `rotl(i, rot)` over the -/// 6-bit word index rather than at `i`: stage `s` then pairs slots -/// `2^((s + rot) % 6)` apart, while the bits it moves are unchanged. +/// Runs the 6 butterfly stages on a BITS x BITS block, lowest first. Stage +/// `s` swaps bit index bit `s` with word index bit `s % log_d`, so +/// `log_d = 6` is the full transpose. See [`bit_transpose`] for smaller +/// `log_d`. #[inline(always)] -fn bit_block_transpose(stages: Range, rot: u32, xs: &mut [Word; Word::BITS as usize]) { - let mut j = (Word::BITS / 2) as usize; - let mut mask = (1 << j) - 1; // Lower half high - - while j >= 1 << stages.end { - j >>= 1; - mask ^= mask << j; - } - - while j >= 1 << stages.start { +fn bit_block_transpose(log_d: u32, xs: &mut [Word; Word::BITS as usize]) { + const STAGES: usize = Word::BITS.trailing_zeros() as usize; + // Generate alternating bit pattern, top stage first + // 00001111 + // 00110011 + // 01010101 + let mut masks = [0; STAGES]; + let mut mask = Word::MAX; + for s in (0..STAGES).rev() { + mask ^= mask << (1 << s); + masks[s] = mask; + } + + for (s, mask) in (0..).zip(masks) { + let j = 1 << s; // Distance between the paired words. - let jw = 1 << ((j.trailing_zeros() + rot) % Word::BITS.trailing_zeros()); + let jw = 1 << (s % log_d); let mut k: usize = 0; while k < Word::BITS as usize { let t = xs[k]; @@ -279,12 +268,6 @@ fn bit_block_transpose(stages: Range, rot: u32, xs: &mut [Word; Word::BITS // (k|jw) count the upper part. +1 advances the upper part. !jw converts it into the lower index again except for when it hits the next round. In that case the | jw was k = ((k | jw) + 1) & !jw; } - j >>= 1; - // Generate alternating bit pattern - // 00001111 - // 00110011 - // 01010101 - mask ^= mask << j; } } @@ -562,7 +545,7 @@ mod tests { } let mut transposed = original; - bit_block_transpose(0..6, 0, &mut transposed); + bit_block_transpose(6, &mut transposed); for (i, &transposed_word) in transposed.iter().enumerate() { for (j, &original_word) in original.iter().enumerate() { From 5b37d62be1b8662873aa169de38548ac240e8046 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Mon, 28 Sep 2026 17:49:26 +0800 Subject: [PATCH 09/22] common: compute butterfly masks inline instead of precomputing them --- crates/common/src/table.rs | 17 ++++------------- 1 file changed, 4 insertions(+), 13 deletions(-) diff --git a/crates/common/src/table.rs b/crates/common/src/table.rs index 54d7ec97..a167e7fa 100644 --- a/crates/common/src/table.rs +++ b/crates/common/src/table.rs @@ -242,20 +242,11 @@ fn bit_transpose_blocks(xs: &[Word], dim1: usize, dim2: usize) /// `log_d`. #[inline(always)] fn bit_block_transpose(log_d: u32, xs: &mut [Word; Word::BITS as usize]) { - const STAGES: usize = Word::BITS.trailing_zeros() as usize; - // Generate alternating bit pattern, top stage first - // 00001111 - // 00110011 - // 01010101 - let mut masks = [0; STAGES]; - let mut mask = Word::MAX; - for s in (0..STAGES).rev() { - mask ^= mask << (1 << s); - masks[s] = mask; - } - - for (s, mask) in (0..).zip(masks) { + for s in 0..Word::BITS.trailing_zeros() { let j = 1 << s; + // Alternating runs of j bits, lowest run set: + // 01010101, 00110011, 00001111 + let mask = Word::MAX / ((1 << j) + 1); // Distance between the paired words. let jw = 1 << (s % log_d); let mut k: usize = 0; From f891e96b84df150cbdd45a43d2183cbebdf4402f Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Mon, 28 Sep 2026 17:55:06 +0800 Subject: [PATCH 10/22] common: add divan benchmark of bit_transpose against the reference --- crates/common/benches/table.rs | 53 ++++++++++++++++++++++++++++++++++ 1 file changed, 53 insertions(+) create mode 100644 crates/common/benches/table.rs diff --git a/crates/common/benches/table.rs b/crates/common/benches/table.rs new file mode 100644 index 00000000..72ff288b --- /dev/null +++ b/crates/common/benches/table.rs @@ -0,0 +1,53 @@ +//! Compare the blocked `bit_transpose` with the bit-at-a-time reference. +//! +//! Run with `cargo bench -p common --features bench --bench table`. +//! Both transpose a `128 x DIM2`-bit matrix (`dim2` contiguous), including +//! output allocation. Input generation runs outside the measured loop. + +use common::table::bench::{bit_transpose, transpose_reference}; +use divan::counter::BytesCount; +use divan::{Bencher, black_box}; + +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(); +} + +/// SplitMix64 for deterministic, well-mixed bits. +fn input() -> Vec { + let mut state = 0u64; + (0..DIM1 * DIM2 / u64::BITS as usize) + .map(|_| { + state = state.wrapping_add(0x9e37_79b9_7f4a_7c15); + let mut word = state; + word = (word ^ (word >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9); + word = (word ^ (word >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); + word ^ (word >> 31) + }) + .collect() +} + +#[divan::bench(sample_count = 20)] +fn blocked(bencher: Bencher) { + let xs = input(); + bencher + .counter(BytesCount::of_slice(&xs)) + .bench(|| bit_transpose(black_box(&xs), DIM1, DIM2)); +} + +#[divan::bench(sample_count = 20)] +fn reference(bencher: Bencher) { + let xs = input(); + bencher + .counter(BytesCount::of_slice(&xs)) + .bench(|| transpose_reference::(black_box(&xs))); +} From 2d7fe1e9a853cc57429091bdba5435ce3f101f8d Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Mon, 28 Sep 2026 18:08:49 +0800 Subject: [PATCH 11/22] fixup! common: add divan benchmark of bit_transpose against the reference --- crates/common/src/table.rs | 8 ++++---- crates/gkr/src/lib.rs | 1 - 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/crates/common/src/table.rs b/crates/common/src/table.rs index a167e7fa..7d32d50c 100644 --- a/crates/common/src/table.rs +++ b/crates/common/src/table.rs @@ -348,7 +348,7 @@ pub struct StepIter<'a, const S: usize> { impl<'a, const S: usize> StepIter<'a, S> { const CHECK_SIZE: () = assert!( - S != 0 && S <= Word::BITS as usize && Word::BITS as usize % S == 0, + S != 0 && S <= Word::BITS as usize && (Word::BITS as usize).is_multiple_of(S), "step size must be nonzero, at most 64, and divide 64 evenly" ); const MASK: u64 = 1u64.unbounded_shl(S as u32).wrapping_sub(1); @@ -362,7 +362,7 @@ impl<'a, const S: usize> StepIter<'a, S> { } } -impl<'a, const S: usize> Iterator for StepIter<'_, S> { +impl Iterator for StepIter<'_, S> { type Item = u64; fn next(&mut self) -> Option { @@ -383,9 +383,9 @@ impl<'a, const S: usize> Iterator for StepIter<'_, S> { } } -impl<'a, const S: usize> ExactSizeIterator for StepIter<'_, S> { +impl ExactSizeIterator for StepIter<'_, S> { fn len(&self) -> usize { - (self.remaining as usize + self.elements.len() * (BitTable::BITS as usize)) / S + (self.remaining as usize + self.elements.len() * BitTable::BITS) / S } } diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs index 818e76fb..1f390e6a 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs @@ -1,6 +1,5 @@ use std::collections::VecDeque; -use common::BitTable; use field::{F128, Wide256}; use num_traits::{ConstOne, ConstZero}; use rayon::prelude::*; From b0340d2a70d7b6a3818ee7abbc98c489531cfea3 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Mon, 28 Sep 2026 18:39:02 +0800 Subject: [PATCH 12/22] common: make BitTable generic over its storage and replace TransposedBitTable with OwnedBitTable --- crates/common/src/lib.rs | 2 +- crates/common/src/table.rs | 88 +++++++++++++++++-------------------- crates/prover/src/reduce.rs | 1 - 3 files changed, 42 insertions(+), 49 deletions(-) diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs index 78fb71b7..95e09694 100644 --- a/crates/common/src/lib.rs +++ b/crates/common/src/lib.rs @@ -19,7 +19,7 @@ pub use fold::{ 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, BitTableBase, OwnedBitTable, TableError, TransposeError}; pub use virtual_map::{ TransposedWeights, VirtualMap, VirtualMapError, VirtualStatement, VirtualStatementError, }; diff --git a/crates/common/src/table.rs b/crates/common/src/table.rs index 7d32d50c..e8aca404 100644 --- a/crates/common/src/table.rs +++ b/crates/common/src/table.rs @@ -1,6 +1,6 @@ //! The committed bit table. -use std::slice::Iter; +use std::{ops::Deref, slice::Iter}; use crate::{Shape, ShapeError}; /// The machine word `packed` is sliced into. [`BitTable::BITS`] is derived @@ -42,19 +42,32 @@ pub enum TransposeError { /// for the monomial basis `beta_v = X^v`. In that basis bit `v` of an `F128`'s /// little-endian `lo || hi` *is* the coefficient of `X^v`, so the committed /// form and the bit form are the same bytes, and the caller holds one copy. +/// +/// Generic over the storage `S` of the packed words; use the [`BitTable`] +/// (borrowed) and [`OwnedBitTable`] aliases rather than naming this directly. #[derive(Debug, Clone, Copy)] -pub struct BitTable<'a> { +pub struct BitTableBase { shape: Shape, - packed: &'a [Word], + packed: S, } -impl<'a> BitTable<'a> { +/// A [`BitTableBase`] borrowing the prover's packed witness. +pub type BitTable<'a> = BitTableBase<&'a [Word]>; + +/// A [`BitTableBase`] owning its packed words, as [`BitTableBase::transpose`] +/// returns. +pub type OwnedBitTable = BitTableBase>; + +impl BitTableBase { pub const BITS: usize = Word::BITS as usize; +} + +impl> BitTableBase { /// 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 [Word]) -> Result { + pub(crate) fn new(shape: Shape, packed: S) -> 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); @@ -92,13 +105,13 @@ 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 [Word] { + pub fn column(&self, column: usize) -> &[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) -> BitsIter<'a> { + pub fn column_bits(&self, column: usize) -> BitsIter<'_> { BitsIter::new(self.column(column).iter()) } @@ -132,7 +145,7 @@ impl<'a> BitTable<'a> { /// (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 { + pub fn transpose(&self) -> Result { let shape = Shape::new(self.shape.log_columns(), self.shape.log_rows()).map_err(|error| { debug_assert_eq!( @@ -144,9 +157,22 @@ impl<'a> BitTable<'a> { TransposeError::ColumnCountTooNarrow })?; - let packed = bit_transpose(self.packed, self.shape.columns(), self.shape.rows()); + let packed = bit_transpose(&self.packed, self.shape.columns(), self.shape.rows()); + + Ok(OwnedBitTable { shape, packed }) + } + + /// Borrows this table as an ordinary [`BitTable`]. + pub fn as_table(&self) -> BitTable<'_> { + BitTable { + shape: self.shape, + packed: &self.packed, + } + } - Ok(TransposedBitTable { shape, packed }) + /// Unwraps the packed bits. + pub fn into_packed(self) -> S { + self.packed } } @@ -297,44 +323,15 @@ fn transpose_reference(tt: &[Word]) -> Vec { #[cfg(feature = "bench")] #[doc(hidden)] pub mod bench { - pub fn bit_transpose(xs: &[u64], dim1: usize, dim2: usize) -> Vec { + pub fn bit_transpose(xs: &[super::Word], dim1: usize, dim2: usize) -> Vec { super::bit_transpose(xs, dim1, dim2) } - pub fn transpose_reference(tt: &[u64]) -> Vec { + pub fn transpose_reference(tt: &[super::Word]) -> Vec { super::transpose_reference::(tt) } } -/// 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 - } -} - pub type BitsIter<'a> = StepIter<'a, 1>; /// Iterator over one column's bits, in ascending row order — see @@ -349,9 +346,9 @@ pub struct StepIter<'a, const S: usize> { 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 64, and divide 64 evenly" + "step size must be nonzero, at most Word::BITS, and divide Word::BITS evenly" ); - const MASK: u64 = 1u64.unbounded_shl(S as u32).wrapping_sub(1); + 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 { @@ -363,7 +360,7 @@ impl<'a, const S: usize> StepIter<'a, S> { } impl Iterator for StepIter<'_, S> { - type Item = u64; + type Item = Word; fn next(&mut self) -> Option { if self.remaining == 0 { @@ -631,8 +628,6 @@ mod tests { 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. @@ -656,7 +651,6 @@ mod tests { let roundtripped = table .transpose() .unwrap() - .as_table() .transpose() .unwrap() .into_packed(); diff --git a/crates/prover/src/reduce.rs b/crates/prover/src/reduce.rs index bf879336..0262a26f 100644 --- a/crates/prover/src/reduce.rs +++ b/crates/prover/src/reduce.rs @@ -24,7 +24,6 @@ fn init_circuit(table: &BitTable, fold: &Fold) -> GrandProductCircuit { 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)) { From 403e34a0f4920cedaa865dd7e03443ff90e08ebb Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Tue, 29 Sep 2026 10:50:13 +0800 Subject: [PATCH 13/22] fixup! common: add divan benchmark of bit_transpose against the reference --- Cargo.lock | 1 + crates/common/Cargo.toml | 1 + crates/common/benches/table.rs | 16 +++++----------- 3 files changed, 7 insertions(+), 11 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index c645d1a1..dc8cedd9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -391,6 +391,7 @@ dependencies = [ "field", "num-traits", "poly", + "rand 0.10.2", "rayon", "spongefish", ] diff --git a/crates/common/Cargo.toml b/crates/common/Cargo.toml index c6bf33a6..260d5963 100644 --- a/crates/common/Cargo.toml +++ b/crates/common/Cargo.toml @@ -22,6 +22,7 @@ num-traits = { workspace = true } [dev-dependencies] divan = { workspace = true } +rand = { workspace = true } [[bench]] name = "table" diff --git a/crates/common/benches/table.rs b/crates/common/benches/table.rs index 72ff288b..8b897352 100644 --- a/crates/common/benches/table.rs +++ b/crates/common/benches/table.rs @@ -7,6 +7,8 @@ use common::table::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. @@ -22,21 +24,14 @@ fn main() { divan::main(); } -/// SplitMix64 for deterministic, well-mixed bits. +/// Seeded so every run transposes the same matrix. fn input() -> Vec { - let mut state = 0u64; + let mut rng = StdRng::seed_from_u64(0); (0..DIM1 * DIM2 / u64::BITS as usize) - .map(|_| { - state = state.wrapping_add(0x9e37_79b9_7f4a_7c15); - let mut word = state; - word = (word ^ (word >> 30)).wrapping_mul(0xbf58_476d_1ce4_e5b9); - word = (word ^ (word >> 27)).wrapping_mul(0x94d0_49bb_1331_11eb); - word ^ (word >> 31) - }) + .map(|_| rng.random()) .collect() } -#[divan::bench(sample_count = 20)] fn blocked(bencher: Bencher) { let xs = input(); bencher @@ -44,7 +39,6 @@ fn blocked(bencher: Bencher) { .bench(|| bit_transpose(black_box(&xs), DIM1, DIM2)); } -#[divan::bench(sample_count = 20)] fn reference(bencher: Bencher) { let xs = input(); bencher From 8f8c5e519adaf336a6e1dc1e446be6fc30720b44 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Tue, 29 Sep 2026 12:43:28 +0800 Subject: [PATCH 14/22] common: switch Word to u128 and rotate bit indices ahead of coarse-to-fine butterflies --- crates/common/benches/table.rs | 6 +- crates/common/src/table.rs | 206 +++++++++++++++++++++++---------- 2 files changed, 149 insertions(+), 63 deletions(-) diff --git a/crates/common/benches/table.rs b/crates/common/benches/table.rs index 8b897352..2aa6a4fc 100644 --- a/crates/common/benches/table.rs +++ b/crates/common/benches/table.rs @@ -25,13 +25,14 @@ fn main() { } /// Seeded so every run transposes the same matrix. -fn input() -> Vec { +fn input() -> Vec { let mut rng = StdRng::seed_from_u64(0); - (0..DIM1 * DIM2 / u64::BITS as usize) + (0..DIM1 * DIM2 / u128::BITS as usize) .map(|_| rng.random()) .collect() } +#[divan::bench] fn blocked(bencher: Bencher) { let xs = input(); bencher @@ -39,6 +40,7 @@ fn blocked(bencher: Bencher) { .bench(|| bit_transpose(black_box(&xs), DIM1, DIM2)); } +#[divan::bench] fn reference(bencher: Bencher) { let xs = input(); bencher diff --git a/crates/common/src/table.rs b/crates/common/src/table.rs index e8aca404..46cd7995 100644 --- a/crates/common/src/table.rs +++ b/crates/common/src/table.rs @@ -3,10 +3,16 @@ use std::{ops::Deref, slice::Iter}; use crate::{Shape, ShapeError}; -/// The machine word `packed` is sliced into. [`BitTable::BITS`] is derived -/// from this, so changing it is the only step needed to repack into a -/// different word size. -type Word = u64; +/// The word `packed` is sliced into: one `F128` element. [`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 `BitTable::transpose`. +type Word = u128; +/// Bits in a [`Word`]'s bit index. +const LOG_BITS: u32 = Word::BITS.trailing_zeros(); /// A witness that does not match its shape. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -185,16 +191,15 @@ impl> BitTableBase { // `r = BITS / d` consecutive words each: word i = c * d + a is row a, word c. // Writing a block's index bits as (word | bit), the input is // ([c | a] | [p_hi | p_lo]) with a and p_hi log2(d) bits wide, and the -// output needs ([c | p_hi] | [p_lo | a]): the bit index rotates by log2(d). +// output needs ([c | p_hi] | [p_lo | a]). // -// A butterfly stage swaps one word index bit with one bit index bit. Stages -// s < log d swap a with the low bits of p: every d x d tile is transposed and -// each word holds correct d-bit pieces. The later stages move whole pieces -// with the same word index bits, s % log d: stage s puts the bit that stage -// s - log d took out of the bit index back in, log d places higher. That is -// one sweep of 6 stages for any d >= 2, and it has to run upwards. The word -// index ends as [c | p_hi] with p_hi's bits rotated by 6 % log d, which the -// indexed write undoes. For d = BITS this is the plain blocked transpose. +// p_lo only moves within its word, up by log2(d), so that is done first: +// each gathered word's bit index is rotated to [p_lo | p_hi], which +// transposes the word read as d rows of r bits. That leaves a and p_hi to +// trade places bit for bit: log2(d) butterfly stages, stage s swapping word +// index bit s with bit index bit s, which transposes every d x d tile of the +// block. The stages commute and run coarse to fine. For d = BITS the +// rotation is the identity and this is the plain blocked transpose. fn bit_transpose(xs: &[Word], dim1: usize, dim2: usize) -> Vec { assert_eq!(xs.len() * BitTable::BITS, dim1 * dim2); assert!(dim1.is_power_of_two(), "dim1 {dim1} is not a power of two"); @@ -210,14 +215,14 @@ fn bit_transpose(xs: &[Word], dim1: usize, dim2: usize) -> Vec { 3 => bit_transpose_blocks::<3>(xs, dim1, dim2), 4 => bit_transpose_blocks::<4>(xs, dim1, dim2), 5 => bit_transpose_blocks::<5>(xs, dim1, dim2), - _ => bit_transpose_blocks::<6>(xs, dim1, dim2), + 6 => bit_transpose_blocks::<6>(xs, dim1, dim2), + _ => bit_transpose_blocks::(xs, dim1, dim2), } } /// [`bit_transpose`] for blocks of `d = 2^LOG_D` rows, `LOG_D >= 1`. fn bit_transpose_blocks(xs: &[Word], dim1: usize, dim2: usize) -> Vec { const BITS: usize = Word::BITS as usize; - const LOG_BITS: u32 = Word::BITS.trailing_zeros(); let d = 1 << LOG_D; // rows per block let r = BITS / d; // words per row per block @@ -229,14 +234,15 @@ fn bit_transpose_blocks(xs: &[Word], dim1: usize, dim2: usize) "dim2 {dim2} does not fill whole {d}-row blocks" ); - // The sweep leaves p_hi in the low LOG_D word index bits rotated left - // by `rot`; the indexed write rotates them back. - let rot = LOG_BITS % LOG_D; - let mut out = bytemuck::zeroed_vec(xs.len()); - for g1 in 0..dim1 / d { - for g2 in 0..row_words / r { + // 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. g1 only takes more than one value + // when d = BITS. + for g2 in 0..row_words / r { + for g1 in 0..dim1 / d { // Word i = [c | a] is row a, word c. Two loops rather than one // over i: splitting i back into a and c measured 4-12% slower // for d < BITS. @@ -247,13 +253,14 @@ fn bit_transpose_blocks(xs: &[Word], dim1: usize, dim2: usize) } } + for word in &mut block { + *word = rotate_bit_index::(*word); + } bit_block_transpose(LOG_D, &mut block); // o is output row g2 * BITS + o, word g1 in [dim2][dim1]; g1 is // only nonzero when d = BITS. - for (i, &word) in block.iter().enumerate() { - let low = i & (d - 1); - let o = (i & !(d - 1)) | ((low >> rot) | (low << (LOG_D - rot))) & (d - 1); + for (o, &word) in block.iter().enumerate() { out[(g2 * BITS + o) * out_stride + g1] = word; } } @@ -262,28 +269,81 @@ fn bit_transpose_blocks(xs: &[Word], dim1: usize, dim2: usize) out } -/// Runs the 6 butterfly stages on a BITS x BITS block, lowest first. Stage -/// `s` swaps bit index bit `s` with word index bit `s % log_d`, so -/// `log_d = 6` is the full transpose. See [`bit_transpose`] for smaller -/// `log_d`. +/// The bits whose index has bit `s` clear: alternating runs of `2^s` bits, +/// lowest run set. 01010101, 00110011, 00001111 for `s` = 0, 1, 2. +const fn index_bit_clear(s: u32) -> Word { + Word::MAX / ((1 << (1 << s)) + 1) +} + +/// Rotates the bit index of every bit in `x` left by `LOG_D`: bit +/// `[p_hi | p_lo]` moves to `[p_lo | p_hi]`, `p_hi` being `LOG_D` bits wide. +/// Read as `d = 2^LOG_D` rows of `r = BITS / d` bits, the word is transposed. +/// A delta swap per entry of [`rotation_swaps`], none for `LOG_D = LOG_BITS`. +#[inline(always)] +fn rotate_bit_index(mut x: Word) -> Word { + // Evaluated at compile time, so every shift and mask is a constant. + let (swaps, n) = const { rotation_swaps(LOG_D) }; + for &(shift, mask) in &swaps[..n] { + let t = ((x >> shift) ^ x) & mask; + x ^= t ^ (t << shift); + } + x +} + +/// The delta swaps that rotate a bit index left by `log_d`, as `(shift, +/// mask)` pairs: each bit in `mask` trades places with the bit `shift` above +/// it. Returns the pairs and how many of them are used. +/// +/// The rotation splits the index bits into `gcd(LOG_BITS, log_d)` cycles. +/// Each cycle is walked from its lowest index bit `low`, which is swapped +/// with the cycle's other bits `i` in turn: `LOG_BITS - gcd` swaps in all. +const fn rotation_swaps(log_d: u32) -> ([(u32, Word); LOG_BITS as usize], usize) { + let (mut cycles, mut rest) = (LOG_BITS, log_d); + while rest != 0 { + (cycles, rest) = (rest, cycles % rest); + } + + let mut swaps = [(0, 0); LOG_BITS as usize]; + let mut n = 0; + let mut low = 0; + while low < cycles { + let mut i = (low + log_d) % LOG_BITS; + while i != low { + // The bits whose index has bit low set and bit i clear. + swaps[n] = ( + (1 << i) - (1 << low), + !index_bit_clear(low) & index_bit_clear(i), + ); + n += 1; + i = (i + log_d) % LOG_BITS; + } + low += 1; + } + (swaps, n) +} + +/// Transposes every `2^log_d x 2^log_d` tile of a BITS x BITS block, the +/// tile's words being those that differ only in their low `log_d` index +/// bits. Stage `s` swaps bit index bit `s` with word index bit `s`, coarse +/// to fine; the stages commute, so the order is free. `log_d = LOG_BITS` is +/// the full transpose. #[inline(always)] fn bit_block_transpose(log_d: u32, xs: &mut [Word; Word::BITS as usize]) { - for s in 0..Word::BITS.trailing_zeros() { - let j = 1 << s; - // Alternating runs of j bits, lowest run set: - // 01010101, 00110011, 00001111 - let mask = Word::MAX / ((1 << j) + 1); - // Distance between the paired words. - let jw = 1 << (s % log_d); + for s in (0..log_d).rev() { + // Distance between the paired bits, and between the paired words. + let j: usize = 1 << s; + let mask = index_bit_clear(s); let mut k: usize = 0; while k < Word::BITS as usize { - let t = xs[k]; - let b = xs[k | jw]; - // Alternative is a delta swap, however this is more readable. - xs[k] = (t & mask) | ((b << j) & !mask); - xs[k | jw] = ((t >> j) & mask) | (b & !mask); - // (k|jw) count the upper part. +1 advances the upper part. !jw converts it into the lower index again except for when it hits the next round. In that case the | jw was - k = ((k | jw) + 1) & !jw; + // A delta swap across the pair: t is where the upper runs of + // xs[k] and the lower runs of xs[k | j] differ, and flipping those + // bits in both swaps the runs. Measured 6-10% faster than + // rebuilding both words from masked halves. + let t = ((xs[k] >> j) ^ xs[k | j]) & mask; + xs[k | j] ^= t; + xs[k] ^= t << j; + // (k|j) count the upper part. +1 advances the upper part. !j converts it into the lower index again except for when it hits the next round. In that case the | j was + k = ((k | j) + 1) & !j; } } } @@ -484,7 +544,7 @@ mod tests { let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); assert_eq!(table.column(3).len(), shape.rows() / BitTable::BITS); - assert_eq!(table.column(3), [0b101u64, 1u64]); + assert_eq!(table.column(3), [0b101 | (1 << 64)]); assert!(table.column(4).iter().all(|&element| element == 0)); } @@ -525,15 +585,20 @@ mod tests { (seed.wrapping_mul(0x9E37_79B9_7F4A_7C15) >> 63) & 1 == 1 } + /// 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 + } + #[test] fn bit_block_transpose_matches_a_brute_force_reference() { let mut original = [0 as Word; BitTable::BITS]; for (i, word) in original.iter_mut().enumerate() { - *word = (i as Word).wrapping_mul(0x2545_F491_4F6C_DD1D) ^ 0xA5A5; + *word = pseudo_random_word(i); } let mut transposed = original; - bit_block_transpose(6, &mut transposed); + bit_block_transpose(LOG_BITS, &mut transposed); for (i, &transposed_word) in transposed.iter().enumerate() { for (j, &original_word) in original.iter().enumerate() { @@ -546,12 +611,33 @@ mod tests { } } + /// Bit `p` must land at `p` rotated left by `LOG_D` within `LOG_BITS` bits. + fn check_rotate_bit_index() { + for p in 0..BitTable::BITS { + let rotated = ((p << LOG_D) | (p >> (LOG_BITS - LOG_D))) % BitTable::BITS; + assert_eq!( + rotate_bit_index::(1 << p), + 1 << rotated, + "LOG_D {LOG_D}, bit {p}" + ); + } + } + + #[test] + fn rotate_bit_index_rotates_every_bit_index() { + check_rotate_bit_index::<1>(); + check_rotate_bit_index::<2>(); + check_rotate_bit_index::<3>(); + check_rotate_bit_index::<4>(); + check_rotate_bit_index::<5>(); + check_rotate_bit_index::<6>(); + check_rotate_bit_index::<7>(); + } + /// `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(|i| (i as Word).wrapping_mul(0x2545_F491_4F6C_DD1D) ^ 0xA5A5) - .collect(); + let tt: Vec = (0..GS * words).map(pseudo_random_word).collect(); assert_eq!( bit_transpose(&tt, GS, words * BitTable::BITS), transpose_reference::(&tt), @@ -563,25 +649,23 @@ mod tests { fn bit_transpose_matches_reference_for_every_row_count() { // Two blocks along dim2 at every row count, so block placement is // exercised as well as the block contents. - check_bit_transpose::<1>(128); - check_bit_transpose::<2>(64); - check_bit_transpose::<4>(32); - check_bit_transpose::<8>(16); - check_bit_transpose::<16>(8); - check_bit_transpose::<32>(4); - check_bit_transpose::<64>(2); + check_bit_transpose::<1>(256); + check_bit_transpose::<2>(128); + check_bit_transpose::<4>(64); + check_bit_transpose::<8>(32); + check_bit_transpose::<16>(16); + check_bit_transpose::<32>(8); + check_bit_transpose::<64>(4); check_bit_transpose::<128>(2); check_bit_transpose::<256>(2); } #[test] fn transpose_reference_matches_bit_transpose() { - // 128 segments of 4 words each: two block groups on both axes. - const GS: usize = 128; - let segment_n = 4; - let tt: Vec = (0..GS * segment_n) - .map(|i| (i as Word).wrapping_mul(0x2545_F491_4F6C_DD1D) ^ 0xA5A5) - .collect(); + // 256 segments of 2 words each: two block groups on both axes. + const GS: usize = 256; + let segment_n = 2; + let tt: Vec = (0..GS * segment_n).map(pseudo_random_word).collect(); let reference = transpose_reference::(&tt); assert_eq!(reference.len(), tt.len()); From 435d7454c03c0b550e0489899a5293685ca52809 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Wed, 30 Sep 2026 19:59:58 +0800 Subject: [PATCH 15/22] common: run in-lane bit index rotation swaps on u64 lanes so they vectorise --- crates/common/src/table.rs | 63 +++++++++++++++++++++++-------------- crates/prover/src/reduce.rs | 1 + 2 files changed, 41 insertions(+), 23 deletions(-) diff --git a/crates/common/src/table.rs b/crates/common/src/table.rs index 46cd7995..21290ed0 100644 --- a/crates/common/src/table.rs +++ b/crates/common/src/table.rs @@ -13,6 +13,9 @@ use crate::{Shape, ShapeError}; type Word = u128; /// Bits in a [`Word`]'s bit index. const LOG_BITS: u32 = Word::BITS.trailing_zeros(); +/// Bits in a `u64`'s bit index. Swapping two bit index bits below this never +/// moves a bit from one of a word's `u64` lanes to another. +const LOG_LANE_BITS: u32 = u64::BITS.trailing_zeros(); /// A witness that does not match its shape. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -253,9 +256,7 @@ fn bit_transpose_blocks(xs: &[Word], dim1: usize, dim2: usize) } } - for word in &mut block { - *word = rotate_bit_index::(*word); - } + rotate_bit_index::(&mut block); bit_block_transpose(LOG_D, &mut block); // o is output row g2 * BITS + o, word g1 in [dim2][dim1]; g1 is @@ -275,35 +276,50 @@ const fn index_bit_clear(s: u32) -> Word { Word::MAX / ((1 << (1 << s)) + 1) } -/// Rotates the bit index of every bit in `x` left by `LOG_D`: bit -/// `[p_hi | p_lo]` moves to `[p_lo | p_hi]`, `p_hi` being `LOG_D` bits wide. -/// Read as `d = 2^LOG_D` rows of `r = BITS / d` bits, the word is transposed. -/// A delta swap per entry of [`rotation_swaps`], none for `LOG_D = LOG_BITS`. +/// Rotates the bit index of every bit in every word of `block` left by +/// `LOG_D`: bit `[p_hi | p_lo]` moves to `[p_lo | p_hi]`, `p_hi` being +/// `LOG_D` bits wide. Read as `d = 2^LOG_D` rows of `r = BITS / d` bits, +/// each word is transposed. One pass over the block per entry of +/// [`rotation_swaps`], none for `LOG_D = LOG_BITS`. +/// +/// A pass that stays inside `u64` lanes runs on the lanes rather than on +/// whole words, which vectorises: on `u128` that made the rotation about +/// 40% cheaper. Only a swap with bit index bit 6 or above needs whole words. #[inline(always)] -fn rotate_bit_index(mut x: Word) -> Word { +fn rotate_bit_index(block: &mut [Word; Word::BITS as usize]) { // Evaluated at compile time, so every shift and mask is a constant. let (swaps, n) = const { rotation_swaps(LOG_D) }; - for &(shift, mask) in &swaps[..n] { - let t = ((x >> shift) ^ x) & mask; - x ^= t ^ (t << shift); + for &(shift, mask, in_lane) in &swaps[..n] { + if in_lane { + let mask = mask as u64; + for x in bytemuck::cast_slice_mut::(block) { + let t = ((*x >> shift) ^ *x) & mask; + *x ^= t ^ (t << shift); + } + } else { + for x in block.iter_mut() { + let t = ((*x >> shift) ^ *x) & mask; + *x ^= t ^ (t << shift); + } + } } - x } /// The delta swaps that rotate a bit index left by `log_d`, as `(shift, -/// mask)` pairs: each bit in `mask` trades places with the bit `shift` above -/// it. Returns the pairs and how many of them are used. +/// mask, in_lane)`: each bit in `mask` trades places with the bit `shift` +/// above it, and `in_lane` says no bit crosses from one `u64` lane to +/// another. Returns the swaps and how many of them are used. /// /// The rotation splits the index bits into `gcd(LOG_BITS, log_d)` cycles. /// Each cycle is walked from its lowest index bit `low`, which is swapped /// with the cycle's other bits `i` in turn: `LOG_BITS - gcd` swaps in all. -const fn rotation_swaps(log_d: u32) -> ([(u32, Word); LOG_BITS as usize], usize) { +const fn rotation_swaps(log_d: u32) -> ([(u32, Word, bool); LOG_BITS as usize], usize) { let (mut cycles, mut rest) = (LOG_BITS, log_d); while rest != 0 { (cycles, rest) = (rest, cycles % rest); } - let mut swaps = [(0, 0); LOG_BITS as usize]; + let mut swaps = [(0, 0, false); LOG_BITS as usize]; let mut n = 0; let mut low = 0; while low < cycles { @@ -313,6 +329,7 @@ const fn rotation_swaps(log_d: u32) -> ([(u32, Word); LOG_BITS as usize], usize) swaps[n] = ( (1 << i) - (1 << low), !index_bit_clear(low) & index_bit_clear(i), + i < LOG_LANE_BITS, ); n += 1; i = (i + log_d) % LOG_BITS; @@ -611,15 +628,15 @@ mod tests { } } - /// Bit `p` must land at `p` rotated left by `LOG_D` within `LOG_BITS` bits. + /// Bit `p` must land at `p` rotated left by `LOG_D` within `LOG_BITS` + /// bits. Word `p` of the block holds only bit `p`, so one call checks + /// every bit. fn check_rotate_bit_index() { - for p in 0..BitTable::BITS { + let mut block: [Word; BitTable::BITS] = std::array::from_fn(|p| 1 << p); + rotate_bit_index::(&mut block); + for (p, &word) in block.iter().enumerate() { let rotated = ((p << LOG_D) | (p >> (LOG_BITS - LOG_D))) % BitTable::BITS; - assert_eq!( - rotate_bit_index::(1 << p), - 1 << rotated, - "LOG_D {LOG_D}, bit {p}" - ); + assert_eq!(word, 1 << rotated, "LOG_D {LOG_D}, bit {p}"); } } diff --git a/crates/prover/src/reduce.rs b/crates/prover/src/reduce.rs index 0262a26f..753e942f 100644 --- a/crates/prover/src/reduce.rs +++ b/crates/prover/src/reduce.rs @@ -31,6 +31,7 @@ fn init_circuit(table: &BitTable, fold: &Fold) -> GrandProductCircuit { } } } + // TODO: probably should be dealt with earlier instead of relying on a fallback Err(TransposeError::ColumnCountTooNarrow) => { // Narrow tables cannot form packed columns after transposition. for (b, &row_image) in fold.row_images.iter().enumerate() { From c56514d2f747e099c0aef73e51764f7ac6013e54 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Thu, 1 Oct 2026 14:19:30 +0800 Subject: [PATCH 16/22] common: replace bit_transpose's derivation comment with a short doc comment --- crates/common/src/table.rs | 22 ++++------------------ 1 file changed, 4 insertions(+), 18 deletions(-) diff --git a/crates/common/src/table.rs b/crates/common/src/table.rs index 21290ed0..db968b91 100644 --- a/crates/common/src/table.rs +++ b/crates/common/src/table.rs @@ -185,24 +185,10 @@ impl> BitTableBase { } } -// Bits packed into machine words -// out of place variant. -// dim1 and dim2 should be given in bits -// dim2 is the axis over which the data is adjacent, think xs[dim1][dim2]. -// -// A block is always BITS words, built from `d = min(dim1, BITS)` rows of -// `r = BITS / d` consecutive words each: word i = c * d + a is row a, word c. -// Writing a block's index bits as (word | bit), the input is -// ([c | a] | [p_hi | p_lo]) with a and p_hi log2(d) bits wide, and the -// output needs ([c | p_hi] | [p_lo | a]). -// -// p_lo only moves within its word, up by log2(d), so that is done first: -// each gathered word's bit index is rotated to [p_lo | p_hi], which -// transposes the word read as d rows of r bits. That leaves a and p_hi to -// trade places bit for bit: log2(d) butterfly stages, stage s swapping word -// index bit s with bit index bit s, which transposes every d x d tile of the -// block. The stages commute and run coarse to fine. For d = BITS the -// rotation is the identity and this is the plain blocked transpose. +/// out of place variant. +/// dim1 and dim2 are in bits +/// dim2 is the axis over which the data is adjacent, think xs[dim1][dim2]. +/// dim1 < WORD::Bits dim2 >= WORD::Bits fn bit_transpose(xs: &[Word], dim1: usize, dim2: usize) -> Vec { assert_eq!(xs.len() * BitTable::BITS, dim1 * dim2); assert!(dim1.is_power_of_two(), "dim1 {dim1} is not a power of two"); From 783a5291067981ba9967107eeed0fade537b6c57 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Mon, 5 Oct 2026 19:57:49 +0800 Subject: [PATCH 17/22] common: move the t >= 7 gate from Shape to BitZParams behind a WordColumns view and let transpose take fewer than 128 columns --- crates/common/src/fold.rs | 27 ++-- crates/common/src/lib.rs | 2 +- crates/common/src/params.rs | 37 ++++- crates/common/src/shape.rs | 15 +- crates/common/src/table.rs | 267 +++++++++++++++++++++----------- crates/prover/benches/prover.rs | 8 +- crates/prover/src/fold.rs | 5 +- crates/prover/src/reduce.rs | 34 ++-- crates/tests/tests/prove.rs | 20 ++- 9 files changed, 271 insertions(+), 144 deletions(-) diff --git a/crates/common/src/fold.rs b/crates/common/src/fold.rs index f29d29ab..d46d69a3 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}; +use crate::{BitTable, LinearClaim, Shape, WordColumns}; /// A round whose parts do not describe the shape they belong to. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -26,7 +26,10 @@ pub enum FoldError { /// `exponents` must be the claim's own — [`LinearClaim::row_exponents`]. /// That is what makes the sum safe: each is below `q` and there are `k_1` of /// them, so it is at most `k_1 (q - 1)`, which admissibility put below `|K|`. -pub fn fold_column(table: &BitTable<'_>, exponents: &[u128], column: usize) -> u128 { +/// +/// Reads the table as whole-element columns, which the `t >= 7` gate in +/// [`crate::BitZParams::new`] guarantees for the table it shapes. +pub fn fold_column(table: &WordColumns<'_>, exponents: &[u128], column: usize) -> u128 { table .column(column) .iter() @@ -46,8 +49,8 @@ pub fn fold_column(table: &BitTable<'_>, exponents: &[u128], column: usize) -> u /// /// Columns are independent and read disjoint slices of the witness, so the /// only sharing is the read-only `exponents`. -pub fn fold_columns(table: &BitTable<'_>, exponents: &[u128]) -> Vec { - let columns = 0..table.shape().columns(); +pub fn fold_columns(table: &WordColumns<'_>, exponents: &[u128]) -> Vec { + let columns = 0..table.columns(); #[cfg(feature = "parallel")] { columns @@ -232,13 +235,17 @@ mod tests { let packed = witness(&shape, &[(1, 0), (5, 0), (127, 0), (64, 4)]); let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); + let columns = table.word_columns().unwrap(); assert_eq!( - fold_column(&table, &claim.row_exponents(), 0), + fold_column(&columns, &claim.row_exponents(), 0), weights[1] + weights[5] + weights[127] ); - assert_eq!(fold_column(&table, &claim.row_exponents(), 4), weights[64]); - assert_eq!(fold_column(&table, &claim.row_exponents(), 9), 0); + assert_eq!( + fold_column(&columns, &claim.row_exponents(), 4), + weights[64] + ); + assert_eq!(fold_column(&columns, &claim.row_exponents(), 9), 0); } #[test] @@ -252,9 +259,10 @@ 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, bytemuck::cast_slice(&packed)).unwrap(); + let columns = table.word_columns().unwrap(); assert_eq!( - fold_column(&table, &claim.row_exponents(), 0), + fold_column(&columns, &claim.row_exponents(), 0), params().fold_bound() ); } @@ -273,12 +281,13 @@ mod tests { .collect(); let packed = witness(&shape, &bits); let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); + let columns = table.word_columns().unwrap(); let expected: u128 = (0..shape.rows()) .filter(|&row| table.bit(6, row)) .map(|row| claim.row_exponents()[row]) .sum(); - assert_eq!(fold_column(&table, &claim.row_exponents(), 6), expected); + assert_eq!(fold_column(&columns, &claim.row_exponents(), 6), expected); } #[test] diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs index 95e09694..eea337fc 100644 --- a/crates/common/src/lib.rs +++ b/crates/common/src/lib.rs @@ -19,7 +19,7 @@ pub use fold::{ pub use opening::OpeningQuery; pub use params::{BitZParams, ParamsError, VirtualParams, VirtualParamsError}; pub use shape::{Shape, ShapeError}; -pub use table::{BitTable, BitTableBase, OwnedBitTable, TableError, TransposeError}; +pub use table::{BitTable, BitTableBase, OwnedBitTable, TableError, WordColumns}; pub use virtual_map::{ TransposedWeights, VirtualMap, VirtualMapError, VirtualStatement, VirtualStatementError, }; diff --git a/crates/common/src/params.rs b/crates/common/src/params.rs index fab4cf5d..72fc0226 100644 --- a/crates/common/src/params.rs +++ b/crates/common/src/params.rs @@ -3,11 +3,15 @@ use field::{F128, gf128::is_generator}; use spongefish::Encoding; -use crate::{BitTable, Shape, TableError, VirtualMap}; +use crate::{BitTable, Shape, TableError, VirtualMap, shape::PACK_BITS}; /// A parameter set one of the pre-claim gates rejects. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ParamsError { + /// Fewer than seven row-index bits: a packed row would not fill one + /// codeword position, and the fold reads a column as whole packed + /// elements. + RowIndexTooNarrow, /// `(k_1 + 1)(Q - 1)` reaches `ord(g)`, so two folds could collide in the /// exponent. FoldBoundExceeded, @@ -35,6 +39,12 @@ impl BitZParams { // `Fq` asserts Q is an odd prime below 2^126 on its own behalf, so // no modulus gate is needed here. + // A table admits any row width; the fold walks the table this shapes + // and is the one that needs whole-element columns. + if shape.log_rows() < PACK_BITS as usize { + return Err(ParamsError::RowIndexTooNarrow); + } + // `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 @@ -305,6 +315,31 @@ mod tests { ); } + #[test] + fn rejects_a_row_index_narrower_than_the_pack_width() { + assert_eq!( + params_at(Shape::new(6, 16).unwrap()).err(), + Some(ParamsError::RowIndexTooNarrow) + ); + assert!(params_at(Shape::new(7, 15).unwrap()).is_ok()); + } + + /// The committed bits only reach the table and the opening, neither of + /// which needs whole-element columns. + #[test] + fn the_committed_shape_may_be_narrower_than_the_pack_width() { + let committed = Shape::new(0, 22).unwrap(); + let params = VirtualParams::new( + params_at(shape()).unwrap(), + committed, + &Dimensions { h_len: 4, f_len: 4 }, + ) + .unwrap(); + let packed = vec![F128::ZERO; (1 << committed.log_bits()) / 128]; + + assert_eq!(params.table(&packed).unwrap().shape(), &committed); + } + #[test] fn rejects_a_generator_of_partial_order() { assert_eq!( diff --git a/crates/common/src/shape.rs b/crates/common/src/shape.rs index 22333b41..12f8d811 100644 --- a/crates/common/src/shape.rs +++ b/crates/common/src/shape.rs @@ -1,6 +1,6 @@ //! The public instance shape and the gates that run before any proof work. -/// The seven low bits of a row index select the basis coefficients that one +/// The seven low bits of a bit index select the basis coefficients that one /// packed field element carries. pub const PACK_BITS: u32 = 7; @@ -16,9 +16,6 @@ pub const MAX_LOG_BITS: usize = 35; /// A shape one of the admissibility constraints rejects. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ShapeError { - /// Fewer than seven row-index bits: a packed row would not fill one - /// codeword position. The paper writes this count `t`. - RowIndexTooNarrow, /// The total bit count falls outside `2^22..=2^35`. CommitmentSizeOutOfRange, } @@ -27,7 +24,6 @@ pub enum ShapeError { /// /// One counts the bits of a column, the other the columns; the paper writes /// them `t` and `s`. Geometry only: the modulus sits in [`crate::BitZParams`], -/// with the one gate that couples the two. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct Shape { log_rows: usize, @@ -37,9 +33,6 @@ pub struct Shape { impl Shape { /// Checks the geometric admissibility constraints and returns the shape. pub fn new(log_rows: usize, log_columns: usize) -> Result { - if log_rows < PACK_BITS as usize { - return Err(ShapeError::RowIndexTooNarrow); - } // The row width alone is bounded first: the total is at least as // large, so an out-of-window row width is already a rejection, and the // later `MAX_LOG_BITS - log_rows` cannot underflow. @@ -111,9 +104,9 @@ mod tests { } #[test] - fn rejects_a_row_index_narrower_than_the_pack_width() { - assert_eq!(Shape::new(6, 16), Err(ShapeError::RowIndexTooNarrow)); - assert!(Shape::new(7, 15).is_ok()); + fn accepts_a_row_index_narrower_than_the_pack_width() { + assert!(Shape::new(6, 16).is_ok()); + assert!(Shape::new(0, MIN_LOG_BITS).is_ok()); } #[test] diff --git a/crates/common/src/table.rs b/crates/common/src/table.rs index db968b91..0f7fdb1a 100644 --- a/crates/common/src/table.rs +++ b/crates/common/src/table.rs @@ -2,14 +2,14 @@ use std::{ops::Deref, slice::Iter}; -use crate::{Shape, ShapeError}; +use crate::Shape; /// The word `packed` is sliced into: one `F128` element. [`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 `BitTable::transpose`. +/// code with `u64` measured 20-55% slower on `WordColumns::transpose`. type Word = u128; /// Bits in a [`Word`]'s bit index. const LOG_BITS: u32 = Word::BITS.trailing_zeros(); @@ -24,21 +24,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 @@ -51,6 +36,8 @@ pub enum TransposeError { /// for the monomial basis `beta_v = X^v`. In that basis bit `v` of an `F128`'s /// little-endian `lo || hi` *is* the coefficient of `X^v`, so the committed /// form and the bit form are the same bytes, and the caller holds one copy. +/// The formula is written for `t >= 7`; in general element `k` holds bit +/// indices `128 k` to `128 k + 127`, which below that is several whole columns. /// /// Generic over the storage `S` of the packed words; use the [`BitTable`] /// (borrowed) and [`OwnedBitTable`] aliases rather than naming this directly. @@ -63,7 +50,7 @@ pub struct BitTableBase { /// A [`BitTableBase`] borrowing the prover's packed witness. pub type BitTable<'a> = BitTableBase<&'a [Word]>; -/// A [`BitTableBase`] owning its packed words, as [`BitTableBase::transpose`] +/// A [`BitTableBase`] owning its packed words, as [`WordColumns::transpose`] /// returns. pub type OwnedBitTable = BitTableBase>; @@ -108,81 +95,101 @@ impl> BitTableBase { ((element >> offset) & 1) == 1 } + /// The table read column by column as whole elements, or `None` when + /// `t < 7` and a column is narrower than one. + pub fn word_columns(&self) -> Option> { + let words = self.shape.rows() / BitTable::BITS; + (words > 0).then(|| WordColumns { + packed: &self.packed, + words, + }) + } + + /// The column's bits, in ascending row order. + pub fn column_bits(&self, column: usize) -> BitsIter<'_> { + if let Some(columns) = self.word_columns() { + return BitsIter::new(columns.column(column).iter()); + } + // A column narrower than an element shares it with its neighbours: + // shift it down to bit 0 and stop after its `2^t` rows. + let start = column << self.shape.log_rows(); + BitsIter::preloaded( + self.packed[start / BitTable::BITS] >> (start % BitTable::BITS), + self.shape.rows(), + ) + } + + /// Unwraps the packed bits. + pub fn into_packed(self) -> S { + self.packed + } +} + +/// A table's columns as whole elements, which needs `t >= 7`; from +/// [`BitTableBase::word_columns`]. The fold and [`WordColumns::transpose`] +/// both read the table this way. +/// +// Refactoring Seam might be removed later. +#[derive(Debug, Clone, Copy)] +pub struct WordColumns<'a> { + packed: &'a [Word], + /// Elements per column, `2^(t-7)`. + words: usize, +} + +impl<'a> WordColumns<'a> { /// The `2^(t-7)` elements holding one column, in ascending row order. /// /// A column starts at bit `c * 2^t` and `t >= 7`, so it begins on an /// 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) -> &[Word] { - let elements = self.shape.rows() / BitTable::BITS; - &self.packed[column * elements..(column + 1) * elements] + pub fn column(&self, column: usize) -> &'a [Word] { + &self.packed[column * self.words..(column + 1) * self.words] } - /// The column's bits, in ascending row order. - pub fn column_bits(&self, column: usize) -> BitsIter<'_> { - BitsIter::new(self.column(column).iter()) + /// The number of columns. + pub fn columns(&self) -> usize { + self.packed.len() / self.words } /// 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`. + /// table.bit(c, b)` for every row `b` and column `c` of the table. /// /// 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 + /// row by row: the packing is column-major (see [`BitTableBase`]), so a + /// row-major walk over the table 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. + /// - **Row count**: `bit_transpose` reads each column as whole words, + /// which is what this view guarantees. The column count is free: with + /// fewer than `128` columns the result's columns share words. /// - **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. + /// sits in `[MIN_LOG_BITS, MAX_LOG_BITS]` whenever the table 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 packed = bit_transpose(&self.packed, self.shape.columns(), self.shape.rows()); - - Ok(OwnedBitTable { shape, packed }) - } + /// the same size as the packed witness (up to 4 GiB at `m = 35`), never + /// a view over the original. + /// - **Involution**: transposing the result returns the table's bits + /// exactly, provided it has at least `128` columns. With fewer, the + /// result has no word view of its own to transpose. + pub fn transpose(&self) -> OwnedBitTable { + let (columns, rows) = (self.columns(), self.words * BitTable::BITS); + let shape = Shape::new(columns.ilog2() as usize, rows.ilog2() as usize).expect( + "log_bits is preserved by swapping row/column axes, and the \ + commitment-size window is all a shape is checked against", + ); - /// Borrows this table as an ordinary [`BitTable`]. - pub fn as_table(&self) -> BitTable<'_> { - BitTable { - shape: self.shape, - packed: &self.packed, + OwnedBitTable { + shape, + packed: bit_transpose(self.packed, columns, rows), } } - - /// Unwraps the packed bits. - pub fn into_packed(self) -> S { - self.packed - } } /// out of place variant. @@ -262,10 +269,13 @@ const fn index_bit_clear(s: u32) -> Word { Word::MAX / ((1 << (1 << s)) + 1) } -/// Rotates the bit index of every bit in every word of `block` left by -/// `LOG_D`: bit `[p_hi | p_lo]` moves to `[p_lo | p_hi]`, `p_hi` being -/// `LOG_D` bits wide. Read as `d = 2^LOG_D` rows of `r = BITS / d` bits, -/// each word is transposed. One pass over the block per entry of +/// Interleaves the `d = 2^LOG_D` chunks of `r = BITS / d` bits in every word +/// of `block`, a `d`-way perfect shuffle: bit `j` of chunk `k` moves to bit +/// `j * d + k`. On bit indices that is a left rotation by `LOG_D`: bit +/// `[p_hi | p_lo] = [k | j]` moves to `[p_lo | p_hi]`, so unlike +/// `rotate_left` each bit moves its own distance. That puts `p_hi` in the low +/// `LOG_D` index bits, where [`bit_block_transpose`] swaps it with the row +/// bits of the word index. One pass over the block per entry of /// [`rotation_swaps`], none for `LOG_D = LOG_BITS`. /// /// A pass that stays inside `u64` lanes runs on the lanes rather than on @@ -296,9 +306,13 @@ fn rotate_bit_index(block: &mut [Word; Word::BITS as usize]) { /// above it, and `in_lane` says no bit crosses from one `u64` lane to /// another. Returns the swaps and how many of them are used. /// -/// The rotation splits the index bits into `gcd(LOG_BITS, log_d)` cycles. -/// Each cycle is walked from its lowest index bit `low`, which is swapped -/// with the cycle's other bits `i` in turn: `LOG_BITS - gcd` swaps in all. +/// Each swap exchanges two bits of the bit index, `low` and `i`: the bits +/// whose index has `low` set and `i` clear move up by `shift`, their partners +/// move down by it, and the other half of the word stays put. The rotation +/// splits the index bits into `gcd(LOG_BITS, log_d)` cycles. Each cycle is +/// walked from its lowest index bit `low`, which is swapped with the cycle's +/// other bits `i` in turn: `LOG_BITS - gcd` swaps in all, the fewest index +/// bit swaps that make the rotation. const fn rotation_swaps(log_d: u32) -> ([(u32, Word, bool); LOG_BITS as usize], usize) { let (mut cycles, mut rest) = (LOG_BITS, log_d); while rest != 0 { @@ -420,6 +434,18 @@ impl<'a, const S: usize> StepIter<'a, S> { remaining: 0, } } + + /// Emits the low `bits` bits of `word` and nothing after; `bits` is a + /// multiple of `S` and at most `Word::BITS`. + fn preloaded(word: Word, bits: usize) -> Self { + let () = Self::CHECK_SIZE; + debug_assert!(bits.is_multiple_of(S) && bits <= Word::BITS as usize); + Self { + elements: [].iter(), + word, + remaining: bits as u8, + } + } } impl Iterator for StepIter<'_, S> { @@ -545,10 +571,22 @@ mod tests { let shape = small_shape(); let packed = with_bits(&shape, &[(0, 3), (2, 3), (64, 3)]); let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); + let columns = table.word_columns().unwrap(); + + assert_eq!(columns.columns(), shape.columns()); + assert_eq!(columns.column(3).len(), shape.rows() / BitTable::BITS); + assert_eq!(columns.column(3), [0b101 | (1 << 64)]); + assert!(columns.column(4).iter().all(|&element| element == 0)); + } + + /// So neither the fold nor `transpose` can be handed it. + #[test] + fn a_column_narrower_than_an_element_has_no_word_view() { + let shape = Shape::new(6, 16).unwrap(); + let packed = vec![F128::ZERO; (1 << shape.log_bits()) / PACKED_BITS]; + let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); - 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)); + assert!(table.word_columns().is_none()); } #[test] @@ -574,7 +612,7 @@ mod tests { let packed = with_bits(&shape, &rows); let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); - for (index, &element) in table.column(2).iter().enumerate() { + for (index, &element) in table.word_columns().unwrap().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}"); @@ -678,6 +716,31 @@ mod tests { ); } + /// Every row width below one element, where columns share elements. + #[test] + fn column_bits_agrees_with_the_bit_accessor_below_one_element() { + for log_rows in 0..PACKED_SHIFT as usize { + let shape = Shape::new(log_rows, 22 - log_rows).unwrap(); + let packed: Vec = (0..(1 << shape.log_bits()) / BitTable::BITS) + .map(pseudo_random_word) + .collect(); + let table = BitTable::new(shape, &packed).unwrap(); + + // Both ends of the first, a middle and the last element. + let per_element = BitTable::BITS / shape.rows(); + let last = shape.columns() - 1; + for column in [0, 1, per_element - 1, per_element, last / 2, last] { + let iter = table.column_bits(column); + assert_eq!(iter.len(), shape.rows()); + let bits: Vec = iter.map(|b| b == 1).collect(); + let expected: Vec = (0..shape.rows()) + .map(|row| table.bit(column, row)) + .collect(); + assert_eq!(bits, expected, "t {log_rows}, column {column}"); + } + } + } + /// 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. @@ -711,7 +774,7 @@ mod tests { let packed = pseudo_random_table(&shape); let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); - let transposed = table.transpose().unwrap(); + let transposed = table.word_columns().unwrap().transpose(); assert_eq!(transposed.shape().log_rows(), shape.log_columns()); assert_eq!(transposed.shape().log_columns(), shape.log_rows()); @@ -736,27 +799,47 @@ mod tests { let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); let roundtripped = table - .transpose() + .word_columns() .unwrap() .transpose() + .word_columns() .unwrap() + .transpose() .into_packed(); assert_eq!(roundtripped, bytemuck::cast_slice::(&packed)); } + /// Every column count below one element, down to a single column: the + /// result's rows share elements, and its columns read back through + /// `column_bits` as the source's rows. #[test] - 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. - 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(); - - assert_eq!( - table.transpose().err(), - Some(TransposeError::ColumnCountTooNarrow) - ); + fn transpose_handles_fewer_columns_than_an_element() { + for log_columns in 0..PACKED_SHIFT as usize { + let shape = Shape::new(22 - log_columns, log_columns).unwrap(); + let packed: Vec = (0..(1 << shape.log_bits()) / BitTable::BITS) + .map(pseudo_random_word) + .collect(); + let table = BitTable::new(shape, &packed).unwrap(); + + let transposed = table.word_columns().unwrap().transpose(); + assert_eq!(transposed.shape().log_rows(), log_columns); + assert_eq!(transposed.shape().log_columns(), shape.log_rows()); + + for row in [ + 0, + 1, + PACKED_BITS - 1, + PACKED_BITS, + shape.rows() / 2, + shape.rows() - 1, + ] { + let bits: Vec = transposed.column_bits(row).map(|b| b == 1).collect(); + let expected: Vec = (0..shape.columns()) + .map(|column| table.bit(column, row)) + .collect(); + assert_eq!(bits, expected, "s {log_columns}, row {row}"); + } + } } } diff --git a/crates/prover/benches/prover.rs b/crates/prover/benches/prover.rs index 6c70a9f4..4f9c9377 100644 --- a/crates/prover/benches/prover.rs +++ b/crates/prover/benches/prover.rs @@ -29,10 +29,10 @@ fn random_table(shape: Shape) -> BitTable<'static> { #[divan::bench] fn gkr(bencher: Bencher) { // `Shape::new(log_rows, log_columns)` rejects fewer than `2^22` committed - // bits total (`MIN_LOG_BITS`) and a row width under 7 (`PACK_BITS`), so - // the two args must sum to at least 22. That floor isn't a security bound - // computed here -- it's the range flock-core ships precomputed Ligerito - // configs for (see the comment on `MIN_LOG_BITS`). + // bits total (`MIN_LOG_BITS`), and `BitZParams::new` a row width under 7 + // (`PACK_BITS`), so the two args must sum to at least 22. That floor + // isn't a security bound computed here -- it's the range flock-core ships + // precomputed Ligerito configs for (see the comment on `MIN_LOG_BITS`). let shape = Shape::new(10, 15).unwrap(); let table = random_table(shape); diff --git a/crates/prover/src/fold.rs b/crates/prover/src/fold.rs index 8c272fe5..c2b27a92 100644 --- a/crates/prover/src/fold.rs +++ b/crates/prover/src/fold.rs @@ -48,7 +48,10 @@ impl BitZProver { // Lifted once: the fold reads it per set bit across every column. let exponents = claim.row_exponents(); - let folds = fold_columns(table, &exponents); + let columns = table + .word_columns() + .expect("the shape matches the parameters, whose gate admits only t >= 7"); + let folds = fold_columns(&columns, &exponents); for fold in &folds { transcript.prover_message(&fold.to_le_bytes()); diff --git a/crates/prover/src/reduce.rs b/crates/prover/src/reduce.rs index 753e942f..00410437 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; @@ -21,29 +21,15 @@ fn init_circuit(table: &BitTable, fold: &Fold) -> GrandProductCircuit { // 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. - 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 == 1 { row_image } else { F128::ONE }; - } - } - } - // TODO: probably should be dealt with earlier instead of relying on a fallback - 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 - }; - } - } + // Transposed so each row can be read sequentially. + let transposed = table + .word_columns() + .expect("the fold's table, whose gate admits only t >= 7") + .transpose(); + 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 == 1 { row_image } else { F128::ONE }; } } diff --git a/crates/tests/tests/prove.rs b/crates/tests/tests/prove.rs index 3c87650f..49160585 100644 --- a/crates/tests/tests/prove.rs +++ b/crates/tests/tests/prove.rs @@ -1,6 +1,6 @@ //! The top-level prove and verify, through the real opening. -use common::{Root, TableError}; +use common::{Root, Shape, TableError}; use field::{F128, Fq}; use num_traits::{ConstOne, ConstZero}; use pcs::{HashKind, LigeritoProfile, Pcs, VerifyError as PcsVerifyError}; @@ -42,6 +42,24 @@ fn an_honest_proof_verifies_on_both_floor_shapes() { } } +/// Fewer columns than one element holds: the leaves are read from a +/// transpose whose rows share elements. +#[test] +fn an_honest_proof_verifies_with_fewer_than_128_columns() { + let instance = Instance::honest(Shape::new(16, 6).unwrap(), 31); + let proof = prove(&instance); + + instance + .verifier + .verify( + &instance.claim, + &instance.pcs, + instance.com, + verifier_transcript(&proof), + ) + .unwrap(); +} + #[test] fn a_proof_replayed_under_a_different_commitment_is_refused() { let instance = Instance::honest(narrow_shape(), 32); From f8405bca6cb2045cea3011ba54999e61392fcd0f Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Tue, 6 Oct 2026 14:37:32 +0800 Subject: [PATCH 18/22] fixup! common: move the t >= 7 gate from Shape to BitZParams behind a WordColumns view and let transpose take fewer than 128 columns --- crates/common/src/params.rs | 16 ---------------- 1 file changed, 16 deletions(-) diff --git a/crates/common/src/params.rs b/crates/common/src/params.rs index 72fc0226..077a50e8 100644 --- a/crates/common/src/params.rs +++ b/crates/common/src/params.rs @@ -324,22 +324,6 @@ mod tests { assert!(params_at(Shape::new(7, 15).unwrap()).is_ok()); } - /// The committed bits only reach the table and the opening, neither of - /// which needs whole-element columns. - #[test] - fn the_committed_shape_may_be_narrower_than_the_pack_width() { - let committed = Shape::new(0, 22).unwrap(); - let params = VirtualParams::new( - params_at(shape()).unwrap(), - committed, - &Dimensions { h_len: 4, f_len: 4 }, - ) - .unwrap(); - let packed = vec![F128::ZERO; (1 << committed.log_bits()) / 128]; - - assert_eq!(params.table(&packed).unwrap().shape(), &committed); - } - #[test] fn rejects_a_generator_of_partial_order() { assert_eq!( From 706a2ae05447c5e98a38d0cc5517b8af03319ad4 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Wed, 7 Oct 2026 16:03:59 +0800 Subject: [PATCH 19/22] gkr: move leaf construction into GrandProductCircuit::new, taking row images and the transposed table --- crates/gkr/src/lib.rs | 30 +++++++++++++++++++++++------- crates/prover/src/reduce.rs | 12 +----------- crates/verifier/src/reduce.rs | 16 +++------------- 3 files changed, 27 insertions(+), 31 deletions(-) diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs index 1f390e6a..1ab35f25 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs @@ -1,5 +1,6 @@ use std::collections::VecDeque; +use common::OwnedBitTable; use field::{F128, Wide256}; use num_traits::{ConstOne, ConstZero}; use rayon::prelude::*; @@ -342,8 +343,23 @@ pub struct GrandProductCircuit { } impl GrandProductCircuit { - // Fails if the leafs are 0. - pub fn new(mut leafs: Vec) -> Self { + pub fn new(row_images: &[F128], table: OwnedBitTable) -> Self { + let columns = table.shape().columns(); + let rows = table.shape().rows(); + let dim = columns * rows; + let mut leafs = F128::zeroed_vec(dim); + + for (b, &row_image) in row_images.iter().enumerate() { + let leafs = &mut leafs[b * rows..(b + 1) * rows]; + for (leaf, bit) in leafs.iter_mut().zip(table.column_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 +444,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 +466,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 +527,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 +573,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 +597,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 00410437..6166b9a7 100644 --- a/crates/prover/src/reduce.rs +++ b/crates/prover/src/reduce.rs @@ -15,10 +15,6 @@ 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. // Transposed so each row can be read sequentially. @@ -26,14 +22,8 @@ fn init_circuit(table: &BitTable, fold: &Fold) -> GrandProductCircuit { .word_columns() .expect("the fold's table, whose gate admits only t >= 7") .transpose(); - 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 == 1 { row_image } else { F128::ONE }; - } - } - GrandProductCircuit::new(leafs) + GrandProductCircuit::new(&fold.row_images, transposed) } /// Reduces the grand-product circuit to a factored claim on the committed bits. diff --git a/crates/verifier/src/reduce.rs b/crates/verifier/src/reduce.rs index eec8a599..a0cc92cb 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,8 @@ 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 transposed = table.word_columns().unwrap().transpose(); + let circuit = GrandProductCircuit::new(&row_images, transposed); let (top_layer, witnesses) = circuit.batched_eval(shape.columns()); // Fold::new evaluates top_layer at zeta to obtain GKR's initial claim. From 336e8097225aca4f8146f35594740b0e5054d839 Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Wed, 7 Oct 2026 19:29:41 +0800 Subject: [PATCH 20/22] common: split BitTable into a shaped BitTable view and a geometry-only BitMatrix, moving the t >= 7 gate back to Shape --- crates/common/Cargo.toml | 2 +- crates/common/benches/{table.rs => matrix.rs} | 4 +- crates/common/src/fold.rs | 27 +- crates/common/src/lib.rs | 4 +- crates/common/src/matrix.rs | 614 ++++++++++++++++ crates/common/src/params.rs | 21 +- crates/common/src/shape.rs | 15 +- crates/common/src/table.rs | 689 +----------------- crates/gkr/src/lib.rs | 13 +- crates/prover/benches/prover.rs | 8 +- crates/prover/src/fold.rs | 5 +- crates/prover/src/reduce.rs | 5 +- crates/verifier/src/reduce.rs | 2 +- 13 files changed, 691 insertions(+), 718 deletions(-) rename crates/common/benches/{table.rs => matrix.rs} (90%) create mode 100644 crates/common/src/matrix.rs diff --git a/crates/common/Cargo.toml b/crates/common/Cargo.toml index 260d5963..b22086da 100644 --- a/crates/common/Cargo.toml +++ b/crates/common/Cargo.toml @@ -25,6 +25,6 @@ divan = { workspace = true } rand = { workspace = true } [[bench]] -name = "table" +name = "matrix" harness = false required-features = ["bench"] diff --git a/crates/common/benches/table.rs b/crates/common/benches/matrix.rs similarity index 90% rename from crates/common/benches/table.rs rename to crates/common/benches/matrix.rs index 2aa6a4fc..7ec9b350 100644 --- a/crates/common/benches/table.rs +++ b/crates/common/benches/matrix.rs @@ -1,10 +1,10 @@ //! Compare the blocked `bit_transpose` with the bit-at-a-time reference. //! -//! Run with `cargo bench -p common --features bench --bench table`. +//! 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::table::bench::{bit_transpose, transpose_reference}; +use common::matrix::bench::{bit_transpose, transpose_reference}; use divan::counter::BytesCount; use divan::{Bencher, black_box}; use rand::rngs::StdRng; diff --git a/crates/common/src/fold.rs b/crates/common/src/fold.rs index d46d69a3..f29d29ab 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, WordColumns}; +use crate::{BitTable, LinearClaim, Shape}; /// A round whose parts do not describe the shape they belong to. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -26,10 +26,7 @@ pub enum FoldError { /// `exponents` must be the claim's own — [`LinearClaim::row_exponents`]. /// That is what makes the sum safe: each is below `q` and there are `k_1` of /// them, so it is at most `k_1 (q - 1)`, which admissibility put below `|K|`. -/// -/// Reads the table as whole-element columns, which the `t >= 7` gate in -/// [`crate::BitZParams::new`] guarantees for the table it shapes. -pub fn fold_column(table: &WordColumns<'_>, exponents: &[u128], column: usize) -> u128 { +pub fn fold_column(table: &BitTable<'_>, exponents: &[u128], column: usize) -> u128 { table .column(column) .iter() @@ -49,8 +46,8 @@ pub fn fold_column(table: &WordColumns<'_>, exponents: &[u128], column: usize) - /// /// Columns are independent and read disjoint slices of the witness, so the /// only sharing is the read-only `exponents`. -pub fn fold_columns(table: &WordColumns<'_>, exponents: &[u128]) -> Vec { - let columns = 0..table.columns(); +pub fn fold_columns(table: &BitTable<'_>, exponents: &[u128]) -> Vec { + let columns = 0..table.shape().columns(); #[cfg(feature = "parallel")] { columns @@ -235,17 +232,13 @@ mod tests { let packed = witness(&shape, &[(1, 0), (5, 0), (127, 0), (64, 4)]); let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); - let columns = table.word_columns().unwrap(); assert_eq!( - fold_column(&columns, &claim.row_exponents(), 0), + fold_column(&table, &claim.row_exponents(), 0), weights[1] + weights[5] + weights[127] ); - assert_eq!( - fold_column(&columns, &claim.row_exponents(), 4), - weights[64] - ); - assert_eq!(fold_column(&columns, &claim.row_exponents(), 9), 0); + assert_eq!(fold_column(&table, &claim.row_exponents(), 4), weights[64]); + assert_eq!(fold_column(&table, &claim.row_exponents(), 9), 0); } #[test] @@ -259,10 +252,9 @@ 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, bytemuck::cast_slice(&packed)).unwrap(); - let columns = table.word_columns().unwrap(); assert_eq!( - fold_column(&columns, &claim.row_exponents(), 0), + fold_column(&table, &claim.row_exponents(), 0), params().fold_bound() ); } @@ -281,13 +273,12 @@ mod tests { .collect(); let packed = witness(&shape, &bits); let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); - let columns = table.word_columns().unwrap(); let expected: u128 = (0..shape.rows()) .filter(|&row| table.bit(6, row)) .map(|row| claim.row_exponents()[row]) .sum(); - assert_eq!(fold_column(&columns, &claim.row_exponents(), 6), expected); + assert_eq!(fold_column(&table, &claim.row_exponents(), 6), expected); } #[test] diff --git a/crates/common/src/lib.rs b/crates/common/src/lib.rs index eea337fc..d5c5e7ec 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, BitTableBase, OwnedBitTable, TableError, WordColumns}; +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 00000000..5af32cb9 --- /dev/null +++ b/crates/common/src/matrix.rs @@ -0,0 +1,614 @@ +//! Bit matrices in memory order, and the transpose between them. + +use std::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; +/// Bits in a [`Word`]'s bit index. +const LOG_BITS: u32 = Word::BITS.trailing_zeros(); +/// Bits in a `u64`'s bit index. Swapping two bit index bits below this never +/// moves a bit from one of a word's `u64` lanes to another. +const LOG_LANE_BITS: u32 = u64::BITS.trailing_zeros(); + +/// 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. +#[derive(Debug, Clone)] +pub struct BitMatrix { + packed: Vec, + log_dim2: usize, +} + +impl BitMatrix { + /// 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: Vec, log_dim2: usize) -> Self { + let bits = packed.len() * BITS; + assert!( + bits.is_power_of_two() && log_dim2 <= bits.ilog2() as usize, + "{bits} bits do not split into rows of 2^{log_dim2}" + ); + Self { packed, log_dim2 } + } + + /// Transposes `xs[dim1][2^log_dim2]` into an owned `[dim2][dim1]`, as + /// [`BitMatrix::transpose`] does, reading the words where they lie. + pub(crate) fn transposed(xs: &[Word], log_dim2: usize) -> Self { + let dim2 = 1 << log_dim2; + let dim1 = xs.len() * BITS / dim2; + Self::new(bit_transpose(xs, dim1, dim2), dim1.ilog2() as usize) + } + + /// 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 bit `xs[i1][i2]`. + pub fn bit(&self, i1: usize, i2: usize) -> bool { + debug_assert!(i1 < self.dim1(), "i1 {i1} is outside dim1"); + debug_assert!(i2 < self.dim2(), "i2 {i2} is outside dim2"); + + let index = (i1 << self.log_dim2) | i2; + (self.packed[index / BITS] >> (index % BITS)) & 1 == 1 + } + + /// The `dim2` bits of `xs[i1]`, in ascending order. + pub fn bits(&self, i1: usize) -> BitsIter<'_> { + let start = i1 << self.log_dim2; + if self.dim2() >= BITS { + let words = self.dim2() / BITS; + return BitsIter::new(self.packed[start / BITS..][..words].iter()); + } + // Shorter than a word, so it shares one with its neighbours: shift it + // down to bit 0 and stop after its `dim2` bits. + BitsIter::preloaded(self.packed[start / BITS] >> (start % BITS), self.dim2()) + } + + /// Swaps the axes, returning an owned copy: `result.bit(i2, i1) == + /// self.bit(i1, i2)` for every `i1` and `i2`. + /// + /// Built for callers that must walk the strided axis: in place that is a + /// scatter, one `bit()` call and cache miss per bit. Transposing once up + /// front turns it into the sequential access [`BitMatrix::bits`] gives. + /// + /// # Constraints + /// + /// - **`dim2`**: at least `128`, since `bit_transpose` reads each + /// `xs[i1]` as whole words. `dim1` is free: below `128` the result's + /// `xs[i2]` share words. + /// - **Allocation**: always a full out-of-place copy -- a fresh buffer + /// the same size as the packed words (up to 4 GiB at `m = 35`), never + /// a view over the original. + /// - **Involution**: transposing the result returns the original bits, + /// provided `dim1` is at least `128` as well. + pub fn transpose(&self) -> BitMatrix { + Self::transposed(&self.packed, self.log_dim2) + } + + /// 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. + pub fn into_packed(self) -> Vec { + self.packed + } +} + +/// out of place variant. +/// dim1 and dim2 are in bits +/// dim2 is the axis over which the data is adjacent, think xs[dim1][dim2]. +/// dim1 < WORD::Bits dim2 >= WORD::Bits +fn bit_transpose(xs: &[Word], dim1: usize, dim2: usize) -> Vec { + assert_eq!(xs.len() * BITS, dim1 * dim2); + assert!(dim1.is_power_of_two(), "dim1 {dim1} is not a power of two"); + + // One body, compiled once per block height. With `log2(d)` a runtime + // value the stage loops cannot be unrolled with constant masks and + // distances, which cost ~30% at d = 64. + match dim1.min(BITS).trailing_zeros() { + // A single row is its own transpose. + 0 => xs.to_vec(), + 1 => bit_transpose_blocks::<1>(xs, dim1, dim2), + 2 => bit_transpose_blocks::<2>(xs, dim1, dim2), + 3 => bit_transpose_blocks::<3>(xs, dim1, dim2), + 4 => bit_transpose_blocks::<4>(xs, dim1, dim2), + 5 => bit_transpose_blocks::<5>(xs, dim1, dim2), + 6 => bit_transpose_blocks::<6>(xs, dim1, dim2), + _ => bit_transpose_blocks::(xs, dim1, dim2), + } +} + +/// [`bit_transpose`] for blocks of `d = 2^LOG_D` rows, `LOG_D >= 1`. +fn bit_transpose_blocks(xs: &[Word], dim1: usize, dim2: usize) -> Vec { + let d = 1 << LOG_D; // rows per block + let r = BITS / d; // words per row per block + let row_words = dim2 / BITS; + // Output words between consecutive words of a block; 1 when d < BITS. + let out_stride = dim1 / d; + assert!( + dim2.is_multiple_of(BITS) && row_words.is_multiple_of(r), + "dim2 {dim2} does not fill whole {d}-row blocks" + ); + + let mut out = bytemuck::zeroed_vec(xs.len()); + + // 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. g1 only takes more than one value + // when d = BITS. + for g2 in 0..row_words / r { + for g1 in 0..dim1 / d { + // Word i = [c | a] is row a, word c. Two loops rather than one + // over i: splitting i back into a and c measured 4-12% slower + // for d < BITS. + let mut block = [0; BITS]; + for c in 0..r { + for a in 0..d { + block[c * d + a] = xs[(g1 * d + a) * row_words + g2 * r + c]; + } + } + + rotate_bit_index::(&mut block); + bit_block_transpose(LOG_D, &mut block); + + // o is output row g2 * BITS + o, word g1 in [dim2][dim1]; g1 is + // only nonzero when d = BITS. + for (o, &word) in block.iter().enumerate() { + out[(g2 * BITS + o) * out_stride + g1] = word; + } + } + } + + out +} + +/// The bits whose index has bit `s` clear: alternating runs of `2^s` bits, +/// lowest run set. 01010101, 00110011, 00001111 for `s` = 0, 1, 2. +const fn index_bit_clear(s: u32) -> Word { + Word::MAX / ((1 << (1 << s)) + 1) +} + +/// Interleaves the `d = 2^LOG_D` chunks of `r = BITS / d` bits in every word +/// of `block`, a `d`-way perfect shuffle: bit `j` of chunk `k` moves to bit +/// `j * d + k`. On bit indices that is a left rotation by `LOG_D`: bit +/// `[p_hi | p_lo] = [k | j]` moves to `[p_lo | p_hi]`, so unlike +/// `rotate_left` each bit moves its own distance. That puts `p_hi` in the low +/// `LOG_D` index bits, where [`bit_block_transpose`] swaps it with the row +/// bits of the word index. One pass over the block per entry of +/// [`rotation_swaps`], none for `LOG_D = LOG_BITS`. +/// +/// A pass that stays inside `u64` lanes runs on the lanes rather than on +/// whole words, which vectorises: on `u128` that made the rotation about +/// 40% cheaper. Only a swap with bit index bit 6 or above needs whole words. +#[inline(always)] +fn rotate_bit_index(block: &mut [Word; Word::BITS as usize]) { + // Evaluated at compile time, so every shift and mask is a constant. + let (swaps, n) = const { rotation_swaps(LOG_D) }; + for &(shift, mask, in_lane) in &swaps[..n] { + if in_lane { + let mask = mask as u64; + for x in bytemuck::cast_slice_mut::(block) { + let t = ((*x >> shift) ^ *x) & mask; + *x ^= t ^ (t << shift); + } + } else { + for x in block.iter_mut() { + let t = ((*x >> shift) ^ *x) & mask; + *x ^= t ^ (t << shift); + } + } + } +} + +/// The delta swaps that rotate a bit index left by `log_d`, as `(shift, +/// mask, in_lane)`: each bit in `mask` trades places with the bit `shift` +/// above it, and `in_lane` says no bit crosses from one `u64` lane to +/// another. Returns the swaps and how many of them are used. +/// +/// Each swap exchanges two bits of the bit index, `low` and `i`: the bits +/// whose index has `low` set and `i` clear move up by `shift`, their partners +/// move down by it, and the other half of the word stays put. The rotation +/// splits the index bits into `gcd(LOG_BITS, log_d)` cycles. Each cycle is +/// walked from its lowest index bit `low`, which is swapped with the cycle's +/// other bits `i` in turn: `LOG_BITS - gcd` swaps in all, the fewest index +/// bit swaps that make the rotation. +const fn rotation_swaps(log_d: u32) -> ([(u32, Word, bool); LOG_BITS as usize], usize) { + let (mut cycles, mut rest) = (LOG_BITS, log_d); + while rest != 0 { + (cycles, rest) = (rest, cycles % rest); + } + + let mut swaps = [(0, 0, false); LOG_BITS as usize]; + let mut n = 0; + let mut low = 0; + while low < cycles { + let mut i = (low + log_d) % LOG_BITS; + while i != low { + // The bits whose index has bit low set and bit i clear. + swaps[n] = ( + (1 << i) - (1 << low), + !index_bit_clear(low) & index_bit_clear(i), + i < LOG_LANE_BITS, + ); + n += 1; + i = (i + log_d) % LOG_BITS; + } + low += 1; + } + (swaps, n) +} + +/// Transposes every `2^log_d x 2^log_d` tile of a BITS x BITS block, the +/// tile's words being those that differ only in their low `log_d` index +/// bits. Stage `s` swaps bit index bit `s` with word index bit `s`, coarse +/// to fine; the stages commute, so the order is free. `log_d = LOG_BITS` is +/// the full transpose. +#[inline(always)] +fn bit_block_transpose(log_d: u32, xs: &mut [Word; Word::BITS as usize]) { + for s in (0..log_d).rev() { + // Distance between the paired bits, and between the paired words. + let j: usize = 1 << s; + let mask = index_bit_clear(s); + let mut k: usize = 0; + while k < Word::BITS as usize { + // A delta swap across the pair: t is where the upper runs of + // xs[k] and the lower runs of xs[k | j] differ, and flipping those + // bits in both swaps the runs. Measured 6-10% faster than + // rebuilding both words from masked halves. + let t = ((xs[k] >> j) ^ xs[k | j]) & mask; + xs[k | j] ^= t; + xs[k] ^= t << j; + // (k|j) count the upper part. +1 advances the upper part. !j converts it into the lower index again except for when it hits the next round. In that case the | j was + k = ((k | j) + 1) & !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 `bit()` and `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, + } + } + + /// Emits the low `bits` bits of `word` and nothing after; `bits` is a + /// multiple of `S` and at most `Word::BITS`. + fn preloaded(word: Word, bits: usize) -> Self { + let () = Self::CHECK_SIZE; + debug_assert!(bits.is_multiple_of(S) && bits <= Word::BITS as usize); + Self { + elements: [].iter(), + word, + remaining: bits as u8, + } + } +} + +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 { + let packed = (0..(1 << (log_dim1 + log_dim2)) / BITS) + .map(pseudo_random_word) + .collect(); + BitMatrix::new(packed, log_dim2) + } + + #[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!(matrix.bit(i1, i2), (i1, i2) == (1, 130), "xs[{i1}][{i2}]"); + } + } + } + + /// `dim2` of whole words, and shorter than one, where neighbours share it. + #[test] + fn bits_agrees_with_the_bit_accessor() { + for log_dim2 in 0..=LOG_BITS as usize + 1 { + 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| matrix.bit(i1, i2)).collect(); + assert_eq!(bits, expected, "log_dim2 {log_dim2}, i1 {i1}"); + } + } + } + + #[test] + fn bit_block_transpose_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; + bit_block_transpose(LOG_BITS, &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 `p` must land at `p` rotated left by `LOG_D` within `LOG_BITS` + /// bits. Word `p` of the block holds only bit `p`, so one call checks + /// every bit. + fn check_rotate_bit_index() { + let mut block: [Word; BITS] = std::array::from_fn(|p| 1 << p); + rotate_bit_index::(&mut block); + for (p, &word) in block.iter().enumerate() { + let rotated = ((p << LOG_D) | (p >> (LOG_BITS - LOG_D))) % BITS; + assert_eq!(word, 1 << rotated, "LOG_D {LOG_D}, bit {p}"); + } + } + + #[test] + fn rotate_bit_index_rotates_every_bit_index() { + check_rotate_bit_index::<1>(); + check_rotate_bit_index::<2>(); + check_rotate_bit_index::<3>(); + check_rotate_bit_index::<4>(); + check_rotate_bit_index::<5>(); + check_rotate_bit_index::<6>(); + check_rotate_bit_index::<7>(); + } + + /// `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_reference_for_every_row_count() { + // Two blocks along dim2 at every row count, so block placement is + // exercised as well as the block contents. + check_bit_transpose::<1>(256); + check_bit_transpose::<2>(128); + check_bit_transpose::<4>(64); + check_bit_transpose::<8>(32); + check_bit_transpose::<16>(16); + check_bit_transpose::<32>(8); + check_bit_transpose::<64>(4); + check_bit_transpose::<128>(2); + check_bit_transpose::<256>(2); + } + + #[test] + fn transpose_reference_matches_bit_transpose() { + // 256 segments of 2 words each: two block groups on both axes. + const GS: usize = 256; + let segment_n = 2; + let tt: Vec = (0..GS * segment_n).map(pseudo_random_word).collect(); + + let reference = transpose_reference::(&tt); + assert_eq!(reference.len(), tt.len()); + assert_eq!(reference, bit_transpose(&tt, GS, segment_n * BITS)); + } + + /// `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!( + transposed.bit(i2, i1), + matrix.bit(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 `dim1` below one word, down to one: the result's `xs[i2]` are + /// shorter than a word and share it. + #[test] + fn transpose_handles_dim1_below_a_word() { + for log_dim1 in 0..LOG_BITS as usize { + let matrix = pseudo_random_matrix(log_dim1, 15 - log_dim1); + + let transposed = matrix.transpose(); + assert_eq!(transposed.dim1(), matrix.dim2()); + assert_eq!(transposed.dim2(), matrix.dim1()); + + for i2 in 0..matrix.dim2() { + let bits: Vec = transposed.bits(i2).map(|b| b == 1).collect(); + let expected: Vec = (0..matrix.dim1()).map(|i1| matrix.bit(i1, i2)).collect(); + assert_eq!(bits, expected, "log_dim1 {log_dim1}, i2 {i2}"); + } + } + } + + /// 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!( + reshaped.bit(index >> log_dim2, index & (reshaped.dim2() - 1)), + matrix.bit(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 077a50e8..fab4cf5d 100644 --- a/crates/common/src/params.rs +++ b/crates/common/src/params.rs @@ -3,15 +3,11 @@ use field::{F128, gf128::is_generator}; use spongefish::Encoding; -use crate::{BitTable, Shape, TableError, VirtualMap, shape::PACK_BITS}; +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 seven row-index bits: a packed row would not fill one - /// codeword position, and the fold reads a column as whole packed - /// elements. - RowIndexTooNarrow, /// `(k_1 + 1)(Q - 1)` reaches `ord(g)`, so two folds could collide in the /// exponent. FoldBoundExceeded, @@ -39,12 +35,6 @@ impl BitZParams { // `Fq` asserts Q is an odd prime below 2^126 on its own behalf, so // no modulus gate is needed here. - // A table admits any row width; the fold walks the table this shapes - // and is the one that needs whole-element columns. - if shape.log_rows() < PACK_BITS as usize { - return Err(ParamsError::RowIndexTooNarrow); - } - // `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 @@ -315,15 +305,6 @@ mod tests { ); } - #[test] - fn rejects_a_row_index_narrower_than_the_pack_width() { - assert_eq!( - params_at(Shape::new(6, 16).unwrap()).err(), - Some(ParamsError::RowIndexTooNarrow) - ); - assert!(params_at(Shape::new(7, 15).unwrap()).is_ok()); - } - #[test] fn rejects_a_generator_of_partial_order() { assert_eq!( diff --git a/crates/common/src/shape.rs b/crates/common/src/shape.rs index 12f8d811..22333b41 100644 --- a/crates/common/src/shape.rs +++ b/crates/common/src/shape.rs @@ -1,6 +1,6 @@ //! The public instance shape and the gates that run before any proof work. -/// The seven low bits of a bit index select the basis coefficients that one +/// The seven low bits of a row index select the basis coefficients that one /// packed field element carries. pub const PACK_BITS: u32 = 7; @@ -16,6 +16,9 @@ pub const MAX_LOG_BITS: usize = 35; /// A shape one of the admissibility constraints rejects. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum ShapeError { + /// Fewer than seven row-index bits: a packed row would not fill one + /// codeword position. The paper writes this count `t`. + RowIndexTooNarrow, /// The total bit count falls outside `2^22..=2^35`. CommitmentSizeOutOfRange, } @@ -24,6 +27,7 @@ pub enum ShapeError { /// /// One counts the bits of a column, the other the columns; the paper writes /// them `t` and `s`. Geometry only: the modulus sits in [`crate::BitZParams`], +/// with the one gate that couples the two. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct Shape { log_rows: usize, @@ -33,6 +37,9 @@ pub struct Shape { impl Shape { /// Checks the geometric admissibility constraints and returns the shape. pub fn new(log_rows: usize, log_columns: usize) -> Result { + if log_rows < PACK_BITS as usize { + return Err(ShapeError::RowIndexTooNarrow); + } // The row width alone is bounded first: the total is at least as // large, so an out-of-window row width is already a rejection, and the // later `MAX_LOG_BITS - log_rows` cannot underflow. @@ -104,9 +111,9 @@ mod tests { } #[test] - fn accepts_a_row_index_narrower_than_the_pack_width() { - assert!(Shape::new(6, 16).is_ok()); - assert!(Shape::new(0, MIN_LOG_BITS).is_ok()); + fn rejects_a_row_index_narrower_than_the_pack_width() { + assert_eq!(Shape::new(6, 16), Err(ShapeError::RowIndexTooNarrow)); + assert!(Shape::new(7, 15).is_ok()); } #[test] diff --git a/crates/common/src/table.rs b/crates/common/src/table.rs index 0f7fdb1a..b8cdc655 100644 --- a/crates/common/src/table.rs +++ b/crates/common/src/table.rs @@ -1,21 +1,9 @@ //! The committed bit table. -use std::{ops::Deref, slice::Iter}; - -use crate::Shape; -/// The word `packed` is sliced into: one `F128` element. [`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 `WordColumns::transpose`. -type Word = u128; -/// Bits in a [`Word`]'s bit index. -const LOG_BITS: u32 = Word::BITS.trailing_zeros(); -/// Bits in a `u64`'s bit index. Swapping two bit index bits below this never -/// moves a bit from one of a word's `u64` lanes to another. -const LOG_LANE_BITS: u32 = u64::BITS.trailing_zeros(); +use crate::{ + Shape, + matrix::{BitMatrix, Word}, +}; /// A witness that does not match its shape. #[derive(Debug, Clone, Copy, PartialEq, Eq)] @@ -36,34 +24,20 @@ pub enum TableError { /// for the monomial basis `beta_v = X^v`. In that basis bit `v` of an `F128`'s /// little-endian `lo || hi` *is* the coefficient of `X^v`, so the committed /// form and the bit form are the same bytes, and the caller holds one copy. -/// The formula is written for `t >= 7`; in general element `k` holds bit -/// indices `128 k` to `128 k + 127`, which below that is several whole columns. -/// -/// Generic over the storage `S` of the packed words; use the [`BitTable`] -/// (borrowed) and [`OwnedBitTable`] aliases rather than naming this directly. #[derive(Debug, Clone, Copy)] -pub struct BitTableBase { +pub struct BitTable<'a> { shape: Shape, - packed: S, + packed: &'a [Word], } -/// A [`BitTableBase`] borrowing the prover's packed witness. -pub type BitTable<'a> = BitTableBase<&'a [Word]>; - -/// A [`BitTableBase`] owning its packed words, as [`WordColumns::transpose`] -/// returns. -pub type OwnedBitTable = BitTableBase>; - -impl BitTableBase { +impl<'a> BitTable<'a> { pub const BITS: usize = Word::BITS as usize; -} -impl> BitTableBase { /// 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: S) -> Result { + 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); @@ -95,49 +69,6 @@ impl> BitTableBase { ((element >> offset) & 1) == 1 } - /// The table read column by column as whole elements, or `None` when - /// `t < 7` and a column is narrower than one. - pub fn word_columns(&self) -> Option> { - let words = self.shape.rows() / BitTable::BITS; - (words > 0).then(|| WordColumns { - packed: &self.packed, - words, - }) - } - - /// The column's bits, in ascending row order. - pub fn column_bits(&self, column: usize) -> BitsIter<'_> { - if let Some(columns) = self.word_columns() { - return BitsIter::new(columns.column(column).iter()); - } - // A column narrower than an element shares it with its neighbours: - // shift it down to bit 0 and stop after its `2^t` rows. - let start = column << self.shape.log_rows(); - BitsIter::preloaded( - self.packed[start / BitTable::BITS] >> (start % BitTable::BITS), - self.shape.rows(), - ) - } - - /// Unwraps the packed bits. - pub fn into_packed(self) -> S { - self.packed - } -} - -/// A table's columns as whole elements, which needs `t >= 7`; from -/// [`BitTableBase::word_columns`]. The fold and [`WordColumns::transpose`] -/// both read the table this way. -/// -// Refactoring Seam might be removed later. -#[derive(Debug, Clone, Copy)] -pub struct WordColumns<'a> { - packed: &'a [Word], - /// Elements per column, `2^(t-7)`. - words: usize, -} - -impl<'a> WordColumns<'a> { /// The `2^(t-7)` elements holding one column, in ascending row order. /// /// A column starts at bit `c * 2^t` and `t >= 7`, so it begins on an @@ -145,333 +76,21 @@ impl<'a> WordColumns<'a> { /// 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 [Word] { - &self.packed[column * self.words..(column + 1) * self.words] + let elements = self.shape.rows() / BitTable::BITS; + &self.packed[column * elements..(column + 1) * elements] } - /// The number of columns. - pub fn columns(&self) -> usize { - self.packed.len() / self.words - } - - /// Swaps rows and columns, returning an owned copy: `result.bit(b, c) == - /// table.bit(c, b)` for every row `b` and column `c` of the table. + /// Swaps rows and columns into an owned [`BitMatrix`] with `dim1` the + /// rows and `dim2` the columns: `result.bit(b, c) == self.bit(c, b)`. /// /// Built for callers like a leaf construction that must walk the table - /// row by row: the packing is column-major (see [`BitTableBase`]), so a - /// row-major walk over the table 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 - /// - /// - **Row count**: `bit_transpose` reads each column as whole words, - /// which is what this view guarantees. The column count is free: with - /// fewer than `128` columns the result's columns share words. - /// - **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 the table did. - /// - **Allocation**: always a full out-of-place copy -- a fresh buffer - /// the same size as the packed witness (up to 4 GiB at `m = 35`), never - /// a view over the original. - /// - **Involution**: transposing the result returns the table's bits - /// exactly, provided it has at least `128` columns. With fewer, the - /// result has no word view of its own to transpose. - pub fn transpose(&self) -> OwnedBitTable { - let (columns, rows) = (self.columns(), self.words * BitTable::BITS); - let shape = Shape::new(columns.ilog2() as usize, rows.ilog2() as usize).expect( - "log_bits is preserved by swapping row/column axes, and the \ - commitment-size window is all a shape is checked against", - ); - - OwnedBitTable { - shape, - packed: bit_transpose(self.packed, columns, rows), - } - } -} - -/// out of place variant. -/// dim1 and dim2 are in bits -/// dim2 is the axis over which the data is adjacent, think xs[dim1][dim2]. -/// dim1 < WORD::Bits dim2 >= WORD::Bits -fn bit_transpose(xs: &[Word], dim1: usize, dim2: usize) -> Vec { - assert_eq!(xs.len() * BitTable::BITS, dim1 * dim2); - assert!(dim1.is_power_of_two(), "dim1 {dim1} is not a power of two"); - - // One body, compiled once per block height. With `log2(d)` a runtime - // value the stage loops cannot be unrolled with constant masks and - // distances, which cost ~30% at d = 64. - match dim1.min(BitTable::BITS).trailing_zeros() { - // A single row is its own transpose. - 0 => xs.to_vec(), - 1 => bit_transpose_blocks::<1>(xs, dim1, dim2), - 2 => bit_transpose_blocks::<2>(xs, dim1, dim2), - 3 => bit_transpose_blocks::<3>(xs, dim1, dim2), - 4 => bit_transpose_blocks::<4>(xs, dim1, dim2), - 5 => bit_transpose_blocks::<5>(xs, dim1, dim2), - 6 => bit_transpose_blocks::<6>(xs, dim1, dim2), - _ => bit_transpose_blocks::(xs, dim1, dim2), - } -} - -/// [`bit_transpose`] for blocks of `d = 2^LOG_D` rows, `LOG_D >= 1`. -fn bit_transpose_blocks(xs: &[Word], dim1: usize, dim2: usize) -> Vec { - const BITS: usize = Word::BITS as usize; - - let d = 1 << LOG_D; // rows per block - let r = BITS / d; // words per row per block - let row_words = dim2 / BITS; - // Output words between consecutive words of a block; 1 when d < BITS. - let out_stride = dim1 / d; - assert!( - dim2.is_multiple_of(BITS) && row_words.is_multiple_of(r), - "dim2 {dim2} does not fill whole {d}-row blocks" - ); - - let mut out = bytemuck::zeroed_vec(xs.len()); - - // 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. g1 only takes more than one value - // when d = BITS. - for g2 in 0..row_words / r { - for g1 in 0..dim1 / d { - // Word i = [c | a] is row a, word c. Two loops rather than one - // over i: splitting i back into a and c measured 4-12% slower - // for d < BITS. - let mut block = [0; BITS]; - for c in 0..r { - for a in 0..d { - block[c * d + a] = xs[(g1 * d + a) * row_words + g2 * r + c]; - } - } - - rotate_bit_index::(&mut block); - bit_block_transpose(LOG_D, &mut block); - - // o is output row g2 * BITS + o, word g1 in [dim2][dim1]; g1 is - // only nonzero when d = BITS. - for (o, &word) in block.iter().enumerate() { - out[(g2 * BITS + o) * out_stride + g1] = word; - } - } - } - - out -} - -/// The bits whose index has bit `s` clear: alternating runs of `2^s` bits, -/// lowest run set. 01010101, 00110011, 00001111 for `s` = 0, 1, 2. -const fn index_bit_clear(s: u32) -> Word { - Word::MAX / ((1 << (1 << s)) + 1) -} - -/// Interleaves the `d = 2^LOG_D` chunks of `r = BITS / d` bits in every word -/// of `block`, a `d`-way perfect shuffle: bit `j` of chunk `k` moves to bit -/// `j * d + k`. On bit indices that is a left rotation by `LOG_D`: bit -/// `[p_hi | p_lo] = [k | j]` moves to `[p_lo | p_hi]`, so unlike -/// `rotate_left` each bit moves its own distance. That puts `p_hi` in the low -/// `LOG_D` index bits, where [`bit_block_transpose`] swaps it with the row -/// bits of the word index. One pass over the block per entry of -/// [`rotation_swaps`], none for `LOG_D = LOG_BITS`. -/// -/// A pass that stays inside `u64` lanes runs on the lanes rather than on -/// whole words, which vectorises: on `u128` that made the rotation about -/// 40% cheaper. Only a swap with bit index bit 6 or above needs whole words. -#[inline(always)] -fn rotate_bit_index(block: &mut [Word; Word::BITS as usize]) { - // Evaluated at compile time, so every shift and mask is a constant. - let (swaps, n) = const { rotation_swaps(LOG_D) }; - for &(shift, mask, in_lane) in &swaps[..n] { - if in_lane { - let mask = mask as u64; - for x in bytemuck::cast_slice_mut::(block) { - let t = ((*x >> shift) ^ *x) & mask; - *x ^= t ^ (t << shift); - } - } else { - for x in block.iter_mut() { - let t = ((*x >> shift) ^ *x) & mask; - *x ^= t ^ (t << shift); - } - } - } -} - -/// The delta swaps that rotate a bit index left by `log_d`, as `(shift, -/// mask, in_lane)`: each bit in `mask` trades places with the bit `shift` -/// above it, and `in_lane` says no bit crosses from one `u64` lane to -/// another. Returns the swaps and how many of them are used. -/// -/// Each swap exchanges two bits of the bit index, `low` and `i`: the bits -/// whose index has `low` set and `i` clear move up by `shift`, their partners -/// move down by it, and the other half of the word stays put. The rotation -/// splits the index bits into `gcd(LOG_BITS, log_d)` cycles. Each cycle is -/// walked from its lowest index bit `low`, which is swapped with the cycle's -/// other bits `i` in turn: `LOG_BITS - gcd` swaps in all, the fewest index -/// bit swaps that make the rotation. -const fn rotation_swaps(log_d: u32) -> ([(u32, Word, bool); LOG_BITS as usize], usize) { - let (mut cycles, mut rest) = (LOG_BITS, log_d); - while rest != 0 { - (cycles, rest) = (rest, cycles % rest); - } - - let mut swaps = [(0, 0, false); LOG_BITS as usize]; - let mut n = 0; - let mut low = 0; - while low < cycles { - let mut i = (low + log_d) % LOG_BITS; - while i != low { - // The bits whose index has bit low set and bit i clear. - swaps[n] = ( - (1 << i) - (1 << low), - !index_bit_clear(low) & index_bit_clear(i), - i < LOG_LANE_BITS, - ); - n += 1; - i = (i + log_d) % LOG_BITS; - } - low += 1; - } - (swaps, n) -} - -/// Transposes every `2^log_d x 2^log_d` tile of a BITS x BITS block, the -/// tile's words being those that differ only in their low `log_d` index -/// bits. Stage `s` swaps bit index bit `s` with word index bit `s`, coarse -/// to fine; the stages commute, so the order is free. `log_d = LOG_BITS` is -/// the full transpose. -#[inline(always)] -fn bit_block_transpose(log_d: u32, xs: &mut [Word; Word::BITS as usize]) { - for s in (0..log_d).rev() { - // Distance between the paired bits, and between the paired words. - let j: usize = 1 << s; - let mask = index_bit_clear(s); - let mut k: usize = 0; - while k < Word::BITS as usize { - // A delta swap across the pair: t is where the upper runs of - // xs[k] and the lower runs of xs[k | j] differ, and flipping those - // bits in both swaps the runs. Measured 6-10% faster than - // rebuilding both words from masked halves. - let t = ((xs[k] >> j) ^ xs[k | j]) & mask; - xs[k | j] ^= t; - xs[k] ^= t << j; - // (k|j) count the upper part. +1 advances the upper part. !j converts it into the lower index again except for when it hits the next round. In that case the | j was - k = ((k | j) + 1) & !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 `bit()` and `StepIter`. - word |= b << offset; - offset += 1; - if offset == BitTable::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 one column's bits, in ascending row order — see -/// [`BitTable::column_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, - } - } - - /// Emits the low `bits` bits of `word` and nothing after; `bits` is a - /// multiple of `S` and at most `Word::BITS`. - fn preloaded(word: Word, bits: usize) -> Self { - let () = Self::CHECK_SIZE; - debug_assert!(bits.is_multiple_of(S) && bits <= Word::BITS as usize); - Self { - elements: [].iter(), - word, - remaining: bits as u8, - } - } -} - -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() * BitTable::BITS) / S + /// row by row: the packing is column major, so a row-major walk over the + /// table is a scatter, one `bit()` call and cache miss per cell. After + /// the transpose each row is contiguous, read by [`BitMatrix::bits`]. See + /// [`BitMatrix::transpose`] for the constraints; the table's own `t >= 7` + /// is what the first one asks for. + pub fn transpose(&self) -> BitMatrix { + BitMatrix::transposed(self.packed, self.shape.log_rows()) } } @@ -482,12 +101,8 @@ mod tests { // We want to use compare against use field::F128; - /// Bits in a packed element, and the shift that divides an index into element - /// and offset. + /// Bits in a packed element. const PACKED_BITS: usize = 128; - const PACKED_SHIFT: u32 = 7; - /// Bits in each half of a packed element. - const HALF_BITS: usize = 64; use super::*; @@ -571,38 +186,10 @@ mod tests { let shape = small_shape(); let packed = with_bits(&shape, &[(0, 3), (2, 3), (64, 3)]); let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); - let columns = table.word_columns().unwrap(); - - assert_eq!(columns.columns(), shape.columns()); - assert_eq!(columns.column(3).len(), shape.rows() / BitTable::BITS); - assert_eq!(columns.column(3), [0b101 | (1 << 64)]); - assert!(columns.column(4).iter().all(|&element| element == 0)); - } - - /// So neither the fold nor `transpose` can be handed it. - #[test] - fn a_column_narrower_than_an_element_has_no_word_view() { - let shape = Shape::new(6, 16).unwrap(); - let packed = vec![F128::ZERO; (1 << shape.log_bits()) / PACKED_BITS]; - let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); - assert!(table.word_columns().is_none()); - } - - #[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, bytemuck::cast_slice(&packed)).unwrap(); - - let iter = table.column_bits(2); - assert_eq!(iter.len(), shape.rows()); - let bits: Vec = iter.map(|b| b == 1).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] @@ -612,7 +199,7 @@ mod tests { let packed = with_bits(&shape, &rows); let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); - for (index, &element) in table.word_columns().unwrap().column(2).iter().enumerate() { + 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}"); @@ -620,226 +207,24 @@ mod tests { } } - /// 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 - } - - /// 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 - } - - #[test] - fn bit_block_transpose_matches_a_brute_force_reference() { - let mut original = [0 as Word; BitTable::BITS]; - for (i, word) in original.iter_mut().enumerate() { - *word = pseudo_random_word(i); - } - - let mut transposed = original; - bit_block_transpose(LOG_BITS, &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 `p` must land at `p` rotated left by `LOG_D` within `LOG_BITS` - /// bits. Word `p` of the block holds only bit `p`, so one call checks - /// every bit. - fn check_rotate_bit_index() { - let mut block: [Word; BitTable::BITS] = std::array::from_fn(|p| 1 << p); - rotate_bit_index::(&mut block); - for (p, &word) in block.iter().enumerate() { - let rotated = ((p << LOG_D) | (p >> (LOG_BITS - LOG_D))) % BitTable::BITS; - assert_eq!(word, 1 << rotated, "LOG_D {LOG_D}, bit {p}"); - } - } - - #[test] - fn rotate_bit_index_rotates_every_bit_index() { - check_rotate_bit_index::<1>(); - check_rotate_bit_index::<2>(); - check_rotate_bit_index::<3>(); - check_rotate_bit_index::<4>(); - check_rotate_bit_index::<5>(); - check_rotate_bit_index::<6>(); - check_rotate_bit_index::<7>(); - } - - /// `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 * BitTable::BITS), - transpose_reference::(&tt), - "dim1 {GS}, {words} words per row" - ); - } - - #[test] - fn bit_transpose_matches_reference_for_every_row_count() { - // Two blocks along dim2 at every row count, so block placement is - // exercised as well as the block contents. - check_bit_transpose::<1>(256); - check_bit_transpose::<2>(128); - check_bit_transpose::<4>(64); - check_bit_transpose::<8>(32); - check_bit_transpose::<16>(16); - check_bit_transpose::<32>(8); - check_bit_transpose::<64>(4); - check_bit_transpose::<128>(2); - check_bit_transpose::<256>(2); - } - - #[test] - fn transpose_reference_matches_bit_transpose() { - // 256 segments of 2 words each: two block groups on both axes. - const GS: usize = 256; - let segment_n = 2; - let tt: Vec = (0..GS * segment_n).map(pseudo_random_word).collect(); - - let reference = transpose_reference::(&tt); - assert_eq!(reference.len(), tt.len()); - assert_eq!( - reference, - bit_transpose(&tt, GS, segment_n * BitTable::BITS) - ); - } - - /// Every row width below one element, where columns share elements. - #[test] - fn column_bits_agrees_with_the_bit_accessor_below_one_element() { - for log_rows in 0..PACKED_SHIFT as usize { - let shape = Shape::new(log_rows, 22 - log_rows).unwrap(); - let packed: Vec = (0..(1 << shape.log_bits()) / BitTable::BITS) - .map(pseudo_random_word) - .collect(); - let table = BitTable::new(shape, &packed).unwrap(); - - // Both ends of the first, a middle and the last element. - let per_element = BitTable::BITS / shape.rows(); - let last = shape.columns() - 1; - for column in [0, 1, per_element - 1, per_element, last / 2, last] { - let iter = table.column_bits(column); - assert_eq!(iter.len(), shape.rows()); - let bits: Vec = iter.map(|b| b == 1).collect(); - let expected: Vec = (0..shape.rows()) - .map(|row| table.bit(column, row)) - .collect(); - assert_eq!(bits, expected, "t {log_rows}, column {column}"); - } - } - } - - /// 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); - } - } - } - } - packed - } - - #[test] - fn transpose_swaps_row_and_column_bits() { - let shape = transpose_shape(); - let packed = pseudo_random_table(&shape); - let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); - - let transposed = table.word_columns().unwrap().transpose(); - assert_eq!(transposed.shape().log_rows(), shape.log_columns()); - assert_eq!(transposed.shape().log_columns(), shape.log_rows()); - - // 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" - ); - } - } - } - + /// The transpose has the table's rows along `dim1`, each one contiguous. #[test] - fn transposing_twice_recovers_the_original_bits() { - let shape = transpose_shape(); - let packed = pseudo_random_table(&shape); + 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 roundtripped = table - .word_columns() - .unwrap() - .transpose() - .word_columns() - .unwrap() - .transpose() - .into_packed(); + let transposed = table.transpose(); + assert_eq!(transposed.dim1(), shape.rows()); + assert_eq!(transposed.dim2(), shape.columns()); - assert_eq!(roundtripped, bytemuck::cast_slice::(&packed)); - } - - /// Every column count below one element, down to a single column: the - /// result's rows share elements, and its columns read back through - /// `column_bits` as the source's rows. - #[test] - fn transpose_handles_fewer_columns_than_an_element() { - for log_columns in 0..PACKED_SHIFT as usize { - let shape = Shape::new(22 - log_columns, log_columns).unwrap(); - let packed: Vec = (0..(1 << shape.log_bits()) / BitTable::BITS) - .map(pseudo_random_word) + 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(); - let table = BitTable::new(shape, &packed).unwrap(); - - let transposed = table.word_columns().unwrap().transpose(); - assert_eq!(transposed.shape().log_rows(), log_columns); - assert_eq!(transposed.shape().log_columns(), shape.log_rows()); - - for row in [ - 0, - 1, - PACKED_BITS - 1, - PACKED_BITS, - shape.rows() / 2, - shape.rows() - 1, - ] { - let bits: Vec = transposed.column_bits(row).map(|b| b == 1).collect(); - let expected: Vec = (0..shape.columns()) - .map(|column| table.bit(column, row)) - .collect(); - assert_eq!(bits, expected, "s {log_columns}, row {row}"); - } + assert_eq!(bits, expected, "row {row}"); } } } diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs index 1ab35f25..f34ffdc4 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs @@ -1,6 +1,6 @@ use std::collections::VecDeque; -use common::OwnedBitTable; +use common::BitMatrix; use field::{F128, Wide256}; use num_traits::{ConstOne, ConstZero}; use rayon::prelude::*; @@ -343,15 +343,14 @@ pub struct GrandProductCircuit { } impl GrandProductCircuit { - pub fn new(row_images: &[F128], table: OwnedBitTable) -> Self { - let columns = table.shape().columns(); - let rows = table.shape().rows(); - let dim = columns * rows; + pub fn new(row_images: &[F128], matrix: BitMatrix) -> Self { + 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 * rows..(b + 1) * rows]; - for (leaf, bit) in leafs.iter_mut().zip(table.column_bits(b)) { + 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 }; } } diff --git a/crates/prover/benches/prover.rs b/crates/prover/benches/prover.rs index 4f9c9377..6c70a9f4 100644 --- a/crates/prover/benches/prover.rs +++ b/crates/prover/benches/prover.rs @@ -29,10 +29,10 @@ fn random_table(shape: Shape) -> BitTable<'static> { #[divan::bench] fn gkr(bencher: Bencher) { // `Shape::new(log_rows, log_columns)` rejects fewer than `2^22` committed - // bits total (`MIN_LOG_BITS`), and `BitZParams::new` a row width under 7 - // (`PACK_BITS`), so the two args must sum to at least 22. That floor - // isn't a security bound computed here -- it's the range flock-core ships - // precomputed Ligerito configs for (see the comment on `MIN_LOG_BITS`). + // bits total (`MIN_LOG_BITS`) and a row width under 7 (`PACK_BITS`), so + // the two args must sum to at least 22. That floor isn't a security bound + // computed here -- it's the range flock-core ships precomputed Ligerito + // configs for (see the comment on `MIN_LOG_BITS`). let shape = Shape::new(10, 15).unwrap(); let table = random_table(shape); diff --git a/crates/prover/src/fold.rs b/crates/prover/src/fold.rs index c2b27a92..8c272fe5 100644 --- a/crates/prover/src/fold.rs +++ b/crates/prover/src/fold.rs @@ -48,10 +48,7 @@ impl BitZProver { // Lifted once: the fold reads it per set bit across every column. let exponents = claim.row_exponents(); - let columns = table - .word_columns() - .expect("the shape matches the parameters, whose gate admits only t >= 7"); - let folds = fold_columns(&columns, &exponents); + let folds = fold_columns(table, &exponents); for fold in &folds { transcript.prover_message(&fold.to_le_bytes()); diff --git a/crates/prover/src/reduce.rs b/crates/prover/src/reduce.rs index 6166b9a7..4fd17779 100644 --- a/crates/prover/src/reduce.rs +++ b/crates/prover/src/reduce.rs @@ -18,10 +18,7 @@ fn init_circuit(table: &BitTable, fold: &Fold) -> GrandProductCircuit { // 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. // Transposed so each row can be read sequentially. - let transposed = table - .word_columns() - .expect("the fold's table, whose gate admits only t >= 7") - .transpose(); + let transposed = table.transpose(); GrandProductCircuit::new(&fold.row_images, transposed) } diff --git a/crates/verifier/src/reduce.rs b/crates/verifier/src/reduce.rs index a0cc92cb..f8aa68cf 100644 --- a/crates/verifier/src/reduce.rs +++ b/crates/verifier/src/reduce.rs @@ -100,7 +100,7 @@ mod round_trip_ai_test { .map(|i| F128::from((i as u128 + 3) * 0xABCDEF12345u128 + 1)) .collect(); - let transposed = table.word_columns().unwrap().transpose(); + let transposed = table.transpose(); let circuit = GrandProductCircuit::new(&row_images, transposed); let (top_layer, witnesses) = circuit.batched_eval(shape.columns()); From e1505fade75cd7c9a75668eac1b0e3a81c1bc71d Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Wed, 7 Oct 2026 20:18:57 +0800 Subject: [PATCH 21/22] common: limit bit_transpose to whole 128x128 blocks and gate s < 7 in BitZParams Production never transposes a table with fewer than 128 columns: the CLI builds t = 7, s = m - 7 shapes and the reference split has s >= 8. Go back to origin/main's 128x128 block kernel and drop the small-dim1 machinery (bit_transpose_blocks, rotate_bit_index, rotation_swaps and StepIter::preloaded) together with its tests. That version lives in 336e8097225aca4f8146f35594740b0e5054d839. The bit_transpose doc comment points there and notes that it extends to dim2 < 128 if needed. BitZParams::new now rejects s < 7 with ParamsError::ColumnCountTooNarrow rather than letting the reduction panic in the transpose. It is not a protocol bound but a limit of the kernel, and restoring that version lifts it. --- crates/common/src/matrix.rs | 310 ++++++++---------------------------- crates/common/src/params.rs | 30 ++++ crates/common/src/table.rs | 25 ++- crates/prover/src/reduce.rs | 5 +- crates/tests/tests/prove.rs | 20 +-- 5 files changed, 123 insertions(+), 267 deletions(-) diff --git a/crates/common/src/matrix.rs b/crates/common/src/matrix.rs index 5af32cb9..78ccb86e 100644 --- a/crates/common/src/matrix.rs +++ b/crates/common/src/matrix.rs @@ -12,11 +12,6 @@ use std::slice::Iter; pub(crate) type Word = u128; /// Bits in a [`Word`]. const BITS: usize = Word::BITS as usize; -/// Bits in a [`Word`]'s bit index. -const LOG_BITS: u32 = Word::BITS.trailing_zeros(); -/// Bits in a `u64`'s bit index. Swapping two bit index bits below this never -/// moves a bit from one of a word's `u64` lanes to another. -const LOG_LANE_BITS: u32 = u64::BITS.trailing_zeros(); /// Packed bits read as `xs[dim1][dim2]`: bit `(i1, i2)` is bit /// `(i1 << log2(dim2)) | i2` of the words, least significant bit first. Each @@ -74,15 +69,19 @@ impl BitMatrix { } /// The `dim2` bits of `xs[i1]`, in ascending order. + /// + /// # Panics + /// + /// If `dim2` is below `128`: `xs[i1]` then shares a word with its + /// neighbours, which only [`BitMatrix::bit`] reads. pub fn bits(&self, i1: usize) -> BitsIter<'_> { - let start = i1 << self.log_dim2; - if self.dim2() >= BITS { - let words = self.dim2() / BITS; - return BitsIter::new(self.packed[start / BITS..][..words].iter()); - } - // Shorter than a word, so it shares one with its neighbours: shift it - // down to bit 0 and stop after its `dim2` bits. - BitsIter::preloaded(self.packed[start / BITS] >> (start % BITS), self.dim2()) + 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()) } /// Swaps the axes, returning an owned copy: `result.bit(i2, i1) == @@ -94,14 +93,13 @@ impl BitMatrix { /// /// # Constraints /// - /// - **`dim2`**: at least `128`, since `bit_transpose` reads each - /// `xs[i1]` as whole words. `dim1` is free: below `128` the result's - /// `xs[i2]` share words. + /// - **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 buffer /// the same size as the packed words (up to 4 GiB at `m = 35`), never /// a view over the original. - /// - **Involution**: transposing the result returns the original bits, - /// provided `dim1` is at least `128` as well. + /// - **Involution**: transposing the result returns the original bits. pub fn transpose(&self) -> BitMatrix { Self::transposed(&self.packed, self.log_dim2) } @@ -120,68 +118,42 @@ impl BitMatrix { } } -/// out of place variant. -/// dim1 and dim2 are in bits -/// dim2 is the axis over which the data is adjacent, think xs[dim1][dim2]. -/// dim1 < WORD::Bits dim2 >= WORD::Bits +/// 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_power_of_two(), "dim1 {dim1} is not a power of two"); - - // One body, compiled once per block height. With `log2(d)` a runtime - // value the stage loops cannot be unrolled with constant masks and - // distances, which cost ~30% at d = 64. - match dim1.min(BITS).trailing_zeros() { - // A single row is its own transpose. - 0 => xs.to_vec(), - 1 => bit_transpose_blocks::<1>(xs, dim1, dim2), - 2 => bit_transpose_blocks::<2>(xs, dim1, dim2), - 3 => bit_transpose_blocks::<3>(xs, dim1, dim2), - 4 => bit_transpose_blocks::<4>(xs, dim1, dim2), - 5 => bit_transpose_blocks::<5>(xs, dim1, dim2), - 6 => bit_transpose_blocks::<6>(xs, dim1, dim2), - _ => bit_transpose_blocks::(xs, dim1, dim2), - } -} - -/// [`bit_transpose`] for blocks of `d = 2^LOG_D` rows, `LOG_D >= 1`. -fn bit_transpose_blocks(xs: &[Word], dim1: usize, dim2: usize) -> Vec { - let d = 1 << LOG_D; // rows per block - let r = BITS / d; // words per row per block - let row_words = dim2 / BITS; - // Output words between consecutive words of a block; 1 when d < BITS. - let out_stride = dim1 / d; assert!( - dim2.is_multiple_of(BITS) && row_words.is_multiple_of(r), - "dim2 {dim2} does not fill whole {d}-row blocks" + 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. g1 only takes more than one value - // when d = BITS. - for g2 in 0..row_words / r { - for g1 in 0..dim1 / d { - // Word i = [c | a] is row a, word c. Two loops rather than one - // over i: splitting i back into a and c measured 4-12% slower - // for d < BITS. - let mut block = [0; BITS]; - for c in 0..r { - for a in 0..d { - block[c * d + a] = xs[(g1 * d + a) * row_words + g2 * r + c]; - } + // ~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]; } - rotate_bit_index::(&mut block); - bit_block_transpose(LOG_D, &mut block); + transpose_bit_block(&mut block); - // o is output row g2 * BITS + o, word g1 in [dim2][dim1]; g1 is - // only nonzero when d = BITS. for (o, &word) in block.iter().enumerate() { - out[(g2 * BITS + o) * out_stride + g1] = word; + out[(g2 * BITS + o) * dim1_words + g1] = word; } } } @@ -189,105 +161,28 @@ fn bit_transpose_blocks(xs: &[Word], dim1: usize, dim2: usize) out } -/// The bits whose index has bit `s` clear: alternating runs of `2^s` bits, -/// lowest run set. 01010101, 00110011, 00001111 for `s` = 0, 1, 2. -const fn index_bit_clear(s: u32) -> Word { - Word::MAX / ((1 << (1 << s)) + 1) -} - -/// Interleaves the `d = 2^LOG_D` chunks of `r = BITS / d` bits in every word -/// of `block`, a `d`-way perfect shuffle: bit `j` of chunk `k` moves to bit -/// `j * d + k`. On bit indices that is a left rotation by `LOG_D`: bit -/// `[p_hi | p_lo] = [k | j]` moves to `[p_lo | p_hi]`, so unlike -/// `rotate_left` each bit moves its own distance. That puts `p_hi` in the low -/// `LOG_D` index bits, where [`bit_block_transpose`] swaps it with the row -/// bits of the word index. One pass over the block per entry of -/// [`rotation_swaps`], none for `LOG_D = LOG_BITS`. +/// 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`. /// -/// A pass that stays inside `u64` lanes runs on the lanes rather than on -/// whole words, which vectorises: on `u128` that made the rotation about -/// 40% cheaper. Only a swap with bit index bit 6 or above needs whole words. -#[inline(always)] -fn rotate_bit_index(block: &mut [Word; Word::BITS as usize]) { - // Evaluated at compile time, so every shift and mask is a constant. - let (swaps, n) = const { rotation_swaps(LOG_D) }; - for &(shift, mask, in_lane) in &swaps[..n] { - if in_lane { - let mask = mask as u64; - for x in bytemuck::cast_slice_mut::(block) { - let t = ((*x >> shift) ^ *x) & mask; - *x ^= t ^ (t << shift); - } - } else { - for x in block.iter_mut() { - let t = ((*x >> shift) ^ *x) & mask; - *x ^= t ^ (t << shift); - } - } - } -} - -/// The delta swaps that rotate a bit index left by `log_d`, as `(shift, -/// mask, in_lane)`: each bit in `mask` trades places with the bit `shift` -/// above it, and `in_lane` says no bit crosses from one `u64` lane to -/// another. Returns the swaps and how many of them are used. -/// -/// Each swap exchanges two bits of the bit index, `low` and `i`: the bits -/// whose index has `low` set and `i` clear move up by `shift`, their partners -/// move down by it, and the other half of the word stays put. The rotation -/// splits the index bits into `gcd(LOG_BITS, log_d)` cycles. Each cycle is -/// walked from its lowest index bit `low`, which is swapped with the cycle's -/// other bits `i` in turn: `LOG_BITS - gcd` swaps in all, the fewest index -/// bit swaps that make the rotation. -const fn rotation_swaps(log_d: u32) -> ([(u32, Word, bool); LOG_BITS as usize], usize) { - let (mut cycles, mut rest) = (LOG_BITS, log_d); - while rest != 0 { - (cycles, rest) = (rest, cycles % rest); - } - - let mut swaps = [(0, 0, false); LOG_BITS as usize]; - let mut n = 0; - let mut low = 0; - while low < cycles { - let mut i = (low + log_d) % LOG_BITS; - while i != low { - // The bits whose index has bit low set and bit i clear. - swaps[n] = ( - (1 << i) - (1 << low), - !index_bit_clear(low) & index_bit_clear(i), - i < LOG_LANE_BITS, - ); - n += 1; - i = (i + log_d) % LOG_BITS; - } - low += 1; - } - (swaps, n) -} - -/// Transposes every `2^log_d x 2^log_d` tile of a BITS x BITS block, the -/// tile's words being those that differ only in their low `log_d` index -/// bits. Stage `s` swaps bit index bit `s` with word index bit `s`, coarse -/// to fine; the stages commute, so the order is free. `log_d = LOG_BITS` is -/// the full transpose. -#[inline(always)] -fn bit_block_transpose(log_d: u32, xs: &mut [Word; Word::BITS as usize]) { - for s in (0..log_d).rev() { - // Distance between the paired bits, and between the paired words. - let j: usize = 1 << s; - let mask = index_bit_clear(s); - let mut k: usize = 0; - while k < Word::BITS as usize { - // A delta swap across the pair: t is where the upper runs of - // xs[k] and the lower runs of xs[k | j] differ, and flipping those - // bits in both swaps the runs. Measured 6-10% faster than - // rebuilding both words from masked halves. - let t = ((xs[k] >> j) ^ xs[k | j]) & mask; - xs[k | j] ^= t; - xs[k] ^= t << j; - // (k|j) count the upper part. +1 advances the upper part. !j converts it into the lower index again except for when it hits the next round. In that case the | j was +/// 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 `BitMatrix::bit` would need to do the same +/// work. +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; } } @@ -360,18 +255,6 @@ impl<'a, const S: usize> StepIter<'a, S> { remaining: 0, } } - - /// Emits the low `bits` bits of `word` and nothing after; `bits` is a - /// multiple of `S` and at most `Word::BITS`. - fn preloaded(word: Word, bits: usize) -> Self { - let () = Self::CHECK_SIZE; - debug_assert!(bits.is_multiple_of(S) && bits <= Word::BITS as usize); - Self { - elements: [].iter(), - word, - remaining: bits as u8, - } - } } impl Iterator for StepIter<'_, S> { @@ -433,10 +316,10 @@ mod tests { } } - /// `dim2` of whole words, and shorter than one, where neighbours share it. + /// `dim2` of one word and of two. #[test] fn bits_agrees_with_the_bit_accessor() { - for log_dim2 in 0..=LOG_BITS as usize + 1 { + for log_dim2 in [7, 8] { let matrix = pseudo_random_matrix(12 - log_dim2, log_dim2); for i1 in 0..matrix.dim1() { @@ -450,14 +333,14 @@ mod tests { } #[test] - fn bit_block_transpose_matches_a_brute_force_reference() { + 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; - bit_block_transpose(LOG_BITS, &mut transposed); + transpose_bit_block(&mut transposed); for (i, &transposed_word) in transposed.iter().enumerate() { for (j, &original_word) in original.iter().enumerate() { @@ -470,29 +353,6 @@ mod tests { } } - /// Bit `p` must land at `p` rotated left by `LOG_D` within `LOG_BITS` - /// bits. Word `p` of the block holds only bit `p`, so one call checks - /// every bit. - fn check_rotate_bit_index() { - let mut block: [Word; BITS] = std::array::from_fn(|p| 1 << p); - rotate_bit_index::(&mut block); - for (p, &word) in block.iter().enumerate() { - let rotated = ((p << LOG_D) | (p >> (LOG_BITS - LOG_D))) % BITS; - assert_eq!(word, 1 << rotated, "LOG_D {LOG_D}, bit {p}"); - } - } - - #[test] - fn rotate_bit_index_rotates_every_bit_index() { - check_rotate_bit_index::<1>(); - check_rotate_bit_index::<2>(); - check_rotate_bit_index::<3>(); - check_rotate_bit_index::<4>(); - check_rotate_bit_index::<5>(); - check_rotate_bit_index::<6>(); - check_rotate_bit_index::<7>(); - } - /// `bit_transpose` of `GS` rows of `words` words each against the /// bit-at-a-time reference. fn check_bit_transpose(words: usize) { @@ -505,32 +365,15 @@ mod tests { } #[test] - fn bit_transpose_matches_reference_for_every_row_count() { - // Two blocks along dim2 at every row count, so block placement is - // exercised as well as the block contents. - check_bit_transpose::<1>(256); - check_bit_transpose::<2>(128); - check_bit_transpose::<4>(64); - check_bit_transpose::<8>(32); - check_bit_transpose::<16>(16); - check_bit_transpose::<32>(8); - check_bit_transpose::<64>(4); + 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); } - #[test] - fn transpose_reference_matches_bit_transpose() { - // 256 segments of 2 words each: two block groups on both axes. - const GS: usize = 256; - let segment_n = 2; - let tt: Vec = (0..GS * segment_n).map(pseudo_random_word).collect(); - - let reference = transpose_reference::(&tt); - assert_eq!(reference.len(), tt.len()); - assert_eq!(reference, bit_transpose(&tt, GS, segment_n * BITS)); - } - /// `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; @@ -565,25 +408,6 @@ mod tests { assert_eq!(roundtripped.into_packed(), matrix.into_packed()); } - /// Every `dim1` below one word, down to one: the result's `xs[i2]` are - /// shorter than a word and share it. - #[test] - fn transpose_handles_dim1_below_a_word() { - for log_dim1 in 0..LOG_BITS as usize { - let matrix = pseudo_random_matrix(log_dim1, 15 - log_dim1); - - let transposed = matrix.transpose(); - assert_eq!(transposed.dim1(), matrix.dim2()); - assert_eq!(transposed.dim2(), matrix.dim1()); - - for i2 in 0..matrix.dim2() { - let bits: Vec = transposed.bits(i2).map(|b| b == 1).collect(); - let expected: Vec = (0..matrix.dim1()).map(|i1| matrix.bit(i1, i2)).collect(); - assert_eq!(bits, expected, "log_dim1 {log_dim1}, i2 {i2}"); - } - } - } - /// Every split of the index, from one long `xs[0]` down to `dim2 = 1`. #[test] fn reshape_leaves_every_bit_in_place() { diff --git a/crates/common/src/params.rs b/crates/common/src/params.rs index fab4cf5d..fad3619a 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 + /// [`BitTable::transpose`], 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 @@ -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 b8cdc655..7e252c58 100644 --- a/crates/common/src/table.rs +++ b/crates/common/src/table.rs @@ -86,9 +86,15 @@ impl<'a> BitTable<'a> { /// Built for callers like a leaf construction that must walk the table /// row by row: the packing is column major, so a row-major walk over the /// table is a scatter, one `bit()` call and cache miss per cell. After - /// the transpose each row is contiguous, read by [`BitMatrix::bits`]. See - /// [`BitMatrix::transpose`] for the constraints; the table's own `t >= 7` - /// is what the first one asks for. + /// the transpose each row is contiguous, read by [`BitMatrix::bits`]. + /// + /// # Panics + /// + /// If the table has fewer than `128` columns (`s < 7`): see + /// [`BitMatrix::transpose`], which needs both axes at least `128`. The + /// table's own `t >= 7` covers the rows, and + /// [`crate::ParamsError::ColumnCountTooNarrow`] keeps such shapes away + /// from the reduction. pub fn transpose(&self) -> BitMatrix { BitMatrix::transposed(self.packed, self.shape.log_rows()) } @@ -227,4 +233,17 @@ mod tests { assert_eq!(bits, expected, "row {row}"); } } + + #[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. + 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(); + + table.transpose(); + } } diff --git a/crates/prover/src/reduce.rs b/crates/prover/src/reduce.rs index 4fd17779..0a42ab29 100644 --- a/crates/prover/src/reduce.rs +++ b/crates/prover/src/reduce.rs @@ -125,8 +125,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/tests/tests/prove.rs b/crates/tests/tests/prove.rs index 49160585..3c87650f 100644 --- a/crates/tests/tests/prove.rs +++ b/crates/tests/tests/prove.rs @@ -1,6 +1,6 @@ //! The top-level prove and verify, through the real opening. -use common::{Root, Shape, TableError}; +use common::{Root, TableError}; use field::{F128, Fq}; use num_traits::{ConstOne, ConstZero}; use pcs::{HashKind, LigeritoProfile, Pcs, VerifyError as PcsVerifyError}; @@ -42,24 +42,6 @@ fn an_honest_proof_verifies_on_both_floor_shapes() { } } -/// Fewer columns than one element holds: the leaves are read from a -/// transpose whose rows share elements. -#[test] -fn an_honest_proof_verifies_with_fewer_than_128_columns() { - let instance = Instance::honest(Shape::new(16, 6).unwrap(), 31); - let proof = prove(&instance); - - instance - .verifier - .verify( - &instance.claim, - &instance.pcs, - instance.com, - verifier_transcript(&proof), - ) - .unwrap(); -} - #[test] fn a_proof_replayed_under_a_different_commitment_is_refused() { let instance = Instance::honest(narrow_shape(), 32); From f80168471beed10d6e244fc97b36b8723f5548be Mon Sep 17 00:00:00 2001 From: Xander van der Goot Date: Thu, 8 Oct 2026 12:34:49 +0800 Subject: [PATCH 22/22] common: replace BitTable::transpose with a borrowed as_matrix view and let GrandProductCircuit::new take the table BitTable::as_matrix wraps the table's words as an [columns][rows] BitMatrix without moving a bit, and transposing goes only through BitMatrix::transpose. To borrow the words, BitMatrix now holds a Cow<'a, [Word]>; transpose always returns an owned BitMatrix<'static>. The crate-private transposed helper is gone, and BitMatrix computes no ilog2. GrandProductCircuit::new takes the BitTable and does as_matrix().transpose() itself, so the prover and the verifier test no longer transpose before building the circuit. BitMatrix::bit was only read by tests, so it moves into the matrix test module as the oracle there; the as_matrix test reads columns through bits() instead. --- crates/common/src/matrix.rs | 111 +++++++++++++++++----------------- crates/common/src/params.rs | 8 +-- crates/common/src/table.rs | 47 ++++++++------ crates/gkr/src/lib.rs | 7 ++- crates/prover/src/reduce.rs | 7 +-- crates/verifier/src/reduce.rs | 3 +- 6 files changed, 95 insertions(+), 88 deletions(-) diff --git a/crates/common/src/matrix.rs b/crates/common/src/matrix.rs index 78ccb86e..2f7f5130 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::slice::Iter; +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 @@ -23,30 +23,46 @@ const BITS: usize = Word::BITS as usize; /// 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 { - packed: Vec, +pub struct BitMatrix<'a> { + packed: Cow<'a, [Word]>, log_dim2: usize, } -impl BitMatrix { +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: Vec, log_dim2: usize) -> Self { + 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() && log_dim2 <= bits.ilog2() as usize, + bits.is_power_of_two() && (1 << log_dim2) <= bits, "{bits} bits do not split into rows of 2^{log_dim2}" ); Self { packed, log_dim2 } } - /// Transposes `xs[dim1][2^log_dim2]` into an owned `[dim2][dim1]`, as - /// [`BitMatrix::transpose`] does, reading the words where they lie. - pub(crate) fn transposed(xs: &[Word], log_dim2: usize) -> Self { - let dim2 = 1 << log_dim2; - let dim1 = xs.len() * BITS / dim2; - Self::new(bit_transpose(xs, dim1, dim2), dim1.ilog2() as usize) + /// 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. @@ -59,21 +75,12 @@ impl BitMatrix { 1 << self.log_dim2 } - /// The bit `xs[i1][i2]`. - pub fn bit(&self, i1: usize, i2: usize) -> bool { - debug_assert!(i1 < self.dim1(), "i1 {i1} is outside dim1"); - debug_assert!(i2 < self.dim2(), "i2 {i2} is outside dim2"); - - let index = (i1 << self.log_dim2) | i2; - (self.packed[index / BITS] >> (index % BITS)) & 1 == 1 - } - /// The `dim2` bits of `xs[i1]`, in ascending order. /// /// # Panics /// /// If `dim2` is below `128`: `xs[i1]` then shares a word with its - /// neighbours, which only [`BitMatrix::bit`] reads. + /// neighbours. pub fn bits(&self, i1: usize) -> BitsIter<'_> { assert!( self.dim2() >= BITS, @@ -84,26 +91,6 @@ impl BitMatrix { BitsIter::new(self.packed[i1 * words..(i1 + 1) * words].iter()) } - /// Swaps the axes, returning an owned copy: `result.bit(i2, i1) == - /// self.bit(i1, i2)` for every `i1` and `i2`. - /// - /// Built for callers that must walk the strided axis: in place that is a - /// scatter, one `bit()` call and cache miss per bit. Transposing once up - /// front turns it into the sequential access [`BitMatrix::bits`] gives. - /// - /// # 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 buffer - /// the same size as the packed words (up to 4 GiB at `m = 35`), never - /// a view over the original. - /// - **Involution**: transposing the result returns the original bits. - pub fn transpose(&self) -> BitMatrix { - Self::transposed(&self.packed, self.log_dim2) - } - /// 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 @@ -112,9 +99,9 @@ impl BitMatrix { Self::new(self.packed, log_dim2) } - /// Unwraps the packed bits. + /// Unwraps the packed bits, copying them if they were borrowed. pub fn into_packed(self) -> Vec { - self.packed + self.packed.into_owned() } } @@ -168,8 +155,7 @@ fn bit_transpose(xs: &[Word], dim1: usize, dim2: usize) -> Vec { /// 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 `BitMatrix::bit` would need to do the same -/// work. +/// `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; @@ -202,7 +188,7 @@ fn transpose_reference(tt: &[Word]) -> Vec { 'outer: loop { for s in &mut segments { let Some(b) = s.next() else { break 'outer }; - // LSB first, like `bit()` and `StepIter`. + // LSB first, like `StepIter`. word |= b << offset; offset += 1; if offset == BITS { @@ -294,13 +280,23 @@ mod tests { } /// A pseudo-random `xs[2^log_dim1][2^log_dim2]`. - fn pseudo_random_matrix(log_dim1: usize, log_dim2: usize) -> BitMatrix { - let packed = (0..(1 << (log_dim1 + log_dim2)) / BITS) + 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. @@ -311,7 +307,7 @@ mod tests { for i1 in 0..matrix.dim1() { for i2 in 0..matrix.dim2() { - assert_eq!(matrix.bit(i1, i2), (i1, i2) == (1, 130), "xs[{i1}][{i2}]"); + assert_eq!(bit(&matrix, i1, i2), (i1, i2) == (1, 130), "xs[{i1}][{i2}]"); } } } @@ -326,7 +322,8 @@ mod tests { 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| matrix.bit(i1, i2)).collect(); + let expected: Vec = + (0..matrix.dim2()).map(|i2| bit(&matrix, i1, i2)).collect(); assert_eq!(bits, expected, "log_dim2 {log_dim2}, i1 {i1}"); } } @@ -390,8 +387,8 @@ mod tests { for i1 in 0..matrix.dim1() { for i2 in 0..matrix.dim2() { assert_eq!( - transposed.bit(i2, i1), - matrix.bit(i1, i2), + bit(&transposed, i2, i1), + bit(&matrix, i1, i2), "xs[{i1}][{i2}] vs its transpose" ); } @@ -421,8 +418,12 @@ mod tests { for index in 0..1 << log_bits { assert_eq!( - reshaped.bit(index >> log_dim2, index & (reshaped.dim2() - 1)), - matrix.bit(index >> TRANSPOSE_LOG_DIM2, index & (matrix.dim2() - 1)), + 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}" ); } diff --git a/crates/common/src/params.rs b/crates/common/src/params.rs index fad3619a..51ee4a07 100644 --- a/crates/common/src/params.rs +++ b/crates/common/src/params.rs @@ -11,10 +11,10 @@ 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 - /// [`BitTable::transpose`], 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 + /// 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, diff --git a/crates/common/src/table.rs b/crates/common/src/table.rs index f3e690d9..04edaada 100644 --- a/crates/common/src/table.rs +++ b/crates/common/src/table.rs @@ -80,23 +80,11 @@ impl<'a> BitTable<'a> { &self.packed[column * elements..(column + 1) * elements] } - /// Swaps rows and columns into an owned [`BitMatrix`] with `dim1` the - /// rows and `dim2` the columns: `result.bit(b, c) == self.bit(c, b)`. - /// - /// Built for callers like a leaf construction that must walk the table - /// row by row: the packing is column major, so a row-major walk over the - /// table is a scatter, one `bit()` call and cache miss per cell. After - /// the transpose each row is contiguous, read by [`BitMatrix::bits`]. - /// - /// # Panics - /// - /// If the table has fewer than `128` columns (`s < 7`): see - /// [`BitMatrix::transpose`], which needs both axes at least `128`. The - /// table's own `t >= 7` covers the rows, and - /// [`crate::ParamsError::ColumnCountTooNarrow`] keeps such shapes away - /// from the reduction. - pub fn transpose(&self) -> BitMatrix { - BitMatrix::transposed(self.packed, self.shape.log_rows()) + /// 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()) } } @@ -213,6 +201,27 @@ mod tests { } } + /// The matrix has the table's columns along `dim1`, each one contiguous. + #[test] + 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}"); + } + } + /// The transpose has the table's rows along `dim1`, each one contiguous. #[test] fn the_transpose_has_the_rows_contiguous() { @@ -221,7 +230,7 @@ mod tests { let packed = with_bits(&shape, &set); let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); - let transposed = table.transpose(); + let transposed = table.as_matrix().transpose(); assert_eq!(transposed.dim1(), shape.rows()); assert_eq!(transposed.dim2(), shape.columns()); @@ -244,6 +253,6 @@ mod tests { let packed = vec![F128::ZERO; (1 << shape.log_bits()) / PACKED_BITS]; let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap(); - table.transpose(); + table.as_matrix().transpose(); } } diff --git a/crates/gkr/src/lib.rs b/crates/gkr/src/lib.rs index f34ffdc4..a00ebfee 100644 --- a/crates/gkr/src/lib.rs +++ b/crates/gkr/src/lib.rs @@ -1,6 +1,6 @@ use std::collections::VecDeque; -use common::BitMatrix; +use common::BitTable; use field::{F128, Wide256}; use num_traits::{ConstOne, ConstZero}; use rayon::prelude::*; @@ -343,7 +343,10 @@ pub struct GrandProductCircuit { } impl GrandProductCircuit { - pub fn new(row_images: &[F128], matrix: BitMatrix) -> 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); diff --git a/crates/prover/src/reduce.rs b/crates/prover/src/reduce.rs index 0a42ab29..17e03b7f 100644 --- a/crates/prover/src/reduce.rs +++ b/crates/prover/src/reduce.rs @@ -15,12 +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 { - // 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. - // Transposed so each row can be read sequentially. - let transposed = table.transpose(); - - GrandProductCircuit::new(&fold.row_images, transposed) + GrandProductCircuit::new(&fold.row_images, table) } /// Reduces the grand-product circuit to a factored claim on the committed bits. diff --git a/crates/verifier/src/reduce.rs b/crates/verifier/src/reduce.rs index f8aa68cf..f684d5b8 100644 --- a/crates/verifier/src/reduce.rs +++ b/crates/verifier/src/reduce.rs @@ -100,8 +100,7 @@ mod round_trip_ai_test { .map(|i| F128::from((i as u128 + 3) * 0xABCDEF12345u128 + 1)) .collect(); - let transposed = table.transpose(); - let circuit = GrandProductCircuit::new(&row_images, transposed); + 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.