Improve benchmarks layout

candle-core/benches/benchmarks/matmul.rs

@@ -0,0 +1,38 @@
+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) {
+    a.matmul(&b.t().unwrap()).unwrap();
+}
+
+fn criterion_benchmark(c: &mut Criterion) {
+    let b = 1;
+    let m = 1;
+    let n = 2048;
+    let k = 2048;
+
+    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(bench_name("matmul"));
+    group.throughput(Throughput::Bytes(flops as u64));
+    group.bench_function("iter", move |b| {
+        b.iter_custom(|iters| {
+            let start = Instant::now();
+            for _i in 0..iters {
+                run(black_box(&lhs), black_box(&rhs));
+            }
+            device.sync().unwrap();
+            start.elapsed()
+        })
+    });
+    group.finish();
+}
+
+criterion_group!(benches, criterion_benchmark);

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"
+    }
+}