Use threadcount for parallelization

nomos-da/kzgrs/Cargo.toml

@@ -34,6 +34,10 @@ harness = false
 name = "fft"
 harness = false
 
+[[bench]]
+name = "fk20"
+harness = false
+
 [features]
 default = ["single"]
 single = []

nomos-da/kzgrs/benches/fft.rs

@@ -1,52 +1,23 @@
-use ark_bls12_381::{Fr, G1Affine};
-use ark_ec::{AffineRepr, CurveGroup};
-use ark_ff::{BigInt, FftField, Field};
+use ark_bls12_381::{Fr, G1Affine, G1Projective};
+use ark_ec::AffineRepr;
+use ark_ff::BigInt;
+use ark_poly::{EvaluationDomain, GeneralEvaluationDomain};
 use divan::counter::ItemsCount;
-use divan::{black_box, counter::BytesCount, AllocProfiler, Bencher};
-use kzgrs::fft::{fft_g1, ifft_g1};
+use divan::{black_box, Bencher};
 fn main() {
     divan::main()
 }
 
 #[divan::bench(args = [16, 32, 64, 128, 256, 512, 1024, 2048, 4096])]
-fn compute_fft_for_size(bencher: Bencher, size: usize) {
+fn compute_ark_fft_for_size(bencher: Bencher, size: usize) {
     bencher
         .with_inputs(|| {
-            let primitive_root = <Fr as FftField>::get_root_of_unity(size as u64).unwrap();
-            let roots_of_unity: Vec<_> = (1..=size)
-                .map(|i| primitive_root.pow::<ark_ff::BigInt<4>>(BigInt::from(i as u64)))
+            let domain = GeneralEvaluationDomain::<Fr>::new(size).unwrap();
+            let buff: Vec<G1Projective> = (0..size)
+                .map(|i| G1Affine::identity().mul_bigint(BigInt::<4>::from(i as u64)))
                 .collect();
-            let buff: Vec<G1Affine> = (0..size)
-                .map(|i| {
-                    G1Affine::identity()
-                        .mul_bigint(BigInt::<4>::from(i as u64))
-                        .into_affine()
-                })
-                .collect();
-            (buff, roots_of_unity)
+            (buff, domain)
         })
         .input_counter(move |_| ItemsCount::new(size))
-        .bench_refs(|(buff, roots_of_unity)| black_box(fft_g1(buff, roots_of_unity)));
-}
-
-#[divan::bench(args = [16, 32, 64, 128, 256, 512, 1024, 2048, 4096])]
-fn compute_ifft_for_size(bencher: Bencher, size: usize) {
-    bencher
-        .with_inputs(|| {
-            let primitive_root = <Fr as FftField>::get_root_of_unity(size as u64).unwrap();
-            let roots_of_unity: Vec<_> = (1..=size)
-                .map(|i| primitive_root.pow::<ark_ff::BigInt<4>>(BigInt::from(i as u64)))
-                .collect();
-            let buff: Vec<G1Affine> = (0..size)
-                .map(|i| {
-                    G1Affine::identity()
-                        .mul_bigint(BigInt::<4>::from(i as u64))
-                        .into_affine()
-                })
-                .collect();
-            let buff = fft_g1(&buff, &roots_of_unity);
-            (buff, roots_of_unity)
-        })
-        .input_counter(move |_| ItemsCount::new(size))
-        .bench_refs(|(buff, roots_of_unity)| black_box(ifft_g1(buff, roots_of_unity)));
+        .bench_refs(|(buff, domain)| black_box(domain.fft(buff)));
 }

nomos-da/kzgrs/benches/fk20.rs

@@ -0,0 +1,67 @@
+use ark_bls12_381::{Bls12_381, Fr, G1Affine, G1Projective};
+use ark_ec::AffineRepr;
+use ark_ff::BigInt;
+use ark_poly::univariate::DensePolynomial;
+use ark_poly::{EvaluationDomain, GeneralEvaluationDomain};
+use ark_poly_commit::kzg10::KZG10;
+use divan::counter::ItemsCount;
+use divan::Bencher;
+use kzgrs::fk20::fk20_batch_generate_elements_proofs;
+use kzgrs::{bytes_to_polynomial, GlobalParameters, BYTES_PER_FIELD_ELEMENT};
+use once_cell::sync::Lazy;
+use rand::SeedableRng;
+#[cfg(feature = "parallel")]
+use rayon::iter::{IntoParallelIterator, ParallelIterator};
+use std::hint::black_box;
+
+fn main() {
+    divan::main()
+}
+
+static GLOBAL_PARAMETERS: Lazy<GlobalParameters> = Lazy::new(|| {
+    let mut rng = rand::rngs::StdRng::seed_from_u64(1987);
+    KZG10::<Bls12_381, DensePolynomial<Fr>>::setup(4096, true, &mut rng).unwrap()
+});
+
+#[divan::bench(args = [16, 32, 64, 128, 256, 512, 1024, 2048, 4096])]
+fn compute_fk20_proofs_for_size(bencher: Bencher, size: usize) {
+    bencher
+        .with_inputs(|| {
+            let buff: Vec<_> = (0..BYTES_PER_FIELD_ELEMENT * size)
+                .map(|i| (i % 255) as u8)
+                .rev()
+                .collect();
+            let domain = GeneralEvaluationDomain::new(size).unwrap();
+            let (_, poly) = bytes_to_polynomial::<BYTES_PER_FIELD_ELEMENT>(&buff, domain).unwrap();
+            poly
+        })
+        .input_counter(move |_| ItemsCount::new(size))
+        .bench_refs(|(poly)| {
+            black_box(fk20_batch_generate_elements_proofs(
+                poly,
+                &GLOBAL_PARAMETERS,
+            ))
+        });
+}
+
+#[cfg(feature = "parallel")]
+#[divan::bench(args = [16, 32, 64, 128, 256, 512, 1024, 2048, 4096])]
+fn compute_parallel_fk20_proofs_for_size(bencher: Bencher, size: usize) {
+    let thread_count: usize = rayon::max_num_threads().min(rayon::current_num_threads());
+    bencher
+        .with_inputs(|| {
+            let buff: Vec<_> = (0..BYTES_PER_FIELD_ELEMENT * size)
+                .map(|i| (i % 255) as u8)
+                .rev()
+                .collect();
+            let domain = GeneralEvaluationDomain::new(size).unwrap();
+            let (_, poly) = bytes_to_polynomial::<BYTES_PER_FIELD_ELEMENT>(&buff, domain).unwrap();
+            poly
+        })
+        .input_counter(move |_| ItemsCount::new(size * thread_count))
+        .bench_refs(|(poly)| {
+            black_box((0..thread_count).into_par_iter().for_each(|_| {
+                fk20_batch_generate_elements_proofs(poly, &GLOBAL_PARAMETERS);
+            }))
+        });
+}

nomos-da/kzgrs/benches/kzg.rs

@@ -3,7 +3,7 @@ use ark_poly::univariate::DensePolynomial;
 use ark_poly::{EvaluationDomain, GeneralEvaluationDomain};
 use ark_poly_commit::kzg10::{UniversalParams, KZG10};
 use divan::counter::ItemsCount;
-use divan::{black_box, counter::BytesCount, AllocProfiler, Bencher};
+use divan::{black_box, Bencher};
 use once_cell::sync::Lazy;
 use rand::RngCore;
 #[cfg(feature = "parallel")]
@@ -61,7 +61,7 @@ fn commit_polynomial_with_element_count_parallelized(bencher: Bencher, element_c
         })
         .input_counter(move |(_evals, _poly)| ItemsCount::new(threads))
         .bench_refs(|(_evals, poly)| {
-            let commitments: Vec<_> = (0..threads)
+            let _commitments: Vec<_> = (0..threads)
                 .into_par_iter()
                 .map(|_| commit_polynomial(poly, &GLOBAL_PARAMETERS))
                 .collect();

nomos-da/kzgrs/src/common.rs

@@ -4,11 +4,11 @@ use std::fmt;
 // crates
 use crate::{FieldElement, BYTES_PER_FIELD_ELEMENT};
 use ark_bls12_381::fr::Fr;
-use ark_ff::{BigInt, FftField, Field, Zero};
+use ark_ff::Zero;
 use ark_poly::domain::general::GeneralEvaluationDomain;
 use ark_poly::evaluations::univariate::Evaluations;
 use ark_poly::univariate::DensePolynomial;
-use ark_poly::{EvaluationDomain, Radix2EvaluationDomain};
+use ark_poly::EvaluationDomain;
 use blst::BLST_ERROR;
 use num_bigint::BigUint;
 use thiserror::Error;

nomos-da/kzgrs/src/fft.rs

@@ -1,216 +0,0 @@
-use ark_bls12_381::{Bls12_381, Fr, G1Affine};
-use ark_ec::pairing::Pairing;
-use ark_ec::{AffineRepr, CurveGroup};
-use ark_ff::{BigInt, BigInteger, FftField, Field, PrimeField};
-#[cfg(feature = "parallel")]
-use rayon::iter::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator};
-use std::ops::Neg;
-
-pub fn fft_g1(vals: &[G1Affine], roots_of_unity: &[Fr]) -> Vec<G1Affine> {
-    debug_assert_eq!(vals.len(), roots_of_unity.len());
-    let original_len = vals.len();
-    if original_len == 1 {
-        return vals.to_vec();
-    }
-    let half_roots: Vec<_> = roots_of_unity.iter().step_by(2).copied().collect();
-
-    let l = || {
-        fft_g1(
-            vals.iter()
-                .step_by(2)
-                .copied()
-                .collect::<Vec<_>>()
-                .as_slice(),
-            half_roots.as_slice(),
-        )
-    };
-
-    let r = || {
-        fft_g1(
-            vals.iter()
-                .skip(1)
-                .step_by(2)
-                .copied()
-                .collect::<Vec<_>>()
-                .as_slice(),
-            half_roots.as_slice(),
-        )
-    };
-
-    let [l, r]: [Vec<G1Affine>; 2] = {
-        #[cfg(feature = "parallel")]
-        {
-            let (l, r) = rayon::join(l, r);
-            [l, r]
-        }
-        #[cfg(not(feature = "parallel"))]
-        {
-            [l(), r()]
-        }
-    };
-    // Double sized so we can use iterator later on
-    let l = l.into_iter().cycle().take(original_len);
-    let r = r.into_iter().cycle().take(original_len);
-
-    let y_times_root = {
-        #[cfg(feature = "parallel")]
-        {
-            r.into_par_iter()
-        }
-        #[cfg(not(feature = "parallel"))]
-        {
-            r.into_iter()
-        }
-    }
-    .enumerate()
-    .map(|(i, y)| (y * roots_of_unity[i % vals.len()]));
-
-    {
-        #[cfg(feature = "parallel")]
-        {
-            l.into_par_iter()
-        }
-        #[cfg(not(feature = "parallel"))]
-        {
-            l.into_iter()
-        }
-    }
-    .zip(y_times_root)
-    .enumerate()
-    .map(|(i, (x, y_times_root))| {
-        if i < vals.len() / 2 {
-            x + y_times_root
-        } else {
-            x + y_times_root.neg()
-        }
-        .into_affine()
-    })
-    .collect()
-}
-
-pub fn ifft_g1(vals: &[G1Affine], roots_of_unity: &[Fr]) -> Vec<G1Affine> {
-    debug_assert_eq!(vals.len(), roots_of_unity.len());
-    let mut mod_min_2 = BigInt::new(<Fr as PrimeField>::MODULUS.0);
-    mod_min_2.sub_with_borrow(&BigInt::<4>::from(2u64));
-    let invlen = Fr::from(vals.len() as u64).pow(mod_min_2).into_bigint();
-    {
-        #[cfg(feature = "parallel")]
-        {
-            fft_g1(vals, roots_of_unity).into_par_iter()
-        }
-        #[cfg(not(feature = "parallel"))]
-        {
-            fft_g1(vals, roots_of_unity).into_iter()
-        }
-    }
-    .map(|g| g.mul_bigint(invlen).into_affine())
-    .collect()
-}
-
-pub fn fft_fr(vals: &[Fr], roots_of_unity: &[Fr]) -> Vec<Fr> {
-    debug_assert_eq!(vals.len(), roots_of_unity.len());
-    let original_len = vals.len();
-    if original_len == 1 {
-        return vals.to_vec();
-    }
-    let half_roots: Vec<_> = roots_of_unity.iter().step_by(2).copied().collect();
-
-    let l = || {
-        crate::fft::fft_fr(
-            vals.iter()
-                .step_by(2)
-                .copied()
-                .collect::<Vec<_>>()
-                .as_slice(),
-            half_roots.as_slice(),
-        )
-    };
-
-    let r = || {
-        crate::fft::fft_fr(
-            vals.iter()
-                .skip(1)
-                .step_by(2)
-                .copied()
-                .collect::<Vec<_>>()
-                .as_slice(),
-            half_roots.as_slice(),
-        )
-    };
-
-    let [l, r]: [Vec<Fr>; 2] = {
-        #[cfg(feature = "parallel")]
-        {
-            let (l, r) = rayon::join(l, r);
-            [l, r]
-        }
-        #[cfg(not(feature = "parallel"))]
-        {
-            [l(), r()]
-        }
-    };
-    // Double sized so we can use iterator later on
-    let l: Vec<_> = l.into_iter().cycle().take(original_len).collect();
-    let r: Vec<_> = r.into_iter().cycle().take(original_len).collect();
-
-    let y_times_root = {
-        #[cfg(feature = "parallel")]
-        {
-            r.into_par_iter()
-        }
-        #[cfg(not(feature = "parallel"))]
-        {
-            r.into_iter()
-        }
-    }
-    .enumerate()
-    .map(|(i, y)| y * roots_of_unity[i % vals.len()]);
-
-    {
-        #[cfg(feature = "parallel")]
-        {
-            l.into_par_iter()
-        }
-        #[cfg(not(feature = "parallel"))]
-        {
-            l.into_iter()
-        }
-    }
-    .zip(y_times_root)
-    .enumerate()
-    .map(|(i, (x, y_times_root))| {
-        if i < vals.len() / 2 {
-            x + y_times_root
-        } else {
-            x - y_times_root
-        }
-    })
-    .collect()
-}
-
-#[cfg(test)]
-mod test {
-    use crate::common::compute_roots_of_unity;
-    use crate::fft::{fft_g1, ifft_g1};
-    use ark_bls12_381::{Fr, G1Affine};
-    use ark_ec::{AffineRepr, CurveGroup};
-    use ark_ff::{BigInt, FftField, Field};
-
-    #[test]
-    fn test_fft_ifft_g1() {
-        for size in [16usize, 32, 64, 128, 256, 512, 1024, 2048, 4096] {
-            let roots_of_unity = compute_roots_of_unity(size);
-            let r: Vec<_> = roots_of_unity.iter().map(|a| a.to_string()).collect();
-            let buff: Vec<G1Affine> = (0..size)
-                .map(|i| {
-                    G1Affine::identity()
-                        .mul_bigint(BigInt::<4>::from(i as u64))
-                        .into_affine()
-                })
-                .collect();
-            let fft = fft_g1(&buff, &roots_of_unity);
-            let ifft = ifft_g1(&fft, &roots_of_unity);
-            assert_eq!(buff, ifft);
-        }
-    }
-}

nomos-da/kzgrs/src/fk20.rs

@@ -1,23 +1,19 @@
-use ark_bls12_381::{Fr, G1Affine, G1Projective};
-use ark_ec::{AffineRepr, CurveGroup};
-use ark_ff::{FftField, Field};
-use ark_poly::{EvaluationDomain, GeneralEvaluationDomain, Polynomial as _};
-use num_traits::Zero;
-
-//use crate::common::compute_roots_of_unity;
-use crate::fft::{fft_fr, fft_g1, ifft_g1};
 use crate::{GlobalParameters, Polynomial, Proof};
-
-
+use ark_bls12_381::{Fr, G1Affine, G1Projective};
+use ark_ec::CurveGroup;
+use ark_ff::Field;
+use ark_poly::{EvaluationDomain, GeneralEvaluationDomain};
+use num_traits::Zero;
 
 fn toeplitz1(global_parameters: &[G1Affine], polynomial_degree: usize) -> Vec<G1Projective> {
     debug_assert_eq!(global_parameters.len(), polynomial_degree);
     debug_assert!(polynomial_degree.is_power_of_two());
-    let domain: GeneralEvaluationDomain<Fr> =
-        GeneralEvaluationDomain::new(polynomial_degree*2).expect("Domain should be able to build");
+    let domain: GeneralEvaluationDomain<Fr> = GeneralEvaluationDomain::new(polynomial_degree * 2)
+        .expect("Domain should be able to build");
     let vector_extended: Vec<G1Projective> = global_parameters
         .iter()
-        .copied().map(|pi| G1Projective::from(pi))
+        .copied()
+        .map(G1Projective::from)
         .chain(std::iter::repeat_with(G1Projective::zero).take(polynomial_degree))
         .collect();
     domain.fft(&vector_extended)
@@ -27,9 +23,7 @@ fn toeplitz2(coefficients: &[Fr], extended_vector: &[G1Projective]) -> Vec<G1Pro
     debug_assert!(coefficients.len().is_power_of_two());
     let domain: GeneralEvaluationDomain<Fr> =
         GeneralEvaluationDomain::new(coefficients.len()).expect("Domain should be able to build");
-    // let toeplitz_coefficients_fft = domain.fft(coefficients);
-    let roots_of_unity = domain.elements().take(coefficients.len()).collect();
-    let toeplitz_coefficients_fft = fft_fr(coefficients, &roots_of_unity);
+    let toeplitz_coefficients_fft = domain.fft(coefficients);
     extended_vector
         .iter()
         .copied()
@@ -38,7 +32,7 @@ fn toeplitz2(coefficients: &[Fr], extended_vector: &[G1Projective]) -> Vec<G1Pro
         .collect()
 }
 
-fn toeplitz3(h_extended_fft: &[G1Projective], polynomial_degree: usize) -> Vec<G1Projective> {
+fn toeplitz3(h_extended_fft: &[G1Projective]) -> Vec<G1Projective> {
     let domain: GeneralEvaluationDomain<Fr> =
         GeneralEvaluationDomain::new(h_extended_fft.len()).expect("Domain should be able to build");
     domain.ifft(h_extended_fft)
@@ -67,8 +61,9 @@ pub fn fk20_batch_generate_elements_proofs(
         .chain(polynomial.coeffs.iter().copied())
         .collect();
     let h_extended_vector = toeplitz2(&toeplitz_coefficients, &extended_vector);
-    let h_vector = toeplitz3(&h_extended_vector, polynomial_degree);
-    domain.fft(&h_vector)
+    let h_vector = toeplitz3(&h_extended_vector);
+    domain
+        .fft(&h_vector)
         .into_iter()
         .map(|g1| Proof {
             w: g1.into_affine(),
@@ -79,14 +74,12 @@ pub fn fk20_batch_generate_elements_proofs(
 
 #[cfg(test)]
 mod test {
-    use crate::common::compute_roots_of_unity;
-    use crate::fk20::{fk20_batch_generate_elements_proofs};
+    use crate::fk20::fk20_batch_generate_elements_proofs;
     use crate::{
         common::bytes_to_polynomial, kzg::generate_element_proof, GlobalParameters, Proof,
         BYTES_PER_FIELD_ELEMENT,
     };
     use ark_bls12_381::{Bls12_381, Fr};
-    use ark_ff::FftField;
     use ark_poly::univariate::DensePolynomial;
     use ark_poly::{EvaluationDomain, GeneralEvaluationDomain};
     use ark_poly_commit::kzg10::KZG10;
@@ -101,7 +94,7 @@ mod test {
     #[test]
     fn test_generate_proofs() {
         for size in [16, 32, 64, 128, 256] {
-            let mut buff: Vec<_> = (0..BYTES_PER_FIELD_ELEMENT * size)
+            let buff: Vec<_> = (0..BYTES_PER_FIELD_ELEMENT * size)
                 .map(|i| (i % 255) as u8)
                 .rev()
                 .collect();
@@ -109,7 +102,6 @@ mod test {
             let (evals, poly) =
                 bytes_to_polynomial::<BYTES_PER_FIELD_ELEMENT>(&buff, domain).unwrap();
             let polynomial_degree = poly.len();
-            //let roots_of_unity: Vec<Fr> = compute_roots_of_unity(size);
             let slow_proofs: Vec<Proof> = (0..polynomial_degree)
                 .map(|i| {
                     generate_element_proof(i, &poly, &evals, &GLOBAL_PARAMETERS, domain).unwrap()

nomos-da/kzgrs/src/lib.rs

@@ -1,5 +1,4 @@
 pub mod common;
-pub mod fft;
 pub mod fk20;
 pub mod global_parameters;
 pub mod kzg;