Add poseidon hash benchmark + optim (#300)

This commit is contained in:
Sydhds
2025-04-18 15:08:43 +02:00
committed by GitHub
parent c576af8e62
commit 84d9799d09
6 changed files with 94 additions and 21 deletions

View File

@@ -23,7 +23,7 @@ static POSEIDON: Lazy<Poseidon<Fr>> = Lazy::new(|| Poseidon::<Fr>::from(&ROUND_P
pub fn poseidon_hash(input: &[Fr]) -> Fr {
POSEIDON
.hash(input.to_vec())
.hash(input)
.expect("hash with fixed input size can't fail")
}

View File

@@ -39,3 +39,7 @@ pmtree-ft = ["pmtree"]
[[bench]]
name = "merkle_tree_benchmark"
harness = false
[[bench]]
name = "poseidon_benchmark"
harness = false

View File

@@ -0,0 +1,65 @@
use ark_bn254::Fr;
use criterion::{
black_box, criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion, Throughput,
};
use zerokit_utils::Poseidon;
const ROUND_PARAMS: [(usize, usize, usize, usize); 8] = [
(2, 8, 56, 0),
(3, 8, 57, 0),
(4, 8, 56, 0),
(5, 8, 60, 0),
(6, 8, 60, 0),
(7, 8, 63, 0),
(8, 8, 64, 0),
(9, 8, 63, 0),
];
pub fn poseidon_benchmark(c: &mut Criterion) {
let hasher = Poseidon::<Fr>::from(&ROUND_PARAMS);
let mut group = c.benchmark_group("poseidon Fr");
for size in [10u32, 100, 1000].iter() {
group.throughput(Throughput::Elements(*size as u64));
group.bench_with_input(BenchmarkId::new("Array hash", size), size, |b, &size| {
b.iter_batched(
// Setup: create values for each benchmark iteration
|| {
let mut values = Vec::with_capacity(size as usize);
for i in 0..size {
values.push([Fr::from(i)]);
}
values
},
// Actual benchmark
|values| {
for v in values.iter() {
let _ = hasher.hash(black_box(&v[..]));
}
},
BatchSize::SmallInput,
)
});
}
// Benchmark single hash operation separately
group.bench_function("Single hash", |b| {
let input = [Fr::from(u64::MAX)];
b.iter(|| {
let _ = hasher.hash(black_box(&input[..]));
})
});
group.finish();
}
criterion_group! {
name = benches;
config = Criterion::default()
.warm_up_time(std::time::Duration::from_millis(500))
.measurement_time(std::time::Duration::from_secs(4))
.sample_size(20);
targets = poseidon_benchmark
}
criterion_main!(benches);

View File

@@ -54,24 +54,24 @@ impl<F: PrimeField> Poseidon<F> {
}
}
pub fn get_parameters(&self) -> Vec<RoundParameters<F>> {
self.round_params.clone()
pub fn get_parameters(&self) -> &Vec<RoundParameters<F>> {
&self.round_params
}
pub fn ark(&self, state: &mut [F], c: &[F], it: usize) {
for i in 0..state.len() {
state[i] += c[it + i];
}
state.iter_mut().enumerate().for_each(|(i, elem)| {
*elem += c[it + i];
});
}
pub fn sbox(&self, n_rounds_f: usize, n_rounds_p: usize, state: &mut [F], i: usize) {
if (i < n_rounds_f / 2) || (i >= n_rounds_f / 2 + n_rounds_p) {
for current_state in &mut state.iter_mut() {
state.iter_mut().for_each(|current_state| {
let aux = *current_state;
*current_state *= *current_state;
*current_state *= *current_state;
*current_state *= aux;
}
})
} else {
let aux = state[0];
state[0] *= state[0];
@@ -82,17 +82,18 @@ impl<F: PrimeField> Poseidon<F> {
pub fn mix_2(&self, state: &[F], m: &[Vec<F>], state_2: &mut [F]) {
for i in 0..state.len() {
state_2[i] = F::ZERO;
for (j, state_item) in state.iter().enumerate() {
let mut mij = m[i][j];
mij *= state_item;
state_2[i] += mij;
// Cache the row reference
let row = &m[i];
let mut acc = F::ZERO;
for j in 0..state.len() {
acc += row[j] * state[j];
}
state_2[i] = acc;
}
}
pub fn hash(&self, inp: Vec<F>) -> Result<F, String> {
// Note that the rate t becomes input length + 1, hence for length N we pick parameters with T = N + 1
pub fn hash(&self, inp: &[F]) -> Result<F, String> {
// Note that the rate t becomes input length + 1; hence for length N we pick parameters with T = N + 1
let t = inp.len() + 1;
// We seek the index (Poseidon's round_params is an ordered vector) for the parameters corresponding to t
@@ -105,8 +106,8 @@ impl<F: PrimeField> Poseidon<F> {
let param_index = param_index.unwrap();
let mut state = vec![F::ZERO; t];
state[1..].clone_from_slice(&inp);
let mut state_2 = vec![F::ZERO; state.len()];
let mut state_2 = state.clone();
state[1..].clone_from_slice(inp);
for i in 0..(self.round_params[param_index].n_rounds_f
+ self.round_params[param_index].n_rounds_p)

View File

@@ -3530,7 +3530,8 @@ mod test {
{
// We check if the round constants and matrices correspond to the one generated when instantiating Poseidon with ROUND_PARAMS
let (loaded_c, loaded_m) = load_constants();
let poseidon_parameters = Poseidon::<Fr>::from(&ROUND_PARAMS).get_parameters();
let poseidon_hasher = Poseidon::<Fr>::from(&ROUND_PARAMS);
let poseidon_parameters = poseidon_hasher.get_parameters();
for i in 0..poseidon_parameters.len() {
assert_eq!(loaded_c[i], poseidon_parameters[i].c);
assert_eq!(loaded_m[i], poseidon_parameters[i].m);

View File

@@ -57,16 +57,17 @@ mod test {
),
]);
// map (key: what to hash, value: expected value)
for (k, v) in map.into_iter() {
let hasher = Poseidon::<Fr>::from(&ROUND_PARAMS);
let h = hasher.hash([k].to_vec());
let h = hasher.hash(&[k]);
assert_eq!(h.unwrap(), v);
}
}
#[test]
fn test_poseidon_hash_multi() {
// All hashes done in a merkle tree (with leaves: [0, 1, 2, 3, 4, 5? 6, 7])
// All hashes done in a merkle tree (with leaves: [0, 1, 2, 3, 4, 5, 6, 7])
// ~ leaves
let fr_0 = Fr::ZERO;
@@ -110,6 +111,7 @@ mod test {
)
.unwrap();
// map (key: what to hash, value: expected value)
let map = HashMap::from([
((fr_0, fr_1), fr_0_1),
((fr_2, fr_3), fr_2_3),
@@ -122,7 +124,7 @@ mod test {
for (k, v) in map.into_iter() {
let hasher = Poseidon::<Fr>::from(&ROUND_PARAMS);
let h = hasher.hash([k.0, k.1].to_vec());
let h = hasher.hash(&[k.0, k.1]);
assert_eq!(h.unwrap(), v);
}
}