burn_dispatch/

device.rs

use burn_backend::{DeviceId, DeviceOps, DeviceSettings};

#[allow(unused_imports)] // Empty in backend-free and capture-only builds.
use crate::devices::*;

#[cfg(feature = "capture")]
use burn_capture::CaptureDevice;
#[cfg(feature = "group")]
use burn_group::{GroupDevice, GroupError};

#[cfg(feature = "autodiff")]
use alloc::boxed::Box;

// Throughput types come from `burn-backend` (which re-exports them from cubecl),
// so `burn-dispatch` needs no direct `cubecl` dependency.
#[cfg(feature = "cubecl")]
use alloc::vec::Vec;
#[cfg(feature = "cubecl")]
use burn_backend::cubecl::{DeviceIdentity, ThroughputError, ThroughputKey, ThroughputValue};
// `cubecl` without a runtime feature gives the throughput *types* but no `Cube` device to
// measure, so the measurement itself follows `cube_backend` rather than the feature.
#[cfg(cube_backend)]
use burn_backend::cubecl::measure_peak_throughput;

/// Represents a device for the [`Dispatch`](crate::Dispatch).
///
/// Each variant corresponds to a backend that the [`Dispatch`](crate::Dispatch) can dispatch operations to.
///
/// # Example
///
/// ```ignore
/// use burn::DispatchDevice;
///
/// // One variant covers every cubecl runtime; the device inside says which.
/// #[cfg(feature = "cuda")]
/// let cuda_device = DispatchDevice::Cube(cubecl::Device::Cuda(Default::default()));
///
/// #[cfg(feature = "flex")]
/// let flex_device = DispatchDevice::Flex(Default::default());
/// ```
#[derive(Clone, Eq)]
pub enum DispatchDevice {
    #[cfg(not(backend_enabled))]
    #[doc(hidden)]
    Unavailable(crate::NoBackend),
    /// A device of the [cubecl backend](crate::backends::Cube): CUDA, ROCm, Metal, Vulkan,
    /// WebGPU, wgpu or the CPU runtime.
    #[cfg(cube_backend)]
    Cube(CubeDevice),

    /// The [Flex backend](crate::backends::Flex) device (CPU-only).
    #[cfg(feature = "flex")]
    Flex(FlexDevice),

    /// The [remote backend](crate::backends::Remote) device, identified by a network address.
    #[cfg(feature = "remote")]
    Remote(RemoteDevice),

    /// A non-executing graph capture device.
    #[cfg(feature = "capture")]
    Capture(CaptureDevice),

    /// A [group of devices](crate::backends::Group) that a tensor can be split over.
    #[cfg(feature = "group")]
    Group(GroupDevice),

    /// The [autodiff enabled backend](crate::backends::Autodiff) device.
    #[cfg(feature = "autodiff")]
    Autodiff(AutodiffDevice),
}

#[cfg(feature = "cubecl")]
impl DispatchDevice {
    /// Who this device is, `None` for a backend that does not report one. An autodiff device
    /// answers for the device it wraps. Opens the device.
    pub fn identity(&self) -> Option<DeviceIdentity> {
        match self {
            #[cfg(cube_backend)]
            DispatchDevice::Cube(device) => Some(device.client().properties().identity.clone()),
            #[cfg(feature = "autodiff")]
            DispatchDevice::Autodiff(device) => device.inner.identity(),
            #[allow(unreachable_patterns)]
            _ => None,
        }
    }

    /// Measure peak throughput for this device against the given `keys`.
    ///
    /// Only cubecl-backed devices can measure throughput; other backends
    /// (remote, ...) return an empty vector. An autodiff
    /// device reports the peaks of the device it wraps. Each returned result
    /// corresponds positionally to the key at the same index, and carries a
    /// [`ThroughputError`] where the device has no peak for that key.
    // With `cubecl` on but no runtime compiled in, every arm below ignores `keys`.
    #[cfg_attr(not(cube_backend), allow(unused_variables))]
    pub fn performance_stats(
        &self,
        keys: &[ThroughputKey],
    ) -> Vec<Result<ThroughputValue, ThroughputError>> {
        // No catch-all arm: a new backend must fail to compile here rather
        // than silently report no peaks.
        match self {
            #[cfg(not(backend_enabled))]
            Self::Unavailable(never) => never.unreachable(),
            #[cfg(cube_backend)]
            DispatchDevice::Cube(device) => {
                let client = device.client();
                keys.iter()
                    .map(|key| measure_peak_throughput(&client, *key))
                    .collect()
            }
            // Autodiff does not change the hardware, so measure the wrapped device.
            #[cfg(feature = "autodiff")]
            DispatchDevice::Autodiff(device) => device.performance_stats(keys),

            // Not cubecl-backed, so there are no kernels to measure.
            #[cfg(feature = "flex")]
            DispatchDevice::Flex(_) => Vec::new(),

            // The kernels run on the server, which this local API cannot reach.
            #[cfg(feature = "remote")]
            DispatchDevice::Remote(_) => Vec::new(),
            // A group's peaks are its members', measured on them.
            #[cfg(feature = "group")]
            DispatchDevice::Group(_) => Vec::new(),
            #[cfg(feature = "capture")]
            DispatchDevice::Capture(_) => Vec::new(),
        }
    }
}

#[cfg(feature = "autodiff")]
// This tuple struct mainly restricts users from creating Autodiff(Autodiff) devices.
/// A wrapper that enables automatic differentiation for a [`DispatchDevice`].
///
/// Use [`DispatchDevice::autodiff`] to construct this type.
#[derive(Debug, Clone)]
pub struct AutodiffDevice {
    pub(crate) inner: Box<DispatchDevice>,
    pub(crate) checkpointing: GradientCheckpointingStrategy,
}

/// Compares on hardware identity only, ignoring the checkpointing strategy, so that this agrees
/// with [`DispatchDevice`]'s own [`PartialEq`] — which has to ignore it, since comparing an
/// `Autodiff` device against a raw one has no strategy to compare against. A derived impl would
/// make `Autodiff(a) == Autodiff(b)` disagree with `DispatchDevice::Autodiff(a) ==
/// DispatchDevice::Autodiff(b)`.
///
/// Use [`gradient_checkpointing_strategy`](Self::gradient_checkpointing_strategy) when the
/// strategy is what you actually need to compare.
#[cfg(feature = "autodiff")]
impl PartialEq for AutodiffDevice {
    fn eq(&self, other: &Self) -> bool {
        self.inner == other.inner
    }
}

#[cfg(feature = "autodiff")]
impl Eq for AutodiffDevice {}

#[cfg(feature = "autodiff")]
impl AutodiffDevice {
    pub(crate) fn new(
        device: DispatchDevice,
        checkpointing: GradientCheckpointingStrategy,
    ) -> Self {
        Self {
            inner: Box::new(device),
            checkpointing,
        }
    }

    /// Returns the underlying device, removing the autodiff capability.
    pub fn inner(self) -> DispatchDevice {
        *self.inner
    }

    /// Returns the gradient checkpointing strategy.
    pub fn gradient_checkpointing_strategy(&self) -> GradientCheckpointingStrategy {
        self.checkpointing
    }
}

#[cfg(feature = "autodiff")]
// Useful for match in dispatch macros
impl core::ops::Deref for AutodiffDevice {
    type Target = DispatchDevice;

    fn deref(&self) -> &Self::Target {
        &self.inner
    }
}

#[allow(missing_docs)]
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
/// Gradient checkpointing strategy for autodiff.
#[repr(u8)]
pub enum GradientCheckpointingStrategy {
    /// Recompute selected activations during backpropagation to reduce peak memory usage.
    Balanced,
    /// Disable gradient checkpointing while retaining autodiff tracking.
    #[default]
    Disabled,
}

impl core::fmt::Debug for DispatchDevice {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            #[cfg(not(backend_enabled))]
            Self::Unavailable(never) => never.unreachable(),
            #[cfg(cube_backend)]
            Self::Cube(device) => f.debug_tuple("Cube").field(device).finish(),
            #[cfg(feature = "flex")]
            Self::Flex(device) => f.debug_tuple("Flex").field(device).finish(),
            #[cfg(feature = "remote")]
            Self::Remote(device) => f.debug_tuple("Remote").field(device).finish(),
            #[cfg(feature = "group")]
            Self::Group(device) => f.debug_tuple("Group").field(device).finish(),
            #[cfg(feature = "capture")]
            Self::Capture(device) => f.debug_tuple("Capture").field(device).finish(),
            #[cfg(feature = "autodiff")]
            // Format without `AutodiffDevice` wrapper
            Self::Autodiff(device) => f
                .debug_struct("Autodiff")
                .field("device", &device.inner)
                .field("checkpointing", &device.checkpointing)
                .finish(),
        }
    }
}

impl Default for DispatchDevice {
    /// Select an enabled backend in this order: CUDA, Metal, ROCm, Vulkan, WebGPU,
    /// wgpu, CPU, Flex, Remote. `BURN_DEVICE` overrides this in
    /// std builds. Capture devices must be constructed explicitly.
    ///
    /// The `metal`, `vulkan`, and `webgpu` overrides require their matching Cargo features
    /// and pin the graphics API. Explicit Metal selection requires native MSL support.
    /// The `wgpu` override keeps automatic API and compiler selection with WGSL fallback.
    ///
    /// Panics when no execution backend is enabled, or when `BURN_DEVICE` names an unknown
    /// backend or one whose Cargo feature is not enabled.
    #[allow(unreachable_code)]
    fn default() -> Self {
        // BURN_DEVICE selects one compiled backend or reports a configuration error.

        #[cfg(feature = "std")]
        {
            if let Ok(device_str) = std::env::var("BURN_DEVICE") {
                match device_str.to_lowercase().as_str() {
                    // Every cubecl runtime is the one `Cube` backend; the name here
                    // picks the runtime and, for named wgpu APIs, pins the graphics API.
                    "cuda" => {
                        #[cfg(feature = "cuda")]
                        return Self::Cube(CubeDevice::Cuda(Default::default()));
                        panic!(
                            "BURN_DEVICE=cuda requested, but the 'cuda' feature is not enabled."
                        );
                    }
                    "rocm" => {
                        #[cfg(feature = "rocm")]
                        return Self::Cube(CubeDevice::Hip(Default::default()));
                        panic!(
                            "BURN_DEVICE=rocm requested, but the 'rocm' feature is not enabled."
                        );
                    }
                    "metal" => {
                        #[cfg(feature = "metal")]
                        return Self::Cube(CubeDevice::Wgpu(
                            WgpuDevice::default().on(WgpuBackend::Metal),
                        ));
                        panic!(
                            "BURN_DEVICE=metal requested, but the 'metal' feature is not enabled."
                        );
                    }
                    "vulkan" => {
                        #[cfg(feature = "vulkan")]
                        return Self::Cube(CubeDevice::Wgpu(
                            WgpuDevice::default().on(WgpuBackend::Vulkan),
                        ));
                        panic!(
                            "BURN_DEVICE=vulkan requested, but the 'vulkan' feature is not enabled."
                        );
                    }
                    "webgpu" => {
                        #[cfg(feature = "webgpu")]
                        return Self::Cube(CubeDevice::Wgpu(
                            WgpuDevice::default().on(WgpuBackend::WebGpu),
                        ));
                        panic!(
                            "BURN_DEVICE=webgpu requested, but the 'webgpu' feature is not enabled."
                        );
                    }
                    "wgpu" => {
                        #[cfg(feature = "wgpu")]
                        return Self::Cube(CubeDevice::Wgpu(Default::default()));
                        panic!(
                            "BURN_DEVICE=wgpu requested, but the 'wgpu' feature is not enabled."
                        );
                    }
                    "cpu" => {
                        #[cfg(feature = "cpu")]
                        return Self::Cube(CubeDevice::Cpu(Default::default()));
                        panic!("BURN_DEVICE=cpu requested, but the 'cpu' feature is not enabled.");
                    }
                    "remote" => {
                        #[cfg(feature = "remote")]
                        return Self::Remote(RemoteDevice::default());
                        panic!(
                            "BURN_DEVICE=remote requested, but the 'remote' feature is not enabled."
                        );
                    }
                    "flex" => {
                        #[cfg(feature = "flex")]
                        return Self::Flex(FlexDevice);
                        panic!(
                            "BURN_DEVICE=flex requested, but the 'flex' feature is not enabled."
                        );
                    }
                    _ => panic!("Unknown BURN_DEVICE override: '{}'.", device_str),
                }
            }
        }

        // Spelled out per feature rather than left to `CubeDevice::default()`: that answers for
        // the runtimes *cubecl* compiled in, and cargo unifies features across a build, so a
        // workspace that also builds `burn-cuda` would hand this crate a CUDA default even when
        // it was built with only `wgpu`. The order is the one a caller who did not choose would
        // want — a discrete accelerator, then the portable path, then the CPU.
        #[cfg(feature = "cuda")]
        return Self::Cube(CubeDevice::Cuda(Default::default()));

        #[cfg(feature = "metal")]
        return Self::Cube(CubeDevice::Wgpu(Default::default()));

        #[cfg(feature = "rocm")]
        return Self::Cube(CubeDevice::Hip(Default::default()));

        #[cfg(feature = "vulkan")]
        return Self::Cube(CubeDevice::Wgpu(Default::default()));

        #[cfg(feature = "webgpu")]
        return Self::Cube(CubeDevice::Wgpu(Default::default()));

        #[cfg(feature = "wgpu")]
        return Self::Cube(CubeDevice::Wgpu(Default::default()));

        #[cfg(feature = "cpu")]
        return Self::Cube(CubeDevice::Cpu(Default::default()));

        #[cfg(feature = "flex")]
        return Self::Flex(FlexDevice);

        #[cfg(feature = "remote")]
        return Self::Remote(RemoteDevice::default());

        panic!(
            "No execution backend is enabled. Enable a Burn backend feature such as `flex`, \
             `wgpu`, or `cuda`. To record a graph without executing it, enable `capture` \
             and use Device::capture()."
        );
    }
}

impl PartialEq for DispatchDevice {
    /// Compares devices based on hardware identity.
    ///
    /// Returns `true` if both devices represent the same compute resource.
    /// Note that this comparison ignores autodiff and checkpointing settings.
    fn eq(&self, other: &Self) -> bool {
        match (self, other) {
            // If both are Autodiff, compare the inner devices
            #[cfg(feature = "autodiff")]
            (DispatchDevice::Autodiff(a), DispatchDevice::Autodiff(b)) => {
                a.inner.as_ref() == b.inner.as_ref()
            }
            // If one is Autodiff, compare it to the raw device
            #[cfg(feature = "autodiff")]
            (DispatchDevice::Autodiff(a), b) => a.inner.as_ref() == b,
            #[cfg(feature = "autodiff")]
            (a, DispatchDevice::Autodiff(b)) => a == b.inner.as_ref(),
            #[cfg(cube_backend)]
            (Self::Cube(a), Self::Cube(b)) => a == b,
            #[cfg(feature = "flex")]
            (Self::Flex(a), Self::Flex(b)) => a == b,
            #[cfg(feature = "remote")]
            (Self::Remote(a), Self::Remote(b)) => a == b,
            #[cfg(feature = "group")]
            (Self::Group(a), Self::Group(b)) => a == b,
            #[cfg(feature = "capture")]
            (Self::Capture(a), Self::Capture(b)) => a == b,
            #[allow(unreachable_patterns)]
            (_, _) => false,
        }
    }
}

const INTERNAL_ID_MASK: u16 = 0x00FF;
const BACKEND_SHIFT: u32 = 8;

impl DispatchDevice {
    /// Create the dispatch representation used by the high-level graph-capture device.
    #[cfg(feature = "capture")]
    #[doc(hidden)]
    pub fn capture() -> Self {
        Self::Capture(CaptureDevice::default())
    }

    /// A new group of `devices`, its members, that a tensor can be split over, apart from every
    /// other group.
    ///
    /// Make a group once and reuse it, as any device: each call registers a new group for the
    /// rest of the process.
    ///
    /// Each of its default dtypes is the first of the members' configured ones, in member order,
    /// that every member supports.
    ///
    /// # Errors
    ///
    /// When `devices` is empty or mixes backends, when they are not Cube or Flex devices, when no
    /// member's default for a dtype runs on every member, or when the members quantize
    /// differently. An autodiff group is a group of plain devices made autodiff afterwards.
    ///
    /// # Panics
    ///
    /// When the process already has 65,536 groups.
    #[cfg(feature = "group")]
    pub fn group(devices: alloc::vec::Vec<DispatchDevice>) -> Result<Self, GroupError> {
        if let Some(error) = devices.iter().find_map(Self::refused_as_member) {
            return Err(error);
        }
        let group = match devices.first() {
            #[cfg(cube_backend)]
            Some(Self::Cube(_)) => {
                let members = group_members!(devices, Cube)?;
                GroupDevice::with_member_settings::<crate::backends::Cube>(&members)
            }
            #[cfg(feature = "flex")]
            Some(Self::Flex(_)) => {
                let members = group_members!(devices, Flex)?;
                GroupDevice::with_member_settings::<crate::backends::Flex>(&members)
            }
            Some(other) => unreachable!("{other:?} was refused as a member"),
            None => Err(GroupError::NoDevices),
        }?;
        Ok(Self::Group(group))
    }

    /// Why this device cannot be a group's member, or `None` when it can.
    #[cfg(feature = "group")]
    fn refused_as_member(&self) -> Option<GroupError> {
        match self {
            #[cfg(cube_backend)]
            Self::Cube(_) => None,
            #[cfg(feature = "flex")]
            Self::Flex(_) => None,
            #[cfg(feature = "autodiff")]
            Self::Autodiff(_) => Some(GroupError::AutodiffMember),
            #[cfg(feature = "remote")]
            Self::Remote(_) => Some(GroupError::RemoteMember),
            #[allow(unreachable_patterns)]
            other => Some(GroupError::UnsupportedBackend {
                device: alloc::format!("{other:?}"),
            }),
        }
    }

    #[cfg(feature = "autodiff")]
    /// Creates a new [`DispatchDevice`] with
    /// [automatic differentiation](crate::backends::Autodiff) enabled.
    pub fn autodiff(device: impl Into<DispatchDevice>) -> DispatchDevice {
        Self::autodiff_with_gradient_checkpointing(device, GradientCheckpointingStrategy::Disabled)
    }
    #[cfg(feature = "autodiff")]
    /// Creates a new [`DispatchDevice`] with automatic differentiation and the provided gradient
    /// checkpointing strategy enabled.
    pub fn autodiff_with_gradient_checkpointing(
        device: impl Into<DispatchDevice>,
        checkpointing: GradientCheckpointingStrategy,
    ) -> DispatchDevice {
        let device = device.into();
        DispatchDevice::Autodiff(AutodiffDevice::new(device, checkpointing))
    }

    /// The id of the worker this device's fusion server runs on, or `None` when its operations do
    /// not go through one.
    ///
    /// This is the runtime device's own id, not the one this device reports through
    /// [`Device::to_id`](burn_backend::Device::to_id), which encodes the dispatch variant into its
    /// type. A [`DeviceHandle`](burn_backend::DeviceHandle) opened with it shares the fusion
    /// server's queue, so a task submitted there runs between the operations registered before
    /// and after it. An autodiff device answers for the device it wraps. A device group answers
    /// `None`: each of its members has a fusion server of its own.
    pub fn fusion_worker(&self) -> Option<DeviceId> {
        match self {
            #[cfg(cube_backend)]
            DispatchDevice::Cube(device) => burn_cubecl::FUSED.then(|| device.to_id()),
            #[cfg(feature = "autodiff")]
            DispatchDevice::Autodiff(device) => device.inner.fusion_worker(),
            #[allow(unreachable_patterns)]
            _ => None,
        }
    }

    /// Returns the inner device, without autodiff (when enabled).
    pub fn inner(self) -> Self {
        #[cfg(feature = "autodiff")]
        if let DispatchDevice::Autodiff(device) = self {
            return *device.inner;
        }

        self
    }

    /// Returns a unique number per variant to encode into type_id.
    fn backend_id(&self) -> DispatchDeviceId {
        match self {
            #[cfg(not(backend_enabled))]
            Self::Unavailable(never) => never.unreachable(),
            #[cfg(cube_backend)]
            Self::Cube(_) => DispatchDeviceId::Cube,
            #[cfg(feature = "flex")]
            Self::Flex(_) => DispatchDeviceId::Flex,
            #[cfg(feature = "remote")]
            Self::Remote(_) => DispatchDeviceId::Remote,
            #[cfg(feature = "group")]
            Self::Group(_) => DispatchDeviceId::Group,
            #[cfg(feature = "capture")]
            Self::Capture(_) => DispatchDeviceId::Capture,
            #[cfg(feature = "autodiff")]
            Self::Autodiff(device) => device.inner.backend_id(),
        }
    }

    /// Encode variant ID and backend type ID into a unique `type_id`.
    fn encode_type_id(&self, backend_type_id: u16) -> u16 {
        // Use the lower 8 bits for the backend's internal type ID
        let internal_type_id = backend_type_id & INTERNAL_ID_MASK;
        // Use the upper 8 bits for the DispatchDevice/DispatchDeviceId
        let backend = u16::from(self.backend_id()) << BACKEND_SHIFT;
        backend | internal_type_id
    }

    /// Decode an encoded `type_id` into variant ID and backend type ID.
    pub(crate) fn decode_type_id(type_id: u16) -> (DispatchDeviceId, u16) {
        let backend_raw = type_id >> BACKEND_SHIFT;
        let internal_type_id = type_id & INTERNAL_ID_MASK;

        let backend = DispatchDeviceId::try_from(backend_raw).expect("Unknown DispatchDevice ID");

        (backend, internal_type_id)
    }
}

#[allow(missing_docs)]
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
#[repr(u16)]
pub enum DispatchDeviceId {
    /// Every cubecl runtime: which one is in the device's own id.
    Cube = 0,
    Flex = 4,
    // 5 was LibTorch; keep the other backend IDs stable.
    // 6 was NdArray; keep the other backend IDs stable.
    Remote = 10,
    Capture = 11,
    Group = 12,
}

impl From<DispatchDeviceId> for u16 {
    fn from(variant: DispatchDeviceId) -> Self {
        variant as u16
    }
}

impl TryFrom<u16> for DispatchDeviceId {
    type Error = ();

    fn try_from(value: u16) -> Result<Self, Self::Error> {
        match value {
            #[cfg(cube_backend)]
            0 => Ok(Self::Cube),
            #[cfg(feature = "flex")]
            4 => Ok(Self::Flex),
            #[cfg(feature = "remote")]
            10 => Ok(Self::Remote),
            #[cfg(feature = "group")]
            12 => Ok(Self::Group),
            #[cfg(feature = "capture")]
            11 => Ok(Self::Capture),
            _ => Err(()),
        }
    }
}

impl DeviceOps for DispatchDevice {
    fn defaults(&self) -> DeviceSettings {
        match self {
            #[cfg(not(backend_enabled))]
            Self::Unavailable(never) => never.unreachable(),
            #[cfg(cube_backend)]
            Self::Cube(device) => device.defaults(),
            #[cfg(feature = "flex")]
            Self::Flex(device) => device.defaults(),
            #[cfg(feature = "remote")]
            Self::Remote(device) => device.defaults(),
            #[cfg(feature = "group")]
            Self::Group(device) => device.defaults(),
            #[cfg(feature = "capture")]
            Self::Capture(device) => device.defaults(),
            #[cfg(feature = "autodiff")]
            Self::Autodiff(device) => device.inner.defaults(),
        }
    }
}

impl burn_backend::Device for DispatchDevice {
    fn from_id(mut device_id: DeviceId) -> Self {
        let (dispatch_id, backend_type_id) = Self::decode_type_id(device_id.type_id);
        device_id.type_id = backend_type_id;

        #[allow(unreachable_patterns)]
        match dispatch_id {
            #[cfg(cube_backend)]
            DispatchDeviceId::Cube => Self::Cube(burn_backend::Device::from_id(device_id)),
            #[cfg(feature = "flex")]
            DispatchDeviceId::Flex => Self::Flex(FlexDevice::from_id(device_id)),
            #[cfg(feature = "remote")]
            DispatchDeviceId::Remote => Self::Remote(RemoteDevice::from_id(device_id)),
            #[cfg(feature = "group")]
            DispatchDeviceId::Group => Self::Group(GroupDevice::from_id(device_id)),
            #[cfg(feature = "capture")]
            DispatchDeviceId::Capture => Self::Capture(CaptureDevice::from_id(device_id)),
            _ => unreachable!("No backend feature enabled."),
        }
    }

    fn to_id(&self) -> DeviceId {
        let mut device_id: DeviceId = match self {
            #[cfg(not(backend_enabled))]
            Self::Unavailable(never) => never.unreachable(),
            #[cfg(cube_backend)]
            Self::Cube(device) => device.to_id(),
            #[cfg(feature = "flex")]
            Self::Flex(device) => device.to_id(),
            #[cfg(feature = "remote")]
            Self::Remote(device) => device.to_id(),
            #[cfg(feature = "group")]
            Self::Group(device) => device.to_id(),
            #[cfg(feature = "capture")]
            Self::Capture(device) => device.to_id(),
            #[cfg(feature = "autodiff")]
            Self::Autodiff(device) => device.inner.to_id(),
        };
        device_id.type_id = self.encode_type_id(device_id.type_id);
        device_id
    }
}

/// Every cubecl device reaches the one cubecl variant.
#[cfg(cube_backend)]
impl From<CubeDevice> for DispatchDevice {
    fn from(device: CubeDevice) -> Self {
        DispatchDevice::Cube(device)
    }
}

// A runtime's own device type converts too, since that is what its crate hands
// out. There is one variant to reach now, so a wgpu device no longer needs a
// priority chain of gates to decide which of four it lands in.
#[cfg(feature = "cpu")]
impl From<CpuDevice> for DispatchDevice {
    fn from(device: CpuDevice) -> Self {
        DispatchDevice::Cube(CubeDevice::Cpu(device))
    }
}

#[cfg(feature = "cuda")]
impl From<CudaDevice> for DispatchDevice {
    fn from(device: CudaDevice) -> Self {
        DispatchDevice::Cube(CubeDevice::Cuda(device))
    }
}

#[cfg(feature = "rocm")]
impl From<RocmDevice> for DispatchDevice {
    fn from(device: RocmDevice) -> Self {
        DispatchDevice::Cube(CubeDevice::Hip(device))
    }
}

#[cfg(any(
    feature = "wgpu",
    feature = "metal",
    feature = "vulkan",
    feature = "webgpu"
))]
impl From<WgpuDevice> for DispatchDevice {
    fn from(device: WgpuDevice) -> Self {
        DispatchDevice::Cube(CubeDevice::Wgpu(device))
    }
}

#[cfg(feature = "flex")]
impl From<FlexDevice> for DispatchDevice {
    fn from(device: FlexDevice) -> Self {
        DispatchDevice::Flex(device)
    }
}

#[cfg(feature = "group")]
impl From<GroupDevice> for DispatchDevice {
    fn from(device: GroupDevice) -> Self {
        DispatchDevice::Group(device)
    }
}

#[cfg(feature = "remote")]
impl From<RemoteDevice> for DispatchDevice {
    fn from(device: RemoteDevice) -> Self {
        DispatchDevice::Remote(device)
    }
}

#[cfg(all(test, not(backend_enabled)))]
mod no_backend_tests {
    #[test]
    #[should_panic(expected = "No execution backend is enabled. Enable a Burn backend feature")]
    fn default_requires_backend() {
        super::DispatchDevice::default();
    }
}

#[cfg(all(test, any(feature = "capture", feature = "remote-websocket")))]
mod tests {
    use super::*;
    use burn_backend::Device;

    #[test]
    #[cfg(feature = "remote-websocket")]
    fn remote_device_id_round_trips_through_dispatch() {
        let device = DispatchDevice::from(RemoteDevice::default());
        let restored = DispatchDevice::from_id(device.to_id());

        assert_eq!(restored, device);
    }

    #[test]
    #[cfg(feature = "capture")]
    fn capture_device_id_round_trips_through_dispatch() {
        let device = DispatchDevice::capture();
        let restored = DispatchDevice::from_id(device.to_id());

        assert_eq!(restored, device);
    }
}