perf: add benchmarks for different components (#273)

evm_arithmetization/Cargo.toml

@@ -69,6 +69,11 @@ required-features = ["asmtools"]
 name = "stack_manipulation"
 harness = false
 
+[[bench]]
+name = "fibonacci_25m_gas"
+harness = false
+
+
 # Display math equations properly in documentation
 [package.metadata.docs.rs]
 rustdoc-args = ["--html-in-header", ".cargo/katex-header.html"]

evm_arithmetization/benches/fibonacci_25m_gas.rs

@@ -0,0 +1,184 @@
+//! Benchmarks the CPU execution of a transaction calling a simple Fibonacci
+//! contract, iterating over and over until reaching the 25M gas limit.
+//!
+//! Total number of user instructions: 7_136_858.
+//! Total number of loops: 2_378_952.
+
+use std::collections::HashMap;
+use std::str::FromStr;
+
+use criterion::{criterion_group, criterion_main, BatchSize, BenchmarkId, Criterion};
+use env_logger::{try_init_from_env, Env, DEFAULT_FILTER_ENV};
+use ethereum_types::{Address, H256, U256};
+use evm_arithmetization::cpu::kernel::aggregator::KERNEL;
+use evm_arithmetization::cpu::kernel::opcodes::{get_opcode, get_push_opcode};
+use evm_arithmetization::generation::mpt::{AccountRlp, LegacyReceiptRlp};
+use evm_arithmetization::generation::{GenerationInputs, TrieInputs};
+use evm_arithmetization::proof::{BlockHashes, BlockMetadata, TrieRoots};
+use evm_arithmetization::prover::testing::simulate_execution;
+use evm_arithmetization::Node;
+use hex_literal::hex;
+use keccak_hash::keccak;
+use mpt_trie::nibbles::Nibbles;
+use mpt_trie::partial_trie::{HashedPartialTrie, PartialTrie};
+use plonky2::field::goldilocks_field::GoldilocksField;
+
+type F = GoldilocksField;
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let inputs = prepare_setup().unwrap();
+
+    // Dummy call to preinitialize the kernel.
+    let _ = KERNEL.hash();
+
+    let mut group = c.benchmark_group("fibonacci_25m_gas");
+    group.sample_size(10);
+    group.bench_function(BenchmarkId::from_parameter(8), |b| {
+        b.iter_batched(
+            || inputs.clone(),
+            |inp| simulate_execution::<F>(inp).unwrap(),
+            BatchSize::LargeInput,
+        )
+    });
+
+    // Last run to print the number of CPU cycles.
+    init_logger();
+    simulate_execution::<F>(inputs).unwrap();
+}
+
+fn prepare_setup() -> anyhow::Result<GenerationInputs> {
+    let sender = hex!("8943545177806ED17B9F23F0a21ee5948eCaa776");
+    let to = hex!("159271B89fea49aF29DFaf8b4eCE7D042D5d6f07");
+
+    let sender_state_key = keccak(sender);
+    let to_state_key = keccak(to);
+
+    let sender_nibbles = Nibbles::from_bytes_be(sender_state_key.as_bytes()).unwrap();
+    let to_nibbles = Nibbles::from_bytes_be(to_state_key.as_bytes()).unwrap();
+
+    let push1 = get_push_opcode(1);
+    let push4 = get_push_opcode(4);
+    let add = get_opcode("ADD");
+    let swap1 = get_opcode("SWAP1");
+    let dup2 = get_opcode("DUP2");
+    let jump = get_opcode("JUMP");
+    let jumpdest = get_opcode("JUMPDEST");
+    let code = [
+        push1, 1, push1, 1, jumpdest, dup2, add, swap1, push4, 0, 0, 0, 4, jump,
+    ];
+    let code_hash = keccak(code);
+
+    let empty_trie_root = HashedPartialTrie::from(Node::Empty).hash();
+
+    let sender_account_before = AccountRlp {
+        nonce: 169.into(),
+        balance: U256::from_dec_str("999999999998417410153631615")?,
+        storage_root: empty_trie_root,
+        code_hash: keccak(vec![]),
+    };
+    let to_account_before = AccountRlp {
+        nonce: 1.into(),
+        balance: 0.into(),
+        storage_root: empty_trie_root,
+        code_hash,
+    };
+
+    let mut state_trie_before = HashedPartialTrie::from(Node::Empty);
+    state_trie_before.insert(sender_nibbles, rlp::encode(&sender_account_before).to_vec())?;
+    state_trie_before.insert(to_nibbles, rlp::encode(&to_account_before).to_vec())?;
+
+    let tries_before = TrieInputs {
+        state_trie: state_trie_before,
+        transactions_trie: Node::Empty.into(),
+        receipts_trie: Node::Empty.into(),
+        storage_tries: vec![
+            (sender_state_key, Node::Empty.into()),
+            (to_state_key, Node::Empty.into()),
+        ],
+    };
+
+    let gas_used = U256::from(0x17d7840_u32);
+
+    let txn = hex!("f86981a9843b9aca1084017d784094159271b89fea49af29dfaf8b4ece7d042d5d6f0780808360306ba00cdea08ac2e8075188b289d779fa84bf86020c2b162bbee11d2785b5225b0ccca00ea9a76f4641955a74ae8c1589914fc7d6c5bfe5940454a89daf5b12d6f06617");
+    let value = U256::zero();
+
+    let block_metadata = BlockMetadata {
+        block_beneficiary: Address::from(sender),
+        block_difficulty: 0x0.into(),
+        block_number: 0x176.into(),
+        block_chain_id: 0x301824.into(),
+        block_timestamp: 0x664e63af.into(),
+        block_gaslimit: 0x1c9c380.into(),
+        block_gas_used: gas_used,
+        block_bloom: [0.into(); 8],
+        block_base_fee: 0x11.into(),
+        block_random: H256(hex!(
+            "388bd2892c01ab13e22f713316cc2b5d3c3d963e1426c25a80c7878a1815f889"
+        )),
+    };
+
+    let mut contract_code = HashMap::new();
+    contract_code.insert(keccak(vec![]), vec![]);
+    contract_code.insert(code_hash, code.to_vec());
+
+    let sender_account_after = AccountRlp {
+        balance: sender_account_before.balance - value - gas_used * block_metadata.block_base_fee,
+        nonce: sender_account_before.nonce + 1,
+        ..sender_account_before
+    };
+    let to_account_after = to_account_before;
+
+    let mut expected_state_trie_after = HashedPartialTrie::from(Node::Empty);
+    expected_state_trie_after
+        .insert(sender_nibbles, rlp::encode(&sender_account_after).to_vec())?;
+    expected_state_trie_after.insert(to_nibbles, rlp::encode(&to_account_after).to_vec())?;
+
+    let receipt_0 = LegacyReceiptRlp {
+        status: false,
+        cum_gas_used: gas_used,
+        bloom: vec![0; 256].into(),
+        logs: vec![],
+    };
+    let mut receipts_trie = HashedPartialTrie::from(Node::Empty);
+    receipts_trie.insert(
+        Nibbles::from_str("0x80").unwrap(),
+        rlp::encode(&receipt_0).to_vec(),
+    )?;
+    let transactions_trie: HashedPartialTrie = Node::Leaf {
+        nibbles: Nibbles::from_str("0x80").unwrap(),
+        value: txn.to_vec(),
+    }
+    .into();
+
+    let trie_roots_after = TrieRoots {
+        state_root: expected_state_trie_after.hash(),
+        transactions_root: transactions_trie.hash(),
+        receipts_root: receipts_trie.hash(),
+    };
+
+    Ok(GenerationInputs {
+        signed_txn: Some(txn.to_vec()),
+        withdrawals: vec![],
+        tries: tries_before,
+        trie_roots_after,
+        contract_code,
+        checkpoint_state_trie_root: H256(hex!(
+            "fe07ff6d1ab215df17884b89112ccf2373597285a56c5902150313ad1a53ee57"
+        )),
+        block_metadata,
+        txn_number_before: 0.into(),
+        gas_used_before: 0.into(),
+        gas_used_after: gas_used,
+        block_hashes: BlockHashes {
+            prev_hashes: vec![H256::default(); 256],
+            cur_hash: H256::default(),
+        },
+    })
+}
+
+fn init_logger() {
+    let _ = try_init_from_env(Env::default().filter_or(DEFAULT_FILTER_ENV, "info"));
+}
+
+criterion_group!(benches, criterion_benchmark);
+criterion_main!(benches);

evm_arithmetization/src/cpu/kernel/asm/mpt/read.asm

@@ -106,7 +106,7 @@ global mpt_read_extension_not_found:
     // Not found; return 0.
     %stack (key_part, future_nibbles, key, node_payload_ptr, retdest) -> (retdest, 0)
     JUMP
-mpt_read_extension_found:
+global mpt_read_extension_found:
     // stack: key_part, future_nibbles, key, node_payload_ptr, retdest
     DUP2 %mul_const(4) SHL // key_part_shifted = (key_part << (future_nibbles * 4))
     // stack: key_part_shifted, future_nibbles, key, node_payload_ptr, retdest
@@ -142,7 +142,7 @@ global mpt_read_leaf_not_found:
     // Not found; return 0.
     %stack (node_payload_ptr, retdest) -> (retdest, 0)
     JUMP
-mpt_read_leaf_found:
+global mpt_read_leaf_found:
     // stack: node_payload_ptr, retdest
     %add_const(2) // The value pointer is located after num_nibbles and the key.
     // stack: value_ptr_ptr, retdest

trace_decoder/Cargo.toml

@@ -31,4 +31,10 @@ mpt_trie = { version = "0.3.0", path = "../mpt_trie" }
 evm_arithmetization = { version = "0.2.0", path = "../evm_arithmetization" }
 
 [dev-dependencies]
+criterion = "0.5.1"
 pretty_env_logger = "0.5.0"
+serde_json = { workspace = true }
+
+[[bench]]
+name = "block_processing"
+harness = false

trace_decoder/benches/block_processing.rs

@@ -0,0 +1,49 @@
+//! Benchmarks the processing by the decoder of a block witness obtained from a
+//! node into a sequence of prover inputs ready to be sent to a prover.
+//!
+//! The block being processed here is the 19240650th Ethereum block
+//! (<https://etherscan.io/block/19240650>) containing 201 transactions and 16 withdrawals.
+
+use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
+use serde::{Deserialize, Serialize};
+use trace_decoder::{
+    processed_block_trace::ProcessingMeta,
+    trace_protocol::BlockTrace,
+    types::{CodeHash, OtherBlockData},
+};
+
+#[derive(Clone, Debug, Deserialize, Serialize)]
+pub struct ProverInput {
+    pub block_trace: BlockTrace,
+    pub other_data: OtherBlockData,
+}
+
+fn resolve_code_hash_fn(_: &CodeHash) -> Vec<u8> {
+    todo!()
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let bytes = std::fs::read("benches/block_input.json").unwrap();
+    let prover_input: ProverInput = serde_json::from_slice(&bytes).unwrap();
+
+    c.bench_function("Block 19240650 processing", |b| {
+        b.iter_batched(
+            || prover_input.clone(),
+            |pi| {
+                pi.block_trace
+                    .into_txn_proof_gen_ir(
+                        &ProcessingMeta::new(resolve_code_hash_fn),
+                        prover_input.other_data.clone(),
+                    )
+                    .unwrap()
+            },
+            BatchSize::LargeInput,
+        )
+    });
+}
+
+criterion_group!(
+    name = benches;
+    config = Criterion::default().sample_size(10);
+    targets = criterion_benchmark);
+criterion_main!(benches);

trace_decoder/src/trace_protocol.rs

@@ -37,7 +37,7 @@ use crate::{
 /// Core payload needed to generate a proof for a block. Note that the scheduler
 /// may need to request some additional data from the client along with this in
 /// order to generate a proof.
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct BlockTrace {
     /// The trie pre-images (state & storage) in multiple possible formats.
     pub trie_pre_images: BlockTraceTriePreImages,
@@ -52,7 +52,7 @@ pub struct BlockTrace {
 }
 
 /// Minimal hashed out tries needed by all txns in the block.
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 #[serde(rename_all = "snake_case")]
 pub enum BlockTraceTriePreImages {
     /// The trie pre-image with separate state/storage tries.
@@ -62,7 +62,7 @@ pub enum BlockTraceTriePreImages {
 }
 
 /// State/Storage trie pre-images that are separate.
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct SeparateTriePreImages {
     /// State trie.
     pub state: SeparateTriePreImage,
@@ -71,7 +71,7 @@ pub struct SeparateTriePreImages {
 }
 
 /// A trie pre-image where state & storage are separate.
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 #[serde(rename_all = "snake_case")]
 pub enum SeparateTriePreImage {
     /// Storage or state trie in a bulkier format, that can be processed faster.
@@ -82,7 +82,7 @@ pub enum SeparateTriePreImage {
 }
 
 /// A trie pre-image where both state & storage are combined into one payload.
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 #[serde(rename_all = "snake_case")]
 pub struct CombinedPreImages {
     /// Compact combined state and storage tries.
@@ -91,23 +91,23 @@ pub struct CombinedPreImages {
 
 // TODO
 /// Bulkier format that is quicker to process.
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct TrieUncompressed {}
 
 // TODO
 #[serde_as]
 /// Compact representation of a trie (will likely be very close to <https://github.com/ledgerwatch/erigon/blob/devel/docs/programmers_guide/witness_formal_spec.md>)
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct TrieCompact(#[serde_as(as = "FromInto<ByteString>")] pub Vec<u8>);
 
 // TODO
 /// Trie format that is in exactly the same format of our internal trie format.
 /// This is the fastest format for us to processes.
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct TrieDirect(pub HashedPartialTrie);
 
 /// A trie pre-image where state and storage are separate.
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 #[serde(rename_all = "snake_case")]
 pub enum SeparateStorageTriesPreImage {
     /// A single hash map that contains all node hashes from all storage tries
@@ -121,7 +121,7 @@ pub enum SeparateStorageTriesPreImage {
 }
 
 /// Info specific to txns in the block.
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct TxnInfo {
     /// Trace data for the txn. This is used by the protocol to:
     /// - Mutate it's own trie state between txns to arrive at the correct trie
@@ -136,7 +136,7 @@ pub struct TxnInfo {
 
 /// Structure holding metadata for one transaction.
 #[serde_as]
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct TxnMeta {
     /// Txn byte code.
     #[serde_as(as = "FromInto<ByteString>")]
@@ -162,7 +162,7 @@ pub struct TxnMeta {
 ///
 /// Specifically, since we can not execute the txn before proof generation, we
 /// rely on a separate EVM to run the txn and supply this data for us.
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 pub struct TxnTrace {
     /// If the balance changed, then the new balance will appear here. Will be
     /// `None` if no change.
@@ -198,7 +198,7 @@ pub struct TxnTrace {
 
 /// Contract code access type. Used by txn traces.
 #[serde_as]
-#[derive(Debug, Deserialize, Serialize)]
+#[derive(Clone, Debug, Deserialize, Serialize)]
 #[serde(rename_all = "snake_case")]
 pub enum ContractCodeUsage {
     /// Contract was read.