From 29472e63a6a9fe0fada09cc751de2b61e038f779 Mon Sep 17 00:00:00 2001 From: Manish Kumar Date: Thu, 1 Aug 2024 14:22:17 +0530 Subject: [PATCH] test for correctness of polygon keccak, separate script for running poseidon hash added --- hash/plonky2/bench/build.sh | 1 - hash/plonky2/bench/run.sh | 4 +- hash/plonky2/bench/run_tree.sh | 15 +++++ .../src/bench/keccak256/keccak_polygon.rs | 61 ++++++++++++++++++- hash/plonky2/bench/src/bench/poseidon.rs | 2 +- hash/plonky2/bench/src/main.rs | 1 - hash/plonky2/external/zk_evm | 2 +- 7 files changed, 78 insertions(+), 8 deletions(-) create mode 100755 hash/plonky2/bench/run_tree.sh diff --git a/hash/plonky2/bench/build.sh b/hash/plonky2/bench/build.sh index 7d27c85..7f06aa8 100755 --- a/hash/plonky2/bench/build.sh +++ b/hash/plonky2/bench/build.sh @@ -4,7 +4,6 @@ rustup override set nightly # Build -# RUSTFLAGS=-Ctarget-cpu=native cargo build RUSTFLAGS=-Ctarget-cpu=native cargo build --release --bin plonky2_hash_benchmarks diff --git a/hash/plonky2/bench/run.sh b/hash/plonky2/bench/run.sh index 8d21b7e..82d817e 100755 --- a/hash/plonky2/bench/run.sh +++ b/hash/plonky2/bench/run.sh @@ -1,10 +1,10 @@ #!/bin/bash if [ -z ${ZKBENCH_HASH_TYPE} ]; then -ZKBENCH_HASH_TYPE="poseidon" +ZKBENCH_HASH_TYPE="sha256" fi if [ -z ${ZKBENCH_INPUT_SIZE_BYTES} ]; then -ZKBENCH_INPUT_SIZE_BYTES=4 +ZKBENCH_INPUT_SIZE_BYTES=256 fi echo "Running benchmarks with the following configurations:" diff --git a/hash/plonky2/bench/run_tree.sh b/hash/plonky2/bench/run_tree.sh new file mode 100755 index 0000000..e1e33d6 --- /dev/null +++ b/hash/plonky2/bench/run_tree.sh @@ -0,0 +1,15 @@ +#!/bin/bash +if [ -z ${ZKBENCH_HASH_TYPE_TREE} ]; then +ZKBENCH_HASH_TYPE_TREE="poseidon" +fi + +if [ -z ${ZKBENCH_TREE_DEPTH} ]; then +ZKBENCH_TREE_DEPTH=4 +fi + +echo "Running benchmarks with the following configurations:" +echo "HASH = $ZKBENCH_HASH_TYPE_TREE" +echo "Tree Depth = $ZKBENCH_TREE_DEPTH" + +# Run the benchmarks +./target/release/plonky2_hash_benchmarks $ZKBENCH_HASH_TYPE_TREE $ZKBENCH_TREE_DEPTH \ No newline at end of file diff --git a/hash/plonky2/bench/src/bench/keccak256/keccak_polygon.rs b/hash/plonky2/bench/src/bench/keccak256/keccak_polygon.rs index 8deed2f..3d3282c 100644 --- a/hash/plonky2/bench/src/bench/keccak256/keccak_polygon.rs +++ b/hash/plonky2/bench/src/bench/keccak256/keccak_polygon.rs @@ -25,9 +25,23 @@ use env_logger::{ try_init_from_env }; -const NUM_INPUTS: usize = 25; +#[allow(dead_code)] +/// Number of rounds in a Keccak permutation. +pub(crate) const NUM_ROUNDS: usize = 24; -pub fn keccak_polygon_bench(num_perms: usize) -> Result<()> { +/// Number of 64-bit elements in the Keccak permutation input. +pub(crate) const NUM_INPUTS: usize = 25; + +fn ceil_div(a: usize, b: usize) -> usize { + (a + b - 1) / b +} + +pub fn keccak_polygon_bench(size: usize) -> Result<()> { + // here input in in terms of block + let mut num_perms = 1; + if size > 25 * 64/8 { + num_perms = ceil_div(size, 25 * 64/8) as usize; + } const D: usize = 2; type C = PoseidonGoldilocksConfig; @@ -99,4 +113,47 @@ pub fn keccak_polygon_bench(num_perms: usize) -> Result<()> { fn init_logger() { let _ = try_init_from_env(Env::default().filter_or(DEFAULT_FILTER_ENV, "debug")); +} + +#[cfg(test)] +mod tests { + use anyhow::Result; + use plonky2::field::types::PrimeField64; + use plonky2::plonk::config::{GenericConfig, PoseidonGoldilocksConfig}; + use tiny_keccak::keccakf; + use evm_arithmetization::keccak::columns::reg_output_limb; + use super::*; + const NUM_ROUNDS: usize = 24; + + #[test] + fn keccak_correctness_test() -> Result<()> { + let input: [u64; NUM_INPUTS] = rand::random(); + + const D: usize = 2; + type C = PoseidonGoldilocksConfig; + type F = >::F; + type S = KeccakStark; + + let stark: KeccakStark<>::F, 2> = S::default(); + + let rows = stark.generate_trace_rows(vec![(input, 0)], 8); + let last_row = rows[NUM_ROUNDS - 1]; + let output = (0..NUM_INPUTS) + .map(|i| { + let hi = last_row[reg_output_limb(2 * i + 1)].to_canonical_u64(); + let lo = last_row[reg_output_limb(2 * i)].to_canonical_u64(); + (hi << 32) | lo + }) + .collect::>(); + + let expected = { + let mut state = input; + keccakf(&mut state); + state + }; + + assert_eq!(output, expected); + + Ok(()) + } } \ No newline at end of file diff --git a/hash/plonky2/bench/src/bench/poseidon.rs b/hash/plonky2/bench/src/bench/poseidon.rs index fed9508..8867ff6 100644 --- a/hash/plonky2/bench/src/bench/poseidon.rs +++ b/hash/plonky2/bench/src/bench/poseidon.rs @@ -41,7 +41,7 @@ pub fn poseidon_bench(depth: usize) -> Result<()> { let hash = builder.hash_or_noop::(initial.clone()); // Public inputs are the initial value (provided below) and the result (which is generated). - // builder.register_public_inputs(initial.clone().as_slice()); + builder.register_public_inputs(initial.clone().as_slice()); builder.register_public_input(hash.elements[0]); builder.register_public_input(hash.elements[1]); builder.register_public_input(hash.elements[2]); diff --git a/hash/plonky2/bench/src/main.rs b/hash/plonky2/bench/src/main.rs index a5853f8..e804c69 100644 --- a/hash/plonky2/bench/src/main.rs +++ b/hash/plonky2/bench/src/main.rs @@ -58,7 +58,6 @@ fn main() { "keccak-polygon" => { println!("Running keccak of plolygon zk_evm: "); - eprintln!("number of permutation: {:?}", size); let _ = keccak_polygon_bench(size); } diff --git a/hash/plonky2/external/zk_evm b/hash/plonky2/external/zk_evm index 29f8620..5511569 160000 --- a/hash/plonky2/external/zk_evm +++ b/hash/plonky2/external/zk_evm @@ -1 +1 @@ -Subproject commit 29f8620338f2b3bb8be62def5666a2d9d7ba5331 +Subproject commit 5511569d166c29cc85211df31bb9485870ca53ab