Reorder R1CS constraints + CompactPolynomial refactor (#561)

* Add feature flag to reorder R1CS constraints in Az, Bz, Cz * Format * temp * temp * temp * temp2 * temp3 * Update instruction lookups primary sumcheck * Update opening proof sumcheck * Fix Quarks errors * Fix hyperkzg/zeromorph errors * temp * fix subtable tests * temp * Use precomputed lookup tables for small scalar -> Fr conversion * debugging * Optimize MSMs * More debugging * tests passing * Interpolate eval at 1 for primary and arbitrary sumcheck * Optimize openings * MultilinearPolynomial::batch_evaluate * optimize Spartan Az Bz Cz computation * temp * temp2 * debug * clippy * Fix test build * more CI fixes * forge fmt * Use new ordering for Spartan matrices by default * Remove boolean_polynomial (it'll be back) * Comments and cleanup * Uncomment Surge * Fix icicle MSM stuff * Add comments * Optimize for 1s and 0s in CompactPolynomial::bind * Fix MSM logic for CompactPolynomial<i64> * delete special_polys.rs --------- Co-authored-by: James Parker <[email protected]>
a16z · Jan 17, 2025 · b8d6f0a · b8d6f0a
1 parent 1145ec4
commit b8d6f0a
Show file tree

Hide file tree

Showing 83 changed files with 3,738 additions and 5,495 deletions.
diff --git a/common/src/lib.rs b/common/src/lib.rs
@@ -1,4 +1,3 @@
 pub mod attributes;
 pub mod constants;
 pub mod rv_trace;
-pub mod serializable;
diff --git a/common/src/serializable.rs b/common/src/serializable.rs
diff --git a/jolt-core/Cargo.toml b/jolt-core/Cargo.toml
@@ -29,7 +29,12 @@ default = [
     "rayon",
 ]
 host = ["dep:reqwest", "dep:tokio"]
-icicle = ["default", "dep:icicle-runtime", "dep:icicle-core", "dep:icicle-bn254"]
+icicle = [
+    "default",
+    "dep:icicle-runtime",
+    "dep:icicle-core",
+    "dep:icicle-bn254",
+]
 
 [dependencies]
 ark-bn254 = "0.4.0"
@@ -119,9 +124,13 @@ name = "jolt_core"
 path = "src/lib.rs"
 
 [target.'cfg(not(target_arch = "wasm32"))'.dependencies]
-icicle-runtime = { git = "https://github.com/ingonyama-zk/icicle-jolt.git", features = ["cuda_backend"], rev = "ed93e21", optional = true }
+icicle-runtime = { git = "https://github.com/ingonyama-zk/icicle-jolt.git", features = [
+    "cuda_backend",
+], rev = "ed93e21", optional = true }
 icicle-core = { git = "https://github.com/ingonyama-zk/icicle-jolt.git", rev = "ed93e21", optional = true }
-icicle-bn254 = { git = "https://github.com/ingonyama-zk/icicle-jolt.git", features = ["cuda_backend"], rev = "ed93e21", optional = true }
+icicle-bn254 = { git = "https://github.com/ingonyama-zk/icicle-jolt.git", features = [
+    "cuda_backend",
+], rev = "ed93e21", optional = true }
 memory-stats = "1.0.0"
 sys-info = "0.9.1"
 tokio = { version = "1.38.0", optional = true, features = ["rt-multi-thread"] }

diff --git a/jolt-core/benches/binding.rs b/jolt-core/benches/binding.rs
@@ -49,7 +49,7 @@ fn benchmark_dense<F: JoltField>(c: &mut Criterion, num_vars: usize) {
                 },
                 |(mut poly, r)| {
                     for i in 0..num_vars {
-                        criterion::black_box(poly.bound_poly_var_top_par(&r[i]));
+                        criterion::black_box(poly.bound_poly_var_top(&r[i]));
                     }
                 },
             );
@@ -157,19 +157,19 @@ fn main() {
     benchmark_sparse_interleaved::<Fr>(&mut criterion, 64, 21, 0.1);
     benchmark_sparse_interleaved::<Fr>(&mut criterion, 128, 21, 0.1);
 
-    // benchmark_dense::<Fr>(&mut criterion, 20);
-    // benchmark_dense::<Fr>(&mut criterion, 22);
-    // benchmark_dense::<Fr>(&mut criterion, 24);
+    benchmark_dense::<Fr>(&mut criterion, 20);
+    benchmark_dense::<Fr>(&mut criterion, 22);
+    benchmark_dense::<Fr>(&mut criterion, 24);
 
-    // benchmark_dense_interleaved::<Fr>(&mut criterion, 22);
+    benchmark_dense_interleaved::<Fr>(&mut criterion, 22);
     // benchmark_dense_interleaved::<Fr>(&mut criterion, 23);
-    // benchmark_dense_interleaved::<Fr>(&mut criterion, 24);
+    benchmark_dense_interleaved::<Fr>(&mut criterion, 24);
     // benchmark_dense_interleaved::<Fr>(&mut criterion, 25);
 
-    // benchmark_dense_batch::<Fr>(&mut criterion, 20, 4);
-    // benchmark_dense_batch::<Fr>(&mut criterion, 20, 8);
-    // benchmark_dense_batch::<Fr>(&mut criterion, 20, 16);
-    // benchmark_dense_batch::<Fr>(&mut criterion, 20, 32);
+    benchmark_dense_batch::<Fr>(&mut criterion, 20, 4);
+    benchmark_dense_batch::<Fr>(&mut criterion, 20, 8);
+    benchmark_dense_batch::<Fr>(&mut criterion, 20, 16);
+    benchmark_dense_batch::<Fr>(&mut criterion, 20, 32);
 
     criterion.final_summary();
 }
diff --git a/jolt-core/benches/commit.rs b/jolt-core/benches/commit.rs
@@ -5,6 +5,7 @@ use jolt_core::poly::commitment::commitment_scheme::{BatchType, CommitShape, Com
 use jolt_core::poly::commitment::hyperkzg::HyperKZG;
 use jolt_core::poly::commitment::kzg::CommitMode;
 use jolt_core::poly::commitment::zeromorph::Zeromorph;
+use jolt_core::poly::multilinear_polynomial::MultilinearPolynomial;
 use jolt_core::utils::transcript::{KeccakTranscript, Transcript};
 use rand_chacha::ChaCha20Rng;
 use rand_core::{RngCore, SeedableRng};
@@ -77,8 +78,8 @@ fn benchmark_commit<PCS, F, ProofTranscript>(
     let (leaves, setup, _) =
         setup_bench::<PCS, F, ProofTranscript>(num_layer, layer_size, threshold);
     let leaves = leaves
-        .iter()
-        .map(|layer| layer.as_slice())
+        .into_iter()
+        .map(|layer| MultilinearPolynomial::from(layer))
         .collect::<Vec<_>>();
     let mode = match batch_type {
         BatchType::GrandProduct => CommitMode::GrandProduct,
@@ -88,7 +89,11 @@ fn benchmark_commit<PCS, F, ProofTranscript>(
         &format!("{} Commit(mode:{:?}): {}% Ones", name, mode, threshold),
         |b| {
             b.iter(|| {
-                PCS::batch_commit(&leaves, &setup, batch_type.clone());
+                PCS::batch_commit(
+                    &leaves.iter().collect::<Vec<_>>(),
+                    &setup,
+                    batch_type.clone(),
+                );
             });
         },
     );

diff --git a/jolt-core/benches/compute_cubic.rs b/jolt-core/benches/compute_cubic.rs
@@ -3,13 +3,11 @@ use ark_std::{rand::Rng, test_rng};
 use criterion::Criterion;
 use jolt_core::field::JoltField;
 use jolt_core::poly::dense_interleaved_poly::DenseInterleavedPolynomial;
-use jolt_core::poly::dense_mlpoly::DensePolynomial;
 use jolt_core::poly::sparse_interleaved_poly::{SparseCoefficient, SparseInterleavedPolynomial};
 use jolt_core::poly::split_eq_poly::SplitEqPolynomial;
-use jolt_core::subprotocols::sumcheck::{BatchedCubicSumcheck, Bindable};
+use jolt_core::subprotocols::sumcheck::BatchedCubicSumcheck;
 use jolt_core::utils::math::Math;
 use jolt_core::utils::transcript::KeccakTranscript;
-use rayon::prelude::*;
 
 fn random_dense_coeffs<F: JoltField>(rng: &mut impl Rng, num_vars: usize) -> Vec<F> {
     std::iter::repeat_with(|| F::random(rng))
@@ -110,11 +108,11 @@ fn main() {
         .configure_from_args()
         .warm_up_time(std::time::Duration::from_secs(5));
 
-    // benchmark_dense_interleaved::<Fr>(&mut criterion, 20);
+    benchmark_dense_interleaved::<Fr>(&mut criterion, 20);
     // benchmark_dense_interleaved::<Fr>(&mut criterion, 21);
-    // benchmark_dense_interleaved::<Fr>(&mut criterion, 22);
+    benchmark_dense_interleaved::<Fr>(&mut criterion, 22);
     // benchmark_dense_interleaved::<Fr>(&mut criterion, 23);
-    // benchmark_dense_interleaved::<Fr>(&mut criterion, 24);
+    benchmark_dense_interleaved::<Fr>(&mut criterion, 24);
     // benchmark_dense_interleaved::<Fr>(&mut criterion, 25);
 
     benchmark_sparse_interleaved::<Fr>(&mut criterion, 64, 20, 0.1);

diff --git a/jolt-core/benches/iai.rs b/jolt-core/benches/iai.rs
@@ -1,20 +1,9 @@
-use ark_bn254::{Fr, G1Projective};
-use ark_ec::CurveGroup;
-use ark_std::{test_rng, UniformRand};
+use ark_bn254::Fr;
+use ark_std::test_rng;
 use iai_callgrind::{library_benchmark, library_benchmark_group, main};
-use jolt_core::{field::JoltField, msm::VariableBaseMSM, poly::dense_mlpoly::DensePolynomial};
+use jolt_core::{field::JoltField, poly::dense_mlpoly::DensePolynomial};
 use std::hint::black_box;
 
-fn msm_setup<G: CurveGroup>(num_points: usize) -> (Vec<G>, Vec<G::ScalarField>) {
-    let mut rng = test_rng();
-
-    // Generate a vector of random affine points on the curve.
-    (
-        vec![G::rand(&mut rng); num_points],
-        vec![G::ScalarField::rand(&mut rng); num_points],
-    )
-}
-
 fn bound_poly_setup<F: JoltField>(size: usize) -> (DensePolynomial<F>, F) {
     let mut rng = test_rng();
 
@@ -32,12 +21,6 @@ fn eval_poly_setup<F: JoltField>(size: usize) -> (DensePolynomial<F>, Vec<F>) {
     (poly, points)
 }
 
-#[library_benchmark]
-#[bench::long(msm_setup::<G1Projective>(4096))]
-fn bench_msm<G: CurveGroup>(input: (Vec<G>, Vec<G::ScalarField>)) -> G {
-    black_box(VariableBaseMSM::msm(&G::normalize_batch(&input.0), None, &input.1).unwrap())
-}
-
 #[library_benchmark]
 #[bench::long(bound_poly_setup::<Fr>(4096))]
 fn bench_polynomial_binding<F: JoltField>(input: (DensePolynomial<F>, F)) {
@@ -54,7 +37,7 @@ fn bench_polynomial_evaluate<F: JoltField>(input: (DensePolynomial<F>, Vec<F>))
 
 library_benchmark_group!(
     name = jolt_core_ops;
-    benchmarks = bench_msm, bench_polynomial_binding, bench_polynomial_evaluate
+    benchmarks = bench_polynomial_binding, bench_polynomial_evaluate
 );
 
 main!(library_benchmark_groups = jolt_core_ops);
diff --git a/jolt-core/benches/msm.rs b/jolt-core/benches/msm.rs
@@ -1,62 +1,65 @@
 use ark_bn254::{Bn254, Fr, G1Affine, G1Projective};
-use ark_ff::{BigInteger, PrimeField};
-use ark_std::rand::Rng;
 use ark_std::UniformRand;
-use ark_std::{One, Zero};
 use criterion::Criterion;
 use jolt_core::field::JoltField;
 #[cfg(feature = "icicle")]
 use jolt_core::msm::Icicle;
-use jolt_core::msm::{icicle_init, GpuBaseType, MsmType, VariableBaseMSM};
+use jolt_core::msm::{icicle_init, GpuBaseType, VariableBaseMSM};
 use jolt_core::poly::commitment::commitment_scheme::CommitmentScheme;
 use jolt_core::poly::commitment::zeromorph::Zeromorph;
+use jolt_core::poly::multilinear_polynomial::MultilinearPolynomial;
 use jolt_core::utils::transcript::{KeccakTranscript, Transcript};
 use rand_chacha::ChaCha20Rng;
 use rand_core::{RngCore, SeedableRng};
-use rayon::prelude::*;
 
 const SRS_SIZE: usize = 1 << 20;
 
 // Sets up the benchmark
 fn setup_bench<PCS, F, ProofTranscript>(
-    msm_type: MsmType,
+    max_num_bits: usize,
 ) -> (
     Vec<G1Affine>,
     Option<Vec<GpuBaseType<G1Projective>>>,
-    Vec<Fr>,
+    MultilinearPolynomial<Fr>,
 )
 where
     F: JoltField,
     PCS: CommitmentScheme<ProofTranscript, Field = F>,
     ProofTranscript: Transcript,
 {
     let mut rng = ChaCha20Rng::seed_from_u64(SRS_SIZE as u64);
-
-    let scalars = match msm_type {
-        MsmType::Zero => {
-            vec![Fr::zero(); SRS_SIZE]
-        }
-        MsmType::One => {
-            vec![Fr::one(); SRS_SIZE]
-        }
-        MsmType::Small(_) => (0..SRS_SIZE)
-            .into_iter()
-            .map(|_| {
-                let i = rng.gen_range(0..(1 << 10));
-                <Fr as JoltField>::from_u64(i).unwrap()
-            })
-            .collect(),
-        MsmType::Medium(_) => (0..SRS_SIZE)
-            .into_iter()
-            .map(|_| {
-                let i = rng.next_u64();
-                <Fr as JoltField>::from_u64(i).unwrap()
-            })
-            .collect(),
-        MsmType::Large(_) => (0..SRS_SIZE)
-            .into_iter()
-            .map(|_| Fr::random(&mut rng))
-            .collect(),
+    let poly = match max_num_bits {
+        0 => MultilinearPolynomial::from(vec![0u8; SRS_SIZE]),
+        1..=8 => MultilinearPolynomial::from(
+            (0..SRS_SIZE)
+                .into_iter()
+                .map(|_| (rng.next_u32() & ((1 << max_num_bits) - 1)) as u8)
+                .collect::<Vec<_>>(),
+        ),
+        9..=16 => MultilinearPolynomial::from(
+            (0..SRS_SIZE)
+                .into_iter()
+                .map(|_| (rng.next_u32() & ((1 << max_num_bits) - 1)) as u16)
+                .collect::<Vec<_>>(),
+        ),
+        17..=32 => MultilinearPolynomial::from(
+            (0..SRS_SIZE)
+                .into_iter()
+                .map(|_| (rng.next_u64() & ((1 << max_num_bits) - 1)) as u32)
+                .collect::<Vec<_>>(),
+        ),
+        33..=64 => MultilinearPolynomial::from(
+            (0..SRS_SIZE)
+                .into_iter()
+                .map(|_| rng.next_u64() & ((1 << max_num_bits) - 1))
+                .collect::<Vec<_>>(),
+        ),
+        _ => MultilinearPolynomial::from(
+            (0..SRS_SIZE)
+                .into_iter()
+                .map(|_| Fr::random(&mut rng))
+                .collect::<Vec<_>>(),
+        ),
     };
 
     let bases: Vec<G1Affine> = std::iter::repeat_with(|| G1Affine::rand(&mut rng))
@@ -72,23 +75,17 @@ where
     #[cfg(not(feature = "icicle"))]
     let gpu_bases = None;
 
-    let max_num_bits = scalars
-        .par_iter()
-        .map(|s| s.clone().into_bigint().num_bits())
-        .max()
-        .unwrap();
-
     println!("Using max num bits: {}", max_num_bits);
-    (bases, gpu_bases, scalars)
+    (bases, gpu_bases, poly)
 }
 
-fn benchmark_msm<PCS, F, ProofTranscript>(c: &mut Criterion, name: &str, msm_type: MsmType)
+fn benchmark_msm<PCS, F, ProofTranscript>(c: &mut Criterion, name: &str, max_num_bits: usize)
 where
     F: JoltField,
     PCS: CommitmentScheme<ProofTranscript, Field = F>,
     ProofTranscript: Transcript,
 {
-    let (bases, gpu_bases, scalars) = setup_bench::<PCS, F, ProofTranscript>(msm_type);
+    let (bases, gpu_bases, poly) = setup_bench::<PCS, F, ProofTranscript>(max_num_bits);
     icicle_init();
     #[cfg(feature = "icicle")]
     let id = format!("{} [mode:Icicle]", name);
@@ -97,7 +94,7 @@ where
     c.bench_function(&id, |b| {
         b.iter(|| {
             let msm =
-                <G1Projective as VariableBaseMSM>::msm(&bases, gpu_bases.as_deref(), &scalars);
+                <G1Projective as VariableBaseMSM>::msm(&bases, gpu_bases.as_deref(), &poly, None);
             let _ = msm.expect("MSM failed");
         });
     });
@@ -110,18 +107,33 @@ fn main() {
         .warm_up_time(std::time::Duration::from_secs(5));
     benchmark_msm::<Zeromorph<Bn254, KeccakTranscript>, Fr, KeccakTranscript>(
         &mut criterion,
-        "VariableBaseMSM::msm(Large)",
-        MsmType::Large(0 /* unused */),
+        "VariableBaseMSM::msm(256 bit scalars)",
+        256,
+    );
+    benchmark_msm::<Zeromorph<Bn254, KeccakTranscript>, Fr, KeccakTranscript>(
+        &mut criterion,
+        "VariableBaseMSM::msm(64 bit scalars)",
+        64,
+    );
+    benchmark_msm::<Zeromorph<Bn254, KeccakTranscript>, Fr, KeccakTranscript>(
+        &mut criterion,
+        "VariableBaseMSM::msm(32 bit scalars)",
+        32,
+    );
+    benchmark_msm::<Zeromorph<Bn254, KeccakTranscript>, Fr, KeccakTranscript>(
+        &mut criterion,
+        "VariableBaseMSM::msm(16 bit scalars)",
+        16,
     );
     benchmark_msm::<Zeromorph<Bn254, KeccakTranscript>, Fr, KeccakTranscript>(
         &mut criterion,
-        "VariableBaseMSM::msm(Medium)",
-        MsmType::Medium(0 /* unused */),
+        "VariableBaseMSM::msm(8 bit scalars)",
+        8,
     );
     benchmark_msm::<Zeromorph<Bn254, KeccakTranscript>, Fr, KeccakTranscript>(
         &mut criterion,
-        "VariableBaseMSM::msm(Small)",
-        MsmType::Small(0 /* unused */),
+        "VariableBaseMSM::msm(1 bit scalars)",
+        1,
     );
     criterion.final_summary();
 }