1use cubecl::prelude::*;
2
3use crate::accelerate::Accelerator;
4
5pub fn sniff_best_accelerator(enable: &[Accelerator]) -> Option<Accelerator> {
10 for accelerator in enable {
11 let is_enabled = match accelerator {
12 #[cfg(feature = "cuda")]
13 Accelerator::Cuda => probe_runtime::<cubecl::cuda::CudaRuntime>("CUDA"),
14 #[cfg(feature = "rocm")]
15 Accelerator::Rocm => probe_runtime::<cubecl::hip::HipRuntime>("ROCM"),
16 #[cfg(feature = "vulkan")]
17 Accelerator::Vulkan => probe_runtime::<cubecl::wgpu::WgpuRuntime>("VULKAN"),
18 #[cfg(feature = "metal")]
19 Accelerator::Metal => probe_runtime::<cubecl::wgpu::WgpuRuntime>("METAL"),
20 #[cfg(feature = "cpu")]
21 Accelerator::Cpu => probe_runtime::<cubecl::cpu::CpuRuntime>("CPU"),
22 #[cfg(docsrs)]
27 #[allow(unreachable_patterns)]
28 _ => unreachable!(),
29 };
30
31 if is_enabled {
32 return Some(*accelerator);
33 }
34 }
35
36 None
37}
38
39fn probe_runtime<R: Runtime>(name: &'static str) -> bool {
40 let device = <R::Device as Default>::default();
41 let client = R::client(&device);
42 match cubecl::future::block_on(client.sync()) {
43 Ok(()) => true,
44 Err(err) => {
45 tracing::debug!(err = ?err, "could not use {name} runtime");
46 false
47 },
48 }
49}