From 126b0b890b779c2398875e3194d6f8b0b642985a Mon Sep 17 00:00:00 2001
From: Daniel Sanchez <sanchez.quiros.daniel@gmail.com>
Date: Mon, 1 Jul 2024 10:54:48 +0200
Subject: [PATCH] Da: Missing parallel features (#669)

* Added missing ark parallel features

* Remove unnecessary copy

* Toeplitz3 in place

* Fix blob size in encoder bench
---
 nomos-da/kzgrs-backend/benches/encoder.rs | 24 +++++++++++------------
 nomos-da/kzgrs/Cargo.toml                 |  5 ++++-
 nomos-da/kzgrs/src/fk20.rs                | 12 +++++++-----
 3 files changed, 23 insertions(+), 18 deletions(-)

diff --git a/nomos-da/kzgrs-backend/benches/encoder.rs b/nomos-da/kzgrs-backend/benches/encoder.rs
index 3ae4d059..e74aa6be 100644
--- a/nomos-da/kzgrs-backend/benches/encoder.rs
+++ b/nomos-da/kzgrs-backend/benches/encoder.rs
@@ -1,7 +1,6 @@
 use divan::counter::BytesCount;
 use divan::Bencher;
 use kzgrs_backend::encoder::{DaEncoder, DaEncoderParams};
-use once_cell::sync::Lazy;
 use rand::RngCore;
 use std::hint::black_box;
 
@@ -9,22 +8,23 @@ fn main() {
     divan::main()
 }
 
-static ENCODER: Lazy<DaEncoder> = Lazy::new(|| {
-    let params = DaEncoderParams::new(4096, true);
-    DaEncoder::new(params)
-});
-
-const KB: usize = 1024;
+const MB1: usize = 1024 * 1024;
 
 pub fn rand_data(elements_count: usize) -> Vec<u8> {
     let mut buff = vec![0; elements_count * DaEncoderParams::MAX_BLS12_381_ENCODING_CHUNK_SIZE];
     rand::thread_rng().fill_bytes(&mut buff);
     buff
 }
-#[divan::bench(args = [32, 64, 128, 256, 512, 1024, 2048], sample_count = 1, sample_size = 1)]
-fn encode(bencher: Bencher, size: usize) {
+#[divan::bench(args = [128, 256, 512, 1024, 2048, 4096], sample_count = 1, sample_size = 1)]
+fn encode(bencher: Bencher, column_size: usize) {
     bencher
-        .with_inputs(|| rand_data(size * KB / DaEncoderParams::MAX_BLS12_381_ENCODING_CHUNK_SIZE))
-        .input_counter(|buff| BytesCount::new(buff.len()))
-        .bench_refs(|buff| black_box(ENCODER.encode(buff)));
+        .with_inputs(|| {
+            let params = DaEncoderParams::new(column_size, true);
+            (
+                DaEncoder::new(params),
+                rand_data(MB1 / DaEncoderParams::MAX_BLS12_381_ENCODING_CHUNK_SIZE),
+            )
+        })
+        .input_counter(|(_, buff)| BytesCount::new(buff.len()))
+        .bench_refs(|(encoder, buff)| black_box(encoder.encode(buff)));
 }
diff --git a/nomos-da/kzgrs/Cargo.toml b/nomos-da/kzgrs/Cargo.toml
index f4690e6b..8d192805 100644
--- a/nomos-da/kzgrs/Cargo.toml
+++ b/nomos-da/kzgrs/Cargo.toml
@@ -49,5 +49,8 @@ parallel = [
     "ark-poly/parallel",
     "ark-poly/rayon",
     "ark-poly-commit/rayon",
-    "ark-poly-commit/parallel"
+    "ark-poly-commit/parallel",
+    "ark-ec/parallel",
+    "ark-ec/rayon",
+    "ark-bls12-381-ext/parallel"
 ]
diff --git a/nomos-da/kzgrs/src/fk20.rs b/nomos-da/kzgrs/src/fk20.rs
index 270c10f2..b7a2a28d 100644
--- a/nomos-da/kzgrs/src/fk20.rs
+++ b/nomos-da/kzgrs/src/fk20.rs
@@ -5,6 +5,7 @@ use ark_ff::Field;
 use ark_poly::{EvaluationDomain, GeneralEvaluationDomain};
 use num_traits::Zero;
 use std::borrow::Cow;
+use std::ops::Mul;
 
 fn toeplitz1(global_parameters: &[G1Affine], polynomial_degree: usize) -> Vec<G1Projective> {
     debug_assert_eq!(global_parameters.len(), polynomial_degree);
@@ -27,16 +28,17 @@ fn toeplitz2(coefficients: &[Fr], extended_vector: &[G1Projective]) -> Vec<G1Pro
     let toeplitz_coefficients_fft = domain.fft(coefficients);
     extended_vector
         .iter()
-        .copied()
         .zip(toeplitz_coefficients_fft)
-        .map(|(v, c)| (v * c))
+        .map(|(v, c)| (v.mul(c)))
         .collect()
 }
 
-fn toeplitz3(h_extended_fft: &[G1Projective]) -> Vec<G1Projective> {
+fn toeplitz3(mut h_extended_fft: Vec<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)
+
+    domain.ifft_in_place(&mut h_extended_fft);
+    h_extended_fft
 }
 
 pub fn fk20_batch_generate_elements_proofs(
@@ -67,7 +69,7 @@ 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);
+    let h_vector = toeplitz3(h_extended_vector);
     domain
         .fft(&h_vector)
         .into_iter()