Skip to main content

av_denoise/
sniff.rs

1use cubecl::prelude::*;
2
3use crate::accelerate::Accelerator;
4
5/// Probe each enabled accelerator and return the first one whose client
6/// can be created and synchronised. cubecl 0.10 kernels are fully
7/// asynchronous, so a successful `client.sync()` is sufficient to
8/// confirm the backend is usable (no test kernel needed).
9pub 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            // docs.rs widens `Accelerator` variants behind `cfg(docsrs)` so
23            // they appear in the rendered enum even when the matching
24            // backend feature is off. This arm makes the match exhaustive
25            // in that configuration; it is never reached at runtime.
26            #[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}