From 9298cc93fe690f2f690770aebb6c738957e512e8 Mon Sep 17 00:00:00 2001
From: Shahar Papini <43779613+spapinistarkware@users.noreply.github.com>
Date: Wed, 3 Apr 2024 09:33:42 +0300
Subject: [PATCH] simple merkle benchmark (#527)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

<!-- Reviewable:start -->
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/starkware-libs/stwo/527)
<!-- Reviewable:end -->
---
 Cargo.toml        |  4 ++++
 benches/merkle.rs | 40 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 44 insertions(+)
 create mode 100644 benches/merkle.rs

diff --git a/Cargo.toml b/Cargo.toml
index 739b49a03..0a2e5d743 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -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);