Updated feature separated benchmarks

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
-

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);

candle-core/benches/benchmarks/mod.rs

@@ -1,4 +1,5 @@
 pub(crate) mod matmul;
+pub(crate) mod random;
 
 use candle_core::{Device, Result};
 

candle-core/benches/random.rs → 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,51 +17,42 @@ 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| {
             let start = Instant::now();
             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| {
             let start = Instant::now();
             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()
         })
     });
     group.finish();
 }
 
 criterion_group!(benches, criterion_benchmark);
-criterion_main!(benches);