Skip to content

Commit

Permalink
optimisations + benchmark setup
Browse files Browse the repository at this point in the history
  • Loading branch information
remybar committed Jan 24, 2025
1 parent ac35492 commit 3728e6d
Show file tree
Hide file tree
Showing 4 changed files with 219 additions and 2 deletions.
4 changes: 4 additions & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

10 changes: 9 additions & 1 deletion crates/katana/core/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -10,7 +10,7 @@ version.workspace = true
katana-db.workspace = true
katana-executor = { workspace = true, features = [ "blockifier" ] }
katana-pool.workspace = true
katana-primitives.workspace = true
katana-primitives = { workspace = true, features = [ "arbitrary" ]}
katana-provider.workspace = true
katana-tasks.workspace = true
katana-trie.workspace = true
Expand All @@ -25,6 +25,7 @@ metrics.workspace = true
num-traits.workspace = true
parking_lot.workspace = true
reqwest.workspace = true
rayon.workspace = true
serde.workspace = true
serde_json.workspace = true
starknet.workspace = true
Expand All @@ -46,8 +47,15 @@ alloy-transport = { workspace = true, default-features = false }

[dev-dependencies]
assert_matches.workspace = true
criterion.workspace = true
hex.workspace = true
tempfile.workspace = true
arbitrary.workspace = true
rand.workspace = true

[features]
starknet-messaging = [ "dep:starknet-crypto" ]

[[bench]]
name = "commit"
harness = false
132 changes: 132 additions & 0 deletions crates/katana/core/benches/commit.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,132 @@
use arbitrary::{Arbitrary, Unstructured};
use criterion::{black_box, criterion_group, criterion_main, BatchSize, Criterion};
use katana_core::backend::UncommittedBlock;
use std::collections::BTreeMap;

use katana_primitives::{
block::{GasPrices, PartialHeader},
da::L1DataAvailabilityMode,
receipt::{Receipt, ReceiptWithTxHash},
state::StateUpdates,
transaction::{Tx, TxWithHash},
version::CURRENT_STARKNET_VERSION,
ContractAddress, Felt,
};
use katana_provider::providers::db::DbProvider;

const NB_OF_TXS: usize = 20;
const NB_OF_RECEIPTS: usize = 20;
const NB_OF_NONCES: usize = 100;
const NB_OF_STORAGE_KEYS: usize = 100;
const NB_OF_STORAGE_VALUES: usize = 100;
const NB_OF_CLASSES: usize = 100;
const NB_OF_CONTRACTS: usize = 100;

pub fn commit(block: UncommittedBlock<'_, DbProvider>) {
let _ = block.commit();
}

pub fn commit_parallel(block: UncommittedBlock<'_, DbProvider>) {
let _ = block.commit_parallel();
}

#[inline(always)]
pub fn random_array(size: usize) -> Vec<u8> {
(0..size).map(|_| rand::random::<u8>()).collect()
}

#[inline(always)]
pub fn random_felt() -> Felt {
Felt::arbitrary(&mut Unstructured::new(&random_array(Felt::size_hint(0).0))).unwrap()
}

#[inline(always)]
pub fn random_tx() -> Tx {
Tx::arbitrary(&mut Unstructured::new(&random_array(Tx::size_hint(0).0))).unwrap()
}

#[inline(always)]
pub fn random_tx_with_hash() -> TxWithHash {
TxWithHash { hash: random_felt(), transaction: random_tx() }
}

#[inline(always)]
pub fn random_receipt() -> Receipt {
Receipt::arbitrary(&mut Unstructured::new(&random_array(Receipt::size_hint(0).0))).unwrap()
}

#[inline(always)]
pub fn random_receipt_with_hash() -> ReceiptWithTxHash {
ReceiptWithTxHash { tx_hash: random_felt(), receipt: random_receipt() }
}

#[inline(always)]
pub fn random_felt_to_felt_map(size: usize) -> BTreeMap<Felt, Felt> {
(0..size).map(|_| (random_felt(), random_felt())).collect()
}

#[inline(always)]
pub fn random_address_to_felt_map(size: usize) -> BTreeMap<ContractAddress, Felt> {
(0..size).map(|_| (ContractAddress::new(random_felt()), random_felt())).collect()
}

pub fn commit_benchmark(c: &mut Criterion) {
let provider = DbProvider::new_ephemeral();

let gas_prices = GasPrices { eth: 100 * u128::pow(10, 9), strk: 100 * u128::pow(10, 9) };
let sequencer_address = ContractAddress(1u64.into());

let header = PartialHeader {
protocol_version: CURRENT_STARKNET_VERSION,
number: 1,
timestamp: 100,
sequencer_address,
parent_hash: 123u64.into(),
l1_gas_prices: gas_prices.clone(),
l1_data_gas_prices: gas_prices.clone(),
l1_da_mode: L1DataAvailabilityMode::Calldata,
};

let transactions: Vec<TxWithHash> = (0..NB_OF_TXS).map(|_| random_tx_with_hash()).collect();
let receipts: Vec<ReceiptWithTxHash> =
(0..NB_OF_RECEIPTS).map(|_| random_receipt_with_hash()).collect();

let nonce_updates: BTreeMap<ContractAddress, Felt> =
(0..NB_OF_NONCES).map(|_| (ContractAddress::new(random_felt()), random_felt())).collect();

let storage_updates: BTreeMap<ContractAddress, BTreeMap<Felt, Felt>> = (0..NB_OF_STORAGE_KEYS)
.map(|_| {
(ContractAddress::new(random_felt()), random_felt_to_felt_map(NB_OF_STORAGE_VALUES))
})
.collect();

let declared_classes: BTreeMap<Felt, Felt> = random_felt_to_felt_map(NB_OF_CLASSES);
let deployed_contracts: BTreeMap<ContractAddress, Felt> =
random_address_to_felt_map(NB_OF_CONTRACTS);

let state_updates = StateUpdates {
nonce_updates,
storage_updates,
declared_classes,
deployed_contracts,
..Default::default()
};

let block =
UncommittedBlock::new(header, transactions, &receipts.as_slice(), &state_updates, provider);

c.bench_function("commit", |b| {
b.iter_batched(|| block.clone(), |input| commit(black_box(input)), BatchSize::SmallInput);
});

c.bench_function("commit_parallel", |b| {
b.iter_batched(
|| block.clone(),
|input| commit_parallel(black_box(input)),
BatchSize::SmallInput,
);
});
}

criterion_group!(benches, commit_benchmark);
criterion_main!(benches);
75 changes: 74 additions & 1 deletion crates/katana/core/src/backend/mod.rs
Original file line number Diff line number Diff line change
@@ -1,3 +1,4 @@
use rayon::prelude::*;
use std::sync::Arc;

use gas_oracle::L1GasOracle;
Expand Down Expand Up @@ -220,6 +221,7 @@ impl<'a, P: TrieWriter> UncommittedBlock<'a, P> {
let transaction_count = self.transactions.len() as u32;
let state_diff_length = self.state_updates.len() as u32;

// optimisation 1
let state_root = self.compute_new_state_root();
let transactions_commitment = self.compute_transaction_commitment();
let events_commitment = self.compute_event_commitment();
Expand Down Expand Up @@ -250,6 +252,51 @@ impl<'a, P: TrieWriter> UncommittedBlock<'a, P> {
SealedBlock { hash, header, body: self.transactions }
}

pub fn commit_parallel(self) -> SealedBlock {
// get the hash of the latest committed block
let parent_hash = self.header.parent_hash;
let events_count = self.receipts.iter().map(|r| r.events().len() as u32).sum::<u32>();
let transaction_count = self.transactions.len() as u32;
let state_diff_length = self.state_updates.len() as u32;

let mut state_root = Felt::default();
let mut transactions_commitment = Felt::default();
let mut events_commitment = Felt::default();
let mut receipts_commitment = Felt::default();
let mut state_diff_commitment = Felt::default();

rayon::scope(|s| {
s.spawn(|_| state_root = self.compute_new_state_root());
s.spawn(|_| transactions_commitment = self.compute_transaction_commitment());
s.spawn(|_| events_commitment = self.compute_event_commitment_parallel());
s.spawn(|_| receipts_commitment = self.compute_receipt_commitment_parallel());
s.spawn(|_| state_diff_commitment = self.compute_state_diff_commitment());
});

let header = Header {
state_root,
parent_hash,
events_count,
state_diff_length,
transaction_count,
events_commitment,
receipts_commitment,
state_diff_commitment,
transactions_commitment,
number: self.header.number,
timestamp: self.header.timestamp,
l1_da_mode: self.header.l1_da_mode,
l1_gas_prices: self.header.l1_gas_prices,
l1_data_gas_prices: self.header.l1_data_gas_prices,
sequencer_address: self.header.sequencer_address,
protocol_version: self.header.protocol_version,
};

let hash = header.compute_hash();

SealedBlock { hash, header, body: self.transactions }
}

fn compute_transaction_commitment(&self) -> Felt {
let tx_hashes = self.transactions.iter().map(|t| t.hash).collect::<Vec<TxHash>>();
compute_merkle_root::<hash::Poseidon>(&tx_hashes).unwrap()
Expand All @@ -260,6 +307,12 @@ impl<'a, P: TrieWriter> UncommittedBlock<'a, P> {
compute_merkle_root::<hash::Poseidon>(&receipt_hashes).unwrap()
}

fn compute_receipt_commitment_parallel(&self) -> Felt {
let receipt_hashes =
self.receipts.par_iter().map(|r| r.compute_hash()).collect::<Vec<Felt>>();
compute_merkle_root::<hash::Poseidon>(&receipt_hashes).unwrap()
}

fn compute_state_diff_commitment(&self) -> Felt {
compute_state_diff_hash(self.state_updates.clone())
}
Expand All @@ -275,7 +328,6 @@ impl<'a, P: TrieWriter> UncommittedBlock<'a, P> {
// the iterator will yield all events from all the receipts, each one paired with the
// transaction hash that emitted it: (tx hash, event).
let events = self.receipts.iter().flat_map(|r| r.events().iter().map(|e| (r.tx_hash, e)));

let mut hashes = Vec::new();
for (tx, event) in events {
let event_hash = event_hash(tx, event);
Expand All @@ -286,6 +338,27 @@ impl<'a, P: TrieWriter> UncommittedBlock<'a, P> {
compute_merkle_root::<hash::Poseidon>(&hashes).unwrap()
}

fn compute_event_commitment_parallel(&self) -> Felt {
// h(emitter_address, tx_hash, h(keys), h(data))
fn event_hash(tx: TxHash, event: &Event) -> Felt {
let keys_hash = hash::Poseidon::hash_array(&event.keys);
let data_hash = hash::Poseidon::hash_array(&event.data);
hash::Poseidon::hash_array(&[tx, event.from_address.into(), keys_hash, data_hash])
}

// the iterator will yield all events from all the receipts, each one paired with the
// transaction hash that emitted it: (tx hash, event).
let events = self.receipts.iter().flat_map(|r| r.events().iter().map(|e| (r.tx_hash, e)));
let hashes = events
.par_bridge()
.into_par_iter()
.map(|(tx, event)| event_hash(tx, event))
.collect::<Vec<_>>();

// compute events commitment
compute_merkle_root::<hash::Poseidon>(&hashes).unwrap()
}

// state_commitment = hPos("STARKNET_STATE_V0", contract_trie_root, class_trie_root)
fn compute_new_state_root(&self) -> Felt {
let class_trie_root = self
Expand Down

0 comments on commit 3728e6d

Please sign in to comment.