Optimize Mersenne31 Field (#921)

* optimize add * save changes. Add, sub and mul checked * fix tests * add new inv * add mult by powers of two * replace inverse * test new inv * modify old algorithm for inv * fix tests extension * add mul for degree 4 extension * add fp4 isField and isSubField operations and benchmarks * new version for fp4 mul based on the paper * add mul of a fp2e by non-residue * change inv using mul_fp2_by_non_resiude * save work * wip fp2 test * add fp2 tests * add 2 * a^2 - 1 function * use karatsuba in fp4 mul version 1 * clean up * fix Fp as subfield of Fp2. Tests Fp plus Fp4 is now correct * fix inv * fix comments * fix comments * fixes * fix clippy * fix cargo check no-std * fix clippy * change zero function of isField to rust default * fix two_square_minus_one function and optimize inv function * fix clippy --------- Co-authored-by: Nicole <[email protected]> Co-authored-by: Joaquin Carletti <[email protected]> Co-authored-by: diegokingston <[email protected]> Co-authored-by: Diego K <[email protected]> Co-authored-by: Joaquin Carletti <[email protected]>
lambdaclass · Oct 1, 2024 · 7a427aa · 7a427aa
1 parent 9617e52
commit 7a427aa
Show file tree

Hide file tree

Showing 9 changed files with 863 additions and 428 deletions.
diff --git a/math/benches/criterion_field.rs b/math/benches/criterion_field.rs
@@ -2,7 +2,7 @@ use criterion::{criterion_group, criterion_main, Criterion};
 use pprof::criterion::{Output, PProfProfiler};
 
 mod fields;
-use fields::mersenne31::mersenne31_ops_benchmarks;
+use fields::mersenne31::{mersenne31_extension_ops_benchmarks, mersenne31_ops_benchmarks};
 use fields::mersenne31_montgomery::mersenne31_mont_ops_benchmarks;
 use fields::{
     stark252::starkfield_ops_benchmarks, u64_goldilocks::u64_goldilocks_ops_benchmarks,
@@ -12,6 +12,6 @@ use fields::{
 criterion_group!(
     name = field_benches;
     config = Criterion::default().with_profiler(PProfProfiler::new(100, Output::Flamegraph(None)));
-    targets = starkfield_ops_benchmarks, mersenne31_ops_benchmarks, mersenne31_mont_ops_benchmarks, u64_goldilocks_ops_benchmarks, u64_goldilocks_montgomery_ops_benchmarks
+    targets = mersenne31_ops_benchmarks, mersenne31_extension_ops_benchmarks, mersenne31_mont_ops_benchmarks, starkfield_ops_benchmarks, u64_goldilocks_ops_benchmarks, u64_goldilocks_montgomery_ops_benchmarks
 );
 criterion_main!(field_benches);
diff --git a/math/benches/fields/mersenne31.rs b/math/benches/fields/mersenne31.rs
@@ -1,10 +1,18 @@
 use std::hint::black_box;
 
 use criterion::Criterion;
-use lambdaworks_math::field::{element::FieldElement, fields::mersenne31::field::Mersenne31Field};
+use lambdaworks_math::field::{
+    element::FieldElement,
+    fields::mersenne31::{
+        extensions::{Degree2ExtensionField, Degree4ExtensionField},
+        field::Mersenne31Field,
+    },
+};
 use rand::random;
 
 pub type F = FieldElement<Mersenne31Field>;
+pub type Fp2E = FieldElement<Degree2ExtensionField>;
+pub type Fp4E = FieldElement<Degree4ExtensionField>;
 
 #[inline(never)]
 #[no_mangle]
@@ -17,6 +25,60 @@ pub fn rand_field_elements(num: usize) -> Vec<(F, F)> {
     result
 }
 
+//TODO: Check if this is the correct way to bench.
+pub fn rand_fp4e(num: usize) -> Vec<(Fp4E, Fp4E)> {
+    let mut result = Vec::with_capacity(num);
+    for _ in 0..result.capacity() {
+        result.push((
+            Fp4E::new([
+                Fp2E::new([F::new(random()), F::new(random())]),
+                Fp2E::new([F::new(random()), F::new(random())]),
+            ]),
+            Fp4E::new([
+                Fp2E::new([F::new(random()), F::new(random())]),
+                Fp2E::new([F::new(random()), F::new(random())]),
+            ]),
+        ));
+    }
+    result
+}
+
+pub fn mersenne31_extension_ops_benchmarks(c: &mut Criterion) {
+    let input: Vec<Vec<(Fp4E, Fp4E)>> = [1000000].into_iter().map(rand_fp4e).collect::<Vec<_>>();
+
+    let mut group = c.benchmark_group("Mersenne31 Fp4 operations");
+
+    for i in input.clone().into_iter() {
+        group.bench_with_input(format!("Mul of Fp4 {:?}", &i.len()), &i, |bench, i| {
+            bench.iter(|| {
+                for (x, y) in i {
+                    black_box(black_box(x) * black_box(y));
+                }
+            });
+        });
+    }
+
+    for i in input.clone().into_iter() {
+        group.bench_with_input(format!("Square of Fp4 {:?}", &i.len()), &i, |bench, i| {
+            bench.iter(|| {
+                for (x, _) in i {
+                    black_box(black_box(x).square());
+                }
+            });
+        });
+    }
+
+    for i in input.clone().into_iter() {
+        group.bench_with_input(format!("Inv of Fp4 {:?}", &i.len()), &i, |bench, i| {
+            bench.iter(|| {
+                for (x, _) in i {
+                    black_box(black_box(x).inv().unwrap());
+                }
+            });
+        });
+    }
+}
+
 pub fn mersenne31_ops_benchmarks(c: &mut Criterion) {
     let input: Vec<Vec<(F, F)>> = [1, 10, 100, 1000, 10000, 100000, 1000000]
         .into_iter()