From 87efb5d8eb6a6c3f17acf326aadcb11ad6900306 Mon Sep 17 00:00:00 2001
From: Ivar Flakstad <69173633+ivarflakstad@users.noreply.github.com>
Date: Tue, 9 Jan 2024 19:04:31 +0100
Subject: [PATCH] Updated feature separated benchmarks

---
 candle-core/Cargo.toml                        |  5 ---
 candle-core/benches/bench_main.rs             |  2 +-
 candle-core/benches/benchmarks/mod.rs         |  1 +
 .../benches/{ => benchmarks}/random.rs        | 32 +++++++------------
 4 files changed, 14 insertions(+), 26 deletions(-)
 rename candle-core/benches/{ => benchmarks}/random.rs (53%)

diff --git a/candle-core/Cargo.toml b/candle-core/Cargo.toml
index 3fae7f07..afdb67cd 100644
--- a/candle-core/Cargo.toml
+++ b/candle-core/Cargo.toml
@@ -48,8 +48,3 @@ metal = ["dep:metal", "dep:candle-metal-kernels"]
 [[bench]]
 name = "bench_main"
 harness = false
-
-[[bench]]
-name = "random"
-harness = false
-
diff --git a/candle-core/benches/bench_main.rs b/candle-core/benches/bench_main.rs
index 4425f2fb..8913df4f 100644
--- a/candle-core/benches/bench_main.rs
+++ b/candle-core/benches/bench_main.rs
@@ -1,4 +1,4 @@
 mod benchmarks;
 
 use criterion::criterion_main;
-criterion_main!(benchmarks::matmul::benches);
+criterion_main!(benchmarks::matmul::benches, benchmarks::random::benches);
diff --git a/candle-core/benches/benchmarks/mod.rs b/candle-core/benches/benchmarks/mod.rs
index 1344770d..6bb37a70 100644
--- a/candle-core/benches/benchmarks/mod.rs
+++ b/candle-core/benches/benchmarks/mod.rs
@@ -1,4 +1,5 @@
 pub(crate) mod matmul;
+pub(crate) mod random;
 
 use candle_core::{Device, Result};
 
diff --git a/candle-core/benches/random.rs b/candle-core/benches/benchmarks/random.rs
similarity index 53%
rename from candle-core/benches/random.rs
rename to candle-core/benches/benchmarks/random.rs
index 781d8b39..e4a4a390 100644
--- a/candle-core/benches/random.rs
+++ b/candle-core/benches/benchmarks/random.rs
@@ -1,9 +1,10 @@
+use crate::benchmarks::{bench_name, device, BenchDevice};
 use candle_core::{DType, Device, Tensor};
-use criterion::{black_box, criterion_group, criterion_main, Criterion, Throughput};
+use criterion::{black_box, criterion_group, Criterion, Throughput};
 use std::time::Instant;
 
 fn rand_uniform(a: &Tensor) {
-    a.rand_like(0.0, 1.0).unwrap();
+    a.rand_like(-1.0, 123.0).unwrap();
 }
 
 fn rand_normal(a: &Tensor) {
@@ -16,14 +17,13 @@ fn criterion_benchmark(c: &mut Criterion) {
     let rows = 2048;
     let cols = 2048;
 
-    let device = Device::new_metal(0).unwrap();
-    let device2 = device.clone();
+    let d = device().unwrap();
     let dtype = DType::F32;
-    let tensor = Tensor::zeros((b, rows, cols), dtype, &device).unwrap();
+    let tensor = Tensor::zeros((b, rows, cols), dtype, &d).unwrap();
 
-    let flops = b * rows * cols;
+    let flops = b * rows * cols * dtype.size_in_bytes();
 
-    let mut group = c.benchmark_group("metal_random_uniform");
+    let mut group = c.benchmark_group(bench_name("random_uniform"));
     group.throughput(Throughput::Bytes(flops as u64));
     group.bench_function("iter", move |benches| {
         benches.iter_custom(|iters| {
@@ -31,19 +31,16 @@ fn criterion_benchmark(c: &mut Criterion) {
             for _i in 0..iters {
                 rand_uniform(black_box(&tensor));
             }
-            if let Device::Metal(device) = &device {
-                device.wait_until_completed().unwrap();
-            } else {
-                panic!("Expected metal device");
-            }
+            d.sync().unwrap();
             start.elapsed()
         })
     });
     group.finish();
 
-    let tensor = Tensor::zeros((b, rows, cols), dtype, &device2).unwrap();
+    let d = device().unwrap();
+    let tensor = Tensor::zeros((b, rows, cols), dtype, &d).unwrap();
 
-    let mut group = c.benchmark_group("metal_random_normal");
+    let mut group = c.benchmark_group(bench_name("random_normal"));
     group.throughput(Throughput::Bytes(flops as u64));
     group.bench_function("iter", move |benches| {
         benches.iter_custom(|iters| {
@@ -51,11 +48,7 @@ fn criterion_benchmark(c: &mut Criterion) {
             for _i in 0..iters {
                 rand_normal(black_box(&tensor));
             }
-            if let Device::Metal(device) = &device2 {
-                device.wait_until_completed().unwrap();
-            } else {
-                panic!("Expected metal device");
-            }
+            d.sync().unwrap();
             start.elapsed()
         })
     });
@@ -63,4 +56,3 @@ fn criterion_benchmark(c: &mut Criterion) {
 }
 
 criterion_group!(benches, criterion_benchmark);
-criterion_main!(benches);