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 { + 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>(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() }) });