burn_dispatch/
device.rs

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

burn_dispatch/
device.rs
