more CI fixes

a16z · Jan 15, 2025 · c691c63 · c691c63
1 parent 4ca1333
commit c691c63
Show file tree

Hide file tree

Showing 10 changed files with 213 additions and 337 deletions.
diff --git a/jolt-core/benches/binding.rs b/jolt-core/benches/binding.rs
@@ -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();
 }