Feat/batched grand products (#481)

* temp

* Passes prover assertions

* temp2

* temp3

* Sparse layers working(?)

* temp

* temp

* Fix timestamp range check opening points

* temp

* Tests passing

* Update profiling command

* Optimize init_final_leaves computation

* Add SplitEqPolynomial

* Integrate SplitEqPolynomial into dense grand product

* SparseInterleavedPolynomial

* Integrate SparseInterleavedPolynomial into sparse grand product

* temp: par_blocks

* Replace bind with bind_par_blocks

* too slow :(

* Revert to paralellizing over batch in SparseInterleavedPoly

* temp

* temp2

* temp

* Use DenseInterlaevedPolynomial in dense grand product

* Fix dense tests

* Add dense polynomial benchmarks

* switch to scratch space approach

* Use DenseInterleavedPolynomial for coalesced SparseInterleavedPolynomial

* temp

* Refactor (move BatchedCubicSumcheck impls into dense_interleaved_poly and sparse_interleaved_poly)

* Fix SparseInterleavedPolynomial::layer_output and add debugging

* temp

* tests passing thank god

* cleanup

* compute_cubic bench

* Switch summation order in DenesInterleavedPolynomial::compute_cubic

* Avoid recomputing E1_evals

* Switch summation order in SparseInterleavedPolynomial::compute_cubic

* Fix SparseInterleavedPolynomial::compute_cubic bench

* Drop bound polynomials in background (in primary sumcheck)

* Switch summation order in BatchedGrandProductToggleLayer::compute_cubic

* Optimize for 1s in SparseInterleavedPolynomial::compute_cubic

* cleanup

* Use simplified RS fingerprint for timestamp range-checks

* Un-comment surge.rs

* Update Quarks grand product

* clippy

* fix benches

* Update grand_product_example script and disable test

* Add comments

README.md

@@ -6,7 +6,7 @@ Just One Lookup Table.
 
 Jolt is a zkVM (zero-knowledge virtual machine) for RISC-V, built to be the simplest, fastest, and most extensible general-purpose of its kind. This repository currently contains an implementation of Jolt for the RISC-V 32-bit Base Integer instruction set (RV32I). _Contributors are welcome!_
 
-The Jolt [paper](https://eprint.iacr.org/2023/1217.pdf) was written by Arasu Arun, Srinath Setty, and Justin Thaler. 
+The Jolt [paper](https://eprint.iacr.org/2023/1217.pdf) was written by Arasu Arun, Srinath Setty, and Justin Thaler.
 
 ## Resources
 
@@ -71,15 +71,15 @@ Examples in the [`examples`](./examples/) directory can be run using e.g.
 
 ## Performance profiling
 
-Jolt uses [tracing_chrome](https://crates.io/crates/tracing-chrome) for performance profiling. 
+Jolt uses [tracing_chrome](https://crates.io/crates/tracing-chrome) for performance profiling.
 
 To generate a trace, run:
 
-```cargo run --profile build-fast -p jolt-core trace --name sha3 --format chrome```
+```cargo run --release -p jolt-core trace --name sha3 --format chrome --pcs hyper-kzg```
 
 Where `--name` can be `sha2`, `sha3`, `sha2-chain`, or `fibonacci`. The corresponding guest programs can be found in the [`examples`](./examples/) directory. The benchmark inputs are provided in [`bench.rs`](./jolt-core/src/benches/bench.rs).
 
-The above command will output a JSON file, e.g. `trace-1712455107389520.json`, which can be viewed in [Perfetto](https://ui.perfetto.dev/). 
+The above command will output a JSON file, e.g. `trace-1712455107389520.json`, which can be viewed in [Perfetto](https://ui.perfetto.dev/).
 
 ## Acknowledgements
 

jolt-core/Cargo.toml

@@ -86,6 +86,14 @@ harness = false
 name = "commit"
 harness = false
 
+[[bench]]
+name = "binding"
+harness = false
+
+[[bench]]
+name = "compute_cubic"
+harness = false
+
 [lib]
 name = "jolt_core"
 path = "src/lib.rs"

jolt-core/benches/binding.rs

@@ -0,0 +1,175 @@
+use ark_bn254::Fr;
+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::subprotocols::sumcheck::Bindable;
+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))
+        .take(1 << num_vars)
+        .collect()
+}
+
+fn random_sparse_coeffs<F: JoltField>(
+    rng: &mut impl Rng,
+    batch_size: usize,
+    num_vars: usize,
+    density: f64,
+) -> Vec<Vec<SparseCoefficient<F>>> {
+    (0..batch_size)
+        .map(|batch_index| {
+            let mut coeffs: Vec<SparseCoefficient<F>> = vec![];
+            for i in 0..(1 << num_vars) {
+                if rng.gen_bool(density) {
+                    coeffs.push((batch_index * (1 << num_vars) + i, F::random(rng)).into())
+                }
+            }
+            coeffs
+        })
+        .collect()
+}
+
+fn benchmark_dense<F: JoltField>(c: &mut Criterion, num_vars: usize) {
+    c.bench_function(
+        &format!("DensePolynomial::bind {} variables", num_vars),
+        |b| {
+            b.iter_with_setup(
+                || {
+                    let mut rng = test_rng();
+                    let coeffs = random_dense_coeffs(&mut rng, num_vars);
+                    let poly = DensePolynomial::new(coeffs);
+                    let r: Vec<F> = std::iter::repeat_with(|| F::random(&mut rng))
+                        .take(num_vars)
+                        .collect();
+                    (poly, r)
+                },
+                |(mut poly, r)| {
+                    for i in 0..num_vars {
+                        criterion::black_box(poly.bound_poly_var_top_par(&r[i]));
+                    }
+                },
+            );
+        },
+    );
+}
+
+fn benchmark_dense_batch<F: JoltField>(c: &mut Criterion, num_vars: usize, batch_size: usize) {
+    c.bench_function(
+        &format!(
+            "DensePolynomial::bind {} x {} variables",
+            batch_size, num_vars
+        ),
+        |b| {
+            b.iter_with_setup(
+                || {
+                    let mut rng = test_rng();
+                    let mut polys = vec![];
+                    for _ in 0..batch_size {
+                        let coeffs = random_dense_coeffs(&mut rng, num_vars);
+                        polys.push(DensePolynomial::new(coeffs));
+                    }
+                    let r: Vec<F> = std::iter::repeat_with(|| F::random(&mut rng))
+                        .take(num_vars)
+                        .collect();
+                    (polys, r)
+                },
+                |(mut polys, r)| {
+                    for i in 0..num_vars {
+                        polys
+                            .par_iter_mut()
+                            .for_each(|poly| poly.bound_poly_var_bot(&r[i]))
+                    }
+                },
+            );
+        },
+    );
+}
+
+fn benchmark_dense_interleaved<F: JoltField>(c: &mut Criterion, num_vars: usize) {
+    c.bench_function(
+        &format!("DenseInterleavedPolynomial::bind {} variables", num_vars),
+        |b| {
+            b.iter_with_setup(
+                || {
+                    let mut rng = test_rng();
+                    let coeffs = random_dense_coeffs(&mut rng, num_vars);
+                    let poly = DenseInterleavedPolynomial::new(coeffs);
+                    let r: Vec<F> = std::iter::repeat_with(|| F::random(&mut rng))
+                        .take(num_vars)
+                        .collect();
+                    (poly, r)
+                },
+                |(mut poly, r)| {
+                    for i in 0..num_vars {
+                        criterion::black_box(poly.bind(r[i]));
+                    }
+                },
+            );
+        },
+    );
+}
+
+fn benchmark_sparse_interleaved<F: JoltField>(
+    c: &mut Criterion,
+    batch_size: usize,
+    num_vars: usize,
+    density: f64,
+) {
+    c.bench_function(
+        &format!(
+            "SparseInterleavedPolynomial::bind {} x {} variables, {}% ones",
+            batch_size,
+            num_vars,
+            (1.0 - density) * 100.0
+        ),
+        |b| {
+            b.iter_with_setup(
+                || {
+                    let mut rng = test_rng();
+                    let coeffs = random_sparse_coeffs(&mut rng, batch_size, num_vars, density);
+                    let poly = SparseInterleavedPolynomial::new(coeffs, batch_size << num_vars);
+                    let r: Vec<F> = std::iter::repeat_with(|| F::random(&mut rng))
+                        .take(num_vars)
+                        .collect();
+                    (poly, r)
+                },
+                |(mut poly, r)| {
+                    for i in 0..num_vars {
+                        criterion::black_box(poly.bind(r[i]));
+                    }
+                },
+            );
+        },
+    );
+}
+
+fn main() {
+    let mut criterion = Criterion::default()
+        .configure_from_args()
+        .warm_up_time(std::time::Duration::from_secs(5));
+
+    benchmark_sparse_interleaved::<Fr>(&mut criterion, 64, 20, 0.1);
+    benchmark_sparse_interleaved::<Fr>(&mut criterion, 128, 20, 0.1);
+    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_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, 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);
+
+    criterion.final_summary();
+}

jolt-core/benches/compute_cubic.rs

@@ -0,0 +1,126 @@
+use ark_bn254::Fr;
+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::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))
+        .take(1 << num_vars)
+        .collect()
+}
+
+fn random_sparse_coeffs<F: JoltField>(
+    rng: &mut impl Rng,
+    batch_size: usize,
+    num_vars: usize,
+    density: f64,
+) -> Vec<Vec<SparseCoefficient<F>>> {
+    (0..batch_size)
+        .map(|batch_index| {
+            let mut coeffs: Vec<SparseCoefficient<F>> = vec![];
+            for i in 0..(1 << num_vars) {
+                if rng.gen_bool(density) {
+                    coeffs.push((batch_index * (1 << num_vars) + i, F::random(rng)).into())
+                }
+            }
+            coeffs
+        })
+        .collect()
+}
+
+fn benchmark_dense_interleaved<F: JoltField>(c: &mut Criterion, num_vars: usize) {
+    c.bench_function(
+        &format!(
+            "DenseInterleavedPolynomial::compute_cubic {} variables",
+            num_vars
+        ),
+        |b| {
+            b.iter_with_setup(
+                || {
+                    let mut rng = test_rng();
+                    let coeffs = random_dense_coeffs(&mut rng, num_vars);
+                    let poly = DenseInterleavedPolynomial::new(coeffs);
+                    let r_eq: Vec<F> = std::iter::repeat_with(|| F::random(&mut rng))
+                        .take(num_vars)
+                        .collect();
+                    let eq_poly = SplitEqPolynomial::new(&r_eq);
+                    let claim = F::random(&mut rng);
+                    (poly, eq_poly, claim)
+                },
+                |(poly, eq_poly, claim)| {
+                    criterion::black_box(
+                        BatchedCubicSumcheck::<F, KeccakTranscript>::compute_cubic(
+                            &poly, &eq_poly, claim,
+                        ),
+                    );
+                },
+            );
+        },
+    );
+}
+
+fn benchmark_sparse_interleaved<F: JoltField>(
+    c: &mut Criterion,
+    batch_size: usize,
+    num_vars: usize,
+    density: f64,
+) {
+    c.bench_function(
+        &format!(
+            "SparseInterleavedPolynomial::compute_cubic {} x {} variables, {}% ones",
+            batch_size,
+            num_vars,
+            (1.0 - density) * 100.0
+        ),
+        |b| {
+            b.iter_with_setup(
+                || {
+                    let mut rng = test_rng();
+                    let coeffs = random_sparse_coeffs(&mut rng, batch_size, num_vars, density);
+                    let poly = SparseInterleavedPolynomial::new(coeffs, batch_size << num_vars);
+                    let r_eq: Vec<F> = std::iter::repeat_with(|| F::random(&mut rng))
+                        .take((batch_size << num_vars).next_power_of_two().log_2())
+                        .collect();
+                    let eq_poly = SplitEqPolynomial::new(&r_eq);
+                    let claim = F::random(&mut rng);
+                    (poly, eq_poly, claim)
+                },
+                |(poly, eq_poly, claim)| {
+                    criterion::black_box(
+                        BatchedCubicSumcheck::<F, KeccakTranscript>::compute_cubic(
+                            &poly, &eq_poly, claim,
+                        ),
+                    );
+                },
+            );
+        },
+    );
+}
+
+fn main() {
+    let mut criterion = Criterion::default()
+        .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, 21);
+    // 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, 25);
+
+    benchmark_sparse_interleaved::<Fr>(&mut criterion, 64, 20, 0.1);
+    benchmark_sparse_interleaved::<Fr>(&mut criterion, 128, 20, 0.1);
+    benchmark_sparse_interleaved::<Fr>(&mut criterion, 64, 21, 0.1);
+    benchmark_sparse_interleaved::<Fr>(&mut criterion, 128, 21, 0.1);
+
+    criterion.final_summary();
+}