optimize poseidon2

plonky2/src/hash/arch/aarch64/poseidon_goldilocks_neon.rs

@@ -915,3 +915,28 @@ pub unsafe fn mds_layer(state: &[GoldilocksField; WIDTH]) -> [GoldilocksField; W
     let state = mds_layer_full(state);
     wrap_state(state)
 }
+
+#[inline(always)]
+#[unroll::unroll_for_loops]
+pub unsafe fn vector_add(a: &[u64; WIDTH], b: &[u64; WIDTH]) -> [u64; WIDTH] {
+    let mut res = [0u64; WIDTH];
+    // Process 2 elements at a time using NEON
+    for i in (0..WIDTH).step_by(2) {
+        let a_vec = vld1q_u64(a[i..].as_ptr());
+        let b_vec = vld1q_u64(b[i..].as_ptr());
+
+        // Add the round constants
+        let sum = vaddq_u64(a_vec, b_vec);
+
+        // Check for overflow (if sum < state_vec, we wrapped around)
+        let overflow_mask = vcltq_u64(sum, a_vec);
+
+        // Add EPSILON (0xffffffff) where overflow occurred
+        let epsilon = vdupq_n_u64(0xffffffff);
+        let adjustment = vandq_u64(overflow_mask, epsilon);
+        let result = vaddq_u64(sum, adjustment);
+
+        vst1q_u64(res[i..].as_mut_ptr(), result);
+    }
+    res
+}

plonky2/src/hash/poseidon2/gate.rs

@@ -1,6 +1,5 @@
 //! Implementation of a Plonky2 gate for an entire Poseidon2 permutation over a
 //! state of width 12
-
 use core::marker::PhantomData;
 
 use anyhow::Result;
@@ -495,9 +494,12 @@ impl<F: RichField + Extendable<D> + Poseidon2, const D: usize> SimpleGenerator<F
 mod tests {
     use anyhow::Result;
 
-    use super::Poseidon2Gate;
+    use super::{Poseidon2Gate, *};
     use crate::field::goldilocks_field::GoldilocksField;
     use crate::gates::gate_testing::{test_eval_fns, test_low_degree};
+    use crate::iop::generator::generate_partial_witness;
+    use crate::iop::witness::PartialWitness;
+    use crate::plonk::circuit_data::CircuitConfig;
     use crate::plonk::config::{GenericConfig, Poseidon2GoldilocksConfig};
 
     #[test]
@@ -512,6 +514,67 @@ mod tests {
         assert_eq!(Gate::WIRE_SWAP, 24);
         assert_eq!(Gate::wire_delta(0), 25);
         assert_eq!(Gate::wire_delta(3), 28);
+        assert_eq!(Gate::wire_full_sbox_0(1, 0), 29);
+        assert_eq!(Gate::wire_full_sbox_0(3, 0), 53);
+        assert_eq!(Gate::wire_full_sbox_0(3, 11), 64);
+        assert_eq!(Gate::wire_partial_sbox(0), 65);
+        assert_eq!(Gate::wire_partial_sbox(21), 86);
+        assert_eq!(Gate::wire_full_sbox_1(0, 0), 87);
+        assert_eq!(Gate::wire_full_sbox_1(3, 0), 123);
+        assert_eq!(Gate::wire_full_sbox_1(3, 11), 134);
+    }
+
+    #[test]
+    fn generated_output() {
+        const D: usize = 2;
+        type C = Poseidon2GoldilocksConfig;
+        type F = <C as GenericConfig<D>>::F;
+
+        let config = CircuitConfig {
+            num_wires: 143,
+            ..CircuitConfig::standard_recursion_config()
+        };
+        let mut builder = CircuitBuilder::new(config);
+        type Gate = Poseidon2Gate<F, D>;
+        let gate = Gate::new();
+        let row = builder.add_gate(gate, vec![]);
+        let circuit = builder.build_prover::<C>();
+
+        let permutation_inputs = (0..WIDTH).map(F::from_canonical_usize).collect::<Vec<_>>();
+
+        let mut inputs = PartialWitness::new();
+        inputs
+            .set_wire(
+                Wire {
+                    row,
+                    column: Gate::WIRE_SWAP,
+                },
+                F::ZERO,
+            )
+            .unwrap();
+        for i in 0..WIDTH {
+            inputs
+                .set_wire(
+                    Wire {
+                        row,
+                        column: Gate::wire_input(i),
+                    },
+                    permutation_inputs[i],
+                )
+                .unwrap();
+        }
+
+        let witness =
+            generate_partial_witness(inputs, &circuit.prover_only, &circuit.common).unwrap();
+
+        let expected_outputs: [F; WIDTH] = F::poseidon2(permutation_inputs.try_into().unwrap());
+        for i in 0..WIDTH {
+            let out = witness.get_wire(Wire {
+                row: 0,
+                column: Gate::wire_output(i),
+            });
+            assert_eq!(out, expected_outputs[i]);
+        }
     }
 
     #[test]