Use cfg to seperate benchmark results based on features

1 files changed, 56 insertions, 0 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"
+    };
+}