Device

pub struct Device {
    blob: device_opaque::Opaque,
}

// Aligned, type-erased storage for `DispatchDevice`. See `crate::macros` for
// why this indirection exists (it keeps the dispatch type tree out of
// downstream MIR).
burn_std::obfuscate!(
    type: DispatchDevice,
    module: device_opaque,
    derives: [Send, Sync]
);

error[E0599]: no associated function or constant named `flex` found for struct `burn::tensor::Device` in the current scope
   --> src/bin/noflex.rs:4:19
    |
  4 |   let _ = Device::flex();
    |                   ^^^^ associated function or constant not found in `burn::tensor::Device`
    |
note: if you're trying to build a new `burn::tensor::Device`, consider using `burn::tensor::Device::new` which returns `burn::tensor::Device`
...

    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() {
                    // ...
                    "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."
                        );
                    }
                    // ...
                    _ => 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()));

        // ... rocm、vulkan、webgpu、wgpu、cpu、tch

        // Preserve the preference for Flex over the deprecated NdArray backend.
        #[cfg(feature = "flex")]
        return Self::Flex(FlexDevice);

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

        #[cfg(feature = "ndarray")]
        return Self::NdArray(NdArrayDevice::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()."
        );
    }

--- default
Device::default() = Device<Cube(Wgpu(WgpuDevice { kind: DefaultDevice, backend: Auto }))>
--- BURN_DEVICE=flex
Device::default() = Device<Flex(Cpu)>
exit=0
--- BURN_DEVICE=metal
Device::default() = Device<Cube(Wgpu(WgpuDevice { kind: DefaultDevice, backend: Metal }))>
exit=0
--- BURN_DEVICE=cuda

thread 'main' (3205663) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/burn-dispatch-0.22.1/src/device.rs:261:25:
BURN_DEVICE=cuda requested, but the 'cuda' feature is not enabled.
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
exit=101
--- BURN_DEVICE=foo

thread 'main' (3205665) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/burn-dispatch-0.22.1/src/device.rs:337:26:
Unknown BURN_DEVICE override: 'foo'.
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
exit=101

fn main() {
  println!("building a device...");
  let device = Device::default();
  let t = Tensor::<1>::zeros([2], &device);
  println!("{t}");
}

building a device...

thread 'main' (3149547) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/burn-dispatch-0.22.1/src/device.rs:381:9:
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().
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace

    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,
            // ... Flex、NdArray、LibTorch、Remote、Capture 同样两两比较
            #[allow(unreachable_patterns)]
            (_, _) => false,
        }
    }

== equality
flex == flex: true
flex == flex.autodiff(): true
flex.autodiff() == flex ckpt: true
flex == ndarray: false
metal == wgpu: false
metal == default: false
wgpu == default: true

    #[cfg(feature = "autodiff")]
    #[must_use]
    pub fn autodiff(self) -> Self {
        match self.into_dispatch() {
            device @ DispatchDevice::Autodiff(_) => Self::new(device),
            other => Self::new(DispatchDevice::autodiff(other)),
        }
    }

== debug
flex: Device<Flex(Cpu)> is_autodiff=false
ndarray: Device<NdArray(Cpu)> is_autodiff=false
metal: Device<Cube(Wgpu(WgpuDevice { kind: DefaultDevice, backend: Metal }))> is_autodiff=false
wgpu: Device<Cube(Wgpu(WgpuDevice { kind: DefaultDevice, backend: Auto }))> is_autodiff=false
default: Device<Cube(Wgpu(WgpuDevice { kind: DefaultDevice, backend: Auto }))> is_autodiff=false
flex.autodiff(): Device<Autodiff { device: Flex(Cpu), checkpointing: Disabled }> is_autodiff=true
flex.autodiff().autodiff(): Device<Autodiff { device: Flex(Cpu), checkpointing: Disabled }> is_autodiff=true
flex.autodiff().gradient_checkpointing(): Device<Autodiff { device: Flex(Cpu), checkpointing: Balanced }> is_autodiff=true
flex_ckpt.inner(): Device<Flex(Cpu)>
flex.without_autodiff(): Device<Flex(Cpu)>

--- checkpoint-plain

thread 'main' (3205667) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/burn-tensor-0.22.1/src/device.rs:555:18:
Gradient checkpointing requires autodiff; use Device::autodiff().gradient_checkpointing()
note: run with `RUST_BACKTRACE=1` environment variable to display a backtrace
exit=101

for (name, d) in [("flex", &flex), ("metal", &metal)] {
  d.seed(42);
  let a = Tensor::<1>::random([4], Distribution::Default, d).into_data();
  d.seed(42);
  let b = Tensor::<1>::random([4], Distribution::Default, d).into_data();
  let c = Tensor::<1>::random([4], Distribution::Default, d).into_data();
  // 打印 a、a == b、b == c
}

let (out, dur) = metal.profile(|| a.clone().matmul(b.clone()).into_data()).unwrap();
let ticks = read_sync(dur.resolve());
let (_, dur_lazy) = metal.profile(|| a.clone().matmul(b.clone())).unwrap();
let (_, dur) = flex.profile(|| fa.clone().matmul(fa.clone())).unwrap();

== sync / flush
flex: sync Ok(()), flush Ok(())
metal: sync Ok(()), flush Ok(())
== seed
flex: a [0.13340974, 0.5265574, 0.24873835, 0.5427252]
  same seed equal: true, next draw equal: false
metal: a [0.33610713, 0.3398847, 0.14931732, 0.5403774]
  same seed equal: true, next draw equal: false
== profile
metal profile: out len 1048576, ticks Some(1.318208ms)
metal profile without read: None
flex profile: Some(5.285959ms)

== default dtypes per device
flex: float F32 int I32 bool Native; supports F64 true F16 true BF16 true
ndarray: float F32 int I64 bool Native; supports F64 true F16 false BF16 false
metal: float F32 int I32 bool U8; supports F64 false F16 true BF16 true

== configure
configure ndarray F64 before use: Ok(())
tensor dtype after configure: F64
second configure: Err(Device NdArray(Cpu) settings have already been initialized and cannot be changed.
Default data types lock on the first tensor operation for a device (or on a previous `set_default_dtypes`/`configure` call). Configure the device before creating any tensor on it, or create tensors with an explicit dtype instead.)

    pub fn to_device(self, device: &Device) -> Self {
        let target_device = device.clone().with_autodiff_context_from(&self.device());
        Self::new(K::to_device(self.primitive, &target_device))
    }

let x = Tensor::<2>::from_floats([[1.0, 2.0], [3.0, 4.0]], &flex);
let on_metal = x.clone().to_device(&metal);
let back = (on_metal * 2.0).to_device(&flex);
let i = Tensor::<1, Int>::from_ints([1, 2, 3], &flex).to_device(&metal);
let xa = Tensor::<2>::from_floats([[1.0, 2.0], [3.0, 4.0]], &flex_ad); // flex.autodiff()
let xa_metal = xa.to_device(&metal); // 目标是不带 autodiff 的 metal

== to_device
to metal: device Device<Cube(Wgpu(WgpuDevice { kind: DefaultDevice, backend: Metal }))>
back to flex: device Device<Flex(Cpu)> data [2.0, 4.0, 6.0, 8.0]
int to metal: dtype I32 device Device<Cube(Wgpu(WgpuDevice { kind: DefaultDevice, backend: Metal }))>
autodiff tensor to plain metal device: is_autodiff true device Device<Autodiff { device: Cube(Wgpu(WgpuDevice { kind: DefaultDevice, backend: Metal })), checkpointing: Disabled }>

metal == default: false
wgpu == default: true