Feat(poly): Add Dense and Sparse Multilinear Polynomials (#726)

* directory

* add ops and rest of multilinear functions

* fmt

* finish test and add error infra

* fmt

* convert asserts to errors

* fix cmt

* add sparse multlinear

* no_std + fmt

* cmts + impl error

* add benches + fmt

* ci

* ci

* ci again

* ci

* change SubAssign to Sub

* remove SubAssign

* fix errors add docs

* switch rayon -> parallel

* add docs to get_bits

---------

Co-authored-by: Mariano A. Nicolini <mariano.nicolini.91@gmail.com>

crypto/src/hash/poseidon/mod.rs

@@ -65,7 +65,7 @@ impl<P: PermutationParameters> Poseidon for P {
             let mut new_e = FE::zero();
             for (j, current_state) in state.iter().enumerate() {
                 let mut mij = P::MDS_MATRIX[i * P::N_MDS_MATRIX_COLS + j].clone();
-                mij = mij * current_state;
+                mij *= current_state;
                 new_e += mij;
             }
             new_state.push(new_e);

math/Cargo.toml

@@ -24,6 +24,7 @@ objc = { version = "0.2.7", optional = true }
 # cuda
 cudarc = { version = "0.9.7", optional = true }
 
+
 lambdaworks-gpu = { workspace = true, optional = true }
 
 [dev-dependencies]

math/benches/criterion_polynomial.rs

@@ -1,83 +1,14 @@
-use const_random::const_random;
-use core::hint::black_box;
+mod polynomials;
 use criterion::{criterion_group, criterion_main, Criterion};
-use lambdaworks_math::polynomial::Polynomial;
-use utils::u64_utils::{rand_field_elements, rand_poly, FE};
-
-mod utils;
-
-pub fn polynomial_benchmarks(c: &mut Criterion) {
-    let mut group = c.benchmark_group("Polynomial");
-    let order = const_random!(u64) % 8;
-
-    group.bench_function("evaluate", |bench| {
-        let poly = rand_poly(order);
-        let x = FE::new(rand::random::<u64>());
-        bench.iter(|| poly.evaluate(black_box(&x)));
-    });
-
-    group.bench_function("evaluate_slice", |bench| {
-        let poly = rand_poly(order);
-        let inputs = rand_field_elements(order);
-        bench.iter(|| poly.evaluate_slice(black_box(&inputs)));
-    });
-
-    group.bench_function("add", |bench| {
-        let x_poly = rand_poly(order);
-        let y_poly = rand_poly(order);
-        bench.iter(|| black_box(&x_poly) + black_box(&y_poly));
-    });
-
-    group.bench_function("neg", |bench| {
-        let x_poly = rand_poly(order);
-        bench.iter(|| -black_box(&x_poly));
-    });
-
-    group.bench_function("sub", |bench| {
-        let x_poly = rand_poly(order);
-        let y_poly = rand_poly(order);
-        bench.iter(|| black_box(&x_poly) - black_box(&y_poly));
-    });
-
-    group.bench_function("mul", |bench| {
-        let x_poly = rand_poly(order);
-        let y_poly = rand_poly(order);
-        bench.iter(|| black_box(&x_poly) * black_box(&y_poly));
-    });
-
-    group.bench_function("div", |bench| {
-        let x_poly = rand_poly(order);
-        let y_poly = rand_poly(order);
-        bench.iter_batched(
-            || (x_poly.clone(), y_poly.clone()),
-            |(x_poly, y_poly)| black_box(x_poly) / black_box(y_poly),
-            criterion::BatchSize::SmallInput,
-        );
-    });
-
-    group.bench_function("div by 'x - b' with generic div", |bench| {
-        let poly = rand_poly(order);
-        let m = Polynomial::new_monomial(FE::one(), 1) - rand_field_elements(1)[0];
-        bench.iter_batched(
-            || (poly.clone(), m.clone()),
-            |(poly, m)| poly / m,
-            criterion::BatchSize::SmallInput,
-        );
-    });
-
-    group.bench_function("div by 'x - b' with Ruffini", |bench| {
-        let poly = rand_poly(order);
-        let b = rand_field_elements(1)[0];
-        bench.iter_batched(
-            || (poly.clone(), b),
-            |(mut poly, b)| {
-                poly.ruffini_division_inplace(&b);
-                poly
-            },
-            criterion::BatchSize::SmallInput,
-        );
-    });
-}
-
-criterion_group!(polynomial, polynomial_benchmarks);
+use polynomials::{
+    dense_multilinear_poly::dense_multilinear_polynomial_benchmarks,
+    polynomial::polynomial_benchmarks,
+    sparse_multilinear_poly::sparse_multilinear_polynomial_benchmarks,
+};
+use pprof::criterion::{Output, PProfProfiler};
+
+criterion_group!(
+    name = polynomial;
+    config = Criterion::default().with_profiler(PProfProfiler::new(100, Output::Flamegraph(None)));
+    targets = polynomial_benchmarks, dense_multilinear_polynomial_benchmarks, sparse_multilinear_polynomial_benchmarks);
 criterion_main!(polynomial);

math/benches/polynomials/dense_multilinear_poly.rs

@@ -0,0 +1,43 @@
+use super::utils::{rand_dense_multilinear_poly, rand_field_elements, FE};
+use const_random::const_random;
+use core::hint::black_box;
+use criterion::Criterion;
+use lambdaworks_math::polynomial::dense_multilinear_poly::DenseMultilinearPolynomial;
+
+pub fn dense_multilinear_polynomial_benchmarks(c: &mut Criterion) {
+    let mut group = c.benchmark_group("Polynomial");
+    let order = const_random!(u64) % 8;
+
+    group.bench_function("evaluate", |bench| {
+        let poly = rand_dense_multilinear_poly(order);
+        let r = rand_field_elements(order);
+        bench.iter(|| poly.evaluate(black_box(r.clone())));
+    });
+
+    group.bench_function("evaluate_with", |bench| {
+        let evals = rand_field_elements(order);
+        let r = rand_field_elements(order);
+
+        bench.iter(|| DenseMultilinearPolynomial::evaluate_with(black_box(&evals), black_box(&r)));
+    });
+
+    group.bench_function("merge", |bench| {
+        let x_poly = rand_dense_multilinear_poly(order);
+        let y_poly = rand_dense_multilinear_poly(order);
+        bench.iter(|| {
+            DenseMultilinearPolynomial::merge(black_box(&[x_poly.clone(), y_poly.clone()]))
+        });
+    });
+
+    group.bench_function("add", |bench| {
+        let x_poly = rand_dense_multilinear_poly(order);
+        let y_poly = rand_dense_multilinear_poly(order);
+        bench.iter(|| black_box(black_box(x_poly.clone()) + black_box(y_poly.clone())));
+    });
+
+    group.bench_function("mul", |bench| {
+        let x_poly = rand_dense_multilinear_poly(order);
+        let y = FE::new(rand::random::<u64>());
+        bench.iter(|| black_box(black_box(x_poly.clone()) * black_box(y)));
+    });
+}

math/benches/polynomials/mod.rs

@@ -0,0 +1,4 @@
+pub mod dense_multilinear_poly;
+pub mod polynomial;
+pub mod sparse_multilinear_poly;
+pub mod utils;

math/benches/polynomials/polynomial.rs

@@ -0,0 +1,78 @@
+use super::utils::{rand_field_elements, rand_poly, FE};
+use const_random::const_random;
+use core::hint::black_box;
+use criterion::Criterion;
+use lambdaworks_math::polynomial::Polynomial;
+
+pub fn polynomial_benchmarks(c: &mut Criterion) {
+    let mut group = c.benchmark_group("Polynomial");
+    let order = const_random!(u64) % 8;
+
+    group.bench_function("evaluate", |bench| {
+        let poly = rand_poly(order);
+        let x = FE::new(rand::random::<u64>());
+        bench.iter(|| poly.evaluate(black_box(&x)));
+    });
+
+    group.bench_function("evaluate_slice", |bench| {
+        let poly = rand_poly(order);
+        let inputs = rand_field_elements(order);
+        bench.iter(|| poly.evaluate_slice(black_box(&inputs)));
+    });
+
+    group.bench_function("add", |bench| {
+        let x_poly = rand_poly(order);
+        let y_poly = rand_poly(order);
+        bench.iter(|| black_box(&x_poly) + black_box(&y_poly));
+    });
+
+    group.bench_function("neg", |bench| {
+        let x_poly = rand_poly(order);
+        bench.iter(|| -black_box(&x_poly));
+    });
+
+    group.bench_function("sub", |bench| {
+        let x_poly = rand_poly(order);
+        let y_poly = rand_poly(order);
+        bench.iter(|| black_box(&x_poly) - black_box(&y_poly));
+    });
+
+    group.bench_function("mul", |bench| {
+        let x_poly = rand_poly(order);
+        let y_poly = rand_poly(order);
+        bench.iter(|| black_box(&x_poly) * black_box(&y_poly));
+    });
+
+    group.bench_function("div", |bench| {
+        let x_poly = rand_poly(order);
+        let y_poly = rand_poly(order);
+        bench.iter_batched(
+            || (x_poly.clone(), y_poly.clone()),
+            |(x_poly, y_poly)| black_box(x_poly) / black_box(y_poly),
+            criterion::BatchSize::SmallInput,
+        );
+    });
+
+    group.bench_function("div by 'x - b' with generic div", |bench| {
+        let poly = rand_poly(order);
+        let m = Polynomial::new_monomial(FE::one(), 1) - rand_field_elements(1)[0];
+        bench.iter_batched(
+            || (poly.clone(), m.clone()),
+            |(poly, m)| poly / m,
+            criterion::BatchSize::SmallInput,
+        );
+    });
+
+    group.bench_function("div by 'x - b' with Ruffini", |bench| {
+        let poly = rand_poly(order);
+        let b = rand_field_elements(1)[0];
+        bench.iter_batched(
+            || (poly.clone(), b),
+            |(mut poly, b)| {
+                poly.ruffini_division_inplace(&b);
+                poly
+            },
+            criterion::BatchSize::SmallInput,
+        );
+    });
+}

math/benches/polynomials/sparse_multilinear_poly.rs

@@ -0,0 +1,45 @@
+use super::utils::{rand_field_elements, rand_sparse_multilinear_poly};
+use const_random::const_random;
+use core::hint::black_box;
+use criterion::Criterion;
+use lambdaworks_math::polynomial::sparse_multilinear_poly::SparseMultilinearPolynomial;
+use rand::random;
+
+pub fn sparse_multilinear_polynomial_benchmarks(c: &mut Criterion) {
+    let mut group = c.benchmark_group("Polynomial");
+    let order = const_random!(u64) % 8;
+    let num_vars = [3, 4, 5, 6, 7, 8, 9, 10];
+
+    for num_var in num_vars.iter() {
+        group.bench_with_input(
+            format!("evaluate {:?}", &num_var),
+            num_var,
+            |bench, num_var| {
+                let poly = rand_sparse_multilinear_poly(*num_var, order);
+                let r = rand_field_elements(order);
+                bench.iter(|| poly.evaluate(black_box(&r)));
+            },
+        );
+    }
+
+    for num_var in num_vars.iter() {
+        group.bench_with_input(
+            format!("evaluate_with {:?}", &num_var),
+            num_var,
+            |bench, num_var| {
+                let evals = rand_field_elements(order)
+                    .into_iter()
+                    .map(|eval| (random(), eval))
+                    .collect::<Vec<_>>();
+                let r = rand_field_elements(order);
+                bench.iter(|| {
+                    SparseMultilinearPolynomial::evaluate_with(
+                        black_box(*num_var),
+                        black_box(&evals),
+                        black_box(&r),
+                    )
+                });
+            },
+        );
+    }
+}

math/benches/polynomials/utils.rs

@@ -0,0 +1,61 @@
+use const_random::const_random;
+use lambdaworks_math::{
+    field::fields::u64_prime_field::{U64FieldElement, U64PrimeField},
+    polynomial::{
+        dense_multilinear_poly::DenseMultilinearPolynomial,
+        sparse_multilinear_poly::SparseMultilinearPolynomial, Polynomial,
+    },
+};
+use rand::random;
+
+// Mersenne prime numbers
+// https://www.math.utah.edu/~pa/math/mersenne.html
+const PRIMES: [u64; 39] = [
+    13, 17, 19, 31, 61, 89, 107, 127, 521, 607, 1279, 2203, 2281, 3217, 4253, 4423, 9689, 9941,
+    11213, 19937, 21701, 23209, 44497, 86243, 110503, 132049, 216091, 756839, 859433, 1257787,
+    1398269, 2976221, 3021377, 6972593, 13466917, 20996011, 24036583, 25964951, 30402457,
+];
+
+const MODULUS: u64 = PRIMES[const_random!(usize) % PRIMES.len()];
+pub type FE = U64FieldElement<MODULUS>;
+
+#[allow(dead_code)]
+#[inline(never)]
+#[export_name = "u64_utils::rand_field_elements"]
+pub fn rand_field_elements(order: u64) -> Vec<FE> {
+    let mut result = Vec::with_capacity(1 << order);
+    for _ in 0..result.capacity() {
+        result.push(FE::new(random()));
+    }
+    result
+}
+
+#[allow(dead_code)]
+#[inline(never)]
+#[export_name = "u64_utils::rand_poly"]
+pub fn rand_poly(order: u64) -> Polynomial<FE> {
+    Polynomial::new(&rand_field_elements(order))
+}
+
+#[allow(dead_code)]
+#[inline(never)]
+#[export_name = "u64_utils::rand_dense_multilinear_poly"]
+pub fn rand_dense_multilinear_poly(
+    order: u64,
+) -> DenseMultilinearPolynomial<U64PrimeField<MODULUS>> {
+    DenseMultilinearPolynomial::new(rand_field_elements(order))
+}
+
+#[allow(dead_code)]
+#[inline(never)]
+#[export_name = "u64_utils::rand_sparse_multilinear_poly"]
+pub fn rand_sparse_multilinear_poly(
+    num_vars: usize,
+    order: u64,
+) -> SparseMultilinearPolynomial<U64PrimeField<MODULUS>> {
+    let evals = rand_field_elements(order)
+        .into_iter()
+        .map(|eval| (random(), eval))
+        .collect::<Vec<_>>();
+    SparseMultilinearPolynomial::new(num_vars, evals)
+}

math/src/elliptic_curve/short_weierstrass/curves/bls12_381/pairing.rs

@@ -39,7 +39,7 @@ impl IsPairing for BLS12381AtePairing {
             if !p.is_neutral_element() && !q.is_neutral_element() {
                 let p = p.to_affine();
                 let q = q.to_affine();
-                result = result * miller(&q, &p);
+                result *= miller(&q, &p);
             }
         }
         Ok(final_exponentiation(&result))

math/src/field/element.rs

@@ -14,7 +14,7 @@ use core::iter::Sum;
     feature = "lambdaworks-serde-string"
 ))]
 use core::marker::PhantomData;
-use core::ops::{Add, AddAssign, Div, Mul, Neg, Sub};
+use core::ops::{Add, AddAssign, Div, Mul, MulAssign, Neg, Sub};
 #[cfg(any(
     feature = "lambdaworks-serde-binary",
     feature = "lambdaworks-serde-string"
@@ -301,6 +301,28 @@ where
     }
 }
 
+/// MulAssign operator overloading for field elements
+impl<F, L> MulAssign<FieldElement<F>> for FieldElement<L>
+where
+    F: IsSubFieldOf<L>,
+    L: IsField,
+{
+    fn mul_assign(&mut self, rhs: FieldElement<F>) {
+        self.value = <F as IsSubFieldOf<L>>::mul(&rhs.value, &self.value);
+    }
+}
+
+/// MulAssign operator overloading for field elements
+impl<F, L> MulAssign<&FieldElement<F>> for FieldElement<L>
+where
+    F: IsSubFieldOf<L>,
+    L: IsField,
+{
+    fn mul_assign(&mut self, rhs: &FieldElement<F>) {
+        self.value = <F as IsSubFieldOf<L>>::mul(&rhs.value, &self.value);
+    }
+}
+
 /// Division operator overloading for field elements*/
 impl<F, L> Div<&FieldElement<L>> for &FieldElement<F>
 where