2 files changed, 64 insertions, 8 deletions

diff --git a/candle-core/benches/bench_utils.rs b/candle-core/benches/bench_utils.rs
new file mode 100644
index 00000000..75800761
--- /dev/null
+++ b/candle-core/benches/bench_utils.rs
@@ -0,0 +1,56 @@
+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)
+            }
+        }
+    }
+}
+
+#[allow(dead_code)]
+pub(crate) fn device() -> Result<Device> {
+    return if cfg!(feature = "metal") {
+        Device::new_metal(0)
+    } else if cfg!(feature = "cuda") {
+        Device::new_cuda(0)
+    } else {
+        Ok(Device::Cpu)
+    };
+}
+
+#[allow(dead_code)]
+pub(crate) fn bench_name<S: Into<String>>(name: S) -> String {
+    format!("{}_{}", device_variant(), name.into())
+}
+
+#[allow(dead_code)]
+const fn device_variant() -> &'static str {
+    return 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"
+    };
+}
diff --git a/candle-core/benches/matmul.rs b/candle-core/benches/matmul.rs
index 83679771..4f7dfa6c 100644
--- a/candle-core/benches/matmul.rs
+++ b/candle-core/benches/matmul.rs
@@ -1,4 +1,8 @@
-use candle_core::{DType, Device, Tensor};
+mod bench_utils;
+
+use crate::bench_utils::bench_name;
+use bench_utils::{device, BenchDevice};
+use candle_core::{DType, Tensor};
 use criterion::{black_box, criterion_group, criterion_main, Criterion, Throughput};
 use std::time::Instant;
 
@@ -12,14 +16,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 +31,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()
         })
     });