Merge branch 'ivarflakstad/seperate-benchmarks-by-feature' into ivarflakstad/metal-prng

candle-core/Cargo.toml

@@ -46,7 +46,7 @@ accelerate = ["dep:libc", "dep:accelerate-src"]
 metal = ["dep:metal", "dep:candle-metal-kernels"]
 
 [[bench]]
-name = "matmul"
+name = "bench_main"
 harness = false
 
 [[bench]]

candle-core/benches/bench_main.rs

@@ -0,0 +1,4 @@
+mod benchmarks;
+
+use criterion::criterion_main;
+criterion_main!(benchmarks::matmul::benches);

candle-core/benches/benchmarks/matmul.rs

@@ -1,5 +1,6 @@
-use candle_core::{DType, Device, Tensor};
-use criterion::{black_box, criterion_group, criterion_main, Criterion, Throughput};
+use crate::benchmarks::{bench_name, device, BenchDevice};
+use candle_core::{DType, Tensor};
+use criterion::{black_box, criterion_group, Criterion, Throughput};
 use std::time::Instant;
 
 fn run(a: &Tensor, b: &Tensor) {
@@ -12,14 +13,14 @@ fn criterion_benchmark(c: &mut Criterion) {
     let n = 2048;
     let k = 2048;
 
-    let device = Device::new_metal(0).unwrap();
+    let device = device().unwrap();
     let dtype = DType::F32;
     let lhs = Tensor::zeros((b, m, k), dtype, &device).unwrap();
     let rhs = Tensor::zeros((b, n, k), dtype, &device).unwrap();
 
     let flops = b * m * n * k;
 
-    let mut group = c.benchmark_group("matmul_metal");
+    let mut group = c.benchmark_group(bench_name("matmul"));
     group.throughput(Throughput::Bytes(flops as u64));
     group.bench_function("iter", move |b| {
         b.iter_custom(|iters| {
@@ -27,11 +28,7 @@ fn criterion_benchmark(c: &mut Criterion) {
             for _i in 0..iters {
                 run(black_box(&lhs), black_box(&rhs));
             }
-            if let Device::Metal(device) = &device {
-                device.wait_until_completed().unwrap();
-            } else {
-                panic!("Expected metal device");
-            }
+            device.sync().unwrap();
             start.elapsed()
         })
     });
@@ -39,4 +36,3 @@ fn criterion_benchmark(c: &mut Criterion) {
 }
 
 criterion_group!(benches, criterion_benchmark);
-criterion_main!(benches);

candle-core/benches/benchmarks/mod.rs

@@ -0,0 +1,55 @@
+pub(crate) mod matmul;
+
+use candle_core::{Device, Result};
+
+pub(crate) trait BenchDevice {
+    fn sync(&self) -> Result<()>;
+}
+
+impl BenchDevice for Device {
+    fn sync(&self) -> Result<()> {
+        match self {
+            Device::Cpu => Ok(()),
+            Device::Cuda(device) => {
+                #[cfg(feature = "cuda")]
+                return Ok(device.synchronize()?);
+                #[cfg(not(feature = "cuda"))]
+                panic!("Cuda device without cuda feature enabled: {:?}", device)
+            }
+            Device::Metal(device) => {
+                #[cfg(feature = "metal")]
+                return Ok(device.wait_until_completed()?);
+                #[cfg(not(feature = "metal"))]
+                panic!("Metal device without metal feature enabled: {:?}", device)
+            }
+        }
+    }
+}
+
+pub(crate) fn device() -> Result<Device> {
+    if cfg!(feature = "metal") {
+        Device::new_metal(0)
+    } else if cfg!(feature = "cuda") {
+        Device::new_cuda(0)
+    } else {
+        Ok(Device::Cpu)
+    }
+}
+
+pub(crate) fn bench_name<S: Into<String>>(name: S) -> String {
+    format!("{}_{}", device_variant(), name.into())
+}
+
+const fn device_variant() -> &'static str {
+    if cfg!(feature = "metal") {
+        "metal"
+    } else if cfg!(feature = "cuda") {
+        "cuda"
+    } else if cfg!(feature = "accelerate") {
+        "accelerate"
+    } else if cfg!(feature = "mkl") {
+        "mkl"
+    } else {
+        "cpu"
+    }
+}