Dispatch

#[derive(Debug, Default, Clone)]
pub struct Dispatch;

/// A tensor that can dispatch operations to any enabled backend at runtime.
#[derive(Clone, Debug)]
pub struct DispatchTensor {
    /// Tensor kind primitive.
    pub kind: DispatchTensorKind,
    /// Autodiff backend association for this tensor.
    pub autodiff: DispatchAutodiffContext,
}

pub enum DispatchTensorKind {
    #[cfg(not(backend_enabled))]
    #[doc(hidden)]
    Unavailable(crate::NoBackend),
    /// A tensor on the [cubecl backend](Cube) — its device says which runtime.
    #[cfg(cube_backend)]
    Cube(BackendTensor<Cube>),

    /// The [Flex backend](Flex) tensor.
    #[cfg(feature = "flex")]
    Flex(BackendTensor<Flex>),

    // ...

    /// The [autodiff enabled backend](Autodiff) tensor.
    #[cfg(feature = "autodiff")]
    Autodiff(Box<DispatchTensorKind>),
}

    let cube = cuda || rocm || cpu || metal || vulkan || webgpu || wgpu;
    let local = cube || flex || ndarray || tch;
    let executing = local || remote;

    // Backend-free builds expose tensor/model APIs without installing an execution backend.
    if executing || capture {
        println!("cargo::rustc-cfg=backend_enabled");
    }

    // Every cubecl-backed feature selects the same backend type, since the device says which
    // runtime, so they share one variant under one cfg rather than a seven-way list at each use.
    if cube {
        println!("cargo::rustc-cfg=cube_backend");
    }

    // Capture records operations but never executes them.
    if executing {
        println!("cargo::rustc-cfg=executing_backend");
    }

#[cfg(not(backend_enabled))]
#[doc(hidden)]
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct NoBackend {
    never: core::convert::Infallible,
}

thread 'main' (3150467) 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().

error[E0599]: no associated function or constant named `flex` found for struct `burn::tensor::Device` in the current scope
   --> src/main.rs:5:24
    |
  5 |   let device = Device::flex();
    |                        ^^^^ associated function or constant not found in `burn::tensor::Device`
...

pub(crate) const BACKENDS: &[BackendSpec] = &[
    BackendSpec {
        // Every cubecl runtime is this one backend; a device says which of them
        // it runs on. `cube_backend` is set when any of their features is.
        name: "Cube",
        cfg: "cube_backend",
        // ...
        distributed: true,
        unidirectional_transfer: false,
    },
    BackendSpec {
        name: "Flex",
        cfg: "feature = \"flex\"",
        distributed: false,
        unidirectional_transfer: false,
    },
    // ... NdArray、LibTorch、Remote、Capture
];

#[backend_dispatch]
impl FloatTensorOps<Self> for Dispatch {
    // ...
    fn float_matmul(lhs: FloatTensor<Self>, rhs: FloatTensor<Self>) -> FloatTensor<Self> {
        B::float_matmul(lhs, rhs)
    }

            fn float_matmul(
                lhs: FloatTensor<Self>,
                rhs: FloatTensor<Self>,
            ) -> FloatTensor<Self> {
                let __burn_selected_context = lhs.autodiff;
                let mut __ad_ctx = crate::DispatchAutodiffContext::Disabled;
                __ad_ctx = __ad_ctx.merge(lhs.autodiff);
                __ad_ctx = __ad_ctx.merge(rhs.autodiff);
                match lhs.kind {
                    crate::DispatchTensorKind::Flex(__burn_selected) => {
                        match __ad_ctx {
                            crate::DispatchAutodiffContext::Disabled => {
                                type B = crate::backends::Flex;
                                // ...（检查 lhs 的 context）
                                let lhs = __burn_selected.float();
                                let rhs = match rhs.kind {
                                    crate::DispatchTensorKind::Flex(__inner) => __inner.float(),
                                    _ => {
                                        ::core::panicking::panic_fmt(
                                            format_args!(
                                                "input tensor `{0}` is on the wrong backend",
                                                "rhs",
                                            ),
                                        );
                                    }
                                };
                                let __output = { B::float_matmul(lhs, rhs) };
                                crate::DispatchTensor {
                                    kind: crate::DispatchTensorKind::Flex(
                                        crate::BackendTensor::Float(__output),
                                    ),
                                    autodiff: __ad_ctx,
                                }
                            }
                            crate::DispatchAutodiffContext::Enabled(_) => {
                                // ... panic: "autodiff context requires the `autodiff` feature"
                            }
                        }
                    }
                    #[allow(unreachable_patterns)]
                    _ => {
                        ::core::panicking::panic_fmt(
                            format_args!("unsupported dispatch backend"),
                        );
                    }
                }
            }

let a = Dispatch::float_from_data(TensorData::from([1.0f32, 2.0]), Device::flex().as_dispatch());
let b = Dispatch::float_from_data(TensorData::from([3.0f32, 4.0]), Device::metal(DeviceKind::DefaultDevice).as_dispatch());
let _ = Dispatch::float_add(a, b);

thread 'main' (3146872) panicked at ~/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/burn-dispatch-0.22.1/src/ops/tensor.rs:11:1:
input tensor `rhs` is on the wrong backend

    #[backend_dispatch(skip)]
    fn float_to_device(tensor: FloatTensor<Self>, device: &DispatchDevice) -> FloatTensor<Self> {
        // `Tensor::to_device` aligns the target's autodiff context with the source tensor before
        // reaching this low-level operation. Direct callers must likewise provide compatible
        // contexts; this layer only routes the transfer between compute resources.
        float_to_device!(
            Float,
            float,
            tensor,
            device,
            float_to_device,
            |inner, device| { super::transfer::float_transfer::<B1, B2>(inner, device) }
        )
    }

pub(crate) fn float_transfer<Src: Backend, Dst: Backend>(
    tensor: FloatTensor<Src>,
    device: &Dst::Device,
) -> FloatTensor<Dst> {
    let data = burn_backend::read_sync(Src::float_into_data(tensor))
        .expect("Failed to read tensor during cross-backend transfer");
    Dst::float_from_data(data, device)
}

let a = Tensor::<1>::from_floats([1.0, 2.0, 3.0], &flex);
let b = a.clone().to_device(&metal);
let c = (b * 2.0).to_device(&flex);

let x = Tensor::<1>::from_floats([1.0, 2.0, 3.0], &flex.clone().autodiff()).require_grad();
let y = x.clone().to_device(&metal.clone().autodiff());
let loss = (y.clone() * y).sum();
let grads = loss.backward();
let g = x.grad(&grads).unwrap();

a.device=Device<Flex(Cpu)>
b.device=Device<Cube(Wgpu(WgpuDevice { kind: DefaultDevice, backend: Metal }))>
c=[2.0, 4.0, 6.0] device=Device<Flex(Cpu)>
loss.device=Device<Autodiff { device: Cube(Wgpu(WgpuDevice { kind: DefaultDevice, backend: Metal })), checkpointing: Disabled }>
x.grad=[2.0, 4.0, 6.0] device=Device<Flex(Cpu)>

    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
    }

// 1. Flex directly, through the burn-backend trait.
black_box(Flex::float_add(black_box(fa.clone()), black_box(fb.clone())));
// 2. The Dispatch backend, same trait, DispatchTensor in and out.
black_box(Dispatch::float_add(black_box(da.clone()), black_box(db.clone())));
// 3. The user-facing Tensor API.
black_box(black_box(ta.clone()) + black_box(tb.clone()));

float add, shape [16], f32, 1000000 ops x 7 reps, release
features: metal=false autodiff=false
Flex::float_add              min  138.8 ns  median  140.3 ns  max  141.1 ns
Dispatch::float_add          min  157.9 ns  median  159.7 ns  max  160.7 ns
Tensor + (flex)              min  203.3 ns  median  205.0 ns  max  206.1 ns

float add, shape [16], f32, 1000000 ops x 7 reps, release
features: metal=true autodiff=true
Flex::float_add              min  138.1 ns  median  138.8 ns  max  139.5 ns
Dispatch::float_add          min  173.6 ns  median  174.9 ns  max  176.5 ns
Tensor + (flex)              min  265.1 ns  median  267.8 ns  max  269.8 ns

flex           Flex(Cpu) -> DeviceId(type=1024, index=0)
flex.autodiff  Autodiff { device: Flex(Cpu), checkpointing: Disabled } -> DeviceId(type=1024, index=0)
ndarray        NdArray(Cpu) -> DeviceId(type=1536, index=0)
metal          Cube(Wgpu(WgpuDevice { kind: DefaultDevice, backend: Metal })) -> DeviceId(type=100, index=16384)