Poseidon2 (#683)

crates/prover/Cargo.toml

@@ -23,7 +23,7 @@ starknet-crypto = "0.6.2"
 starknet-ff = "0.3.7"
 thiserror.workspace = true
 tracing.workspace = true
-rayon = {version = "1.10.0", optional = true}
+rayon = { version = "1.10.0", optional = true }
 
 [dev-dependencies]
 aligned = "0.4.2"
@@ -82,3 +82,7 @@ harness = false
 [[bench]]
 name = "quotients"
 harness = false
+
+[[bench]]
+name = "poseidon"
+harness = false

crates/prover/benches/poseidon.rs

@@ -0,0 +1,32 @@
+use criterion::{criterion_group, criterion_main, Criterion, Throughput};
+use stwo_prover::core::backend::simd::SimdBackend;
+use stwo_prover::core::channel::{Blake2sChannel, Channel};
+use stwo_prover::core::fields::m31::BaseField;
+use stwo_prover::core::fields::IntoSlice;
+use stwo_prover::core::prover::prove;
+use stwo_prover::core::vcs::blake2_hash::Blake2sHasher;
+use stwo_prover::core::vcs::hasher::Hasher;
+use stwo_prover::examples::poseidon::{gen_trace, PoseidonAir, PoseidonComponent};
+
+pub fn simd_poseidon(c: &mut Criterion) {
+    const LOG_N_ROWS: u32 = 15;
+    let mut group = c.benchmark_group("poseidon2");
+    group.throughput(Throughput::Elements(1u64 << (LOG_N_ROWS + 3)));
+    group.bench_function(format!("poseidon2 2^{} instances", LOG_N_ROWS + 3), |b| {
+        b.iter(|| {
+            let component = PoseidonComponent {
+                log_n_instances: LOG_N_ROWS,
+            };
+            let trace = gen_trace(component.log_column_size());
+            let channel = &mut Blake2sChannel::new(Blake2sHasher::hash(BaseField::into_slice(&[])));
+            let air = PoseidonAir { component };
+            prove::<SimdBackend>(&air, channel, trace).unwrap()
+        });
+    });
+}
+
+criterion_group!(
+    name = bit_rev;
+    config = Criterion::default().sample_size(10);
+    targets = simd_poseidon);
+criterion_main!(bit_rev);

crates/prover/src/core/backend/simd/m31.rs

@@ -119,6 +119,13 @@ impl AddAssign for PackedM31 {
     }
 }
 
+impl AddAssign<M31> for PackedM31 {
+    #[inline(always)]
+    fn add_assign(&mut self, rhs: M31) {
+        *self = *self + PackedM31::broadcast(rhs);
+    }
+}
+
 impl Mul for PackedM31 {
     type Output = Self;
 
@@ -142,6 +149,15 @@ impl Mul for PackedM31 {
     }
 }
 
+impl Mul<BaseField> for PackedM31 {
+    type Output = Self;
+
+    #[inline(always)]
+    fn mul(self, rhs: BaseField) -> Self::Output {
+        self * PackedM31::broadcast(rhs)
+    }
+}
+
 impl MulAssign for PackedM31 {
     #[inline(always)]
     fn mul_assign(&mut self, rhs: Self) {

crates/prover/src/core/fft.rs

@@ -1,13 +1,20 @@
+use std::ops::{Add, AddAssign, Mul, Sub};
+
 use super::fields::m31::BaseField;
-use super::fields::ExtensionOf;
 
-pub fn butterfly<F: ExtensionOf<BaseField>>(v0: &mut F, v1: &mut F, twid: BaseField) {
+pub fn butterfly<F>(v0: &mut F, v1: &mut F, twid: BaseField)
+where
+    F: Copy + AddAssign<F> + Sub<F, Output = F> + Mul<BaseField, Output = F>,
+{
     let tmp = *v1 * twid;
     *v1 = *v0 - tmp;
     *v0 += tmp;
 }
 
-pub fn ibutterfly<F: ExtensionOf<BaseField>>(v0: &mut F, v1: &mut F, itwid: BaseField) {
+pub fn ibutterfly<F>(v0: &mut F, v1: &mut F, itwid: BaseField)
+where
+    F: Copy + AddAssign<F> + Add<F, Output = F> + Sub<F, Output = F> + Mul<BaseField, Output = F>,
+{
     let tmp = *v0;
     *v0 = tmp + *v1;
     *v1 = (tmp - *v1) * itwid;

crates/prover/src/examples/mod.rs

@@ -1,2 +1,3 @@
 pub mod fibonacci;
+pub mod poseidon;
 pub mod wide_fibonacci;