simd: apply packing for tree leaves

Author: tcoratger

Cargo.toml

@@ -39,10 +39,10 @@ dashmap = "6.1.0"
 serde = { version = "1.0", features = ["derive", "alloc"] }
 thiserror = "2.0"
 
-p3-field = { git = "https://github.com/Plonky3/Plonky3.git", rev = "2117e4b" }
-p3-baby-bear = { git = "https://github.com/Plonky3/Plonky3.git", rev = "2117e4b" }
-p3-koala-bear = { git = "https://github.com/Plonky3/Plonky3.git", rev = "2117e4b" }
-p3-symmetric = { git = "https://github.com/Plonky3/Plonky3.git", rev = "2117e4b" }
+p3-field = { git = "https://github.com/Plonky3/Plonky3.git", rev = "a33a312" }
+p3-baby-bear = { git = "https://github.com/Plonky3/Plonky3.git", rev = "a33a312" }
+p3-koala-bear = { git = "https://github.com/Plonky3/Plonky3.git", rev = "a33a312" }
+p3-symmetric = { git = "https://github.com/Plonky3/Plonky3.git", rev = "a33a312" }
 
 [dev-dependencies]
 criterion = "0.7"

src/lib.rs

@@ -1,3 +1,4 @@
+use p3_field::Field;
 use p3_koala_bear::{
     KoalaBear, Poseidon2KoalaBear, default_koalabear_poseidon2_16, default_koalabear_poseidon2_24,
 };
@@ -11,6 +12,7 @@ pub const TWEAK_SEPARATOR_FOR_TREE_HASH: u8 = 0x01;
 pub const TWEAK_SEPARATOR_FOR_CHAIN_HASH: u8 = 0x00;
 
 type F = KoalaBear;
+pub(crate) type PackedF = <F as Field>::Packing;
 
 pub(crate) mod hypercube;
 pub(crate) mod inc_encoding;

src/signature/generalized_xmss.rs

@@ -205,38 +205,15 @@ where
     let chain_length = IE::BASE;
 
     // the range of epochs covered by that bottom tree
-    let epoch_range_start = bottom_tree_index * leafs_per_bottom_tree;
-    let epoch_range_end = epoch_range_start + leafs_per_bottom_tree;
-    let epoch_range = epoch_range_start..epoch_range_end;
-
-    // parallelize the chain ends hash computation for each epoch in the interval for that bottom tree
-    let chain_ends_hashes = epoch_range
-        .into_par_iter()
-        .map(|epoch| {
-            // each epoch has a number of chains
-            // parallelize the chain ends computation for each chain
-            let chain_ends = (0..num_chains)
-                .into_par_iter()
-                .map(|chain_index| {
-                    // each chain start is just a PRF evaluation
-                    let start =
-                        PRF::get_domain_element(prf_key, epoch as u32, chain_index as u64).into();
-                    // walk the chain to get the public chain end
-                    chain::<TH>(
-                        parameter,
-                        epoch as u32,
-                        chain_index as u8,
-                        0,
-                        chain_length - 1,
-                        &start,
-                    )
-                })
-                .collect::<Vec<_>>();
-            // build hash of chain ends / public keys
-            TH::apply(parameter, &TH::tree_tweak(0, epoch as u32), &chain_ends)
-        })
+    let epoch_start = bottom_tree_index * leafs_per_bottom_tree;
+    let epochs: Vec<u32> = (epoch_start..epoch_start + leafs_per_bottom_tree)
+        .map(|e| e as u32)
         .collect();
 
+    // Compute chain ends for all epochs.
+    let chain_ends_hashes =
+        TH::compute_tree_leaves::<PRF>(prf_key, parameter, &epochs, num_chains, chain_length);
+
     // now that we have the hashes of all chain ends (= leafs of our tree), we can compute the bottom tree
     HashSubTree::new_bottom_tree(
         LOG_LIFETIME,

src/symmetric/message_hash/poseidon.rs

@@ -164,7 +164,7 @@ where
             .copied()
             .collect();
 
-        let hash_fe = poseidon_compress::<_, 24, HASH_LEN_FE>(&perm, &combined_input_vec);
+        let hash_fe = poseidon_compress::<F, _, 24, HASH_LEN_FE>(&perm, &combined_input_vec);
 
         // decode field elements into chunks and return them
         decode_to_chunks::<DIMENSION, BASE, HASH_LEN_FE>(&hash_fe).to_vec()

src/symmetric/message_hash/top_level_poseidon.rs

@@ -159,7 +159,7 @@ where
                 .collect();
 
             let iteration_pos_output =
-                poseidon_compress::<_, 24, POS_OUTPUT_LEN_PER_INV_FE>(&perm, &combined_input);
+                poseidon_compress::<F, _, 24, POS_OUTPUT_LEN_PER_INV_FE>(&perm, &combined_input);
 
             pos_outputs[i * POS_OUTPUT_LEN_PER_INV_FE..(i + 1) * POS_OUTPUT_LEN_PER_INV_FE]
                 .copy_from_slice(&iteration_pos_output);

src/symmetric/tweak_hash.rs

@@ -1,6 +1,9 @@
 use rand::Rng;
+use rayon::prelude::*;
 use serde::{Serialize, de::DeserializeOwned};
 
+use crate::symmetric::prf::Pseudorandom;
+
 /// Trait to model a tweakable hash function.
 /// Such a function takes a public parameter, a tweak, and a
 /// message to be hashed. The tweak should be understood as an
@@ -14,8 +17,13 @@ use serde::{Serialize, de::DeserializeOwned};
 /// to obtain distinct tweaks for applications in chains and
 /// applications in Merkle trees.
 pub trait TweakableHash {
+    /// Public parameter type for the hash function
     type Parameter: Copy + Sized + Send + Sync + Serialize + DeserializeOwned;
+
+    /// Tweak type for domain separation
     type Tweak;
+
+    /// Domain element type (defines output and input types to the hash)
     type Domain: Copy + PartialEq + Sized + Send + Sync + Serialize + DeserializeOwned;
 
     /// Generates a random public parameter.
@@ -39,8 +47,50 @@ pub trait TweakableHash {
         message: &[Self::Domain],
     ) -> Self::Domain;
 
-    /// Function to check internal consistency of any given parameters
-    /// For testing only, and expected to panic if something is wrong.
+    /// Computes bottom tree leaves by walking hash chains for multiple epochs.
+    ///
+    /// This method has a default scalar implementation that processes epochs in parallel.
+    fn compute_tree_leaves<PRF>(
+        prf_key: &PRF::Key,
+        parameter: &Self::Parameter,
+        epochs: &[u32],
+        num_chains: usize,
+        chain_length: usize,
+    ) -> Vec<Self::Domain>
+    where
+        PRF: Pseudorandom,
+        PRF::Domain: Into<Self::Domain>,
+        Self: Sized,
+    {
+        // Default scalar implementation: process each epoch in parallel
+        epochs
+            .par_iter()
+            .map(|&epoch| {
+                // For each epoch, walk all chains in parallel
+                let chain_ends: Vec<_> = (0..num_chains)
+                    .into_par_iter()
+                    .map(|chain_index| {
+                        let start =
+                            PRF::get_domain_element(prf_key, epoch, chain_index as u64).into();
+                        chain::<Self>(
+                            parameter,
+                            epoch,
+                            chain_index as u8,
+                            0,
+                            chain_length - 1,
+                            &start,
+                        )
+                    })
+                    .collect();
+                // Hash all chain ends together to get the leaf
+                Self::apply(parameter, &Self::tree_tweak(0, epoch), &chain_ends)
+            })
+            .collect()
+    }
+
+    /// Function to check internal consistency of any given parameters.
+    ///
+    /// This is for testing only and is expected to panic if something is wrong.
     #[cfg(test)]
     fn internal_consistency_check();
 }
@@ -77,7 +127,6 @@ pub mod poseidon;
 
 #[cfg(test)]
 mod tests {
-
     use crate::symmetric::tweak_hash::poseidon::PoseidonTweak44;
 
     use super::*;