From 6bb4dbd3de45dcda3890d7c7aed87b89333fa556 Mon Sep 17 00:00:00 2001 From: Xinding Wei Date: Sat, 30 Dec 2023 14:21:48 -0700 Subject: [PATCH] Performance improvements: - Cache Poseidon spec - Cache dummy circuit output - Multiple threads in multi_inputs_to_circuit_outputs --- hashes/zkevm/Cargo.toml | 1 + .../src/keccak/component/circuit/shard.rs | 12 ++--- hashes/zkevm/src/keccak/component/encode.rs | 10 +---- hashes/zkevm/src/keccak/component/mod.rs | 45 +++++++++++++++++++ hashes/zkevm/src/keccak/component/output.rs | 41 ++++++++++++----- 5 files changed, 82 insertions(+), 27 deletions(-) diff --git a/hashes/zkevm/Cargo.toml b/hashes/zkevm/Cargo.toml index 169bf018..c53248ba 100644 --- a/hashes/zkevm/Cargo.toml +++ b/hashes/zkevm/Cargo.toml @@ -19,6 +19,7 @@ sha3 = "0.10.8" # always included but without features to use Native poseidon and get CircuitExt trait snark-verifier-sdk = { git = "https://github.com/axiom-crypto/snark-verifier.git", branch = "release-0.1.7-rc", default-features = false } getset = "0.1.2" +type-map = "0.5.0" [dev-dependencies] ethers-signers = "2.0.8" diff --git a/hashes/zkevm/src/keccak/component/circuit/shard.rs b/hashes/zkevm/src/keccak/component/circuit/shard.rs index ef2c99b3..604da477 100644 --- a/hashes/zkevm/src/keccak/component/circuit/shard.rs +++ b/hashes/zkevm/src/keccak/component/circuit/shard.rs @@ -7,6 +7,7 @@ use crate::{ get_words_to_witness_multipliers, num_poseidon_absorb_per_keccak_f, num_word_per_witness, }, + get_poseidon_spec, output::{ calculate_circuit_outputs_commit, dummy_circuit_output, multi_inputs_to_circuit_outputs, KeccakCircuitOutput, @@ -31,10 +32,7 @@ use halo2_base::{ circuit::{Layouter, SimpleFloorPlanner}, plonk::{Circuit, ConstraintSystem, Error}, }, - poseidon::hasher::{ - spec::OptimizedPoseidonSpec, PoseidonCompactChunkInput, PoseidonCompactOutput, - PoseidonHasher, - }, + poseidon::hasher::{PoseidonCompactChunkInput, PoseidonCompactOutput, PoseidonHasher}, safe_types::{SafeBool, SafeTypeChip}, virtual_region::copy_constraints::SharedCopyConstraintManager, AssignedValue, Context, @@ -405,11 +403,7 @@ impl KeccakComponentShardCircuit { pub(crate) fn create_hasher() -> PoseidonHasher { // Construct in-circuit Poseidon hasher. - let spec = OptimizedPoseidonSpec::::new::< - POSEIDON_R_F, - POSEIDON_R_P, - POSEIDON_SECURE_MDS, - >(); + let spec = get_poseidon_spec(); PoseidonHasher::::new(spec) } diff --git a/hashes/zkevm/src/keccak/component/encode.rs b/hashes/zkevm/src/keccak/component/encode.rs index 8767c404..1bed38be 100644 --- a/hashes/zkevm/src/keccak/component/encode.rs +++ b/hashes/zkevm/src/keccak/component/encode.rs @@ -8,14 +8,13 @@ use halo2_base::{ }; use itertools::Itertools; use num_bigint::BigUint; -use snark_verifier_sdk::{snark_verifier, NativeLoader}; use crate::{ keccak::vanilla::{keccak_packed_multi::get_num_keccak_f, param::*}, util::eth_types::Field, }; -use super::param::*; +use super::{create_native_poseidon_sponge, param::*}; // TODO: Abstract this module into a trait for all component circuits. @@ -26,12 +25,7 @@ use super::param::*; pub fn encode_native_input(bytes: &[u8]) -> F { let witnesses_per_keccak_f = pack_native_input(bytes); // Absorb witnesses keccak_f by keccak_f. - let mut native_poseidon_sponge = - snark_verifier::util::hash::Poseidon::::new::< - POSEIDON_R_F, - POSEIDON_R_P, - POSEIDON_SECURE_MDS, - >(&NativeLoader); + let mut native_poseidon_sponge = create_native_poseidon_sponge(); for witnesses in witnesses_per_keccak_f { for absorbing in witnesses.chunks(POSEIDON_RATE) { // To avoid absorbing witnesses crossing keccak_fs together, pad 0s to make sure absorb.len() == RATE. diff --git a/hashes/zkevm/src/keccak/component/mod.rs b/hashes/zkevm/src/keccak/component/mod.rs index 13bbd303..d1a6477e 100644 --- a/hashes/zkevm/src/keccak/component/mod.rs +++ b/hashes/zkevm/src/keccak/component/mod.rs @@ -1,3 +1,14 @@ +use std::sync::RwLock; + +use halo2_base::poseidon::hasher::spec::OptimizedPoseidonSpec; +use lazy_static::lazy_static; +use snark_verifier_sdk::{snark_verifier, NativeLoader}; +use type_map::concurrent::TypeMap; + +use crate::util::eth_types::Field; + +use self::param::{POSEIDON_RATE, POSEIDON_R_F, POSEIDON_R_P, POSEIDON_SECURE_MDS, POSEIDON_T}; + /// Module of Keccak component circuit(s). pub mod circuit; /// Module of encoding raw inputs to component circuit lookup keys. @@ -10,3 +21,37 @@ pub mod output; pub mod param; #[cfg(test)] mod tests; + +lazy_static! { + static ref POSEIDON_SPEC_CACHE: RwLock = Default::default(); +} + +pub(crate) fn get_poseidon_spec() -> OptimizedPoseidonSpec { + let spec = POSEIDON_SPEC_CACHE + .read() + .unwrap_or_else(|e| e.into_inner()) + .get::>() + .cloned(); + if let Some(spec) = spec { + return spec; + } + let spec = { + let mut to_write = POSEIDON_SPEC_CACHE.write().unwrap_or_else(|e| e.into_inner()); + let spec = OptimizedPoseidonSpec::::new::< + POSEIDON_R_F, + POSEIDON_R_P, + POSEIDON_SECURE_MDS, + >(); + to_write.insert(spec.clone()); + spec + }; + spec +} + +pub(crate) fn create_native_poseidon_sponge( +) -> snark_verifier::util::hash::Poseidon { + snark_verifier::util::hash::Poseidon::::from_spec( + &NativeLoader, + get_poseidon_spec(), + ) +} diff --git a/hashes/zkevm/src/keccak/component/output.rs b/hashes/zkevm/src/keccak/component/output.rs index 2fe46ecb..22688b5f 100644 --- a/hashes/zkevm/src/keccak/component/output.rs +++ b/hashes/zkevm/src/keccak/component/output.rs @@ -1,8 +1,12 @@ -use super::{encode::encode_native_input, param::*}; +use std::sync::RwLock; + +use super::{create_native_poseidon_sponge, encode::encode_native_input}; use crate::{keccak::vanilla::keccak_packed_multi::get_num_keccak_f, util::eth_types::Field}; use itertools::Itertools; +use lazy_static::lazy_static; +use rayon::iter::{IntoParallelRefIterator, ParallelIterator}; use sha3::{Digest, Keccak256}; -use snark_verifier_sdk::{snark_verifier, NativeLoader}; +use type_map::concurrent::TypeMap; /// Witnesses to be exposed as circuit outputs. #[derive(Clone, Copy, PartialEq, Debug)] @@ -21,8 +25,8 @@ pub fn multi_inputs_to_circuit_outputs( capacity: usize, ) -> Vec> { assert!(u128::BITS <= F::CAPACITY); - let mut outputs = - inputs.iter().flat_map(|input| input_to_circuit_outputs::(input)).collect_vec(); + let mut outputs: Vec<_> = + inputs.par_iter().flat_map(|input| input_to_circuit_outputs::(input)).collect(); assert!(outputs.len() <= capacity); outputs.resize(capacity, dummy_circuit_output()); outputs @@ -48,8 +52,30 @@ pub fn input_to_circuit_outputs(bytes: &[u8]) -> Vec = Default::default(); +} + /// Return the dummy circuit output for padding. pub fn dummy_circuit_output() -> KeccakCircuitOutput { + let output = DUMMY_CIRCUIT_OUTPUT_CACHE + .read() + .unwrap_or_else(|e| e.into_inner()) + .get::>() + .cloned(); + if let Some(output) = output { + return output; + } + let output = { + let mut to_write = DUMMY_CIRCUIT_OUTPUT_CACHE.write().unwrap_or_else(|e| e.into_inner()); + let output = dummy_circuit_output_impl(); + to_write.insert(output); + output + }; + output +} + +fn dummy_circuit_output_impl() -> KeccakCircuitOutput { assert!(u128::BITS <= F::CAPACITY); let key = encode_native_input(&[]); // Output of Keccak256::digest is big endian. @@ -61,12 +87,7 @@ pub fn dummy_circuit_output() -> KeccakCircuitOutput { /// Calculate the commitment of circuit outputs. pub fn calculate_circuit_outputs_commit(outputs: &[KeccakCircuitOutput]) -> F { - let mut native_poseidon_sponge = - snark_verifier::util::hash::Poseidon::::new::< - POSEIDON_R_F, - POSEIDON_R_P, - POSEIDON_SECURE_MDS, - >(&NativeLoader); + let mut native_poseidon_sponge = create_native_poseidon_sponge(); native_poseidon_sponge.update( &outputs .iter()