av_denoise/
device.rs

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

av_denoise/
device.rs
