simplify poseidon permutation calls

Author: TomWambsgans

Cargo.lock

@@ -21,7 +21,7 @@ Details on how to prove AIR constraints in the multilinear settings are describe
 
 ## Benchmarks
 
-cpu: cpu: i9-12900H, ram: 32 gb
+cpu: i9-12900H, ram: 32 gb
 
 > TLDR: Slow, **but there is hope** (cf [TODO](TODO.md))
 

README.md

@@ -15,7 +15,6 @@ p3-uni-stark.workspace = true
 p3-matrix.workspace = true
 p3-util.workspace = true
 multilinear-toolkit.workspace = true
-p3-symmetric.workspace = true
 p3-koala-bear.workspace = true
 rand.workspace = true
 whir-p3.workspace = true

crates/air/Cargo.toml

@@ -3,7 +3,6 @@ use multilinear_toolkit::prelude::*;
 use p3_air::BaseAir;
 use p3_field::PrimeField64;
 use p3_koala_bear::{KoalaBear, QuinticExtensionFieldKB};
-use p3_symmetric::Permutation;
 use packed_pcs::{
     ColDims, packed_pcs_commit, packed_pcs_global_statements_for_prover,
     packed_pcs_global_statements_for_verifier, packed_pcs_parse_commitment,
@@ -16,7 +15,7 @@ use utils::{
     MyChallenger, Poseidon16Air, Poseidon24Air, build_poseidon_16_air,
     build_poseidon_16_air_packed, build_poseidon_24_air, build_poseidon_24_air_packed,
     build_prover_state, build_verifier_state, generate_trace_poseidon_16,
-    generate_trace_poseidon_24, get_poseidon16, get_poseidon24, init_tracing,
+    generate_trace_poseidon_24, init_tracing, poseidon16_permute, poseidon24_permute,
 };
 use whir_p3::{
     FoldingFactor, SecurityAssumption, WhirConfig, WhirConfigBuilder, precompute_dft_twiddles,
@@ -111,11 +110,11 @@ fn prepare_poseidon(config: &Poseidon2Config) -> PoseidonSetup {
 
     assert_eq!(
         &witness_matrix_16.values[n_columns_16 - 16..n_columns_16 - 8],
-        &get_poseidon16().permute(witness_matrix_16.values[0..16].try_into().unwrap())[..8]
+        &poseidon16_permute(witness_matrix_16.values[0..16].try_into().unwrap())[..8]
     );
     assert_eq!(
         &witness_matrix_24.values[n_columns_24 - 8..n_columns_24],
-        &get_poseidon24().permute(witness_matrix_24.values[0..24].try_into().unwrap())[16..]
+        &poseidon24_permute(witness_matrix_24.values[0..24].try_into().unwrap())[16..]
     );
 
     let witness_matrix_16_transposed = witness_matrix_16.transpose();

crates/air/src/examples/prove_poseidon2.rs

@@ -1,7 +1,6 @@
 use lean_compiler::*;
 use lean_vm::*;
-use p3_symmetric::Permutation;
-use utils::{get_poseidon16, get_poseidon24};
+use utils::{poseidon16_permute, poseidon24_permute};
 
 const DEFAULT_NO_VEC_RUNTIME_MEMORY: usize = 1 << 15;
 
@@ -327,16 +326,15 @@ fn test_mini_program_3() {
         return;
     }
    "#;
-    let mut public_input: [F; 16] = (0..16).map(F::new).collect::<Vec<F>>().try_into().unwrap();
+    let public_input: [F; 16] = (0..16).map(F::new).collect::<Vec<F>>().try_into().unwrap();
     compile_and_run(
         program.to_string(),
         (&public_input, &[]),
         DEFAULT_NO_VEC_RUNTIME_MEMORY,
         false,
     );
 
-    get_poseidon16().permute_mut(&mut public_input);
-    let _ = public_input;
+    let _ = dbg!(poseidon16_permute(public_input));
 }
 
 #[test]
@@ -355,16 +353,15 @@ fn test_mini_program_4() {
         return;
     }
    "#;
-    let mut public_input: [F; 24] = (0..24).map(F::new).collect::<Vec<F>>().try_into().unwrap();
+    let public_input: [F; 24] = (0..24).map(F::new).collect::<Vec<F>>().try_into().unwrap();
     compile_and_run(
         program.to_string(),
         (&public_input, &[]),
         DEFAULT_NO_VEC_RUNTIME_MEMORY,
         false,
     );
 
-    get_poseidon24().permute_mut(&mut public_input);
-    dbg!(&public_input[16..]);
+    dbg!(&poseidon24_permute(public_input)[16..]);
 }
 
 #[test]

crates/lean_compiler/tests/test_compiler.rs

@@ -13,9 +13,8 @@ use crate::witness::{
 };
 use crate::{CodeAddress, HintExecutionContext, SourceLineNumber};
 use p3_field::PrimeCharacteristicRing;
-use p3_symmetric::Permutation;
 use std::collections::BTreeMap;
-use utils::{get_poseidon16, get_poseidon24, pretty_integer};
+use utils::{poseidon16_permute, poseidon24_permute, pretty_integer};
 use xmss::{Poseidon16History, Poseidon24History};
 
 /// Number of instructions to show in stack trace
@@ -51,10 +50,10 @@ pub fn build_public_memory(public_input: &[F]) -> Vec<F> {
 
     public_memory
         [POSEIDON_16_NULL_HASH_PTR * VECTOR_LEN..(POSEIDON_16_NULL_HASH_PTR + 2) * VECTOR_LEN]
-        .copy_from_slice(&get_poseidon16().permute([F::ZERO; 16]));
+        .copy_from_slice(&poseidon16_permute([F::ZERO; 16]));
     public_memory
         [POSEIDON_24_NULL_HASH_PTR * VECTOR_LEN..(POSEIDON_24_NULL_HASH_PTR + 1) * VECTOR_LEN]
-        .copy_from_slice(&get_poseidon24().permute([F::ZERO; 24])[16..]);
+        .copy_from_slice(&poseidon24_permute([F::ZERO; 24])[16..]);
     public_memory
 }
 

crates/lean_vm/src/execution/runner.rs

@@ -12,10 +12,9 @@ use crate::witness::{
 };
 use multilinear_toolkit::prelude::*;
 use p3_field::{BasedVectorSpace, PrimeCharacteristicRing, dot_product};
-use p3_symmetric::Permutation;
 use p3_util::log2_ceil_usize;
 use std::fmt::{Display, Formatter};
-use utils::{ToUsize, get_poseidon16, get_poseidon24};
+use utils::{ToUsize, poseidon16_permute, poseidon24_permute};
 
 /// Complete set of VM instruction types with comprehensive operation support
 #[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
@@ -236,11 +235,7 @@ impl Instruction {
                         *ctx.n_poseidon16_precomputed_used += 1;
                         precomputed.1
                     }
-                    _ => {
-                        let mut output = input;
-                        get_poseidon16().permute_mut(&mut output);
-                        output
-                    }
+                    _ => poseidon16_permute(input),
                 };
 
                 let res0: [F; VECTOR_LEN] = output[..VECTOR_LEN].try_into().unwrap();
@@ -292,8 +287,7 @@ impl Instruction {
                         precomputed.1
                     }
                     _ => {
-                        let mut output = input;
-                        get_poseidon24().permute_mut(&mut output);
+                        let output = poseidon24_permute(input);
                         output[2 * VECTOR_LEN..].try_into().unwrap()
                     }
                 };

crates/lean_vm/src/isa/instruction.rs

@@ -14,6 +14,7 @@ p3-matrix.workspace = true
 p3-challenger.workspace = true
 p3-koala-bear.workspace = true
 tracing-forest.workspace = true
+p3-symmetric.workspace = true
 p3-poseidon2.workspace = true
 p3-poseidon2-air.workspace = true
 tracing-subscriber.workspace = true

crates/utils/Cargo.toml

@@ -19,6 +19,7 @@ use p3_poseidon2_air::p24::Poseidon2Air24;
 use p3_poseidon2_air::p24::RoundConstants24;
 use p3_poseidon2_air::p24::generate_trace_rows_24;
 use p3_poseidon2_air::{p16, p24};
+use p3_symmetric::Permutation;
 
 pub type Poseidon16 = Poseidon2KoalaBear<16>;
 pub type Poseidon24 = Poseidon2KoalaBear<24>;
@@ -81,7 +82,8 @@ pub type MyRoundConstants24<F> = RoundConstants24<F, 24, HALF_FULL_ROUNDS_24, PA
 
 static POSEIDON16_INSTANCE: OnceLock<Poseidon16> = OnceLock::new();
 
-pub fn get_poseidon16() -> &'static Poseidon16 {
+#[inline(always)]
+pub(crate) fn get_poseidon16() -> &'static Poseidon16 {
     POSEIDON16_INSTANCE.get_or_init(|| {
         let round_constants = build_poseidon16_constants();
         let external_constants = ExternalLayerConstants::new(
@@ -95,9 +97,15 @@ pub fn get_poseidon16() -> &'static Poseidon16 {
     })
 }
 
+#[inline(always)]
+pub fn poseidon16_permute(input: [KoalaBear; 16]) -> [KoalaBear; 16] {
+    get_poseidon16().permute(input)
+}
+
 static POSEIDON24_INSTANCE: OnceLock<Poseidon24> = OnceLock::new();
 
-pub fn get_poseidon24() -> &'static Poseidon24 {
+#[inline(always)]
+pub(crate) fn get_poseidon24() -> &'static Poseidon24 {
     POSEIDON24_INSTANCE.get_or_init(|| {
         let round_constants = build_poseidon24_constants();
         let external_constants = ExternalLayerConstants::new(
@@ -111,6 +119,11 @@ pub fn get_poseidon24() -> &'static Poseidon24 {
     })
 }
 
+#[inline(always)]
+pub fn poseidon24_permute(input: [KoalaBear; 24]) -> [KoalaBear; 24] {
+    get_poseidon24().permute(input)
+}
+
 pub fn build_poseidon_16_air() -> Poseidon16Air<KoalaBear> {
     Poseidon16Air::new(build_poseidon16_constants())
 }

crates/utils/src/poseidon2.rs

@@ -11,5 +11,4 @@ p3-koala-bear.workspace = true
 p3-field.workspace = true
 rand.workspace = true
 utils.workspace = true
-p3-symmetric.workspace = true
 p3-util.workspace = true