impl poseidon2 with plonky3

.github/workflows/continuous-integration-workflow.yml

@@ -64,9 +64,8 @@ jobs:
         uses: actions/checkout@v4
 
       - name: Install nightly toolchain
-        uses: dtolnay/rust-toolchain@master
+        uses: dtolnay/rust-toolchain@nightly
         with:
-          toolchain: nightly-2024-06-23
           components: rustfmt, clippy
 
       - name: Rust Cache

field/src/goldilocks_field.rs

@@ -160,6 +160,10 @@ impl Field for GoldilocksField {
         reduce96((n_lo, n_hi))
     }
 
+    fn from_noncanonical_u128_with_96_bits(n: u128) -> Self {
+        reduce128_with_96_bits(n)
+    }
+
     fn from_noncanonical_u128(n: u128) -> Self {
         reduce128(n)
     }
@@ -396,6 +400,14 @@ fn reduce96((x_lo, x_hi): (u64, u32)) -> GoldilocksField {
     GoldilocksField(t2)
 }
 
+#[inline]
+fn reduce128_with_96_bits(x: u128) -> GoldilocksField {
+    let (x_lo, x_hi) = split(x); // This is a no-op
+    let t1 = x_hi * EPSILON;
+    let t2 = unsafe { add_no_canonicalize_trashing_input(x_lo, t1) };
+    GoldilocksField(t2)
+}
+
 /// Reduces to a 64-bit value. The result might not be in canonical form; it could be in between the
 /// field order and `2^64`.
 #[inline]

field/src/lib.rs

@@ -4,7 +4,6 @@
 #![deny(rustdoc::broken_intra_doc_links)]
 #![deny(missing_debug_implementations)]
 #![feature(specialization)]
-#![cfg_attr(target_arch = "x86_64", feature(stdarch_x86_avx512))]
 #![cfg_attr(not(test), no_std)]
 
 extern crate alloc;

field/src/types.rs

@@ -356,6 +356,11 @@ pub trait Field:
         Self::from_noncanonical_u128(n)
     }
 
+    fn from_noncanonical_u128_with_96_bits(n: u128) -> Self {
+        // Default implementation.
+        Self::from_noncanonical_u128(n)
+    }
+
     fn exp_power_of_2(&self, power_log: usize) -> Self {
         let mut res = *self;
         for _ in 0..power_log {

plonky2/Cargo.toml

@@ -38,6 +38,11 @@ plonky2_field = { version = "1.0.0", path = "../field", default-features = false
 plonky2_maybe_rayon = { version = "1.0.0", path = "../maybe_rayon", default-features = false }
 plonky2_util = { version = "1.0.0", path = "../util", default-features = false }
 
+# Plonky3 dependencies for Poseidon2
+p3-poseidon2 = { git = "https://github.com/Plonky3/Plonky3.git", rev = "eeb4e37b20127c4daa871b2bad0df30a7c7380db" }
+p3-goldilocks = { git = "https://github.com/Plonky3/Plonky3.git", rev = "eeb4e37b20127c4daa871b2bad0df30a7c7380db" }
+p3-field = { git = "https://github.com/Plonky3/Plonky3.git", rev = "eeb4e37b20127c4daa871b2bad0df30a7c7380db" }
+p3-symmetric = { git = "https://github.com/Plonky3/Plonky3.git", rev = "eeb4e37b20127c4daa871b2bad0df30a7c7380db" }
 
 [target.'cfg(all(target_arch = "wasm32", target_os = "unknown"))'.dependencies]
 getrandom = { version = "0.2", default-features = false, features = ["js"] }

plonky2/benches/hashing.rs

@@ -1,11 +1,16 @@
 mod allocator;
 
 use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
+use p3_field::AbstractField;
+use p3_goldilocks::{DiffusionMatrixGoldilocks, Goldilocks};
+use p3_poseidon2::{Poseidon2 as P3Poseidon2, Poseidon2ExternalMatrixGeneral};
+use p3_symmetric::Permutation;
 use plonky2::field::goldilocks_field::GoldilocksField;
 use plonky2::field::types::Sample;
 use plonky2::hash::hash_types::{BytesHash, RichField};
 use plonky2::hash::keccak::KeccakHash;
 use plonky2::hash::poseidon::{Poseidon, SPONGE_WIDTH};
+use plonky2::hash::poseidon2::hash::Poseidon2;
 use plonky2::plonk::config::Hasher;
 use tynm::type_name;
 
@@ -32,8 +37,82 @@ pub(crate) fn bench_poseidon<F: Poseidon>(c: &mut Criterion) {
     );
 }
 
+pub(crate) fn bench_poseidon2<F: Poseidon2>(c: &mut Criterion) {
+    c.bench_function(
+        &format!("optimized poseidon2<{}, {SPONGE_WIDTH}>", type_name::<F>()),
+        |b| {
+            b.iter_batched(
+                || F::rand_array::<SPONGE_WIDTH>(),
+                |state| F::poseidon2(state),
+                BatchSize::SmallInput,
+            )
+        },
+    );
+}
+
+pub(crate) fn bench_p3_poseidon2(c: &mut Criterion) {
+    const WIDTH: usize = 12;
+    const D: u64 = 7;
+    const ROUNDS_F: usize = 8;
+    const ROUNDS_P: usize = 22;
+
+    // Create the Poseidon2 instance
+    let external_linear_layer = Poseidon2ExternalMatrixGeneral;
+    let internal_linear_layer = DiffusionMatrixGoldilocks;
+
+    let external_constants = plonky2::hash::poseidon2::config::EXTERNAL_CONSTANTS
+        .iter()
+        .map(|v| {
+            v.iter()
+                .map(|&x| Goldilocks::from_canonical_u64(x))
+                .collect::<Vec<Goldilocks>>()
+                .try_into()
+                .unwrap()
+        })
+        .collect::<Vec<[Goldilocks; WIDTH]>>();
+
+    let internal_constants = plonky2::hash::poseidon2::config::INTERNAL_CONSTANTS
+        .iter()
+        .map(|&x| Goldilocks::from_canonical_u64(x))
+        .collect::<Vec<Goldilocks>>();
+
+    let poseidon = P3Poseidon2::<
+        Goldilocks,
+        Poseidon2ExternalMatrixGeneral,
+        DiffusionMatrixGoldilocks,
+        WIDTH,
+        D,
+    >::new(
+        ROUNDS_F,
+        external_constants,
+        external_linear_layer,
+        ROUNDS_P,
+        internal_constants,
+        internal_linear_layer,
+    );
+
+    c.bench_function("plonky3's poseidon2", |b| {
+        b.iter_batched(
+            || {
+                let mut state = [Goldilocks::zero(); WIDTH];
+                state.iter_mut().for_each(|item| {
+                    *item = Goldilocks::from_canonical_u64(rand::random::<u64>());
+                });
+                state
+            },
+            |mut state| {
+                poseidon.permute_mut(&mut state);
+                state
+            },
+            BatchSize::SmallInput,
+        )
+    });
+}
+
 fn criterion_benchmark(c: &mut Criterion) {
     bench_poseidon::<GoldilocksField>(c);
+    bench_poseidon2::<GoldilocksField>(c);
+    bench_p3_poseidon2(c);
     bench_keccak::<GoldilocksField>(c);
 }
 

plonky2/examples/fibonacci.rs

@@ -1,27 +1,38 @@
+use std::time::Instant;
+
 use anyhow::Result;
+use log::Level;
 use plonky2::field::types::Field;
 use plonky2::iop::witness::{PartialWitness, WitnessWrite};
 use plonky2::plonk::circuit_builder::CircuitBuilder;
 use plonky2::plonk::circuit_data::CircuitConfig;
-use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig};
+use plonky2::plonk::config::{GenericConfig, Poseidon2GoldilocksConfig, PoseidonGoldilocksConfig};
+use plonky2::plonk::prover::prove;
+use plonky2::util::timing::TimingTree;
 
 /// An example of using Plonky2 to prove a statement of the form
 /// "I know the 100th element of the Fibonacci sequence, starting with constants a and b."
 /// When a == 0 and b == 1, this is proving knowledge of the 100th (standard) Fibonacci number.
 fn main() -> Result<()> {
+    env_logger::Builder::from_default_env()
+        .filter_level(log::LevelFilter::Debug)
+        .init();
+    work::<PoseidonGoldilocksConfig>()?;
+    work::<Poseidon2GoldilocksConfig>()
+}
+
+fn work<C: GenericConfig<2>>() -> Result<()> {
     const D: usize = 2;
-    type C = PoseidonGoldilocksConfig;
-    type F = <C as GenericConfig<D>>::F;
 
     let config = CircuitConfig::standard_recursion_config();
-    let mut builder = CircuitBuilder::<F, D>::new(config);
+    let mut builder = CircuitBuilder::<C::F, D>::new(config);
 
     // The arithmetic circuit.
     let initial_a = builder.add_virtual_target();
     let initial_b = builder.add_virtual_target();
     let mut prev_target = initial_a;
     let mut cur_target = initial_b;
-    for _ in 0..99 {
+    for _ in 0..999999 {
         let temp = builder.add(prev_target, cur_target);
         prev_target = cur_target;
         cur_target = temp;
@@ -33,17 +44,29 @@ fn main() -> Result<()> {
     builder.register_public_input(cur_target);
 
     // Provide initial values.
+    let timer1 = Instant::now();
     let mut pw = PartialWitness::new();
-    pw.set_target(initial_a, F::ZERO)?;
-    pw.set_target(initial_b, F::ONE)?;
+    pw.set_target(initial_a, C::F::ZERO)?;
+    pw.set_target(initial_b, C::F::ONE)?;
 
     let data = builder.build::<C>();
-    let proof = data.prove(pw)?;
+    let timer2 = Instant::now();
+
+    // Create a TimingTree to track detailed timing information
+    let mut timing = TimingTree::new("prove", Level::Debug);
+    let proof = prove::<C::F, C, D>(&data.prover_only, &data.common, pw, &mut timing)?;
+    let timer3 = Instant::now();
+
+    // Print the timing tree
+    timing.print();
 
     println!(
         "100th Fibonacci number mod |F| (starting with {}, {}) is: {}",
         proof.public_inputs[0], proof.public_inputs[1], proof.public_inputs[2]
     );
 
+    println!("Build time: {:?}", timer2.duration_since(timer1));
+    println!("Prove time: {:?}", timer3.duration_since(timer2));
+
     data.verify(proof)
 }

plonky2/src/batch_fri/oracle.rs

@@ -460,7 +460,7 @@ mod test {
             &fri_instances,
             &fri_openings,
             &fri_challenges,
-            &[merkle_cap.clone()],
+            std::slice::from_ref(&merkle_cap),
             &proof,
             &fri_params,
         )?;

plonky2/src/fri/validate_shape.rs

@@ -20,7 +20,7 @@ where
     F: RichField + Extendable<D>,
     C: GenericConfig<D, F = F>,
 {
-    validate_batch_fri_proof_shape::<F, C, D>(proof, &[instance.clone()], params)
+    validate_batch_fri_proof_shape::<F, C, D>(proof, std::slice::from_ref(instance), params)
 }
 
 pub(crate) fn validate_batch_fri_proof_shape<F, C, const D: usize>(

plonky2/src/gadgets/interpolation.rs

@@ -86,7 +86,7 @@ mod tests {
 
         let value_targets = values
             .iter()
-            .map(|&v| (builder.constant_extension(v)))
+            .map(|&v| builder.constant_extension(v))
             .collect::<Vec<_>>();
 
         let zt = builder.constant_extension(z);

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

@@ -58,7 +58,7 @@ const_assert!(check_mds_matrix());
 
 /// Addition modulo ORDER accounting for wraparound. Correct only when a + b < 2**64 + ORDER.
 #[inline(always)]
-unsafe fn add_with_wraparound(a: u64, b: u64) -> u64 {
+pub unsafe fn add_with_wraparound(a: u64, b: u64) -> u64 {
     let res: u64;
     let adj: u64;
     asm!(
@@ -872,12 +872,12 @@ unsafe fn partial_rounds(
 */
 
 #[inline(always)]
-fn unwrap_state(state: [GoldilocksField; 12]) -> [u64; 12] {
+pub(crate) fn unwrap_state(state: [GoldilocksField; 12]) -> [u64; 12] {
     state.map(|s| s.0)
 }
 
 #[inline(always)]
-fn wrap_state(state: [u64; 12]) -> [GoldilocksField; 12] {
+pub(crate) fn wrap_state(state: [u64; 12]) -> [GoldilocksField; 12] {
     state.map(GoldilocksField)
 }
 

plonky2/src/hash/merkle_proofs.rs

@@ -59,7 +59,7 @@ pub fn verify_merkle_proof_to_cap<F: RichField, H: Hasher<F>>(
     proof: &MerkleProof<F, H>,
 ) -> Result<()> {
     verify_batch_merkle_proof_to_cap(
-        &[leaf_data.clone()],
+        &[leaf_data],
         &[proof.siblings.len()],
         leaf_index,
         merkle_cap,

plonky2/src/hash/mod.rs

@@ -10,4 +10,5 @@ pub mod merkle_proofs;
 pub mod merkle_tree;
 pub mod path_compression;
 pub mod poseidon;
+pub mod poseidon2;
 pub mod poseidon_goldilocks;