av_denoise/
device.rs

1use std::str::FromStr;
2
3/// Where to run the compute. The library maps each variant onto the
4/// concrete `Device` type of whichever cubecl runtime was selected.
5///
6/// Some variants are only meaningful for certain backends. `Integrated`
7/// and `Virtual` are wgpu-only; `Cpu` is a no-op on the `cpu` runtime
8/// and selects `WgpuDevice::Cpu` on wgpu. Asking for a variant on a
9/// runtime that can't honour it returns an error from the relevant
10/// `to_*` conversion.
11#[derive(Debug, Default, Clone, PartialEq, Eq)]
12pub enum Device {
13    /// Backend-chosen default device.
14    #[default]
15    Default,
16    /// Discrete GPU at ordinal `index`.
17    ///
18    /// Maps to `CudaDevice { index }`, `AmdDevice { index }`, or
19    /// `WgpuDevice::DiscreteGpu(index)`.
20    Discrete { index: usize },
21    /// Integrated GPU at ordinal `index`. wgpu-only.
22    Integrated { index: usize },
23    /// Virtual GPU at ordinal `index`. wgpu-only.
24    Virtual { index: usize },
25    /// Software/CPU device. Valid on the `cpu` runtime and on wgpu
26    /// (where it picks the lavapipe / software adapter).
27    Cpu,
28}
29
30impl FromStr for Device {
31    type Err = String;
32
33    /// `FromStr` accepts the same syntax as the bench CLI:
34    ///
35    /// - `default`
36    /// - `discrete[:N]`, `integrated[:N]`, `virtual[:N]` (default `N = 0`)
37    /// - `cpu`
38    fn from_str(s: &str) -> Result<Self, Self::Err> {
39        let (kind, idx) = s.split_once(':').unwrap_or((s, "0"));
40        let index: usize = idx
41            .parse()
42            .map_err(|_| format!("invalid device index '{idx}' in '{s}'"))?;
43        match kind {
44            "default" => Ok(Device::Default),
45            "discrete" => Ok(Device::Discrete { index }),
46            "integrated" => Ok(Device::Integrated { index }),
47            "virtual" => Ok(Device::Virtual { index }),
48            "cpu" => Ok(Device::Cpu),
49            other => Err(format!(
50                "unknown device kind '{other}'; expected default, discrete[:N], integrated[:N], virtual[:N], or cpu"
51            )),
52        }
53    }
54}
55
56#[cfg(feature = "cuda")]
57impl Device {
58    pub fn to_cuda(&self) -> Result<cubecl::cuda::CudaDevice, anyhow::Error> {
59        match self {
60            Device::Default => Ok(cubecl::cuda::CudaDevice { index: 0 }),
61            Device::Discrete { index } => Ok(cubecl::cuda::CudaDevice { index: *index }),
62            other => Err(anyhow::anyhow!(
63                "device {other:?} is not supported on the CUDA runtime; use `default` or `discrete[:N]`"
64            )),
65        }
66    }
67}
68
69#[cfg(feature = "rocm")]
70impl Device {
71    pub fn to_amd(&self) -> Result<cubecl::hip::AmdDevice, anyhow::Error> {
72        match self {
73            Device::Default => Ok(cubecl::hip::AmdDevice { index: 0 }),
74            Device::Discrete { index } => Ok(cubecl::hip::AmdDevice { index: *index }),
75            other => Err(anyhow::anyhow!(
76                "device {other:?} is not supported on the ROCm runtime; use `default` or `discrete[:N]`"
77            )),
78        }
79    }
80}
81
82#[cfg(any(feature = "vulkan", feature = "metal"))]
83impl Device {
84    pub fn to_wgpu(&self) -> Result<cubecl::wgpu::WgpuDevice, anyhow::Error> {
85        use cubecl::wgpu::WgpuDevice;
86        Ok(match self {
87            Device::Default => WgpuDevice::DefaultDevice,
88            Device::Discrete { index } => WgpuDevice::DiscreteGpu(*index),
89            Device::Integrated { index } => WgpuDevice::IntegratedGpu(*index),
90            Device::Virtual { index } => WgpuDevice::VirtualGpu(*index),
91            Device::Cpu => WgpuDevice::Cpu,
92        })
93    }
94}
95
96#[cfg(feature = "cpu")]
97impl Device {
98    pub fn to_cpu(&self) -> Result<cubecl::cpu::CpuDevice, anyhow::Error> {
99        match self {
100            Device::Default | Device::Cpu => Ok(cubecl::cpu::CpuDevice),
101            other => Err(anyhow::anyhow!(
102                "device {other:?} is not supported on the CPU runtime; use `default` or `cpu`"
103            )),
104        }
105    }
106}
107
108#[cfg(test)]
109mod tests {
110    use super::*;
111
112    #[test]
113    fn parse_default() {
114        assert_eq!("default".parse::<Device>().unwrap(), Device::Default);
115    }
116
117    #[test]
118    fn parse_discrete_with_and_without_index() {
119        assert_eq!(
120            "discrete".parse::<Device>().unwrap(),
121            Device::Discrete { index: 0 },
122        );
123        assert_eq!(
124            "discrete:3".parse::<Device>().unwrap(),
125            Device::Discrete { index: 3 },
126        );
127    }
128
129    #[test]
130    fn parse_integrated_virtual_cpu() {
131        assert_eq!(
132            "integrated:1".parse::<Device>().unwrap(),
133            Device::Integrated { index: 1 },
134        );
135        assert_eq!(
136            "virtual:2".parse::<Device>().unwrap(),
137            Device::Virtual { index: 2 },
138        );
139        assert_eq!("cpu".parse::<Device>().unwrap(), Device::Cpu);
140    }
141
142    #[test]
143    fn parse_rejects_unknown_kind() {
144        assert!("unicorn".parse::<Device>().is_err());
145    }
146
147    #[test]
148    fn parse_rejects_non_numeric_index() {
149        assert!("discrete:abc".parse::<Device>().is_err());
150    }
151
152    #[test]
153    fn default_is_default_variant() {
154        assert_eq!(Device::default(), Device::Default);
155    }
156
157    #[cfg(feature = "cpu")]
158    #[test]
159    fn cpu_runtime_rejects_gpu_variants() {
160        assert!(Device::Default.to_cpu().is_ok());
161        assert!(Device::Cpu.to_cpu().is_ok());
162        assert!(Device::Discrete { index: 0 }.to_cpu().is_err());
163    }
164}
av_denoise/device.rs

av_denoise/
device.rs
