Another attempt to use rayon in poseidon

o1-labs · Sep 16, 2024 · 91ce607 · 91ce607
1 parent fd5d15d
commit 91ce607
Showing 1 changed file with 19 additions and 18 deletions.
diff --git a/poseidon/src/permutation.rs b/poseidon/src/permutation.rs
@@ -3,6 +3,7 @@
 use crate::constants::SpongeConstants;
 use crate::poseidon::{sbox, ArithmeticSpongeParams};
 use ark_ff::Field;
+use rayon::join;
 
 fn apply_mds_matrix<F: Field, SC: SpongeConstants>(
     params: &ArithmeticSpongeParams<F>,
@@ -35,24 +36,24 @@ pub fn full_round<F: Field, SC: SpongeConstants>(
     r: usize,
 ) {
     if SC::PERM_FULL_MDS && state.len() == 3 {
-        state[0] = sbox::<F, SC>(state[0]);
-        state[1] = sbox::<F, SC>(state[1]);
-        state[2] = sbox::<F, SC>(state[2]);
-        *state = vec![
-            // Manually unrolled loops for multiplying each row by the vector
-            params.mds[0][0] * state[0]
-                + params.mds[0][1] * state[1]
-                + params.mds[0][2] * state[2]
-                + params.round_constants[r][0],
-            params.mds[1][0] * state[0]
-                + params.mds[1][1] * state[1]
-                + params.mds[1][2] * state[2]
-                + params.round_constants[r][1],
-            params.mds[2][0] * state[0]
-                + params.mds[2][1] * state[1]
-                + params.mds[2][2] * state[2]
-                + params.round_constants[r][2],
-        ];
+        let (result0, result1) = join(|| sbox::<F, SC>(state[0]), || sbox::<F, SC>(state[1]));
+        let (result1, result2) = join(
+            || result1, // Passed from the first join
+            || sbox::<F, SC>(state[2]),
+        );
+        // Manually unrolled loops for multiplying each row by the vector
+        state[0] = params.mds[0][0] * result0
+            + params.mds[0][1] * result1
+            + params.mds[0][2] * result2
+            + params.round_constants[r][0];
+        state[1] = params.mds[1][0] * result0
+            + params.mds[1][1] * result1
+            + params.mds[1][2] * result2
+            + params.round_constants[r][1];
+        state[2] = params.mds[2][0] * result0
+            + params.mds[2][1] * result1
+            + params.mds[2][2] * result2
+            + params.round_constants[r][2];
     } else {
         for state_i in state.iter_mut() {
             *state_i = sbox::<F, SC>(*state_i);