burn_dispatch/
device.rs

1use burn_backend::{DeviceId, DeviceOps, DeviceSettings};
2
3use crate::backends::*;
4
5#[cfg(feature = "autodiff")]
6use alloc::boxed::Box;
7
8// Throughput types come from `burn-backend` (which re-exports them from cubecl),
9// so `burn-dispatch` needs no direct `cubecl` dependency.
10#[cfg(feature = "cubecl")]
11use alloc::vec::Vec;
12#[cfg(feature = "cubecl")]
13use burn_backend::cubecl::{ThroughputKey, ThroughputValue};
14
15/// Represents a device for the [`Dispatch`](crate::Dispatch).
16///
17/// Each variant corresponds to a backend that the [`Dispatch`](crate::Dispatch) can dispatch operations to.
18///
19/// # Example
20///
21/// ```ignore
22/// use burn::DispatchDevice;
23///
24/// #[cfg(feature = "cpu")]
25/// let cpu_device = DispatchDevice::Cpu(Default::default());
26///
27/// #[cfg(feature = "cuda")]
28/// let cuda_device = DispatchDevice::Cuda(Default::default());
29/// ```
30#[derive(Clone, Eq)]
31pub enum DispatchDevice {
32    /// The [CPU backend](Cpu) device.
33    #[cfg(feature = "cpu")]
34    Cpu(CpuDevice),
35
36    /// The [CUDA backend](Cuda) device.
37    #[cfg(feature = "cuda")]
38    Cuda(CudaDevice),
39
40    /// The [Metal backend](Metal) device (via WGPU runtime).
41    #[cfg(feature = "metal")]
42    Metal(WgpuDevice),
43
44    /// The [ROCm backend](Rocm) device.
45    #[cfg(feature = "rocm")]
46    Rocm(RocmDevice),
47
48    /// The [Vulkan backend](Vulkan) device.
49    #[cfg(feature = "vulkan")]
50    Vulkan(WgpuDevice),
51
52    /// The [Wgpu backend](Wgpu) device (via WGPU runtime with auto-selected compiler).
53    #[cfg(feature = "wgpu")]
54    Wgpu(WgpuDevice),
55
56    /// The [WebGPU backend](WebGpu) device (via WGPU runtime).
57    #[cfg(feature = "webgpu")]
58    WebGpu(WgpuDevice),
59
60    /// The [Flex backend](Flex) device (CPU-only).
61    #[cfg(any(feature = "flex", default_backend))]
62    Flex(FlexDevice),
63
64    /// The [NdArray backend](NdArray) device (CPU-only).
65    #[cfg(feature = "ndarray")]
66    NdArray(NdArrayDevice),
67
68    /// The [LibTorch backend](LibTorch) device.
69    #[cfg(feature = "tch")]
70    LibTorch(LibTorchDevice),
71
72    /// The [remote backend](Remote) device, identified by a network address.
73    #[cfg(feature = "remote")]
74    Remote(RemoteDevice),
75
76    /// The [autodiff enabled backend](Autodiff) device.
77    #[cfg(feature = "autodiff")]
78    Autodiff(AutodiffDevice),
79}
80
81#[cfg(feature = "cubecl")]
82impl DispatchDevice {
83    /// Measure peak throughput for this device against the given `keys`.
84    ///
85    /// Only cubecl-backed devices can measure throughput; other backends
86    /// (ndarray, libtorch, remote, ...) return an empty vector. An autodiff
87    /// device reports the peaks of the device it wraps. Each returned
88    /// [`ThroughputValue`](burn_backend::cubecl::ThroughputValue) corresponds
89    /// positionally to the key at the same index.
90    pub fn performance_stats(&self, keys: &[ThroughputKey]) -> Vec<ThroughputValue> {
91        // No catch-all arm: a new backend must fail to compile here rather
92        // than silently report no peaks.
93        match self {
94            #[cfg(feature = "cpu")]
95            DispatchDevice::Cpu(device) => burn_cpu::device_throughput(device, keys),
96            #[cfg(feature = "cuda")]
97            DispatchDevice::Cuda(device) => burn_cuda::device_throughput(device, keys),
98            #[cfg(feature = "rocm")]
99            DispatchDevice::Rocm(device) => burn_rocm::device_throughput(device, keys),
100            #[cfg(feature = "wgpu")]
101            DispatchDevice::Wgpu(device) => burn_wgpu::device_throughput(device, keys),
102            #[cfg(feature = "vulkan")]
103            DispatchDevice::Vulkan(device) => burn_wgpu::device_throughput(device, keys),
104            #[cfg(feature = "metal")]
105            DispatchDevice::Metal(device) => burn_wgpu::device_throughput(device, keys),
106            #[cfg(feature = "webgpu")]
107            DispatchDevice::WebGpu(device) => burn_wgpu::device_throughput(device, keys),
108            // Autodiff does not change the hardware, so measure the wrapped device.
109            #[cfg(feature = "autodiff")]
110            DispatchDevice::Autodiff(device) => device.performance_stats(keys),
111
112            // Not cubecl-backed, so there are no kernels to measure.
113            #[cfg(any(feature = "flex", default_backend))]
114            DispatchDevice::Flex(_) => Vec::new(),
115            #[cfg(feature = "ndarray")]
116            DispatchDevice::NdArray(_) => Vec::new(),
117            #[cfg(feature = "tch")]
118            DispatchDevice::LibTorch(_) => Vec::new(),
119
120            // The kernels run on the server, which this local API cannot reach.
121            #[cfg(feature = "remote")]
122            DispatchDevice::Remote(_) => Vec::new(),
123        }
124    }
125}
126
127#[cfg(feature = "autodiff")]
128// This tuple struct mainly restricts users from creating Autodiff(Autodiff) devices.
129/// A wrapper that enables automatic differentiation for a [`DispatchDevice`].
130///
131/// Use [`DispatchDevice::autodiff`] to construct this type.
132#[derive(Debug, Clone, PartialEq, Eq)]
133pub struct AutodiffDevice {
134    pub(crate) inner: Box<DispatchDevice>,
135    pub(crate) checkpointing: CheckpointingStrategy,
136}
137
138#[cfg(feature = "autodiff")]
139impl AutodiffDevice {
140    pub(crate) fn new(device: DispatchDevice, checkpointing: CheckpointingStrategy) -> Self {
141        Self {
142            inner: Box::new(device),
143            checkpointing,
144        }
145    }
146
147    /// Returns the underlying device, removing the autodiff capability.
148    pub fn inner(self) -> DispatchDevice {
149        *self.inner
150    }
151}
152
153#[cfg(feature = "autodiff")]
154// Useful for match in dispatch macros
155impl core::ops::Deref for AutodiffDevice {
156    type Target = DispatchDevice;
157
158    fn deref(&self) -> &Self::Target {
159        &self.inner
160    }
161}
162
163#[allow(missing_docs)]
164#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
165/// Checkpointing strategy for autodiff.
166#[repr(u8)]
167pub enum CheckpointingStrategy {
168    Balanced,
169    #[default]
170    None,
171}
172
173#[cfg(feature = "autodiff")]
174pub(crate) fn validate_checkpointing(
175    lhs: Option<crate::CheckpointingStrategy>,
176    rhs: Option<crate::CheckpointingStrategy>,
177) -> Option<crate::CheckpointingStrategy> {
178    match (lhs, rhs) {
179        (Some(lhs), Some(rhs)) => {
180            assert_eq!(
181                lhs, rhs,
182                "Autodiff strategy mismatch: {lhs:?} vs {rhs:?}. Tensors in the same operation must share a strategy."
183            );
184            Some(lhs)
185        }
186        (None, None) => None,
187        // When tensors are created on non-autodiff device there is no checkpointing, but
188        // tensor created with autodiff which moved out (`tensor.inner()`) will still carry the state.
189        // In such cases, we can "promote" the checkpointing.
190        (None, rhs) => rhs,
191        (lhs, None) => lhs,
192    }
193}
194
195impl core::fmt::Debug for DispatchDevice {
196    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
197        match self {
198            #[cfg(feature = "cpu")]
199            Self::Cpu(device) => f.debug_tuple("Cpu").field(device).finish(),
200            #[cfg(feature = "cuda")]
201            Self::Cuda(device) => f.debug_tuple("Cuda").field(device).finish(),
202            #[cfg(feature = "metal")]
203            Self::Metal(device) => f.debug_tuple("Metal").field(device).finish(),
204            #[cfg(feature = "rocm")]
205            Self::Rocm(device) => f.debug_tuple("Rocm").field(device).finish(),
206            #[cfg(feature = "vulkan")]
207            Self::Vulkan(device) => f.debug_tuple("Vulkan").field(device).finish(),
208            #[cfg(feature = "wgpu")]
209            Self::Wgpu(device) => f.debug_tuple("Wgpu").field(device).finish(),
210            #[cfg(feature = "webgpu")]
211            Self::WebGpu(device) => f.debug_tuple("WebGpu").field(device).finish(),
212            #[cfg(any(feature = "flex", default_backend))]
213            Self::Flex(device) => f.debug_tuple("Flex").field(device).finish(),
214            #[cfg(feature = "ndarray")]
215            Self::NdArray(device) => f.debug_tuple("NdArray").field(device).finish(),
216            #[cfg(feature = "tch")]
217            Self::LibTorch(device) => f.debug_tuple("LibTorch").field(device).finish(),
218            #[cfg(feature = "remote")]
219            Self::Remote(device) => f.debug_tuple("Remote").field(device).finish(),
220            #[cfg(feature = "autodiff")]
221            // Format without `AutodiffDevice` wrapper
222            Self::Autodiff(device) => f.debug_tuple("Autodiff").field(&device.inner).finish(),
223        }
224    }
225}
226
227impl Default for DispatchDevice {
228    #[allow(unreachable_code)]
229    fn default() -> Self {
230        // TODO: which priority?
231        // Single override e.g. `BURN_DEVICE=vulkan` forces Vulkan or panics if not available.
232        // Priority list e.g. `BURN_DEVICE_PRIORITY=cuda,vulkan,cpu` sets the order.
233        // Both could be tied into `burn.toml` config
234        // For now we just use `BURN_DEVICE` on CI to force a single device
235
236        #[cfg(feature = "std")]
237        {
238            if let Ok(device_str) = std::env::var("BURN_DEVICE") {
239                match device_str.to_lowercase().as_str() {
240                    "cuda" => {
241                        #[cfg(feature = "cuda")]
242                        return Self::Cuda(CudaDevice::default());
243                        panic!(
244                            "BURN_DEVICE=cuda requested, but the 'cuda' feature is not enabled."
245                        );
246                    }
247                    "metal" => {
248                        #[cfg(feature = "metal")]
249                        return Self::Metal(burn_wgpu::WgpuDevice::default());
250                        panic!(
251                            "BURN_DEVICE=metal requested, but the 'metal' feature is not enabled."
252                        );
253                    }
254                    "rocm" => {
255                        #[cfg(feature = "rocm")]
256                        return Self::Rocm(RocmDevice::default());
257                        panic!(
258                            "BURN_DEVICE=rocm requested, but the 'rocm' feature is not enabled."
259                        );
260                    }
261                    "vulkan" => {
262                        #[cfg(feature = "vulkan")]
263                        return Self::Vulkan(burn_wgpu::WgpuDevice::default());
264                        panic!(
265                            "BURN_DEVICE=vulkan requested, but the 'vulkan' feature is not enabled."
266                        );
267                    }
268                    "webgpu" => {
269                        #[cfg(feature = "webgpu")]
270                        return Self::WebGpu(burn_wgpu::WgpuDevice::default());
271                        panic!(
272                            "BURN_DEVICE=webgpu requested, but the 'webgpu' feature is not enabled."
273                        );
274                    }
275                    "wgpu" => {
276                        #[cfg(feature = "wgpu")]
277                        return Self::Wgpu(burn_wgpu::WgpuDevice::default());
278                        panic!(
279                            "BURN_DEVICE=wgpu requested, but the 'wgpu' feature is not enabled."
280                        );
281                    }
282                    "cpu" => {
283                        #[cfg(feature = "cpu")]
284                        return Self::Cpu(CpuDevice);
285                        panic!("BURN_DEVICE=cpu requested, but the 'cpu' feature is not enabled.");
286                    }
287                    "tch" => {
288                        #[cfg(feature = "tch")]
289                        return Self::LibTorch(LibTorchDevice::default());
290                        panic!("BURN_DEVICE=tch requested, but the 'tch' feature is not enabled.");
291                    }
292                    "remote" => {
293                        #[cfg(feature = "remote")]
294                        return Self::Remote(RemoteDevice::default());
295                        panic!(
296                            "BURN_DEVICE=remote requested, but the 'remote' feature is not enabled."
297                        );
298                    }
299                    "flex" => {
300                        #[cfg(any(feature = "flex", default_backend))]
301                        return Self::Flex(FlexDevice);
302                        panic!(
303                            "BURN_DEVICE=flex requested, but the 'flex' feature is not enabled."
304                        );
305                    }
306                    "ndarray" => {
307                        #[cfg(feature = "ndarray")]
308                        return Self::NdArray(NdArrayDevice::default());
309                        panic!(
310                            "BURN_DEVICE=ndarray requested, but the 'ndarray' feature is not enabled."
311                        );
312                    }
313                    _ => panic!("Unknown BURN_DEVICE override: '{}'.", device_str),
314                }
315            }
316        }
317
318        #[cfg(feature = "cuda")]
319        return Self::Cuda(CudaDevice::default());
320
321        #[cfg(feature = "metal")]
322        return Self::Metal(burn_wgpu::WgpuDevice::default());
323
324        #[cfg(feature = "rocm")]
325        return Self::Rocm(RocmDevice::default());
326
327        #[cfg(feature = "vulkan")]
328        return Self::Vulkan(burn_wgpu::WgpuDevice::default());
329
330        #[cfg(feature = "webgpu")]
331        return Self::WebGpu(burn_wgpu::WgpuDevice::default());
332
333        #[cfg(feature = "wgpu")]
334        return Self::Wgpu(burn_wgpu::WgpuDevice::default());
335
336        #[cfg(feature = "cpu")]
337        return Self::Cpu(CpuDevice);
338
339        #[cfg(feature = "tch")]
340        return Self::LibTorch(LibTorchDevice::default());
341
342        // Prefer Flex over NdArray when both are enabled: Flex is the long-term
343        // CPU backend replacement and should win the default tie.
344        #[cfg(any(feature = "flex", default_backend))]
345        return Self::Flex(FlexDevice);
346
347        #[cfg(feature = "remote")]
348        return Self::Remote(RemoteDevice::default());
349
350        #[cfg(feature = "ndarray")]
351        return Self::NdArray(NdArrayDevice::default());
352    }
353}
354
355impl PartialEq for DispatchDevice {
356    /// Compares devices based on hardware identity.
357    ///
358    /// Returns `true` if both devices represent the same compute resource.
359    /// Note that this comparison ignores autodiff and checkpointing settings.
360    fn eq(&self, other: &Self) -> bool {
361        match (self, other) {
362            // If both are Autodiff, compare the inner devices
363            #[cfg(feature = "autodiff")]
364            (DispatchDevice::Autodiff(a), DispatchDevice::Autodiff(b)) => a == b,
365            // If one is Autodiff, compare it to the raw device
366            #[cfg(feature = "autodiff")]
367            (DispatchDevice::Autodiff(a), b) => a.inner.as_ref() == b,
368            #[cfg(feature = "autodiff")]
369            (a, DispatchDevice::Autodiff(b)) => a == b.inner.as_ref(),
370            #[cfg(feature = "cpu")]
371            (Self::Cpu(a), Self::Cpu(b)) => a == b,
372            #[cfg(feature = "cuda")]
373            (Self::Cuda(a), Self::Cuda(b)) => a == b,
374            #[cfg(feature = "metal")]
375            (Self::Metal(a), Self::Metal(b)) => a == b,
376            #[cfg(feature = "rocm")]
377            (Self::Rocm(a), Self::Rocm(b)) => a == b,
378            #[cfg(feature = "vulkan")]
379            (Self::Vulkan(a), Self::Vulkan(b)) => a == b,
380            #[cfg(feature = "wgpu")]
381            (Self::Wgpu(a), Self::Wgpu(b)) => a == b,
382            #[cfg(feature = "webgpu")]
383            (Self::WebGpu(a), Self::WebGpu(b)) => a == b,
384            #[cfg(any(feature = "flex", default_backend))]
385            (Self::Flex(a), Self::Flex(b)) => a == b,
386            #[cfg(feature = "ndarray")]
387            (Self::NdArray(a), Self::NdArray(b)) => a == b,
388            #[cfg(feature = "tch")]
389            (Self::LibTorch(a), Self::LibTorch(b)) => a == b,
390            #[cfg(feature = "remote")]
391            (Self::Remote(a), Self::Remote(b)) => a == b,
392            #[allow(unreachable_patterns)]
393            (_, _) => false,
394        }
395    }
396}
397
398const INTERNAL_ID_MASK: u16 = 0x00FF;
399const BACKEND_SHIFT: u32 = 8;
400
401impl DispatchDevice {
402    #[cfg(feature = "autodiff")]
403    /// Creates a new [`DispatchDevice`] with [automatic differentiation](Autodiff) enabled.
404    pub fn autodiff(device: impl Into<DispatchDevice>) -> DispatchDevice {
405        Self::autodiff_checkpointed(device, CheckpointingStrategy::None)
406    }
407    #[cfg(feature = "autodiff")]
408    /// Creates a new [`DispatchDevice`] with [automatic differentiation](Autodiff) enabled.
409    pub fn autodiff_checkpointed(
410        device: impl Into<DispatchDevice>,
411        checkpointing: CheckpointingStrategy,
412    ) -> DispatchDevice {
413        let device = device.into();
414        DispatchDevice::Autodiff(AutodiffDevice::new(device, checkpointing))
415    }
416
417    /// Returns the inner device, without autodiff (when enabled).
418    pub fn inner(self) -> Self {
419        #[cfg(feature = "autodiff")]
420        if let DispatchDevice::Autodiff(device) = self {
421            return *device.inner;
422        }
423
424        self
425    }
426
427    /// Returns a unique number per variant to encode into type_id.
428    fn backend_id(&self) -> DispatchDeviceId {
429        match self {
430            #[cfg(feature = "cpu")]
431            Self::Cpu(_) => DispatchDeviceId::Cpu,
432            #[cfg(feature = "cuda")]
433            Self::Cuda(_) => DispatchDeviceId::Cuda,
434            #[cfg(feature = "metal")]
435            Self::Metal(_) => DispatchDeviceId::Metal,
436            #[cfg(feature = "rocm")]
437            Self::Rocm(_) => DispatchDeviceId::Rocm,
438            #[cfg(feature = "vulkan")]
439            Self::Vulkan(_) => DispatchDeviceId::Vulkan,
440            #[cfg(feature = "wgpu")]
441            Self::Wgpu(_) => DispatchDeviceId::Wgpu,
442            #[cfg(feature = "webgpu")]
443            Self::WebGpu(_) => DispatchDeviceId::WebGpu,
444            #[cfg(any(feature = "flex", default_backend))]
445            Self::Flex(_) => DispatchDeviceId::Flex,
446            #[cfg(feature = "ndarray")]
447            Self::NdArray(_) => DispatchDeviceId::NdArray,
448            #[cfg(feature = "tch")]
449            Self::LibTorch(_) => DispatchDeviceId::LibTorch,
450            #[cfg(feature = "remote")]
451            Self::Remote(_) => DispatchDeviceId::Remote,
452            #[cfg(feature = "autodiff")]
453            Self::Autodiff(device) => device.inner.backend_id(),
454        }
455    }
456
457    /// Encode variant ID and backend type ID into a unique `type_id`.
458    fn encode_type_id(&self, backend_type_id: u16) -> u16 {
459        // Use the lower 8 bits for the backend's internal type ID
460        let internal_type_id = backend_type_id & INTERNAL_ID_MASK;
461        // Use the upper 8 bits for the DispatchDevice/DispatchDeviceId
462        let backend = u16::from(self.backend_id()) << BACKEND_SHIFT;
463        backend | internal_type_id
464    }
465
466    /// Decode an encoded `type_id` into variant ID and backend type ID.
467    pub(crate) fn decode_type_id(type_id: u16) -> (DispatchDeviceId, u16) {
468        let backend_raw = type_id >> BACKEND_SHIFT;
469        let internal_type_id = type_id & INTERNAL_ID_MASK;
470
471        let backend = DispatchDeviceId::try_from(backend_raw).expect("Unknown DispatchDevice ID");
472
473        (backend, internal_type_id)
474    }
475}
476
477#[allow(missing_docs)]
478#[derive(Debug, Clone, Copy, PartialEq, Eq)]
479#[repr(u16)]
480pub enum DispatchDeviceId {
481    Cpu = 0,
482    Cuda = 1,
483    Wgpu = 2,
484    Rocm = 3,
485    Flex = 4,
486    LibTorch = 5,
487    NdArray = 6,
488    Metal = 7,
489    Vulkan = 8,
490    WebGpu = 9,
491    Remote = 10,
492}
493
494impl From<DispatchDeviceId> for u16 {
495    fn from(variant: DispatchDeviceId) -> Self {
496        variant as u16
497    }
498}
499
500impl TryFrom<u16> for DispatchDeviceId {
501    type Error = ();
502
503    fn try_from(value: u16) -> Result<Self, Self::Error> {
504        match value {
505            #[cfg(feature = "cpu")]
506            0 => Ok(Self::Cpu),
507            #[cfg(feature = "cuda")]
508            1 => Ok(Self::Cuda),
509            #[cfg(feature = "wgpu")]
510            2 => Ok(Self::Wgpu),
511            #[cfg(feature = "rocm")]
512            3 => Ok(Self::Rocm),
513            #[cfg(any(feature = "flex", default_backend))]
514            4 => Ok(Self::Flex),
515            #[cfg(feature = "tch")]
516            5 => Ok(Self::LibTorch),
517            #[cfg(feature = "ndarray")]
518            6 => Ok(Self::NdArray),
519            #[cfg(feature = "metal")]
520            7 => Ok(Self::Metal),
521            #[cfg(feature = "vulkan")]
522            8 => Ok(Self::Vulkan),
523            #[cfg(feature = "webgpu")]
524            9 => Ok(Self::WebGpu),
525            #[cfg(feature = "remote")]
526            10 => Ok(Self::Remote),
527            _ => Err(()),
528        }
529    }
530}
531
532impl DeviceOps for DispatchDevice {
533    fn defaults(&self) -> DeviceSettings {
534        match self {
535            #[cfg(feature = "cpu")]
536            Self::Cpu(device) => device.defaults(),
537            #[cfg(feature = "cuda")]
538            Self::Cuda(device) => device.defaults(),
539            #[cfg(feature = "metal")]
540            Self::Metal(device) => device.defaults(),
541            #[cfg(feature = "rocm")]
542            Self::Rocm(device) => device.defaults(),
543            #[cfg(feature = "vulkan")]
544            Self::Vulkan(device) => device.defaults(),
545            #[cfg(feature = "wgpu")]
546            Self::Wgpu(device) => device.defaults(),
547            #[cfg(feature = "webgpu")]
548            Self::WebGpu(device) => device.defaults(),
549            #[cfg(any(feature = "flex", default_backend))]
550            Self::Flex(device) => device.defaults(),
551            #[cfg(feature = "ndarray")]
552            Self::NdArray(device) => device.defaults(),
553            #[cfg(feature = "tch")]
554            Self::LibTorch(device) => device.defaults(),
555            #[cfg(feature = "remote")]
556            Self::Remote(device) => device.defaults(),
557            #[cfg(feature = "autodiff")]
558            Self::Autodiff(device) => device.inner.defaults(),
559        }
560    }
561}
562
563impl burn_backend::Device for DispatchDevice {
564    fn from_id(mut device_id: DeviceId) -> Self {
565        let (dispatch_id, backend_type_id) = Self::decode_type_id(device_id.type_id);
566        device_id.type_id = backend_type_id;
567
568        match dispatch_id {
569            #[cfg(feature = "cpu")]
570            DispatchDeviceId::Cpu => Self::Cpu(CpuDevice::from_id(device_id)),
571            #[cfg(feature = "cuda")]
572            DispatchDeviceId::Cuda => Self::Cuda(CudaDevice::from_id(device_id)),
573            #[cfg(feature = "metal")]
574            DispatchDeviceId::Metal => Self::Metal(WgpuDevice::from_id(device_id)),
575            #[cfg(feature = "rocm")]
576            DispatchDeviceId::Rocm => Self::Rocm(RocmDevice::from_id(device_id)),
577            #[cfg(feature = "vulkan")]
578            DispatchDeviceId::Vulkan => Self::Vulkan(WgpuDevice::from_id(device_id)),
579            #[cfg(feature = "wgpu")]
580            DispatchDeviceId::Wgpu => Self::Wgpu(WgpuDevice::from_id(device_id)),
581            #[cfg(feature = "webgpu")]
582            DispatchDeviceId::WebGpu => Self::WebGpu(WgpuDevice::from_id(device_id)),
583            #[cfg(any(feature = "flex", default_backend))]
584            DispatchDeviceId::Flex => Self::Flex(FlexDevice::from_id(device_id)),
585            #[cfg(feature = "ndarray")]
586            DispatchDeviceId::NdArray => Self::NdArray(NdArrayDevice::from_id(device_id)),
587            #[cfg(feature = "tch")]
588            DispatchDeviceId::LibTorch => Self::LibTorch(LibTorchDevice::from_id(device_id)),
589            #[cfg(feature = "remote")]
590            DispatchDeviceId::Remote => Self::Remote(RemoteDevice::from_id(device_id)),
591            _ => unreachable!("No backend feature enabled."),
592        }
593    }
594
595    fn to_id(&self) -> DeviceId {
596        let mut device_id = match self {
597            #[cfg(feature = "cpu")]
598            Self::Cpu(device) => device.to_id(),
599            #[cfg(feature = "cuda")]
600            Self::Cuda(device) => device.to_id(),
601            #[cfg(feature = "metal")]
602            Self::Metal(device) => device.to_id(),
603            #[cfg(feature = "rocm")]
604            Self::Rocm(device) => device.to_id(),
605            #[cfg(feature = "vulkan")]
606            Self::Vulkan(device) => device.to_id(),
607            #[cfg(feature = "wgpu")]
608            Self::Wgpu(device) => device.to_id(),
609            #[cfg(feature = "webgpu")]
610            Self::WebGpu(device) => device.to_id(),
611            #[cfg(any(feature = "flex", default_backend))]
612            Self::Flex(device) => device.to_id(),
613            #[cfg(feature = "ndarray")]
614            Self::NdArray(device) => device.to_id(),
615            #[cfg(feature = "tch")]
616            Self::LibTorch(device) => device.to_id(),
617            #[cfg(feature = "remote")]
618            Self::Remote(device) => device.to_id(),
619            #[cfg(feature = "autodiff")]
620            Self::Autodiff(device) => device.inner.to_id(),
621        };
622        device_id.type_id = self.encode_type_id(device_id.type_id);
623        device_id
624    }
625}
626
627#[cfg(feature = "cpu")]
628impl From<CpuDevice> for DispatchDevice {
629    fn from(device: CpuDevice) -> Self {
630        DispatchDevice::Cpu(device)
631    }
632}
633
634#[cfg(feature = "cuda")]
635impl From<CudaDevice> for DispatchDevice {
636    fn from(device: CudaDevice) -> Self {
637        DispatchDevice::Cuda(device)
638    }
639}
640
641#[cfg(feature = "rocm")]
642impl From<RocmDevice> for DispatchDevice {
643    fn from(device: RocmDevice) -> Self {
644        DispatchDevice::Rocm(device)
645    }
646}
647
648// A bare `WgpuDevice` maps to the auto-compiler [`DispatchDevice::Wgpu`] variant. To target a
649// specific wgpu specialization (Metal, Vulkan, WebGpu) construct the variant explicitly.
650//
651// The gates form a priority chain (metal, then vulkan, then webgpu) rather than mutually
652// exclusive conditions: cargo unifies features across a workspace, so several specializations
653// can be on at once and exclusive gates would leave the conversion with no impl at all.
654#[cfg(all(
655    feature = "wgpu",
656    not(any(feature = "metal", feature = "vulkan", feature = "webgpu"))
657))]
658impl From<WgpuDevice> for DispatchDevice {
659    fn from(device: WgpuDevice) -> Self {
660        DispatchDevice::Wgpu(device)
661    }
662}
663
664#[cfg(feature = "metal")]
665impl From<WgpuDevice> for DispatchDevice {
666    fn from(device: WgpuDevice) -> Self {
667        DispatchDevice::Metal(device)
668    }
669}
670
671#[cfg(all(feature = "vulkan", not(feature = "metal")))]
672impl From<WgpuDevice> for DispatchDevice {
673    fn from(device: WgpuDevice) -> Self {
674        DispatchDevice::Vulkan(device)
675    }
676}
677
678#[cfg(all(feature = "webgpu", not(any(feature = "metal", feature = "vulkan"))))]
679impl From<WgpuDevice> for DispatchDevice {
680    fn from(device: WgpuDevice) -> Self {
681        DispatchDevice::WebGpu(device)
682    }
683}
684
685#[cfg(any(feature = "flex", default_backend))]
686impl From<FlexDevice> for DispatchDevice {
687    fn from(device: FlexDevice) -> Self {
688        DispatchDevice::Flex(device)
689    }
690}
691
692#[cfg(feature = "ndarray")]
693impl From<NdArrayDevice> for DispatchDevice {
694    fn from(device: NdArrayDevice) -> Self {
695        DispatchDevice::NdArray(device)
696    }
697}
698
699#[cfg(feature = "tch")]
700impl From<LibTorchDevice> for DispatchDevice {
701    fn from(device: LibTorchDevice) -> Self {
702        DispatchDevice::LibTorch(device)
703    }
704}
705
706#[cfg(feature = "remote")]
707impl From<RemoteDevice> for DispatchDevice {
708    fn from(device: RemoteDevice) -> Self {
709        DispatchDevice::Remote(device)
710    }
711}
burn_dispatch/device.rs

burn_dispatch/
device.rs
