From 6d9078eb7ea3519785f00ad6b0f0e8514460f221 Mon Sep 17 00:00:00 2001 From: mmagician Date: Tue, 23 Jan 2024 10:26:09 +0100 Subject: [PATCH 1/5] Replace calls to `merge_dual` with `merge` on the concatenated slices --- jolt-core/src/jolt/vm/instruction_lookups.rs | 2 +- jolt-core/src/lasso/surge.rs | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/jolt-core/src/jolt/vm/instruction_lookups.rs b/jolt-core/src/jolt/vm/instruction_lookups.rs index af5f2ea2d..18aa67c71 100644 --- a/jolt-core/src/jolt/vm/instruction_lookups.rs +++ b/jolt-core/src/jolt/vm/instruction_lookups.rs @@ -122,7 +122,7 @@ where fn batch(&self) -> Self::BatchedPolynomials { use rayon::prelude::*; let (batched_dim_read, (batched_final, batched_E, batched_flag)) = rayon::join( - || DensePolynomial::merge_dual(self.dim.as_ref(), self.read_cts.as_ref()), + || DensePolynomial::merge(&[self.dim.as_slice(), self.read_cts.as_slice()].concat()), || { let batched_final = DensePolynomial::merge(&self.final_cts); let (batched_E, batched_flag) = rayon::join( diff --git a/jolt-core/src/lasso/surge.rs b/jolt-core/src/lasso/surge.rs index 941e9dd23..e060ecb20 100644 --- a/jolt-core/src/lasso/surge.rs +++ b/jolt-core/src/lasso/surge.rs @@ -65,7 +65,7 @@ where #[tracing::instrument(skip_all, name = "SurgePolys::batch")] fn batch(&self) -> Self::BatchedPolynomials { let (batched_dim_read, (batched_final, batched_E)) = rayon::join( - || DensePolynomial::merge_dual(self.dim.as_ref(), self.read_cts.as_ref()), + || DensePolynomial::merge(&[self.dim.as_slice(), self.read_cts.as_slice()].concat()), || { rayon::join( || DensePolynomial::merge(&self.final_cts), From 83cb8450792000e460909315fb05b3970460ab49 Mon Sep 17 00:00:00 2001 From: mmagician Date: Tue, 23 Jan 2024 10:27:22 +0100 Subject: [PATCH 2/5] remove function `merge_dual` --- jolt-core/src/poly/dense_mlpoly.rs | 19 ------------------- 1 file changed, 19 deletions(-) diff --git a/jolt-core/src/poly/dense_mlpoly.rs b/jolt-core/src/poly/dense_mlpoly.rs index e4da0ac33..236b636af 100644 --- a/jolt-core/src/poly/dense_mlpoly.rs +++ b/jolt-core/src/poly/dense_mlpoly.rs @@ -469,25 +469,6 @@ impl DensePolynomial { DensePolynomial::new(Z) } - #[tracing::instrument(skip_all, name = "DensePoly.merge_dual")] - pub fn merge_dual(polys_a: &[T], polys_b: &[T]) -> DensePolynomial - where - T: AsRef>, - { - let total_len_a: usize = polys_a.iter().map(|poly| poly.as_ref().len()).sum(); - let total_len_b: usize = polys_b.iter().map(|poly| poly.as_ref().len()).sum(); - let total_len = total_len_a + total_len_b; - - let mut Z: Vec = Vec::with_capacity(total_len.next_power_of_two()); - polys_a.iter().for_each(|poly| Z.extend_from_slice(poly.as_ref().vec())); - polys_b.iter().for_each(|poly| Z.extend_from_slice(poly.as_ref().vec())); - - // pad the polynomial with zero polynomial at the end - Z.resize(Z.capacity(), F::zero()); - - DensePolynomial::new(Z) - } - pub fn combined_commit( &self, label: &'static [u8], From 21a59c6df7220aa92b7d0e8b0b3729a0db149d9a Mon Sep 17 00:00:00 2001 From: mmagician Date: Tue, 23 Jan 2024 10:29:11 +0100 Subject: [PATCH 3/5] `extend` on `DensePolynomial` was unused --- jolt-core/src/poly/dense_mlpoly.rs | 10 ---------- 1 file changed, 10 deletions(-) diff --git a/jolt-core/src/poly/dense_mlpoly.rs b/jolt-core/src/poly/dense_mlpoly.rs index 236b636af..751b1e89b 100644 --- a/jolt-core/src/poly/dense_mlpoly.rs +++ b/jolt-core/src/poly/dense_mlpoly.rs @@ -442,16 +442,6 @@ impl DensePolynomial { self.Z.as_ref() } - pub fn extend(&mut self, other: &DensePolynomial) { - assert_eq!(self.Z.len(), self.len); - let other_vec = other.vec(); - assert_eq!(other_vec.len(), self.len); - self.Z.extend(other_vec); - self.num_vars += 1; - self.len *= 2; - assert_eq!(self.Z.len(), self.len); - } - #[tracing::instrument(skip_all, name = "DensePoly.merge")] pub fn merge(polys: &[T]) -> DensePolynomial where From 570831e6d9347c500e8397227647c4d71190ceaa Mon Sep 17 00:00:00 2001 From: mmagician Date: Tue, 23 Jan 2024 10:38:05 +0100 Subject: [PATCH 4/5] Move msm-related functionality out of `DensePolynomial` and into msm mod --- jolt-core/src/msm/mod.rs | 83 ++++++++++++++++++++++++++++++ jolt-core/src/poly/dense_mlpoly.rs | 77 ++------------------------- 2 files changed, 86 insertions(+), 74 deletions(-) diff --git a/jolt-core/src/msm/mod.rs b/jolt-core/src/msm/mod.rs index 6a76ec523..50c40c1c9 100644 --- a/jolt-core/src/msm/mod.rs +++ b/jolt-core/src/msm/mod.rs @@ -328,3 +328,86 @@ fn ln_without_floats(a: usize) -> usize { // log2(a) * ln(2) (ark_std::log2(a) * 69 / 100) as usize } + +/// Special MSM where all scalar values are 0 / 1 – does not verify. +pub(crate) fn flags_msm(scalars: &[G::ScalarField], bases: &[G::Affine]) -> G { + assert_eq!(scalars.len(), bases.len()); + let result = scalars + .into_iter() + .enumerate() + .filter(|(_index, scalar)| !scalar.is_zero()) + .map(|(index, scalar)| bases[index]) + .sum(); + + result +} + +pub(crate) fn sm_msm( + scalars: &[::BigInt], + bases: &[V::MulBase], +) -> V { + assert_eq!(scalars.len(), bases.len()); + let num_buckets: usize = 1 << 16; // TODO(sragss): This should be passed in / dependent on M = N^{1/C} + + // #[cfg(test)] + // scalars.for_each(|scalar| { + // assert!(scalar < V::ScalarField::from(num_buckets as u64).into_bigint()) + // }); + + // Assign things to buckets based on the scalar + let mut buckets: Vec = vec![V::zero(); num_buckets]; + scalars.into_iter().enumerate().for_each(|(index, scalar)| { + let bucket_index: u64 = scalar.as_ref()[0]; + buckets[bucket_index as usize] += bases[index]; + }); + + let mut result = V::zero(); + let mut running_sum = V::zero(); + buckets + .into_iter() + .skip(1) + .enumerate() + .rev() + .for_each(|(index, bucket)| { + running_sum += bucket; + result += running_sum; + }); + result +} + +#[cfg(test)] +mod tests { + + use ark_std::test_rng; + + use crate::poly::dense_mlpoly::DensePolynomial; + + use super::*; + + #[test] + fn sm_msm_parity() { + use ark_curve25519::{EdwardsAffine as G1Affine, EdwardsProjective as G1Projective, Fr}; + let mut rng = test_rng(); + let bases = vec![ + G1Affine::rand(&mut rng), + G1Affine::rand(&mut rng), + G1Affine::rand(&mut rng), + ]; + let scalars = vec![Fr::from(3), Fr::from(2), Fr::from(1)]; + let expected_result = bases[0] + bases[0] + bases[0] + bases[1] + bases[1] + bases[2]; + assert_eq!(bases[0] + bases[0] + bases[0], bases[0] * scalars[0]); + let expected_result_b = + bases[0] * scalars[0] + bases[1] * scalars[1] + bases[2] * scalars[2]; + assert_eq!(expected_result, expected_result_b); + + let calc_result_a: G1Projective = VariableBaseMSM::msm(&bases, &scalars).unwrap(); + assert_eq!(calc_result_a, expected_result); + + let scalars_bigint: Vec<_> = scalars + .into_iter() + .map(|scalar| scalar.into_bigint()) + .collect(); + let calc_result_b: G1Projective = sm_msm(&scalars_bigint, &bases); + assert_eq!(calc_result_b, expected_result); + } +} diff --git a/jolt-core/src/poly/dense_mlpoly.rs b/jolt-core/src/poly/dense_mlpoly.rs index 751b1e89b..1bae6a4be 100644 --- a/jolt-core/src/poly/dense_mlpoly.rs +++ b/jolt-core/src/poly/dense_mlpoly.rs @@ -1,4 +1,5 @@ #![allow(clippy::too_many_arguments)] +use crate::msm::{flags_msm, sm_msm}; use crate::poly::eq_poly::EqPolynomial; use crate::utils::{self, compute_dotproduct, compute_dotproduct_low_optimized, mul_0_1_optimized}; @@ -219,10 +220,10 @@ impl DensePolynomial { let scalars = self.Z[R_size * i..R_size * (i + 1)].as_ref(); match hint { CommitHint::Normal => Commitments::batch_commit(scalars, &gens), - CommitHint::Flags => Self::flags_msm(scalars, &gens), + CommitHint::Flags => flags_msm(scalars, &gens), CommitHint::Small => { let bigints: Vec<_> = scalars.iter().map(|s| s.into_bigint()).collect(); - Self::sm_msm(&bigints, &gens) + sm_msm(&bigints, &gens) } } }) @@ -230,52 +231,6 @@ impl DensePolynomial { PolyCommitment { C } } - /// Special MSM where all scalar values are 0 / 1 – does not verify. - fn flags_msm(scalars: &[G::ScalarField], bases: &[G::Affine]) -> G { - assert_eq!(scalars.len(), bases.len()); - let result = scalars - .into_iter() - .enumerate() - .filter(|(_index, scalar)| !scalar.is_zero()) - .map(|(index, scalar)| bases[index]) - .sum(); - - result - } - - pub fn sm_msm( - scalars: &[::BigInt], - bases: &[V::MulBase], - ) -> V { - assert_eq!(scalars.len(), bases.len()); - let num_buckets: usize = 1 << 16; // TODO(sragss): This should be passed in / dependent on M = N^{1/C} - - // #[cfg(test)] - // scalars.for_each(|scalar| { - // assert!(scalar < V::ScalarField::from(num_buckets as u64).into_bigint()) - // }); - - // Assign things to buckets based on the scalar - let mut buckets: Vec = vec![V::zero(); num_buckets]; - scalars.into_iter().enumerate().for_each(|(index, scalar)| { - let bucket_index: u64 = scalar.as_ref()[0]; - buckets[bucket_index as usize] += bases[index]; - }); - - let mut result = V::zero(); - let mut running_sum = V::zero(); - buckets - .into_iter() - .skip(1) - .enumerate() - .rev() - .for_each(|(index, bucket)| { - running_sum += bucket; - result += running_sum; - }); - result - } - #[tracing::instrument(skip_all, name = "DensePolynomial.bound")] pub fn bound(&self, L: &[F]) -> Vec { let (left_num_vars, right_num_vars) = @@ -995,32 +950,6 @@ mod tests { ); } - #[test] - fn sm_msm_parity() { - use ark_curve25519::{EdwardsAffine as G1Affine, EdwardsProjective as G1Projective, Fr}; - let mut rng = test_rng(); - let bases = vec![ - G1Affine::rand(&mut rng), - G1Affine::rand(&mut rng), - G1Affine::rand(&mut rng), - ]; - let scalars = vec![Fr::from(3), Fr::from(2), Fr::from(1)]; - let expected_result = bases[0] + bases[0] + bases[0] + bases[1] + bases[1] + bases[2]; - assert_eq!(bases[0] + bases[0] + bases[0], bases[0] * scalars[0]); - let expected_result_b = - bases[0] * scalars[0] + bases[1] * scalars[1] + bases[2] * scalars[2]; - assert_eq!(expected_result, expected_result_b); - - let calc_result_a: G1Projective = VariableBaseMSM::msm(&bases, &scalars).unwrap(); - assert_eq!(calc_result_a, expected_result); - - let scalars_bigint: Vec<_> = scalars - .into_iter() - .map(|scalar| scalar.into_bigint()) - .collect(); - let calc_result_b: G1Projective = DensePolynomial::::sm_msm(&scalars_bigint, &bases); - assert_eq!(calc_result_b, expected_result); - } #[test] fn commit_with_hint_parity() { From 2307aa69adb36f48e0183d5ae929b69418ccf4f8 Mon Sep 17 00:00:00 2001 From: mmagician Date: Thu, 25 Jan 2024 13:09:24 +0100 Subject: [PATCH 5/5] use iterators instead of slices for looping over polys in `merge` --- jolt-core/src/jolt/vm/instruction_lookups.rs | 2 +- jolt-core/src/lasso/surge.rs | 2 +- jolt-core/src/poly/dense_mlpoly.rs | 13 +++++++------ 3 files changed, 9 insertions(+), 8 deletions(-) diff --git a/jolt-core/src/jolt/vm/instruction_lookups.rs b/jolt-core/src/jolt/vm/instruction_lookups.rs index 18aa67c71..ab5a3026f 100644 --- a/jolt-core/src/jolt/vm/instruction_lookups.rs +++ b/jolt-core/src/jolt/vm/instruction_lookups.rs @@ -122,7 +122,7 @@ where fn batch(&self) -> Self::BatchedPolynomials { use rayon::prelude::*; let (batched_dim_read, (batched_final, batched_E, batched_flag)) = rayon::join( - || DensePolynomial::merge(&[self.dim.as_slice(), self.read_cts.as_slice()].concat()), + || DensePolynomial::merge(self.dim.iter().chain(&self.read_cts)), || { let batched_final = DensePolynomial::merge(&self.final_cts); let (batched_E, batched_flag) = rayon::join( diff --git a/jolt-core/src/lasso/surge.rs b/jolt-core/src/lasso/surge.rs index e060ecb20..1769e833f 100644 --- a/jolt-core/src/lasso/surge.rs +++ b/jolt-core/src/lasso/surge.rs @@ -65,7 +65,7 @@ where #[tracing::instrument(skip_all, name = "SurgePolys::batch")] fn batch(&self) -> Self::BatchedPolynomials { let (batched_dim_read, (batched_final, batched_E)) = rayon::join( - || DensePolynomial::merge(&[self.dim.as_slice(), self.read_cts.as_slice()].concat()), + || DensePolynomial::merge(self.dim.iter().chain(&self.read_cts)), || { rayon::join( || DensePolynomial::merge(&self.final_cts), diff --git a/jolt-core/src/poly/dense_mlpoly.rs b/jolt-core/src/poly/dense_mlpoly.rs index 1bae6a4be..884c34e15 100644 --- a/jolt-core/src/poly/dense_mlpoly.rs +++ b/jolt-core/src/poly/dense_mlpoly.rs @@ -398,13 +398,14 @@ impl DensePolynomial { } #[tracing::instrument(skip_all, name = "DensePoly.merge")] - pub fn merge(polys: &[T]) -> DensePolynomial - where - T: AsRef>, - { - let total_len: usize = polys.iter().map(|poly| poly.as_ref().vec().len()).sum(); + pub fn merge(polys: impl IntoIterator> + Clone) -> DensePolynomial { + let polys_iter_cloned = polys.clone().into_iter(); + let total_len: usize = polys + .into_iter() + .map(|poly| poly.as_ref().vec().len()) + .sum(); let mut Z: Vec = Vec::with_capacity(total_len.next_power_of_two()); - for poly in polys { + for poly in polys_iter_cloned { Z.extend_from_slice(poly.as_ref().vec()); }