implement M31Ext3

.github/workflows/ci.yml

@@ -54,8 +54,8 @@ jobs:
       - run: cargo test -p bi-kzg --release
       - run: cargo test --no-default-features --release      
       - run: cargo test --all-features --release
-      - run: cargo run --release -- -t 4 -f 31 
-      - run: cargo run --release -- -t 4 -f 254 
+      - run: cargo run --release -- -t 4 -f m31 
+      - run: cargo run --release -- -t 4 -f fr 
 
   test-linux:
     name: Test-linux
@@ -76,5 +76,5 @@ jobs:
       - run: RUSTFLAGS="-C target-feature=+avx2" cargo test -p bi-kzg --release
       - run: RUSTFLAGS="-C target-feature=+avx2" cargo test --no-default-features --release
       - run: RUSTFLAGS="-C target-feature=+avx2" cargo test --all-features --release
-      - run: RUSTFLAGS="-C target-feature=+avx2" cargo run --release -- -t 4 -f 31 
-      - run: RUSTFLAGS="-C target-feature=+avx2" cargo run --release -- -t 4 -f 254 
+      - run: RUSTFLAGS="-C target-feature=+avx2" cargo run --release -- -t 4 -f m31 
+      - run: RUSTFLAGS="-C target-feature=+avx2" cargo run --release -- -t 4 -f fr 

.gitignore

@@ -21,6 +21,9 @@ Cargo.lock
 /data/circuit8.txt
 /data/compiler_out
 /data/keccak*
+/data/bn254_circuit.txt
+/data/poseidon_witness.txt
+/data/poseidon_100_circuit.txt
 notes.md
 
 # Programming env

Cargo.toml

@@ -43,3 +43,4 @@ log = "0.4"
 rand = "0.8.5"
 rayon = "1.10"
 sha2 = "0.10.8"
+tynm = { version = "0.1.6", default-features = false }

arith/Cargo.toml

@@ -10,5 +10,12 @@ log.workspace = true
 rand.workspace = true
 sha2.workspace = true
 
+[dev-dependencies]
+tynm.workspace = true
+criterion.workspace = true
 
-[features]
+[[bench]]
+name = "field"
+harness = false
+
+[features]

arith/benches/field.rs

@@ -0,0 +1,144 @@
+// this module benchmarks the performance of different field operations
+
+use arith::{Field, M31Ext3, M31};
+use ark_std::test_rng;
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
+use halo2curves::bn256::Fr;
+use tynm::type_name;
+
+#[cfg(target_arch = "x86_64")]
+use arith::PackedM31Ext3;
+
+fn random_element<F: Field>() -> F {
+    let mut rng = test_rng();
+    F::random_unsafe(&mut rng)
+}
+
+pub(crate) fn bench_field_avx512<F: Field>(c: &mut Criterion) {
+    c.bench_function(&format!("mul-throughput<{}>", type_name::<F>()), |b| {
+        b.iter_batched(
+            || {
+                (
+                    random_element::<F>(),
+                    random_element::<F>(),
+                    random_element::<F>(),
+                    random_element::<F>(),
+                )
+            },
+            |(mut x, mut y, mut z, mut w)| {
+                for _ in 0..25 {
+                    (x, y, z, w) = (x * y, y * z, z * w, w * x);
+                }
+                (x, y, z, w)
+            },
+            BatchSize::SmallInput,
+        )
+    });
+
+    c.bench_function(&format!("mul-latency<{}>", type_name::<F>()), |b| {
+        b.iter_batched(
+            || random_element::<F>(),
+            |mut x| {
+                for _ in 0..100 {
+                    x = x * x;
+                }
+                x
+            },
+            BatchSize::SmallInput,
+        )
+    });
+
+    c.bench_function(&format!("sqr-throughput<{}>", type_name::<F>()), |b| {
+        b.iter_batched(
+            || {
+                (
+                    random_element::<F>(),
+                    random_element::<F>(),
+                    random_element::<F>(),
+                    random_element::<F>(),
+                )
+            },
+            |(mut x, mut y, mut z, mut w)| {
+                for _ in 0..25 {
+                    (x, y, z, w) = (x.square(), y.square(), z.square(), w.square());
+                }
+                (x, y, z, w)
+            },
+            BatchSize::SmallInput,
+        )
+    });
+
+    c.bench_function(&format!("sqr-latency<{}>", type_name::<F>()), |b| {
+        b.iter_batched(
+            || random_element::<F>(),
+            |mut x| {
+                for _ in 0..100 {
+                    x = x.square();
+                }
+                x
+            },
+            BatchSize::SmallInput,
+        )
+    });
+
+    c.bench_function(&format!("add-throughput<{}>", type_name::<F>()), |b| {
+        b.iter_batched(
+            || {
+                (
+                    random_element::<F>(),
+                    random_element::<F>(),
+                    random_element::<F>(),
+                    random_element::<F>(),
+                    random_element::<F>(),
+                    random_element::<F>(),
+                    random_element::<F>(),
+                    random_element::<F>(),
+                    random_element::<F>(),
+                    random_element::<F>(),
+                )
+            },
+            |(mut a, mut b, mut c, mut d, mut e, mut f, mut g, mut h, mut i, mut j)| {
+                for _ in 0..10 {
+                    (a, b, c, d, e, f, g, h, i, j) = (
+                        a + b,
+                        b + c,
+                        c + d,
+                        d + e,
+                        e + f,
+                        f + g,
+                        g + h,
+                        h + i,
+                        i + j,
+                        j + a,
+                    );
+                }
+                (a, b, c, d, e, f, g, h, i, j)
+            },
+            BatchSize::SmallInput,
+        )
+    });
+
+    c.bench_function(&format!("add-latency<{}>", type_name::<F>()), |b| {
+        b.iter_batched(
+            || random_element::<F>(),
+            |mut x| {
+                for _ in 0..100 {
+                    x = x + x;
+                }
+                x
+            },
+            BatchSize::SmallInput,
+        )
+    });
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    bench_field_avx512::<M31>(c);
+    bench_field_avx512::<M31Ext3>(c);
+    #[cfg(target_arch = "x86_64")]
+    bench_field_avx512::<PackedM31Ext3>(c);
+    bench_field_avx512::<Fr>(c);
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

arith/src/field.rs

@@ -2,6 +2,8 @@ mod m31;
 pub use m31::*;
 mod bn254;
 pub use bn254::*;
+mod m31_ext;
+pub use m31_ext::*;
 
 use rand::RngCore;
 
@@ -138,5 +140,11 @@ pub trait FieldSerde {
     fn deserialize_from(buffer: &[u8]) -> Self;
 
     /// deserialize bytes into field following ecc format
-    fn deserialize_from_ecc_format(bytes: &[u8; 32]) -> Self;
+    fn deserialize_from_ecc_format(_bytes: &[u8; 32]) -> Self
+    where
+        Self: Sized,
+    {
+        // add default implementation to avoid duplications when this isn't required
+        unimplemented!()
+    }
 }

arith/src/field/m31.rs

@@ -23,7 +23,7 @@ use std::{
 pub const M31_MOD: i32 = 2147483647;
 
 #[inline]
-fn mod_reduce_i32(x: i32) -> i32 {
+pub(crate) fn mod_reduce_i32(x: i32) -> i32 {
     (x & M31_MOD) + (x >> 31)
 }
 
@@ -65,6 +65,20 @@ impl FieldSerde for M31 {
     }
 }
 
+impl M31 {
+    // Add two M31 without mod reduction
+    #[inline(always)]
+    pub fn unsafe_add(&self, rhs: &Self) -> Self {
+        Self { v: self.v + rhs.v }
+    }
+
+    // Double an M31 without mod reduction
+    #[inline(always)]
+    pub fn unsafe_double(&self) -> Self {
+        Self { v: self.v << 1 }
+    }
+}
+
 impl Field for M31 {
     const NAME: &'static str = "Mersenne 31";
 

arith/src/field/m31/m31_avx.rs

@@ -6,7 +6,7 @@ use std::{
     ops::{Add, AddAssign, Mul, MulAssign, Neg, Sub, SubAssign},
 };
 
-use crate::{Field, M31, M31_MOD};
+use crate::{Field, FieldSerde, M31, M31_MOD};
 
 type PackedDataType = __m256i;
 pub(super) const M31_PACK_SIZE: usize = 8;
@@ -50,6 +50,27 @@ impl PackedM31 {
     }
 }
 
+impl FieldSerde for PackedM31 {
+    /// serialize self into bytes
+    fn serialize_into(&self, buffer: &mut [u8]) {
+        unsafe {
+            let data = transmute::<PackedDataType, [u8; 32]>(self.v);
+            buffer[..32].copy_from_slice(&data);
+        }
+    }
+
+    /// deserialize bytes into field
+    fn deserialize_from(buffer: &[u8]) -> Self {
+        let mut data = [0; 32];
+        data.copy_from_slice(buffer);
+        unsafe {
+            PackedM31 {
+                v: transmute::<[u8; 32], PackedDataType>(data),
+            }
+        }
+    }
+}
+
 impl Field for PackedM31 {
     const NAME: &'static str = "AVX Packed Mersenne 31";