Separate Bit Transpose from Bittable

crates/common/Cargo.toml

@@ -8,11 +8,23 @@ license.workspace = true
 [features]
 default = ["parallel"]
 parallel = ["dep:rayon"]
+# Exposes private kernels to `benches/`; not part of the public API.
+bench = []
 
 [dependencies]
 crypto-primitives = { workspace = true }
+bytemuck = {workspace = true}
 field = { workspace = true, features = ["spongefish"] }
 spongefish = { workspace = true }
 poly = { workspace = true }
 rayon = { workspace = true, optional = true }
 num-traits = { workspace = true }
+
+[dev-dependencies]
+divan = { workspace = true }
+rand = { workspace = true }
+
+[[bench]]
+name = "matrix"
+harness = false
+required-features = ["bench"]

crates/common/benches/matrix.rs

@@ -0,0 +1,49 @@
+//! Compare the blocked `bit_transpose` with the bit-at-a-time reference.
+//!
+//! Run with `cargo bench -p common --features bench --bench matrix`.
+//! Both transpose a `128 x DIM2`-bit matrix (`dim2` contiguous), including
+//! output allocation. Input generation runs outside the measured loop.
+
+use common::matrix::bench::{bit_transpose, transpose_reference};
+use divan::counter::BytesCount;
+use divan::{Bencher, black_box};
+use rand::rngs::StdRng;
+use rand::{RngExt, SeedableRng};
+
+const DIM1: usize = 128;
+/// `2^20` bits per row: a 16 MiB matrix, well past the last-level cache.
+const DIM2: usize = 1 << 20;
+
+fn main() {
+    let xs = input();
+    assert_eq!(
+        bit_transpose(&xs, DIM1, DIM2),
+        transpose_reference::<DIM1>(&xs),
+        "bit_transpose disagrees with the reference"
+    );
+    divan::main();
+}
+
+/// Seeded so every run transposes the same matrix.
+fn input() -> Vec<u128> {
+    let mut rng = StdRng::seed_from_u64(0);
+    (0..DIM1 * DIM2 / u128::BITS as usize)
+        .map(|_| rng.random())
+        .collect()
+}
+
+#[divan::bench]
+fn blocked(bencher: Bencher) {
+    let xs = input();
+    bencher
+        .counter(BytesCount::of_slice(&xs))
+        .bench(|| bit_transpose(black_box(&xs), DIM1, DIM2));
+}
+
+#[divan::bench]
+fn reference(bencher: Bencher) {
+    let xs = input();
+    bencher
+        .counter(BytesCount::of_slice(&xs))
+        .bench(|| transpose_reference::<DIM1>(black_box(&xs)));
+}

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::<u128>()
+        .fold((0usize, 0u128), |(base, mut total), &element| {
+            let mut remaining = element;
+            while remaining != 0 {
+                total += exponents[base + remaining.trailing_zeros() as usize];
+                // Clears the lowest set bit.
+                remaining &= remaining - 1;
+            }
+            (base + BitTable::BITS, total)
         })
-        .sum()
+        .1
 }
 
 /// Every column's fold, in column order.
@@ -190,6 +183,7 @@ mod tests {
     use num_traits::{ConstOne, ConstZero};
 
     use super::*;
+    const PACKED_BITS: usize = 128;
     use crate::{BitZParams, Shape};
 
     const Q114: u128 = (1 << 114) - 11;
@@ -237,7 +231,7 @@ mod tests {
         let claim = claim(field_weights, vec![Fq::ONE; shape.columns()]);
 
         let packed = witness(&shape, &[(1, 0), (5, 0), (127, 0), (64, 4)]);
-        let table = BitTable::new(shape, &packed).unwrap();
+        let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap();
 
         assert_eq!(
             fold_column(&table, &claim.row_exponents(), 0),
@@ -257,7 +251,7 @@ mod tests {
         );
         let all: Vec<(usize, usize)> = (0..shape.rows()).map(|row| (row, 0)).collect();
         let packed = witness(&shape, &all);
-        let table = BitTable::new(shape, &packed).unwrap();
+        let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap();
 
         assert_eq!(
             fold_column(&table, &claim.row_exponents(), 0),
@@ -278,7 +272,7 @@ mod tests {
             .map(|row| (row, 6))
             .collect();
         let packed = witness(&shape, &bits);
-        let table = BitTable::new(shape, &packed).unwrap();
+        let table = BitTable::new(shape, bytemuck::cast_slice(&packed)).unwrap();
 
         let expected: u128 = (0..shape.rows())
             .filter(|&row| table.bit(6, row))

crates/common/src/lib.rs

@@ -6,6 +6,7 @@
 
 pub mod claim;
 pub mod fold;
+pub mod matrix;
 pub mod opening;
 pub mod params;
 pub mod shape;
@@ -16,10 +17,11 @@ pub use claim::{ClaimError, LinearClaim, Root};
 pub use fold::{
     Fold, FoldError, column_images, fold_column, fold_columns, reconstruct, row_images,
 };
+pub use matrix::BitMatrix;
 pub use opening::OpeningQuery;
 pub use params::{BitZParams, ParamsError, VirtualParams, VirtualParamsError};
 pub use shape::{Shape, ShapeError};
-pub use table::{BitTable, TableError, TransposeError, TransposedBitTable};
+pub use table::{BitTable, TableError};
 pub use virtual_map::{
     TransposedWeights, VirtualMap, VirtualMapError, VirtualStatement, VirtualStatementError,
 };