Merge grumpkin branch (PR a16z#1211): GLV scalar decomposition + generic MSM crate

Author: quangvdao

Cargo.lock

@@ -36,6 +36,7 @@ members = [
     "jolt-inlines/blake3",
     "jolt-inlines/bigint",
     "jolt-inlines/secp256k1",
+    "jolt-inlines/grumpkin",
     "examples/btreemap/host",
     "examples/btreemap/guest",
     "examples/collatz",
@@ -76,6 +77,8 @@ members = [
     "examples/merkle-tree/guest",
     "examples/hash-bench",
     "examples/hash-bench/guest",
+    "examples/msm",
+    "examples/msm/guest",
     "zklean-extractor",
     "z3-verifier",
 ]
@@ -254,3 +257,4 @@ jolt-inlines-blake2 = { path = "./jolt-inlines/blake2", default-features = false
 jolt-inlines-blake3 = { path = "./jolt-inlines/blake3", default-features = false }
 jolt-inlines-bigint = { path = "./jolt-inlines/bigint", default-features = false }
 jolt-inlines-secp256k1 = { path = "./jolt-inlines/secp256k1", default-features = false }
+jolt-inlines-grumpkin = { path = "./jolt-inlines/grumpkin", default-features = false }

Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "msm-bench"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+jolt-sdk = { workspace = true, features = ["host"] }
+tracing-subscriber.workspace = true
+tracing.workspace = true
+jolt-inlines-grumpkin = { workspace = true, features = ["host"] }
+guest = { package = "msm-bench-guest", path = "./guest" }
+ark-bn254.workspace = true

examples/msm/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "msm-bench-guest"
+version = "0.1.0"
+edition = "2021"
+
+[features]
+guest = []
+
+[dependencies]
+jolt = { package = "jolt-sdk", path = "../../../jolt-sdk", features = [] }
+jolt-inlines-grumpkin.workspace = true

examples/msm/guest/Cargo.toml

@@ -0,0 +1,186 @@
+use crate::fixed_base::FixedBaseTable as GenericFixedBaseTable;
+use crate::traits::{GlvCapable, MsmGroup};
+use jolt_inlines_grumpkin::{GrumpkinFr, GrumpkinPoint, UnwrapOrSpoilProof};
+
+// ============================================================
+// Curve Parameters
+// ============================================================
+
+pub const SCALAR_BITS: usize = 256;
+pub const GLV_SCALAR_BITS: usize = 128;
+
+// Pippenger parameters for baseline (256-bit scalars).
+pub const BASELINE_WINDOW: usize = 12;
+pub const BASELINE_BUCKETS: usize = 1 << BASELINE_WINDOW;
+pub const BASELINE_WINDOWS: usize = SCALAR_BITS.div_ceil(BASELINE_WINDOW);
+
+// Pippenger parameters for GLV (128-bit scalars).
+pub const GLV_WINDOW: usize = 8;
+pub const GLV_BUCKETS: usize = 1 << GLV_WINDOW;
+pub const GLV_WINDOWS: usize = GLV_SCALAR_BITS.div_ceil(GLV_WINDOW);
+
+// Fixed-base (generator) windowed multiplication parameters (256-bit scalars).
+pub const FIXED_BASE_WINDOW: usize = 14;
+pub const FIXED_BASE_BUCKETS: usize = 1 << FIXED_BASE_WINDOW;
+pub const FIXED_BASE_WINDOWS: usize = SCALAR_BITS.div_ceil(FIXED_BASE_WINDOW);
+
+pub type FixedBaseTable =
+    GenericFixedBaseTable<GrumpkinPoint, FIXED_BASE_WINDOWS, FIXED_BASE_BUCKETS>;
+
+// ============================================================
+// Trait Implementations
+// ============================================================
+
+impl MsmGroup for GrumpkinPoint {
+    #[inline(always)]
+    fn identity() -> Self {
+        GrumpkinPoint::infinity()
+    }
+
+    #[inline(always)]
+    fn is_identity(&self) -> bool {
+        self.is_infinity()
+    }
+
+    #[inline(always)]
+    fn add(&self, other: &Self) -> Self {
+        GrumpkinPoint::add(self, other)
+    }
+
+    #[inline(always)]
+    fn neg(&self) -> Self {
+        GrumpkinPoint::neg(self)
+    }
+
+    #[inline(always)]
+    fn double(&self) -> Self {
+        GrumpkinPoint::double(self)
+    }
+
+    #[inline(always)]
+    fn double_and_add(&self, other: &Self) -> Self {
+        GrumpkinPoint::double_and_add(self, other)
+    }
+}
+
+impl GlvCapable for GrumpkinPoint {
+    type HalfScalar = u128;
+    type FullScalar = GrumpkinFr;
+
+    #[inline(always)]
+    fn endomorphism(&self) -> Self {
+        GrumpkinPoint::endomorphism(self)
+    }
+
+    #[inline(always)]
+    fn decompose_scalar(k: &GrumpkinFr) -> [(bool, u128); 2] {
+        GrumpkinPoint::decompose_scalar(k)
+    }
+}
+
+// ============================================================
+// Benchmark Helpers
+// ============================================================
+
+/// Grumpkin Fr modulus limbs for scalar reduction.
+const FR_MODULUS_LIMBS: [u64; 4] = [
+    4332616871279656263,
+    10917124144477883021,
+    13281191951274694749,
+    3486998266802970665,
+];
+
+#[inline(always)]
+fn is_ge_modulus(x: &[u64; 4]) -> bool {
+    let m = FR_MODULUS_LIMBS;
+    if x[3] > m[3] {
+        return true;
+    }
+    if x[3] < m[3] {
+        return false;
+    }
+    if x[2] > m[2] {
+        return true;
+    }
+    if x[2] < m[2] {
+        return false;
+    }
+    if x[1] > m[1] {
+        return true;
+    }
+    if x[1] < m[1] {
+        return false;
+    }
+    x[0] >= m[0]
+}
+
+#[inline(always)]
+fn sub_modulus(x: [u64; 4]) -> [u64; 4] {
+    let m = FR_MODULUS_LIMBS;
+    let mut out = [0u64; 4];
+    let mut borrow = 0u128;
+    let mut i = 0;
+    while i < 4 {
+        let xi = x[i] as u128;
+        let mi = m[i] as u128 + borrow;
+        if xi >= mi {
+            out[i] = (xi - mi) as u64;
+            borrow = 0;
+        } else {
+            out[i] = ((1u128 << 64) + xi - mi) as u64;
+            borrow = 1;
+        }
+        i += 1;
+    }
+    out
+}
+
+/// Reduce scalar modulo Fr.
+#[inline(always)]
+pub fn reduce_scalar(mut scalar: [u64; 4]) -> [u64; 4] {
+    while is_ge_modulus(&scalar) {
+        scalar = sub_modulus(scalar);
+    }
+    scalar
+}
+
+/// Generate deterministic test scalars using a simple LCG.
+#[inline(always)]
+pub fn generate_scalars<const N: usize>(seed: u64) -> [[u64; 4]; N] {
+    let mut scalars = [[0u64; 4]; N];
+    let (a, c) = (6364136223846793005u64, 1442695040888963407u64);
+    let mut state = seed;
+
+    for scalar in scalars.iter_mut() {
+        for limb in scalar.iter_mut() {
+            state = state.wrapping_mul(a).wrapping_add(c);
+            *limb = state;
+        }
+        *scalar = reduce_scalar(*scalar);
+    }
+    scalars
+}
+
+/// Generate deterministic test points by fixed-base scalar multiplication of generator.
+#[inline(always)]
+pub fn generate_points_fixed_base<const N: usize>(
+    seed: u64,
+    table_g: &FixedBaseTable,
+) -> [GrumpkinPoint; N] {
+    let mut points: [GrumpkinPoint; N] = core::array::from_fn(|_| GrumpkinPoint::infinity());
+    let (a, c) = (6364136223846793005u64, 1442695040888963407u64);
+    let mut state = seed;
+
+    for point in points.iter_mut() {
+        state = state.wrapping_mul(a).wrapping_add(c);
+        let small_scalar = [state, 0, 0, 0];
+        *point = table_g.scalar_mul(&small_scalar);
+    }
+    points
+}
+
+/// Convert [u64; 4] to GrumpkinFr.
+#[inline(always)]
+pub fn scalar_to_fr(scalar: &[u64; 4]) -> GrumpkinFr {
+    GrumpkinFr::from_u64_arr(scalar).unwrap_or_spoil_proof()
+}

examples/msm/guest/src/curves/grumpkin.rs

@@ -0,0 +1 @@
+pub mod grumpkin;

examples/msm/guest/src/curves/mod.rs

@@ -0,0 +1,87 @@
+use alloc::boxed::Box;
+
+use crate::traits::{MsmGroup, WindowedScalar};
+
+/// Fixed-base precomputed table for a single base point.
+/// table[window][digit] = digit · 2^(window·width) · base.
+///
+/// Type parameters:
+/// - `G`: The group/point type
+/// - `WINDOWS`: Number of windows (= ceil(scalar_bits / window_bits))
+/// - `BUCKETS`: Number of buckets per window (= 2^window_bits)
+pub struct FixedBaseTable<G, const WINDOWS: usize, const BUCKETS: usize> {
+    table: Box<[[G; BUCKETS]; WINDOWS]>,
+}
+
+impl<G: MsmGroup, const WINDOWS: usize, const BUCKETS: usize> FixedBaseTable<G, WINDOWS, BUCKETS> {
+    /// Window size derived from BUCKETS at compile time.
+    pub const WINDOW_BITS: usize = BUCKETS.trailing_zeros() as usize;
+
+    /// Precompute table for a given base point.
+    /// Window size is derived from BUCKETS const generic.
+    #[inline(always)]
+    pub fn new(base: &G) -> Self {
+        const {
+            assert!(BUCKETS > 1, "BUCKETS must be > 1");
+            assert!(BUCKETS.is_power_of_two(), "BUCKETS must be a power of two");
+            assert!(
+                BUCKETS.trailing_zeros() as usize <= 16,
+                "BUCKETS (window size) must be <= 2^16"
+            );
+        }
+
+        let mut table: Box<[[G; BUCKETS]; WINDOWS]> = Box::new(core::array::from_fn(|_| {
+            core::array::from_fn(|_| G::identity())
+        }));
+
+        let mut window_base = base.clone();
+        for window_table in table.iter_mut() {
+            window_table[0] = G::identity();
+            window_table[1] = window_base.clone();
+            let mut digit = 2;
+            while digit < BUCKETS {
+                window_table[digit] = window_table[digit - 1].add(&window_base);
+                digit += 1;
+            }
+
+            let mut shift = 0;
+            while shift < Self::WINDOW_BITS {
+                window_base = window_base.double();
+                shift += 1;
+            }
+        }
+
+        Self { table }
+    }
+
+    /// Scalar multiplication using table lookup + addition.
+    #[inline(always)]
+    pub fn scalar_mul<S: WindowedScalar>(&self, scalar: &S) -> G {
+        let mut result = G::identity();
+        for (window_idx, window_table) in self.table.iter().enumerate() {
+            let digit = scalar.window(window_idx * Self::WINDOW_BITS, Self::WINDOW_BITS) as usize;
+            if digit != 0 {
+                result = result.add(&window_table[digit]);
+            }
+        }
+        result
+    }
+}
+
+/// Fixed-base MSM: Σ(scalar_i · base) using precomputed table.
+#[inline(always)]
+pub fn msm_fixed_base<G, S, const WINDOWS: usize, const BUCKETS: usize>(
+    scalars: &[S],
+    table: &FixedBaseTable<G, WINDOWS, BUCKETS>,
+) -> G
+where
+    G: MsmGroup,
+    S: WindowedScalar,
+{
+    let mut result = G::identity();
+    for scalar in scalars {
+        let term = table.scalar_mul(scalar);
+        result = result.add(&term);
+    }
+    result
+}