cubecl_runtime/device/wgpu.rs
1use cubecl_common::device::{Device, DeviceId};
2
3/// The graphics API a [`WgpuDevice`] comes up on, which also settles the
4/// shader compiler it is fed: SPIR-V on Vulkan, MSL on Metal, WGSL elsewhere.
5///
6/// [`Auto`](WgpuBackend::Auto) leaves the choice to the runtime, which tries
7/// what this machine offers in the order it would want. Naming one instead
8/// pins it: the device is only reachable on that API, and stays that device
9/// wherever its id travels.
10#[derive(Clone, Copy, Debug, Hash, PartialEq, Eq, Default)]
11pub enum WgpuBackend {
12 /// Whichever of the machine's APIs the runtime would pick.
13 #[default]
14 Auto,
15 /// Vulkan, compiling to `SPIR-V` where the build and the device allow it.
16 Vulkan,
17 /// Metal, compiling to MSL where the build allows it.
18 Metal,
19 /// `DirectX` 12.
20 Dx12,
21 /// `OpenGL`.
22 Gl,
23 /// `WebGPU`, the browser's own.
24 WebGpu,
25}
26
27impl WgpuBackend {
28 /// The three bits this rides in at the top of a [`DeviceId`]'s index.
29 const SHIFT: u32 = 13;
30 const MASK: u16 = 0x7;
31
32 fn from_bits(bits: u16) -> Self {
33 match bits {
34 1 => Self::Vulkan,
35 2 => Self::Metal,
36 3 => Self::Dx12,
37 4 => Self::Gl,
38 5 => Self::WebGpu,
39 _ => Self::Auto,
40 }
41 }
42}
43
44/// Which of a machine's wgpu devices, leaving the graphics API aside.
45#[derive(Clone, Debug, Hash, PartialEq, Eq, Default)]
46pub enum WgpuDeviceKind {
47 /// Discrete GPU with the given index. The index is the index of the discrete GPU in the list
48 /// of all discrete GPUs found on the system.
49 DiscreteGpu(usize),
50
51 /// Integrated GPU with the given index. The index is the index of the integrated GPU in the
52 /// list of all integrated GPUs found on the system.
53 IntegratedGpu(usize),
54
55 /// Virtual GPU with the given index. The index is the index of the virtual GPU in the list of
56 /// all virtual GPUs found on the system.
57 VirtualGpu(usize),
58
59 /// CPU — a software rasterizer such as lavapipe, llvmpipe or WARP.
60 Cpu,
61
62 /// The best available device found with the current graphics API.
63 ///
64 /// This will prioritize GPUs wgpu recognizes as "high power". Additionally, you can override this using
65 /// the `CUBECL_WGPU_DEFAULT_DEVICE` environment variable. This variable is spelled as if i was a `WgpuDevice`,
66 /// so for example `CUBECL_WGPU_DEFAULT_DEVICE=IntegratedGpu(1)` or `CUBECL_WGPU_DEFAULT_DEVICE=Cpu`
67 #[default]
68 DefaultDevice,
69
70 /// Use an externally created, existing, wgpu setup. This is helpful when using `CubeCL` in conjunction
71 /// with some existing wgpu setup (eg. egui or bevy), as resources can be transferred in & out of `CubeCL`.
72 ///
73 /// The graphics API came up with the setup, so pinning one on top of this
74 /// means nothing and is not encoded.
75 ///
76 /// # Notes
77 ///
78 /// This can be initialized with `init_device` in the wgpu runtime.
79 Existing(u32),
80
81 /// An adapter wgpu recognizes but does not classify — what it reports as
82 /// `DeviceType::Other`, which `OpenGL` drivers often do. Indexed among
83 /// its own kind like the rest.
84 Other(usize),
85}
86
87impl WgpuDeviceKind {
88 /// The largest index that survives a round trip through a [`DeviceId`].
89 ///
90 /// The graphics API rides in the top three bits of the index, leaving
91 /// thirteen — far more adapters of one kind than a machine has. A larger
92 /// index is kept at this one rather than wrapped: wrapped, it lands on a
93 /// device that exists; kept here, it names one no machine has, and naming
94 /// it is refused.
95 pub const MAX_INDEX: usize = (1 << WgpuBackend::SHIFT) - 1;
96
97 fn type_id(&self) -> u16 {
98 match *self {
99 Self::DiscreteGpu(_) => 0,
100 Self::IntegratedGpu(_) => 1,
101 Self::VirtualGpu(_) => 2,
102 Self::Cpu => 3,
103 Self::DefaultDevice => 4,
104 Self::Existing(_) => 5,
105 Self::Other(_) => 6,
106 }
107 }
108
109 fn index(&self) -> u16 {
110 match *self {
111 Self::DiscreteGpu(index)
112 | Self::IntegratedGpu(index)
113 | Self::VirtualGpu(index)
114 | Self::Other(index) => index.min(Self::MAX_INDEX) as u16,
115 Self::Cpu | Self::DefaultDevice => 0,
116 Self::Existing(id) => id as u16,
117 }
118 }
119}
120
121/// The device struct when using the `wgpu` backend.
122///
123/// Which device, and which graphics API to reach it on — two things a machine
124/// answers separately, so they are two fields rather than one flattened list.
125///
126/// # Example
127///
128/// ```ignore
129/// use cubecl_wgpu::{WgpuDevice, WgpuDeviceKind};
130///
131/// let first = WgpuDevice::from(WgpuDeviceKind::DiscreteGpu(0)); // First discrete GPU found.
132/// let vulkan = first.clone().on(WgpuBackend::Vulkan); // The same one, on Vulkan.
133/// ```
134#[derive(Clone, Debug, Hash, PartialEq, Eq, Default)]
135pub struct WgpuDevice {
136 /// Which of the machine's devices.
137 pub kind: WgpuDeviceKind,
138 /// The graphics API to reach it on.
139 pub backend: WgpuBackend,
140}
141
142impl WgpuDevice {
143 /// This device on whichever graphics API the runtime would pick.
144 pub fn new(kind: WgpuDeviceKind) -> Self {
145 Self {
146 kind,
147 backend: WgpuBackend::Auto,
148 }
149 }
150
151 /// The same device, pinned to `backend`.
152 ///
153 /// An [externally created](WgpuDeviceKind::Existing) device came up on an
154 /// API of its own, so pinning one there changes nothing.
155 pub fn on(self, backend: WgpuBackend) -> Self {
156 match self.kind {
157 WgpuDeviceKind::Existing(_) => self,
158 kind => Self { kind, backend },
159 }
160 }
161}
162
163impl From<WgpuDeviceKind> for WgpuDevice {
164 fn from(kind: WgpuDeviceKind) -> Self {
165 Self::new(kind)
166 }
167}
168
169impl Device for WgpuDevice {
170 fn from_id(device_id: DeviceId) -> Self {
171 let index = (device_id.index_id & WgpuDeviceKind::MAX_INDEX as u16) as usize;
172
173 let kind = match device_id.type_id {
174 0 => WgpuDeviceKind::DiscreteGpu(index),
175 1 => WgpuDeviceKind::IntegratedGpu(index),
176 2 => WgpuDeviceKind::VirtualGpu(index),
177 3 => WgpuDeviceKind::Cpu,
178 5 => return Self::new(WgpuDeviceKind::Existing(device_id.index_id as u32)),
179 6 => WgpuDeviceKind::Other(index),
180 _ => WgpuDeviceKind::DefaultDevice,
181 };
182
183 Self {
184 kind,
185 backend: WgpuBackend::from_bits(
186 (device_id.index_id >> WgpuBackend::SHIFT) & WgpuBackend::MASK,
187 ),
188 }
189 }
190
191 fn to_id(&self) -> DeviceId {
192 let index = self.kind.index();
193
194 // An external setup keeps the whole index to itself: there is no API
195 // to pin on it, so there are no bits to spend on one.
196 let index_id = match self.kind {
197 WgpuDeviceKind::Existing(_) => index,
198 _ => ((self.backend as u16) << WgpuBackend::SHIFT) | index,
199 };
200
201 DeviceId::new(self.kind.type_id(), index_id)
202 }
203}
204
205#[cfg(test)]
206mod tests {
207 use super::*;
208
209 /// Both halves have to survive the id, or a device pinned to one API comes
210 /// back on another.
211 #[test]
212 fn a_device_round_trips_its_kind_and_its_backend() {
213 for kind in [
214 WgpuDeviceKind::DiscreteGpu(2),
215 WgpuDeviceKind::IntegratedGpu(1),
216 WgpuDeviceKind::VirtualGpu(0),
217 WgpuDeviceKind::Cpu,
218 WgpuDeviceKind::DefaultDevice,
219 WgpuDeviceKind::Other(3),
220 ] {
221 for backend in [
222 WgpuBackend::Auto,
223 WgpuBackend::Vulkan,
224 WgpuBackend::Metal,
225 WgpuBackend::Dx12,
226 WgpuBackend::Gl,
227 WgpuBackend::WebGpu,
228 ] {
229 let device = WgpuDevice::new(kind.clone()).on(backend);
230
231 assert_eq!(WgpuDevice::from_id(device.to_id()), device);
232 }
233 }
234 }
235
236 /// The same device on two APIs must not share an id, or one client serves
237 /// both and the second caller silently gets the first one's backend.
238 #[test]
239 fn pinning_a_backend_changes_the_id() {
240 let auto = WgpuDevice::new(WgpuDeviceKind::DiscreteGpu(0));
241
242 let vulkan = auto.clone().on(WgpuBackend::Vulkan);
243
244 assert_ne!(auto.to_id(), vulkan.to_id());
245 }
246
247 /// An index too wide for the id must not wrap onto a device that exists,
248 /// which is what masking it did: the 8192nd discrete GPU became the first.
249 #[test]
250 fn an_index_too_wide_for_the_id_is_kept_at_the_largest() {
251 let too_wide = WgpuDevice::new(WgpuDeviceKind::DiscreteGpu(WgpuDeviceKind::MAX_INDEX + 1));
252
253 assert_eq!(
254 WgpuDevice::from_id(too_wide.to_id()),
255 WgpuDevice::new(WgpuDeviceKind::DiscreteGpu(WgpuDeviceKind::MAX_INDEX))
256 );
257 }
258
259 /// An external setup spends its whole index on the id it was registered
260 /// with, so pinning an API on it is refused rather than silently truncating.
261 #[test]
262 fn an_existing_device_keeps_its_whole_index() {
263 let existing = WgpuDevice::new(WgpuDeviceKind::Existing(60_000));
264
265 let pinned = existing.clone().on(WgpuBackend::Vulkan);
266
267 assert_eq!(pinned, existing);
268 assert_eq!(WgpuDevice::from_id(existing.to_id()), existing);
269 }
270}