[chore] Keccak Output Computation Optimization (#242)

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"

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<F: Field> KeccakComponentShardCircuit<F> {
 
 pub(crate) fn create_hasher<F: Field>() -> PoseidonHasher<F, POSEIDON_T, POSEIDON_RATE> {
     // Construct in-circuit Poseidon hasher.
-    let spec = OptimizedPoseidonSpec::<F, POSEIDON_T, POSEIDON_RATE>::new::<
-        POSEIDON_R_F,
-        POSEIDON_R_P,
-        POSEIDON_SECURE_MDS,
-    >();
+    let spec = get_poseidon_spec();
     PoseidonHasher::<F, POSEIDON_T, POSEIDON_RATE>::new(spec)
 }
 

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<F: Field>(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::<F, F, POSEIDON_T, POSEIDON_RATE>::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.

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<TypeMap> = Default::default();
+}
+
+pub(crate) fn get_poseidon_spec<F: Field>() -> OptimizedPoseidonSpec<F, POSEIDON_T, POSEIDON_RATE> {
+    let spec = POSEIDON_SPEC_CACHE
+        .read()
+        .unwrap_or_else(|e| e.into_inner())
+        .get::<OptimizedPoseidonSpec<F, POSEIDON_T, POSEIDON_RATE>>()
+        .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::<F, POSEIDON_T, POSEIDON_RATE>::new::<
+            POSEIDON_R_F,
+            POSEIDON_R_P,
+            POSEIDON_SECURE_MDS,
+        >();
+        to_write.insert(spec.clone());
+        spec
+    };
+    spec
+}
+
+pub(crate) fn create_native_poseidon_sponge<F: Field>(
+) -> snark_verifier::util::hash::Poseidon<F, F, POSEIDON_T, POSEIDON_RATE> {
+    snark_verifier::util::hash::Poseidon::<F, F, POSEIDON_T, POSEIDON_RATE>::from_spec(
+        &NativeLoader,
+        get_poseidon_spec(),
+    )
+}

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<F: Field>(
     capacity: usize,
 ) -> Vec<KeccakCircuitOutput<F>> {
     assert!(u128::BITS <= F::CAPACITY);
-    let mut outputs =
-        inputs.iter().flat_map(|input| input_to_circuit_outputs::<F>(input)).collect_vec();
+    let mut outputs: Vec<_> =
+        inputs.par_iter().flat_map(|input| input_to_circuit_outputs::<F>(input)).collect();
     assert!(outputs.len() <= capacity);
     outputs.resize(capacity, dummy_circuit_output());
     outputs
@@ -48,8 +52,30 @@ pub fn input_to_circuit_outputs<F: Field>(bytes: &[u8]) -> Vec<KeccakCircuitOutp
     output
 }
 
+lazy_static! {
+    static ref DUMMY_CIRCUIT_OUTPUT_CACHE: RwLock<TypeMap> = Default::default();
+}
+
 /// Return the dummy circuit output for padding.
 pub fn dummy_circuit_output<F: Field>() -> KeccakCircuitOutput<F> {
+    let output = DUMMY_CIRCUIT_OUTPUT_CACHE
+        .read()
+        .unwrap_or_else(|e| e.into_inner())
+        .get::<KeccakCircuitOutput<F>>()
+        .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<F: Field>() -> KeccakCircuitOutput<F> {
     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<F: Field>() -> KeccakCircuitOutput<F> {
 
 /// Calculate the commitment of circuit outputs.
 pub fn calculate_circuit_outputs_commit<F: Field>(outputs: &[KeccakCircuitOutput<F>]) -> F {
-    let mut native_poseidon_sponge =
-        snark_verifier::util::hash::Poseidon::<F, F, POSEIDON_T, POSEIDON_RATE>::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()