From 1a950fc98b2093995ac6d86f49903e954c816f0f Mon Sep 17 00:00:00 2001
From: Shahar Papini <spapini@starkware.co>
Date: Fri, 22 Mar 2024 15:28:55 +0200
Subject: [PATCH] simple merkle benchmark

---
 Cargo.toml              |  6 +++++-
 benches/merkle.rs       | 40 ++++++++++++++++++++++++++++++++++++++++
 benches/merkle_bench.rs | 21 +--------------------
 3 files changed, 46 insertions(+), 21 deletions(-)
 create mode 100644 benches/merkle.rs

diff --git a/Cargo.toml b/Cargo.toml
index 43bad56e4..27dd1507a 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -6,7 +6,6 @@ edition = "2021"
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-blake2 = "0.10.6"
 blake3 = "1.5.0"
 hex = "0.4.3"
 itertools = "0.12.0"
@@ -14,6 +13,7 @@ num-traits = "0.2.17"
 thiserror = "1.0.56"
 merging-iterator = "1.3.0"
 bytemuck = { version = "1.14.3", features = ["derive"] }
+blake2 = "0.10.6"
 
 [dev-dependencies]
 criterion = { version = "0.5.1", features = ["html_reports"] }
@@ -56,6 +56,10 @@ name = "matrix"
 name = "merkle_bench"
 harness = false
 
+[[bench]]
+name = "merkle"
+harness = false
+
 [[bench]]
 name = "fri"
 harness = false
diff --git a/benches/merkle.rs b/benches/merkle.rs
new file mode 100644
index 000000000..c465f4c93
--- /dev/null
+++ b/benches/merkle.rs
@@ -0,0 +1,40 @@
+#![feature(iter_array_chunks)]
+
+use criterion::Criterion;
+
+#[cfg(target_arch = "x86_64")]
+pub fn cpu_merkle(c: &mut criterion::Criterion) {
+    use itertools::Itertools;
+    use num_traits::Zero;
+    use stwo::commitment_scheme::ops::MerkleOps;
+    use stwo::core::backend::CPUBackend;
+    use stwo::core::fields::m31::BaseField;
+
+    const N_COLS: usize = 1 << 8;
+    const LOG_SIZE: u32 = 20;
+    let cols = (0..N_COLS)
+        .map(|_| {
+            (0..(1 << LOG_SIZE))
+                .map(|_| BaseField::zero())
+                .collect::<Vec<_>>()
+        })
+        .collect::<Vec<_>>();
+
+    let mut group = c.benchmark_group("merkle throughput");
+    group.throughput(criterion::Throughput::Elements((N_COLS << LOG_SIZE) as u64));
+    group.throughput(criterion::Throughput::Bytes(
+        (N_COLS << (LOG_SIZE + 2)) as u64,
+    ));
+    group.bench_function("cpu merkle", |b| {
+        b.iter(|| {
+            CPUBackend::commit_on_layer(LOG_SIZE, None, &cols.iter().collect_vec());
+        })
+    });
+}
+
+#[cfg(target_arch = "x86_64")]
+criterion::criterion_group!(
+    name=merkle;
+    config = Criterion::default().sample_size(10);
+    targets=cpu_merkle);
+criterion::criterion_main!(merkle);
diff --git a/benches/merkle_bench.rs b/benches/merkle_bench.rs
index d85bc8e8f..2a442aa44 100644
--- a/benches/merkle_bench.rs
+++ b/benches/merkle_bench.rs
@@ -1,5 +1,4 @@
 // TODO(Ohad): write better benchmarks. Reduce the variance in sample size.
-use blake2::{Blake2s256, Digest};
 use criterion::measurement::WallTime;
 use criterion::{
     black_box, criterion_group, criterion_main, BatchSize, BenchmarkGroup, BenchmarkId, Criterion,
@@ -82,20 +81,6 @@ fn compare_blakes(c: &mut Criterion) {
     group.finish();
 }
 
-fn single_blake2s_hash_benchmark(c: &mut Criterion) {
-    let input = [0u8; 1];
-    c.bench_function("Single blake2s hash", |b| {
-        b.iter_batched(
-            || -> Blake2s256 { Blake2s256::new() },
-            |mut h| {
-                h.update(&input[..]);
-                h.finalize()
-            },
-            BatchSize::SmallInput,
-        )
-    });
-}
-
 fn single_blake3_hash_benchmark(c: &mut Criterion) {
     let input = [0u8; 1];
     c.bench_function("Single blake3 hash", |b| b.iter(|| blake3::hash(&input)));
@@ -109,10 +94,6 @@ criterion_group!(
 
 criterion_group!(comparisons, compare_blakes,);
 
-criterion_group!(
-    single_hash,
-    single_blake2s_hash_benchmark,
-    single_blake3_hash_benchmark,
-);
+criterion_group!(single_hash, single_blake3_hash_benchmark,);
 
 criterion_main!(comparisons);