Patch plonky3 to use Powdr accelerated Poseidon2

plonky3/Cargo.toml

@@ -22,6 +22,9 @@ nightly-features = [
   "p3-mersenne-31/nightly-features",
   "p3-poseidon2/nightly-features",
 ]
+# As a guest in powdr, use accelerated operations.
+# Silently ignored if target is not zkvm.
+powdr-accel = []
 
 [dependencies]
 powdr-ast.workspace = true
@@ -67,3 +70,12 @@ serde = { version = "1.0", default-features = false, features = [
   "alloc",
 ] }
 
+[target.'cfg(all(target_os = "zkvm", target_arch = "riscv32"))'.dependencies]
+powdr-riscv-runtime = { path = "../riscv-runtime", features = [
+  "std",
+  "getrandom",
+  "allow_fake_rand",
+] }
+indexmap = { version = "1.9.3", features = [
+  "std",
+] }

plonky3/src/bin/gen_poseidon_consts.rs

@@ -8,7 +8,7 @@ use p3_monty_31::DiffusionMatrixParameters;
 use p3_poseidon::Poseidon;
 use p3_poseidon2::Poseidon2ExternalMatrixGeneral;
 use p3_symmetric::Permutation;
-use powdr_plonky3::{goldilocks, poseidon2};
+use powdr_plonky3::poseidon2;
 use rand::{distributions::Standard, Rng, SeedableRng};
 
 fn main() {
@@ -194,7 +194,7 @@ fn poseidon2_bb_consts() {
 
 fn poseidon2_gl_consts() {
     use p3_field::PrimeField64;
-    use powdr_plonky3::goldilocks::{ROUNDS_F, ROUNDS_P, WIDTH};
+    use powdr_plonky3::poseidon2::goldilocks::{PERM, ROUNDS_F, ROUNDS_P, WIDTH};
 
     println!("EXTERNAL_CONSTANTS = [");
     let ec = poseidon2::external_constants::<Goldilocks, WIDTH>(*ROUNDS_F);
@@ -230,15 +230,6 @@ fn poseidon2_gl_consts() {
     );
 
     println!("\n\nTESTS:");
-    let poseidon2 = goldilocks::Perm::new(
-        *goldilocks::ROUNDS_F,
-        poseidon2::external_constants::<Goldilocks, WIDTH>(*ROUNDS_F),
-        Poseidon2ExternalMatrixGeneral,
-        *goldilocks::ROUNDS_P,
-        poseidon2::internal_constants::<Goldilocks>(*ROUNDS_P),
-        p3_goldilocks::DiffusionMatrixGoldilocks,
-    );
-
     let test_vectors = test_vectors::<Goldilocks, WIDTH>();
 
     let mut first_test_vector = test_vectors[0];
@@ -255,7 +246,7 @@ fn poseidon2_gl_consts() {
             );
         }
         println!("\n        poseidon2 0, 0;\n");
-        poseidon2.permute_mut(&mut test_vector);
+        PERM.permute_mut(&mut test_vector);
         for (i, val) in test_vector[..8].iter().enumerate() {
             println!("        assert_eq {}, {val};", i * 8);
         }
@@ -274,7 +265,7 @@ fn poseidon2_gl_consts() {
         );
     }
     println!("\n        poseidon2 100, 104;\n");
-    poseidon2.permute_mut(&mut first_test_vector);
+    PERM.permute_mut(&mut first_test_vector);
     for (i, val) in first_test_vector[..8].iter().enumerate() {
         println!("        assert_eq {}, {val};", 104 + i * 8);
     }

plonky3/src/params/goldilocks.rs

@@ -3,37 +3,31 @@
 //! Inspired from [this example](https://github.com/Plonky3/Plonky3/blob/6a1b0710fdf85136d0fdd645b92933615867740a/keccak-air/examples/prove_goldilocks_keccak.rs#L57)
 //! (But using Poseidon2 instead of Poseidon)
 
-use lazy_static::lazy_static;
-use p3_poseidon2::{poseidon2_round_numbers_128, Poseidon2, Poseidon2ExternalMatrixGeneral};
-
-use crate::params::{poseidon2, Challenger, FieldElementMap, Plonky3Field};
+use crate::{
+    params::{Challenger, FieldElementMap, Plonky3Field},
+    poseidon2::goldilocks::{Permutation, PERM, WIDTH},
+};
 use p3_challenger::DuplexChallenger;
 use p3_commit::ExtensionMmcs;
 use p3_dft::Radix2DitParallel;
 use p3_field::{extension::BinomialExtensionField, AbstractField, Field, PrimeField64};
 use p3_fri::{FriConfig, TwoAdicFriPcs};
-use p3_goldilocks::{DiffusionMatrixGoldilocks, Goldilocks};
+use p3_goldilocks::Goldilocks;
 use p3_merkle_tree::MerkleTreeMmcs;
 use p3_symmetric::{PaddingFreeSponge, TruncatedPermutation};
 use p3_uni_stark::StarkConfig;
 use powdr_number::{FieldElement, GoldilocksField, LargeInt};
 
-// From: https://github.com/Plonky3/Plonky3/blob/64e79fe28c51ab35b509c68242256f253b61d612/poseidon2/benches/poseidon2.rs#L31
-const D: u64 = 7;
-pub const WIDTH: usize = 8;
-pub type Perm =
-    Poseidon2<Goldilocks, Poseidon2ExternalMatrixGeneral, DiffusionMatrixGoldilocks, WIDTH, D>;
-
 const DEGREE: usize = 2;
 type FriChallenge = BinomialExtensionField<Goldilocks, DEGREE>;
 
 const RATE: usize = 4;
 const OUT: usize = 4;
-type Hash = PaddingFreeSponge<Perm, WIDTH, RATE, OUT>;
+type Hash = PaddingFreeSponge<Permutation, WIDTH, RATE, OUT>;
 
 const N: usize = 2;
 const CHUNK: usize = 4;
-type Compress = TruncatedPermutation<Perm, N, CHUNK, WIDTH>;
+type Compress = TruncatedPermutation<Permutation, N, CHUNK, WIDTH>;
 
 const DIGEST_ELEMS: usize = 4;
 type ValMmcs = MerkleTreeMmcs<
@@ -44,7 +38,7 @@ type ValMmcs = MerkleTreeMmcs<
     DIGEST_ELEMS,
 >;
 
-pub type FriChallenger = DuplexChallenger<Goldilocks, Perm, WIDTH, RATE>;
+pub type FriChallenger = DuplexChallenger<Goldilocks, Permutation, WIDTH, RATE>;
 type ChallengeMmcs = ExtensionMmcs<Goldilocks, FriChallenge, ValMmcs>;
 type Dft = Radix2DitParallel<Goldilocks>;
 type MyPcs = TwoAdicFriPcs<Goldilocks, Dft, ValMmcs, ChallengeMmcs>;
@@ -53,20 +47,6 @@ const FRI_LOG_BLOWUP: usize = 1;
 const FRI_NUM_QUERIES: usize = 100;
 const FRI_PROOF_OF_WORK_BITS: usize = 16;
 
-lazy_static! {
-    static ref ROUNDS: (usize, usize) = poseidon2_round_numbers_128::<Goldilocks>(WIDTH, D);
-    pub static ref ROUNDS_F: usize = ROUNDS.0;
-    pub static ref ROUNDS_P: usize = ROUNDS.1;
-    pub static ref PERM_GL: Perm = Perm::new(
-        *ROUNDS_F,
-        poseidon2::external_constants(*ROUNDS_F),
-        Poseidon2ExternalMatrixGeneral,
-        *ROUNDS_P,
-        poseidon2::internal_constants(*ROUNDS_P),
-        DiffusionMatrixGoldilocks
-    );
-}
-
 impl FieldElementMap for GoldilocksField {
     type Config = StarkConfig<MyPcs, FriChallenge, FriChallenger>;
 
@@ -79,13 +59,13 @@ impl FieldElementMap for GoldilocksField {
     }
 
     fn get_challenger() -> Challenger<Self> {
-        FriChallenger::new(PERM_GL.clone())
+        FriChallenger::new(PERM.clone())
     }
 
     fn get_config() -> Self::Config {
-        let hash = Hash::new(PERM_GL.clone());
+        let hash = Hash::new(PERM.clone());
 
-        let compress = Compress::new(PERM_GL.clone());
+        let compress = Compress::new(PERM.clone());
 
         let val_mmcs = ValMmcs::new(hash, compress);
 

plonky3/src/params/poseidon2/goldilocks/mod.rs

@@ -0,0 +1,23 @@
+#[cfg(not(all(feature = "powdr-accel", target_os = "zkvm", target_arch = "riscv32")))]
+mod software_impl;
+#[cfg(not(all(feature = "powdr-accel", target_os = "zkvm", target_arch = "riscv32")))]
+pub use software_impl::*;
+
+#[cfg(all(feature = "powdr-accel", target_os = "zkvm", target_arch = "riscv32"))]
+mod powdr_accel_impl;
+#[cfg(all(feature = "powdr-accel", target_os = "zkvm", target_arch = "riscv32"))]
+pub use powdr_accel_impl::*;
+
+use lazy_static::lazy_static;
+use p3_goldilocks::Goldilocks;
+use p3_poseidon2::poseidon2_round_numbers_128;
+
+// From: https://github.com/Plonky3/Plonky3/blob/64e79fe28c51ab35b509c68242256f253b61d612/poseidon2/benches/poseidon2.rs#L31
+pub const D: u64 = 7;
+pub const WIDTH: usize = 8;
+
+lazy_static! {
+    static ref ROUNDS: (usize, usize) = poseidon2_round_numbers_128::<Goldilocks>(WIDTH, D);
+    pub static ref ROUNDS_F: usize = ROUNDS.0;
+    pub static ref ROUNDS_P: usize = ROUNDS.1;
+}

plonky3/src/params/poseidon2/goldilocks/powdr_accel_impl.rs

@@ -0,0 +1,37 @@
+use lazy_static::lazy_static;
+use p3_field::AbstractField;
+use p3_goldilocks::Goldilocks;
+use p3_symmetric::CryptographicPermutation;
+use powdr_riscv_runtime::{
+    goldilocks::Goldilocks as PowdrGoldilocks,
+    hash::{poseidon2_gl, poseidon2_gl_inplace},
+};
+
+#[derive(Clone, Copy, Debug)]
+pub struct Permutation;
+
+impl p3_symmetric::Permutation<[Goldilocks; 8]> for Permutation {
+    // Both Goldilocks and PowdrGoldilocks are repr(transparent), and both use
+    // canonical representation internally, so it is safe to cast between their
+    // array's references.
+    //
+    // TODO: We are relying on implementation detail. So, ideally, we should
+    // static assert that std::mem::transmute(Goldilocks::one()) == 1u64.
+
+    fn permute(&self, input: [Goldilocks; 8]) -> [Goldilocks; 8] {
+        let input = unsafe { &*(&input as *const _ as *const [PowdrGoldilocks; 8]) };
+        let output = poseidon2_gl(input);
+        // Let's hope the compiler optimizes this into a no-op.
+        output.map(|x| Goldilocks::from_canonical_u64(u64::from(x)))
+    }
+
+    fn permute_mut(&self, data: &mut [Goldilocks; 8]) {
+        let data = unsafe { &mut *(data as *mut _ as *mut [PowdrGoldilocks; 8]) };
+        poseidon2_gl_inplace(data);
+    }
+}
+impl CryptographicPermutation<[Goldilocks; 8]> for Permutation {}
+
+lazy_static! {
+    pub static ref PERM: Permutation = Permutation;
+}

plonky3/src/params/poseidon2/goldilocks/software_impl.rs

@@ -0,0 +1,21 @@
+use crate::poseidon2::{external_constants, internal_constants};
+
+use super::{D, ROUNDS_F, ROUNDS_P, WIDTH};
+use lazy_static::lazy_static;
+use p3_goldilocks::{DiffusionMatrixGoldilocks, Goldilocks};
+use p3_poseidon2::{Poseidon2, Poseidon2ExternalMatrixGeneral};
+
+/// The Poseidon2 permutation type.
+pub type Permutation =
+    Poseidon2<Goldilocks, Poseidon2ExternalMatrixGeneral, DiffusionMatrixGoldilocks, WIDTH, D>;
+
+lazy_static! {
+    pub static ref PERM: Permutation = Permutation::new(
+        *ROUNDS_F,
+        external_constants(*ROUNDS_F),
+        Poseidon2ExternalMatrixGeneral,
+        *ROUNDS_P,
+        internal_constants(*ROUNDS_P),
+        DiffusionMatrixGoldilocks,
+    );
+}

plonky3/src/params/poseidon2.rs → plonky3/src/params/poseidon2/mod.rs

@@ -1,5 +1,7 @@
 //! Poseidon2 constants generation.
 
+pub mod goldilocks;
+
 use alloc::vec::Vec;
 use rand::{distributions::Standard, prelude::Distribution, Rng, SeedableRng};
 

riscv-executor/src/poseidon2_gl.rs

@@ -3,6 +3,6 @@ use p3_symmetric::Permutation;
 
 pub fn poseidon2_gl(input: &[u64; 8]) -> [u64; 8] {
     let mut state = input.map(p3_goldilocks::Goldilocks::from_canonical_u64);
-    powdr_plonky3::goldilocks::PERM_GL.permute_mut(&mut state);
+    powdr_plonky3::poseidon2::goldilocks::PERM.permute_mut(&mut state);
     state.map(|v| v.as_canonical_u64())
 }

riscv/tests/riscv.rs

@@ -209,6 +209,19 @@ fn std_hello_world() {
     verify_riscv_crate(case, Default::default());
 }
 
+#[test]
+fn plonky3_verify() {
+    let case = "plonky3_verify";
+    // Just compile for now
+    // TODO: make it execute like the other tests, and mark it as ignore = "Too slow".
+    let temp_dir = Temp::new_dir().unwrap();
+    powdr_riscv::compile_rust_crate_to_riscv(
+        &format!("tests/riscv_data/{case}/Cargo.toml"),
+        &temp_dir,
+        None,
+    );
+}
+
 #[test]
 #[ignore = "Too slow"]
 fn function_pointer() {

riscv/tests/riscv_data/plonky3_verify/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "plonky3_verify"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+powdr-riscv-syscalls = { path = "../../../../riscv-syscalls" }
+powdr-riscv-runtime = { path = "../../../../riscv-runtime", features = ["std"]}
+powdr-plonky3 = { path = "../../../../plonky3", features = ["powdr-accel"]}
+
+[workspace]

riscv/tests/riscv_data/plonky3_verify/rust-toolchain.toml

@@ -0,0 +1,4 @@
+[toolchain]
+channel = "nightly-2024-08-01"
+targets = ["riscv32im-risc0-zkvm-elf"]
+profile = "minimal"

riscv/tests/riscv_data/plonky3_verify/src/main.rs

@@ -0,0 +1,7 @@
+extern crate powdr_plonky3;
+
+fn main() {
+    // This is just a stub to ensure that the `powdr_plonky3` crate is built and linked correctly.
+    // TODO: perform an actual test.
+    assert_eq!(true, !false);
+}