From 88945f2c227d6a76e8f133d8b72c701b180c4f4c Mon Sep 17 00:00:00 2001
From: Ivar Flakstad <69173633+ivarflakstad@users.noreply.github.com>
Date: Tue, 9 Jan 2024 18:31:28 +0100
Subject: [PATCH] Improve benchmarks layout

---
 candle-core/Cargo.toml                              | 2 +-
 candle-core/benches/bench_main.rs                   | 4 ++++
 candle-core/benches/{ => benchmarks}/matmul.rs      | 7 ++-----
 candle-core/benches/{utils.rs => benchmarks/mod.rs} | 2 ++
 4 files changed, 9 insertions(+), 6 deletions(-)
 create mode 100644 candle-core/benches/bench_main.rs
 rename candle-core/benches/{ => benchmarks}/matmul.rs (85%)
 rename candle-core/benches/{utils.rs => benchmarks/mod.rs} (98%)

diff --git a/candle-core/Cargo.toml b/candle-core/Cargo.toml
index 91655f57..93b718a3 100644
--- a/candle-core/Cargo.toml
+++ b/candle-core/Cargo.toml
@@ -46,6 +46,6 @@ accelerate = ["dep:libc", "dep:accelerate-src"]
 metal = ["dep:metal", "dep:candle-metal-kernels"]
 
 [[bench]]
-name = "matmul"
+name = "bench_main"
 harness = false
 
diff --git a/candle-core/benches/bench_main.rs b/candle-core/benches/bench_main.rs
new file mode 100644
index 00000000..4425f2fb
--- /dev/null
+++ b/candle-core/benches/bench_main.rs
@@ -0,0 +1,4 @@
+mod benchmarks;
+
+use criterion::criterion_main;
+criterion_main!(benchmarks::matmul::benches);
diff --git a/candle-core/benches/matmul.rs b/candle-core/benches/benchmarks/matmul.rs
similarity index 85%
rename from candle-core/benches/matmul.rs
rename to candle-core/benches/benchmarks/matmul.rs
index a5dba9cc..fb173f04 100644
--- a/candle-core/benches/matmul.rs
+++ b/candle-core/benches/benchmarks/matmul.rs
@@ -1,9 +1,7 @@
-mod utils;
-
+use crate::benchmarks::{bench_name, device, BenchDevice};
 use candle_core::{DType, Tensor};
-use criterion::{black_box, criterion_group, criterion_main, Criterion, Throughput};
+use criterion::{black_box, criterion_group, Criterion, Throughput};
 use std::time::Instant;
-use utils::{bench_name, device, BenchDevice};
 
 fn run(a: &Tensor, b: &Tensor) {
     a.matmul(&b.t().unwrap()).unwrap();
@@ -38,4 +36,3 @@ fn criterion_benchmark(c: &mut Criterion) {
 }
 
 criterion_group!(benches, criterion_benchmark);
-criterion_main!(benches);
diff --git a/candle-core/benches/utils.rs b/candle-core/benches/benchmarks/mod.rs
similarity index 98%
rename from candle-core/benches/utils.rs
rename to candle-core/benches/benchmarks/mod.rs
index 3e8b3c57..1344770d 100644
--- a/candle-core/benches/utils.rs
+++ b/candle-core/benches/benchmarks/mod.rs
@@ -1,3 +1,5 @@
+pub(crate) mod matmul;
+
 use candle_core::{Device, Result};
 
 pub(crate) trait BenchDevice {