Add Montgomery u32 backend for BabyBear (#948)

* WIP * add benches. Change cios for mul and then reduction * try another version of inv. Serialization tests not working. * fix from_hex behaviour when working with larger values * fix serialization tests * try plonky3 algorithm for inv * add new mul function * add const functions for mu and r2 parameters * remove commented code and refactor functions * tests big hex and more than 4 bytes failing * add fuzzer for babybear * fix tests from_hex for numbers bigger than u64 * remove comments and refactor some functions * fix cargo clippy * fix clippy * fix clippy metal * rename function * Fix overflow in shift ops --------- Co-authored-by: Nicole <[email protected]> Co-authored-by: Nicole <[email protected]> Co-authored-by: Diego K <[email protected]>
lambdaclass · Jan 6, 2025 · 4ee0600 · 4ee0600
1 parent 2ce8b2e
commit 4ee0600
Show file tree

Hide file tree

Showing 10 changed files with 1,048 additions and 5 deletions.
diff --git a/fuzz/no_gpu_fuzz/Cargo.toml b/fuzz/no_gpu_fuzz/Cargo.toml
@@ -16,6 +16,8 @@ num-traits = "0.2"
 ibig = "0.3.6"
 p3-goldilocks = { git = "https://github.com/Plonky3/Plonky3", rev = "41cd843" }
 p3-mersenne-31 = { git = "https://github.com/Plonky3/Plonky3", rev = "41cd843" }
+p3-field = { git = "https://github.com/Plonky3/Plonky3" }
+p3-baby-bear = { git = "https://github.com/Plonky3/Plonky3" }
 
 [[bin]]
 name = "curve_bls12_381"
@@ -53,6 +55,12 @@ path = "fuzz_targets/field/mersenne31.rs"
 test = false
 doc = false
 
+[[bin]]
+name = "babybear"
+path = "fuzz_targets/field/babybear.rs"
+test = false
+doc = false
+
 [[bin]]
 name = "mini_goldilocks"
 path = "fuzz_targets/field/mini_goldilocks.rs"
@@ -83,5 +91,3 @@ name = "deserialize_stark_proof"
 path = "fuzz_targets/deserialize_stark_proof.rs"
 test = false
 doc = false
-
-
diff --git a/fuzz/no_gpu_fuzz/fuzz_targets/field/babybear.rs b/fuzz/no_gpu_fuzz/fuzz_targets/field/babybear.rs
@@ -0,0 +1,79 @@
+#![no_main]
+
+use lambdaworks_math::field::{
+    element::FieldElement,
+    fields::u32_montgomery_backend_prime_field::U32MontgomeryBackendPrimeField,
+};
+use libfuzzer_sys::fuzz_target;
+use p3_baby_bear::BabyBear;
+use p3_field::{Field, FieldAlgebra, PrimeField32};
+
+pub type U32Babybear31PrimeField = U32MontgomeryBackendPrimeField<2013265921>;
+pub type F = FieldElement<U32Babybear31PrimeField>;
+
+fuzz_target!(|values: (u32, u32)| {
+    // Note: we filter values outside of order as it triggers an assert within plonky3 disallowing values n >= Self::Order
+    let (value_u32_a, value_u32_b) = values;
+
+    if value_u32_a >= 2013265921 || value_u32_b >= 2013265921 {
+        return;
+    }
+    let a = F::from(value_u32_a as u64);
+    let b = F::from(value_u32_b as u64);
+
+    // Note: if we parse using from_canonical_u32 fails due to check that n < Self::Order
+    let a_expected = BabyBear::from_canonical_u32(value_u32_a);
+    let b_expected = BabyBear::from_canonical_u32(value_u32_b);
+
+    let add_u32 = &a + &b;
+    let addition = a_expected + b_expected;
+    assert_eq!(add_u32.representative(), addition.as_canonical_u32());
+
+    let sub_u32 = &a - &b;
+    let substraction = a_expected - b_expected;
+    assert_eq!(sub_u32.representative(), substraction.as_canonical_u32());
+
+    let mul_u32 = &a * &b;
+    let multiplication = a_expected * b_expected;
+    assert_eq!(mul_u32.representative(), multiplication.as_canonical_u32());
+
+    // Axioms soundness
+    let one = F::one();
+    let zero = F::zero();
+
+    assert_eq!(&a + &zero, a, "Neutral add element a failed");
+    assert_eq!(&b + &zero, b, "Neutral mul element b failed");
+    assert_eq!(&a * &one, a, "Neutral add element a failed");
+    assert_eq!(&b * &one, b, "Neutral mul element b failed");
+
+    assert_eq!(&a + &b, &b + &a, "Commutative add property failed");
+    assert_eq!(&a * &b, &b * &a, "Commutative mul property failed");
+
+    let c = &a * &b;
+    assert_eq!(
+        (&a + &b) + &c,
+        &a + (&b + &c),
+        "Associative add property failed"
+    );
+    assert_eq!(
+        (&a * &b) * &c,
+        &a * (&b * &c),
+        "Associative mul property failed"
+    );
+
+    assert_eq!(
+        &a * (&b + &c),
+        &a * &b + &a * &c,
+        "Distributive property failed"
+    );
+
+    assert_eq!(&a - &a, zero, "Inverse add a failed");
+    assert_eq!(&b - &b, zero, "Inverse add b failed");
+
+    if a != zero {
+        assert_eq!(&a * a.inv().unwrap(), one, "Inverse mul a failed");
+    }
+    if b != zero {
+        assert_eq!(&b * b.inv().unwrap(), one, "Inverse mul b failed");
+    }
+});
diff --git a/math/Cargo.toml b/math/Cargo.toml
@@ -39,7 +39,8 @@ const-random = "0.1.15"
 iai-callgrind.workspace = true
 proptest = "1.1.0"
 pprof = { version = "0.13.0", features = ["criterion", "flamegraph"] }
-
+p3-baby-bear = { git = "https://github.com/Plonky3/Plonky3" }
+p3-field = { git = "https://github.com/Plonky3/Plonky3" }
 [features]
 default = ["parallel", "std"]
 std = ["alloc", "serde?/std", "serde_json?/std"]

diff --git a/math/benches/criterion_field.rs b/math/benches/criterion_field.rs
@@ -5,13 +5,16 @@ mod fields;
 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,
+    baby_bear::{babybear_ops_benchmarks, babybear_ops_benchmarks_f64, babybear_p3_ops_benchmarks},
+    stark252::starkfield_ops_benchmarks,
+    u64_goldilocks::u64_goldilocks_ops_benchmarks,
     u64_goldilocks_montgomery::u64_goldilocks_montgomery_ops_benchmarks,
 };
 
 criterion_group!(
     name = field_benches;
     config = Criterion::default().with_profiler(PProfProfiler::new(100, Output::Flamegraph(None)));
-    targets = mersenne31_ops_benchmarks, mersenne31_extension_ops_benchmarks, mersenne31_mont_ops_benchmarks, starkfield_ops_benchmarks, u64_goldilocks_ops_benchmarks, u64_goldilocks_montgomery_ops_benchmarks
+    targets =babybear_ops_benchmarks,babybear_ops_benchmarks_f64, babybear_p3_ops_benchmarks,mersenne31_extension_ops_benchmarks,mersenne31_ops_benchmarks,
+    starkfield_ops_benchmarks,u64_goldilocks_ops_benchmarks,u64_goldilocks_montgomery_ops_benchmarks,mersenne31_mont_ops_benchmarks
 );
 criterion_main!(field_benches);
diff --git a/math/benches/fields/baby_bear.rs b/math/benches/fields/baby_bear.rs
@@ -0,0 +1,216 @@
+use criterion::Criterion;
+use std::hint::black_box;
+
+use lambdaworks_math::field::{
+    element::FieldElement,
+    fields::{
+        fft_friendly::babybear::Babybear31PrimeField,
+        u32_montgomery_backend_prime_field::U32MontgomeryBackendPrimeField,
+    },
+};
+
+use p3_baby_bear::BabyBear;
+use p3_field::{Field, FieldAlgebra};
+
+use rand::random;
+use rand::Rng;
+
+pub type U32Babybear31PrimeField = U32MontgomeryBackendPrimeField<2013265921>;
+pub type F = FieldElement<U32Babybear31PrimeField>;
+pub type F64 = FieldElement<Babybear31PrimeField>;
+
+pub fn rand_field_elements(num: usize) -> Vec<(F, F)> {
+    let mut result = Vec::with_capacity(num);
+    for _ in 0..result.capacity() {
+        result.push((F::from(random::<u64>()), F::from(random::<u64>())));
+    }
+    result
+}
+
+fn rand_babybear_elements_p3(num: usize) -> Vec<(BabyBear, BabyBear)> {
+    let mut rng = rand::thread_rng();
+    (0..num)
+        .map(|_| (rng.gen::<BabyBear>(), rng.gen::<BabyBear>()))
+        .collect()
+}
+
+pub fn babybear_ops_benchmarks(c: &mut Criterion) {
+    let input: Vec<Vec<(F, F)>> = [1000000]
+        .into_iter()
+        .map(rand_field_elements)
+        .collect::<Vec<_>>();
+    let mut group = c.benchmark_group("BabyBear operations using Lambdaworks u32");
+
+    for i in input.clone().into_iter() {
+        group.bench_with_input(format!("Addition {:?}", &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!("Multiplication {:?}", &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 {:?}", &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!("Inverse {:?}", &i.len()), &i, |bench, i| {
+            bench.iter(|| {
+                for (x, _) in i {
+                    black_box(black_box(x).inv().unwrap());
+                }
+            });
+        });
+    }
+
+    for i in input.clone().into_iter() {
+        group.bench_with_input(format!("Division {:?}", &i.len()), &i, |bench, i| {
+            bench.iter(|| {
+                for (x, y) in i {
+                    black_box(black_box(x) / black_box(y));
+                }
+            });
+        });
+    }
+}
+
+pub fn rand_field_elements_u64(num: usize) -> Vec<(F64, F64)> {
+    let mut result = Vec::with_capacity(num);
+    for _ in 0..result.capacity() {
+        result.push((F64::from(random::<u64>()), F64::from(random::<u64>())));
+    }
+    result
+}
+pub fn babybear_ops_benchmarks_f64(c: &mut Criterion) {
+    let input: Vec<Vec<(F64, F64)>> = [1000000]
+        .into_iter()
+        .map(rand_field_elements_u64)
+        .collect::<Vec<_>>();
+    let mut group = c.benchmark_group("BabyBear operations using Lambdaworks u64");
+
+    for i in input.clone().into_iter() {
+        group.bench_with_input(format!("Addition {:?}", &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!("Multiplication {:?}", &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 {:?}", &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!("Inverse {:?}", &i.len()), &i, |bench, i| {
+            bench.iter(|| {
+                for (x, _) in i {
+                    black_box(black_box(x).inv().unwrap());
+                }
+            });
+        });
+    }
+
+    for i in input.clone().into_iter() {
+        group.bench_with_input(format!("Division {:?}", &i.len()), &i, |bench, i| {
+            bench.iter(|| {
+                for (x, y) in i {
+                    black_box(black_box(x) / black_box(y));
+                }
+            });
+        });
+    }
+}
+
+pub fn babybear_p3_ops_benchmarks(c: &mut Criterion) {
+    let input: Vec<Vec<(BabyBear, BabyBear)>> = [1000000]
+        .into_iter()
+        .map(rand_babybear_elements_p3)
+        .collect::<Vec<_>>();
+
+    let mut group = c.benchmark_group("BabyBear operations using Plonky3");
+
+    for i in input.clone().into_iter() {
+        group.bench_with_input(format!("Addition {:?}", &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!("Multiplication {:?}", &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 {:?}", &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!("Inverse {:?}", &i.len()), &i, |bench, i| {
+            bench.iter(|| {
+                for (x, _) in i {
+                    black_box(black_box(x).inverse());
+                }
+            });
+        });
+    }
+
+    for i in input.clone().into_iter() {
+        group.bench_with_input(format!("Division {:?}", &i.len()), &i, |bench, i| {
+            bench.iter(|| {
+                for (x, y) in i {
+                    black_box(black_box(*x) / black_box(*y));
+                }
+            });
+        });
+    }
+}
diff --git a/math/benches/fields/mod.rs b/math/benches/fields/mod.rs
@@ -1,3 +1,4 @@
+pub mod baby_bear;
 pub mod mersenne31;
 pub mod mersenne31_montgomery;
 pub mod stark252;