chore: VID-GPU batch_commit now takes polynomials of different degr…

…ees (#559)

* padding vid polys

* temporarily downgrade tagged-base64

* pad scalars in load_polys_to_gpu

* fix bug

* bug fix

* update test

* use bn254

* hotfix & bug expose

* update icicle to 1.5.1

* update `tagged-base64` back

* bump version

* fix trivial bug

* rename var

* comments

CHANGELOG.md

@@ -19,6 +19,8 @@ We take inspiration from [keep changelog](https://keepachangelog.com/en/1.0.0/)
 
 ### Changed
 
+- [#559](https://github.com/EspressoSystems/jellyfish/pull/559) (`jf-primitives`) GPU-accelerated `batch_commit()` now doesn't require the input polynomials to be of the same degree.
+
 ### Removed
 
 ### Deprecated

Cargo.toml

@@ -3,7 +3,7 @@ members = ["plonk", "primitives", "relation", "utilities"]
 resolver = "2"
 
 [workspace.package]
-version = "0.4.3"
+version = "0.4.4"
 authors = ["Espresso Systems <[email protected]>"]
 edition = "2021"
 license = "MIT"
@@ -32,5 +32,5 @@ rand_chacha = { version = "0.3.1", default-features = false }
 serde = { version = "1.0", default-features = false, features = [ "derive", "rc" ] }
 sha2 = { version = "0.10", default-features = false }
 sha3 = { version = "0.10", default-features = false }
-itertools = { version = "0.12", default-features = false }
-tagged-base64 = { version = "0.4" }
+itertools = { version = "0.12", default-features = false } 
+tagged-base64 = "0.4"

primitives/Cargo.toml

@@ -39,11 +39,9 @@ generic-array = { version = "0", features = [
         "serde",
 ] } # not a direct dependency, but we need serde
 hashbrown = { workspace = true }
-icicle-bls12-377 = { git = "https://github.com/ingonyama-zk/icicle.git", tag = "v1.5.0", optional = true, features = ["arkworks"] }
-icicle-bls12-381 = { git = "https://github.com/ingonyama-zk/icicle.git", tag = "v1.5.0", optional = true, features = ["arkworks"] }
-icicle-bn254 = { git = "https://github.com/ingonyama-zk/icicle.git", tag = "v1.5.0", optional = true, features = ["arkworks"] }
-icicle-core = { git = "https://github.com/ingonyama-zk/icicle.git", tag = "v1.5.0", optional = true }
-icicle-cuda-runtime = { git = "https://github.com/ingonyama-zk/icicle.git", tag = "v1.5.0", optional = true }
+icicle-bn254 = { git = "https://github.com/ingonyama-zk/icicle.git", tag = "v1.5.1", optional = true, features = ["arkworks"] }
+icicle-core = { git = "https://github.com/ingonyama-zk/icicle.git", tag = "v1.5.1", optional = true }
+icicle-cuda-runtime = { git = "https://github.com/ingonyama-zk/icicle.git", tag = "v1.5.1", optional = true }
 itertools = { workspace = true, features = ["use_alloc"] }
 jf-relation = { path = "../relation", default-features = false }
 jf-utils = { path = "../utilities", default-features = false }

primitives/src/pcs/univariate_kzg/mod.rs

@@ -655,6 +655,7 @@ where
 pub(crate) mod icicle {
     use super::*;
     use crate::icicle_deps::{curves::*, *};
+    use itertools::Itertools;
 
     /// Trait for GPU-accelerated PCS.commit APIs
     pub trait GPUCommittable<E: Pairing> {
@@ -699,8 +700,7 @@ pub(crate) mod icicle {
             Ok(comm)
         }
 
-        /// Similar to [`Self::gpu_commit()`] but for a batch of poly of
-        /// the same degree
+        /// Similar to [`Self::gpu_commit()`] but for a batch of polys
         fn gpu_batch_commit(
             prover_param: impl Borrow<UnivariateProverParam<E>>,
             polys: &[DensePolynomial<E::ScalarField>],
@@ -711,12 +711,7 @@ pub(crate) mod icicle {
 
             let stream = warmup_new_stream().unwrap();
 
-            let degree = polys[0].degree();
-            if polys.iter().any(|p| p.degree() != degree) {
-                return Err(PCSError::InvalidParameters(
-                    "all polys should have the same degree".to_string(),
-                ));
-            }
+            let degree = polys.iter().map(|poly| poly.degree()).max().unwrap_or(0);
 
             #[cfg(feature = "kzg-print-trace")]
             let commit_time = start_timer!(|| format!(
@@ -836,8 +831,8 @@ pub(crate) mod icicle {
             Ok(scalars_on_device)
         }
 
-        /// Similar to [`Self::load_poly_to_gpu()`] but handling a batch of poly
-        /// of the same degree at once
+        /// Similar to [`Self::load_poly_to_gpu()`] but handling a batch of
+        /// polys at once
         fn load_batch_poly_to_gpu<'poly>(
             polys: &[DensePolynomial<E::ScalarField>],
         ) -> Result<HostOrDeviceSlice<'poly, <Self::IC as IcicleCurve>::ScalarField>, PCSError>
@@ -848,23 +843,27 @@ pub(crate) mod icicle {
                 ));
             }
 
-            let size = polys[0].degree() + 1;
-            if polys.iter().any(|p| p.degree() + 1 != size) {
-                return Err(PCSError::InvalidParameters(
-                    "all polys should have the same degree".to_string(),
-                ));
-            }
+            let num_coeffs = polys
+                .iter()
+                .map(|poly| poly.degree() + 1)
+                .max()
+                .unwrap_or(1);
 
             let mut scalars_on_device = HostOrDeviceSlice::<
                 '_,
                 <Self::IC as IcicleCurve>::ScalarField,
-            >::cuda_malloc(size * polys.len())?;
+            >::cuda_malloc(num_coeffs * polys.len())?;
 
             #[cfg(feature = "kzg-print-trace")]
             let conv_time = start_timer!(|| "Type Conversion: ark->ICICLE: Scalar");
+            let zero_for_padding = E::ScalarField::zero();
             let scalars: Vec<<Self::IC as IcicleCurve>::ScalarField> = polys
                 .iter()
-                .flat_map(|poly| poly.coeffs())
+                .flat_map(|poly| {
+                    poly.coeffs()
+                        .iter()
+                        .pad_using(num_coeffs, |_| &zero_for_padding)
+                })
                 .collect::<Vec<_>>()
                 .into_par_iter()
                 .map(|&s| Self::ark_field_to_icicle(s))
@@ -1389,17 +1388,24 @@ mod tests {
                 );
 
                 // batch commit
-                let batch_size = rng.gen_range(10..100);
-                let polys: Vec<_> = (0..batch_size).map(|_| <DensePolynomial<E::ScalarField> as DenseUVPolynomial<E::ScalarField>>::rand(
-                        degree, rng,
-                    )).collect();
-                let comms_gpu =
-                    <UnivariateKzgPCS<E> as GPUCommittable<E>>::gpu_batch_commit(&ck, &polys)?;
-                let comms_cpu = UnivariateKzgPCS::<E>::batch_commit(&ck, &polys)?;
-                assert_eq!(comms_gpu, comms_cpu);
-                assert!(
-                    <UnivariateKzgPCS<E> as GPUCommittable<E>>::gpu_batch_commit(&ck, &[]).is_ok()
-                );
+                for i in 0..5 {
+                    let batch_size = 10 + i;
+                    let polys: Vec<_> = (0..batch_size)
+                        .map(|_| {
+                            <DensePolynomial<E::ScalarField> as DenseUVPolynomial<
+                                    E::ScalarField,
+                                >>::rand(degree, rng)
+                        })
+                        .collect();
+                    let comms_gpu =
+                        <UnivariateKzgPCS<E> as GPUCommittable<E>>::gpu_batch_commit(&ck, &polys)?;
+                    let comms_cpu = UnivariateKzgPCS::<E>::batch_commit(&ck, &polys)?;
+                    assert_eq!(comms_gpu, comms_cpu);
+                    assert!(
+                        <UnivariateKzgPCS<E> as GPUCommittable<E>>::gpu_batch_commit(&ck, &[])
+                            .is_ok()
+                    );
+                }
             }
             Ok(())
         }

primitives/src/vid/advz.rs

@@ -378,6 +378,9 @@ where
     ) -> VidResult<Vec<KzgCommit<E>>> {
         // let mut srs_on_gpu = self.srs_on_gpu_and_cuda_stream.as_mut().unwrap().0;
         // let stream = &self.srs_on_gpu_and_cuda_stream.as_ref().unwrap().1;
+        if polys.is_empty() {
+            return Ok(vec![]);
+        }
         let (srs_on_gpu, stream) = self.srs_on_gpu_and_cuda_stream.as_mut().unwrap(); // safe by construction
         <UnivariateKzgPCS<E> as GPUCommittable<E>>::gpu_batch_commit_with_loaded_prover_param(
             srs_on_gpu, polys, stream,

primitives/tests/advz.rs

@@ -1,15 +1,22 @@
 #![cfg(feature = "test-srs")]
-use ark_bls12_381::Bls12_381;
+use ark_bn254::Bn254;
 use ark_ff::{Field, PrimeField};
 use ark_std::rand::seq::SliceRandom;
 use jf_primitives::{
     pcs::{checked_fft_size, prelude::UnivariateKzgPCS, PolynomialCommitmentScheme},
-    vid::advz::Advz,
+    vid::advz,
 };
 use sha2::Sha256;
 
 mod vid;
 
+#[cfg(not(feature = "gpu-vid"))]
+/// Internal Jellyfish VID scheme
+type Advz<E, H> = advz::Advz<E, H>;
+#[cfg(feature = "gpu-vid")]
+/// Internal Jellyfish VID scheme
+type Advz<E, H> = advz::AdvzGPU<'static, E, H>;
+
 #[test]
 fn round_trip() {
     // play with these items
@@ -21,7 +28,7 @@ fn round_trip() {
     let supported_degree = vid_sizes.iter().max_by_key(|v| v.0).unwrap().0 - 1;
     let mut rng = jf_utils::test_rng();
     multiplicities.shuffle(&mut rng);
-    let srs = UnivariateKzgPCS::<Bls12_381>::gen_srs_for_testing(
+    let srs = UnivariateKzgPCS::<Bn254>::gen_srs_for_testing(
         &mut rng,
         checked_fft_size(supported_degree as usize).unwrap()
             * *multiplicities.iter().max().unwrap() as usize,
@@ -30,13 +37,13 @@ fn round_trip() {
 
     println!(
             "modulus byte len: {}",
-            (<<UnivariateKzgPCS<Bls12_381> as PolynomialCommitmentScheme>::Evaluation as Field>::BasePrimeField
+            (<<UnivariateKzgPCS<Bn254> as PolynomialCommitmentScheme>::Evaluation as Field>::BasePrimeField
                 ::MODULUS_BIT_SIZE - 7)/8 + 1
         );
 
     vid::round_trip(
         |recovery_threshold, num_storage_nodes, multiplicity| {
-            Advz::<Bls12_381, Sha256>::with_multiplicity(
+            Advz::<Bn254, Sha256>::with_multiplicity(
                 num_storage_nodes,
                 recovery_threshold,
                 multiplicity,