Binius counts

jolt-core/Cargo.toml

@@ -69,6 +69,7 @@ common = { path = "../common" }
 tracer = { path = "../tracer" }
 bincode = "1.3.3"
 bytemuck = "1.15.0"
+once_cell = "1.19.0"
 
 [build-dependencies]
 common = { path = "../common" }

jolt-core/src/field/ark.rs

@@ -37,4 +37,8 @@ impl JoltField for ark_bn254::Fr {
         assert_eq!(bytes.len(), Self::NUM_BYTES);
         ark_bn254::Fr::from_le_bytes_mod_order(bytes)
     }
+
+    fn from_count_index(index: u64) -> Self {
+        <Self as JoltField>::from_u64(index).unwrap()
+    }
 }

jolt-core/src/field/binius.rs

@@ -1,19 +1,43 @@
 use ark_serialize::{CanonicalDeserialize, CanonicalSerialize};
-use binius_field::{BinaryField128b, BinaryField128bPolyval};
+use binius_field::{BinaryField, BinaryField128b, BinaryField128bPolyval};
 use std::ops::{Add, AddAssign, Div, Mul, MulAssign, Neg, Sub, SubAssign};
 
-use super::JoltField;
+use super::precomputed;
+use super::{
+    precomputed::{
+        PRECOMPUTED_HIGH_128B, PRECOMPUTED_HIGH_128B_POLYVAL, PRECOMPUTED_LOW_128B,
+        PRECOMPUTED_LOW_128B_POLYVAL,
+    },
+    JoltField,
+};
 
 impl BiniusConstructable for BinaryField128b {
     fn new(n: u64) -> Self {
         Self::new(n as u128)
     }
+
+    fn precomputed_generator_multiples() -> (
+        &'static [Self; precomputed::TABLE_SIZE],
+        &'static [Self; precomputed::TABLE_SIZE],
+    ) {
+        (&PRECOMPUTED_LOW_128B, &PRECOMPUTED_HIGH_128B)
+    }
 }
 
 impl BiniusConstructable for BinaryField128bPolyval {
     fn new(n: u64) -> Self {
         Self::new(n as u128)
     }
+
+    fn precomputed_generator_multiples() -> (
+        &'static [Self; precomputed::TABLE_SIZE],
+        &'static [Self; precomputed::TABLE_SIZE],
+    ) {
+        (
+            &PRECOMPUTED_LOW_128B_POLYVAL,
+            &PRECOMPUTED_HIGH_128B_POLYVAL,
+        )
+    }
 }
 
 impl BiniusSpecific for BinaryField128b {}
@@ -22,8 +46,41 @@ impl BiniusSpecific for BinaryField128bPolyval {}
 /// Trait for BiniusField functionality specific to each impl.
 pub trait BiniusSpecific: binius_field::TowerField + BiniusConstructable + bytemuck::Pod {}
 
-pub trait BiniusConstructable {
+pub trait BiniusConstructable: BinaryField {
     fn new(n: u64) -> Self;
+
+    /// Binius counts are constructed from multiplicities of a Binary Field multiplicative generator.
+    /// Precomputing all required counts [0, 2^32] is prohibitively expensive and using iterative multiplication
+    /// or square and multiply is still excessively costly.
+    /// Utilizes a two-table lookup method to handle counts up to `2^32` efficiently:
+    /// - Decompose count `x` as `x = a * 2^16 + b` where `a` and `b` are within the range `[0, 2^16]`.
+    /// - Two Precomptued Lookup Tables:
+    ///   - One table stores `2^16` powers of `g^{2^16}`.
+    ///   - Another table stores `2^16` powers of `g`.
+    /// - Computes any count up to `2^32` using: `g^x = (g^{2^16})^a * g^b`
+    /// This is achieved with two lookups (one from each table) and a single multiplication.
+    ///
+    /// `precomputed_generator_multiples() -> (&low_table, &high_table)`
+    fn precomputed_generator_multiples() -> (
+        &'static [Self; precomputed::TABLE_SIZE],
+        &'static [Self; precomputed::TABLE_SIZE],
+    );
+
+    fn from_count_index(n: u64) -> Self {
+        const MAX_COUNTS: usize = 1 << (2 * precomputed::LOG_TABLE_SIZE);
+        let n = (n) as usize;
+        assert!(n <= MAX_COUNTS);
+        assert!(n != 0);
+
+        let high_index = (n >> precomputed::LOG_TABLE_SIZE) as usize;
+        let low_index = (n & ((1 << precomputed::LOG_TABLE_SIZE) - 1)) as usize;
+        let (precomputed_low, precomputed_high) = Self::precomputed_generator_multiples();
+        if n < precomputed::TABLE_SIZE {
+            precomputed_low[n]
+        } else {
+            precomputed_low[low_index] * precomputed_high[high_index]
+        }
+    }
 }
 
 #[derive(Default, Debug, Copy, Clone, Eq, PartialEq)]
@@ -67,6 +124,10 @@ impl<F: BiniusSpecific> JoltField for BiniusField<F> {
         let field_element = bytemuck::try_from_bytes::<F>(bytes).unwrap();
         Self(field_element.to_owned())
     }
+
+    fn from_count_index(index: u64) -> Self {
+        Self(<F as BiniusConstructable>::from_count_index(index))
+    }
 }
 
 impl<F: BiniusSpecific> Neg for BiniusField<F> {
@@ -183,3 +244,51 @@ impl<F: BiniusSpecific> ark_serialize::Valid for BiniusField<F> {
         todo!()
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use binius_field::Field;
+
+    #[test]
+    fn test_from_count_index() {
+        let actual = BiniusField::<BinaryField128b>::from_count_index(1);
+        let expected = BinaryField128b::MULTIPLICATIVE_GENERATOR;
+        assert_eq!(actual.0, expected);
+
+        let actual = BiniusField::<BinaryField128b>::from_count_index(2);
+        let expected =
+            BinaryField128b::MULTIPLICATIVE_GENERATOR * BinaryField128b::MULTIPLICATIVE_GENERATOR;
+        assert_eq!(actual.0, expected);
+
+        let actual = BiniusField::<BinaryField128b>::from_count_index(3);
+        let expected = BinaryField128b::MULTIPLICATIVE_GENERATOR
+            * BinaryField128b::MULTIPLICATIVE_GENERATOR
+            * BinaryField128b::MULTIPLICATIVE_GENERATOR;
+        assert_eq!(actual.0, expected);
+
+        let actual = BiniusField::<BinaryField128b>::from_count_index(1 << 17);
+        let mut expected = BinaryField128b::MULTIPLICATIVE_GENERATOR;
+        for _ in 0..17 {
+            expected = expected.square();
+        }
+        assert_eq!(actual.0, expected);
+
+        let actual = BiniusField::<BinaryField128b>::from_count_index((1 << 17) + 1);
+        let mut expected = BinaryField128b::MULTIPLICATIVE_GENERATOR;
+        for _ in 0..17 {
+            expected = expected.square();
+        }
+        expected *= BinaryField128b::MULTIPLICATIVE_GENERATOR;
+        assert_eq!(actual.0, expected);
+
+        let actual = BiniusField::<BinaryField128b>::from_count_index((1 << 18) + 2);
+        let mut expected = BinaryField128b::MULTIPLICATIVE_GENERATOR;
+        for _ in 0..18 {
+            expected = expected.square();
+        }
+        expected *= BinaryField128b::MULTIPLICATIVE_GENERATOR;
+        expected *= BinaryField128b::MULTIPLICATIVE_GENERATOR;
+        assert_eq!(actual.0, expected);
+    }
+}

jolt-core/src/field/mod.rs

@@ -37,7 +37,9 @@ pub trait JoltField:
     fn from_u64(n: u64) -> Option<Self>;
     fn square(&self) -> Self;
     fn from_bytes(bytes: &[u8]) -> Self;
+    fn from_count_index(index: u64) -> Self;
 }
 
 pub mod ark;
 pub mod binius;
+mod precomputed;

jolt-core/src/field/precomputed.rs

@@ -0,0 +1,97 @@
+use binius_field::{BinaryField, BinaryField128b, BinaryField128bPolyval};
+use once_cell::sync::Lazy;
+
+pub const LOG_TABLE_SIZE: usize = 16;
+pub const TABLE_SIZE: usize = 1 << LOG_TABLE_SIZE;
+
+// BinaryField128b
+pub static PRECOMPUTED_LOW_128B: Lazy<[BinaryField128b; TABLE_SIZE]> = Lazy::new(|| {
+    compute_powers::<BinaryField128b, TABLE_SIZE>(BinaryField128b::MULTIPLICATIVE_GENERATOR)
+});
+pub static PRECOMPUTED_HIGH_128B: Lazy<[BinaryField128b; TABLE_SIZE]> = Lazy::new(|| {
+    compute_powers_starting_from::<BinaryField128b, TABLE_SIZE>(
+        BinaryField128b::MULTIPLICATIVE_GENERATOR,
+        16,
+    )
+});
+
+// BinaryField128bPolyval
+pub static PRECOMPUTED_LOW_128B_POLYVAL: Lazy<[BinaryField128bPolyval; TABLE_SIZE]> =
+    Lazy::new(|| {
+        compute_powers::<BinaryField128bPolyval, TABLE_SIZE>(
+            BinaryField128bPolyval::MULTIPLICATIVE_GENERATOR,
+        )
+    });
+pub static PRECOMPUTED_HIGH_128B_POLYVAL: Lazy<[BinaryField128bPolyval; TABLE_SIZE]> =
+    Lazy::new(|| {
+        compute_powers_starting_from::<BinaryField128bPolyval, TABLE_SIZE>(
+            BinaryField128bPolyval::MULTIPLICATIVE_GENERATOR,
+            16,
+        )
+    });
+
+/// Computes `N` sequential powers of base^{[0, ... N]}
+const fn compute_powers<F: binius_field::BinaryField, const N: usize>(base: F) -> [F; N] {
+    let mut powers = [F::ZERO; N];
+    powers[0] = F::ONE;
+    powers[1] = base;
+    let mut i = 2;
+    while i < N {
+        powers[i] = powers[i - 1] * base;
+        i += 1;
+    }
+    powers
+}
+
+/// Computes `N` sequential powers of base^{2^starting_power}: {base^{2^starting_power}}^{[1, ... N]}
+const fn compute_powers_starting_from<F: binius_field::BinaryField, const N: usize>(
+    base: F,
+    starting_power: usize,
+) -> [F; N] {
+    let mut starting = base;
+    let mut count = 0;
+    // // Repeated squaring
+    while count < starting_power {
+        starting = starting.mul(starting);
+        count += 1;
+    }
+    compute_powers::<F, N>(starting)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use binius_field::{BinaryField, Field};
+
+    #[test]
+    fn test_compute_powers() {
+        let base = BinaryField128b::MULTIPLICATIVE_GENERATOR;
+        let powers = compute_powers::<BinaryField128b, 5>(base);
+        let expected_powers = [
+            BinaryField128b::ONE,
+            base,
+            base * base,
+            base * base * base,
+            base * base * base * base,
+        ];
+        assert_eq!(powers, expected_powers);
+    }
+
+    #[test]
+    fn test_compute_powers_starting_from() {
+        let base = BinaryField128b::MULTIPLICATIVE_GENERATOR;
+        let powers = compute_powers_starting_from::<BinaryField128b, 5>(base, 2);
+        let expected_starting_base = base * base * base * base;
+        let expected_powers = [
+            BinaryField128b::ONE,
+            expected_starting_base,
+            expected_starting_base * expected_starting_base,
+            expected_starting_base * expected_starting_base * expected_starting_base,
+            expected_starting_base
+                * expected_starting_base
+                * expected_starting_base
+                * expected_starting_base,
+        ];
+        assert_eq!(powers, expected_powers);
+    }
+}

jolt-core/src/lib.rs

@@ -7,6 +7,11 @@
 #![feature(generic_const_exprs)]
 #![feature(iter_next_chunk)]
 #![allow(long_running_const_eval)]
+// Note: Used exclusively by const fn BiniusConstructable::compute_powers.
+// Can be removed with a manual const fn for BinaryField multiplication.
+#![feature(const_trait_impl)]
+#![feature(effects)]
+#![feature(const_refs_to_cell)]
 
 pub mod benches;
 pub mod field;