improve p2 time by 2x

EspressoSystems · Jan 14, 2025 · d7a8321 · d7a8321
1 parent 9a781de
commit d7a8321
Show file tree

Hide file tree

Showing 4 changed files with 17 additions and 5 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -34,3 +34,7 @@ sha3 = { version = "0.10", default-features = false }
 itertools = { version = "0.12", default-features = false }
 tagged-base64 = "0.4"
 zeroize = { version = "^1.8" }
+
+[profile.profiling]
+inherits = "release"
+debug = true
diff --git a/poseidon2/src/external.rs b/poseidon2/src/external.rs
@@ -88,9 +88,8 @@ pub(crate) fn permute_state<F: PrimeField, const T: usize>(
     rc: &'static [F; T],
     d: usize,
 ) {
-    state
-        .iter_mut()
-        .zip(rc.iter())
-        .for_each(|(s, &rc)| add_rc_and_sbox(s, rc, d));
+    for i in 0..T {
+        add_rc_and_sbox(&mut state[i], rc[i], d)
+    }
     matmul_external(state);
 }
diff --git a/poseidon2/src/internal.rs b/poseidon2/src/internal.rs
@@ -22,6 +22,7 @@ fn matmul_internal<F: PrimeField, const T: usize>(
 
 /// One internal round
 // @credit `internal_permute_state()` in plonky3
+#[inline(always)]
 pub(crate) fn permute_state<F: PrimeField, const T: usize>(
     state: &mut [F; T],
     rc: F,

diff --git a/poseidon2/src/lib.rs b/poseidon2/src/lib.rs
@@ -130,7 +130,15 @@ impl<F: PrimeField> Poseidon2<F> {
 #[inline(always)]
 pub(crate) fn add_rc_and_sbox<F: PrimeField>(val: &mut F, rc: F, d: usize) {
     *val += rc;
-    *val = val.pow([d as u64]);
+    if d == 5 {
+        // Perform unrolled computation for val^5, faster
+        let original = *val;
+        val.square_in_place();
+        val.square_in_place();
+        *val *= &original;
+    } else {
+        *val = val.pow([d as u64]);
+    }
 }
 
 /// Poseidon2 Error type