Merge branch 'main' into ivarflakstad/metal-fill

candle-core/benches/benchmarks/fill.rs

@@ -1,4 +1,4 @@
-use crate::benchmarks::{bench_name, device, BenchDevice};
+use crate::benchmarks::{BenchDevice, BenchDeviceHandler};
 use candle_core::{DType, Device, Tensor};
 use criterion::{black_box, criterion_group, Criterion, Throughput};
 use std::time::Instant;
@@ -7,16 +7,14 @@ fn run(shape: (usize, usize, usize), dtype: DType, device: &Device) {
     Tensor::ones(shape, dtype, device).unwrap();
 }
 
-fn run_fill_benchmark(c: &mut Criterion, name: &str, dtype: DType) {
+fn run_fill_benchmark(c: &mut Criterion, device: &Device, name: &str, dtype: DType) {
     let b = 1;
-    let rows = 4096;
-    let columns = 4096;
+    let rows = 1024;
+    let columns = 1024;
 
     let flops = b * rows * columns * dtype.size_in_bytes();
 
-    let device = device().unwrap();
-
-    let mut group = c.benchmark_group(bench_name(name));
+    let mut group = c.benchmark_group(device.bench_name(name));
     group.throughput(Throughput::Bytes(flops as u64));
     group.bench_function("iter", move |bencher| {
         bencher.iter_custom(|iters| {
@@ -36,8 +34,11 @@ fn run_fill_benchmark(c: &mut Criterion, name: &str, dtype: DType) {
 }
 
 fn criterion_benchmark(c: &mut Criterion) {
-    run_fill_benchmark(c, "fill_u8", DType::U8);
-    run_fill_benchmark(c, "fill_f32", DType::F32);
+    let handler = BenchDeviceHandler::new().unwrap();
+    for device in handler.devices {
+        run_fill_benchmark(c, &device, "fill_u8", DType::U8);
+        run_fill_benchmark(c, &device, "fill_f32", DType::F32);
+    }
 }
 
 criterion_group!(benches, criterion_benchmark);

candle-core/benches/benchmarks/matmul.rs

@@ -1,5 +1,5 @@
-use crate::benchmarks::{bench_name, device, BenchDevice};
-use candle_core::{DType, Tensor};
+use crate::benchmarks::{BenchDevice, BenchDeviceHandler};
+use candle_core::{DType, Device, Tensor};
 use criterion::{black_box, criterion_group, Criterion, Throughput};
 use std::time::Instant;
 
@@ -7,20 +7,19 @@ fn run(a: &Tensor, b: &Tensor) {
     a.matmul(&b.t().unwrap()).unwrap();
 }
 
-fn criterion_benchmark(c: &mut Criterion) {
+fn run_bench(c: &mut Criterion, device: &Device) {
     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 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"));
+    let mut group = c.benchmark_group(device.bench_name("matmul"));
     group.throughput(Throughput::Bytes(flops as u64));
     group.bench_function("iter", move |b| {
         b.iter_custom(|iters| {
@@ -35,4 +34,11 @@ fn criterion_benchmark(c: &mut Criterion) {
     group.finish();
 }
 
+fn criterion_benchmark(c: &mut Criterion) {
+    let handler = BenchDeviceHandler::new().unwrap();
+    for device in handler.devices {
+        run_bench(c, &device);
+    }
+}
+
 criterion_group!(benches, criterion_benchmark);

candle-core/benches/benchmarks/mod.rs

@@ -5,6 +5,8 @@ use candle_core::{Device, Result};
 
 pub(crate) trait BenchDevice {
     fn sync(&self) -> Result<()>;
+
+    fn bench_name<S: Into<String>>(&self, name: S) -> String;
 }
 
 impl BenchDevice for Device {
@@ -25,32 +27,38 @@ impl BenchDevice for 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)
+    fn bench_name<S: Into<String>>(&self, name: S) -> String {
+        match self {
+            Device::Cpu => {
+                let cpu_type = if cfg!(feature = "accelerate") {
+                    "accelerate"
+                } else if cfg!(feature = "mkl") {
+                    "mkl"
+                } else {
+                    "cpu"
+                };
+                format!("{}_{}", cpu_type, name.into())
+            }
+            Device::Cuda(_) => format!("cuda_{}", name.into()),
+            Device::Metal(_) => format!("metal_{}", name.into()),
+        }
     }
 }
 
-pub(crate) fn bench_name<S: Into<String>>(name: S) -> String {
-    format!("{}_{}", device_variant(), name.into())
+struct BenchDeviceHandler {
+    devices: Vec<Device>,
 }
 
-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"
+impl BenchDeviceHandler {
+    pub fn new() -> Result<Self> {
+        let mut devices = Vec::new();
+        if cfg!(feature = "metal") {
+            devices.push(Device::new_metal(0)?);
+        } else if cfg!(feature = "cuda") {
+            devices.push(Device::new_cuda(0)?);
+        }
+        devices.push(Device::Cpu);
+        Ok(Self { devices })
     }
 }