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

burn_dispatch/
device.rs
