Backend trait

主线位置：算子主线的「契约」。Tensor 的方法最后都变成 Dispatch 上的某个 FloatTensorOps / IntTensorOps / … 方法，Dispatch 再把同名方法转给 Flex、Cube（GPU）或 Autodiff<B>。这些方法的签名就是本页讲的 trait，定义在 burn-backend。Dispatch 怎么转见 Dispatch，CPU 上的实现见 Flex。

flowchart TB
  bt["BackendTypes"] --> dev["Device: DeviceOps"]
  bt --> prim["四种张量原语"]
  bt --> gprim["GraphPrimitive"]
  be["Backend"] --> bt
  be --> fops["FloatTensorOps"]
  be --> iops["IntTensorOps"]
  be --> bops["BoolTensorOps"]
  be --> qops["QTensorOps"]
  be --> mops["ModuleOps"]
  be --> aops["ActivationOps"]
  be --> tops["Transaction 和 Distributed"]
  ad["AutodiffBackend"] --> be

burn-backend 是 0.20 开发中从 burn-tensor 拆出来的（PR #4111「Move backend traits and types to burn-backend」）。它不依赖 burn-tensor，不认识 Tensor：后端作者实现它，burn-dispatch 和 burn-autodiff 用它，用户代码在 0.22 里一般碰不到它。本页的代码都在 crates/burn-backend/src/backend/ 下，base.rs 589 行，ops/ 下的各个 trait 加起来约 1.1 万行，大部分是文档注释。

BackendTypes：后端的类型表

BackendTypes 只有关联类型，没有方法：

pub trait BackendTypes: Clone + Send + Sync + core::fmt::Debug + 'static {
    /// Device type.
    type Device: DeviceOps;

    /// Tensor primitive to be used for all float operations.
    type FloatTensorPrimitive: TensorMetadata<Device = Self::Device> + 'static;

    /// Tensor primitive to be used for all int operations.
    type IntTensorPrimitive: TensorMetadata<Device = Self::Device> + 'static;

    /// Tensor primitive to be used for all bool operations.
    type BoolTensorPrimitive: TensorMetadata<Device = Self::Device> + 'static;

    /// Tensor primitive to be used for all quantized operations.
    type QuantizedTensorPrimitive: TensorMetadata<Device = Self::Device> + 'static;

    // ...
    type GraphPrimitive: Clone + Send + Sync + core::fmt::Debug + 'static;
}

四种张量原语分开声明，但没要求它们不同。三个后端的取法：

Flex 用一个按字节存储、带 dtype 字段的张量装所有种类，所以浮点、整数、布尔是同一个类型；Dispatch 更进一步，连量化张量也是同一个枚举。这里的「种类」只是 trait 方法签名上的约定，类型系统并不区分：把一个整数 DispatchTensor 传给 float_add 能通过编译，错误要等运行期才暴露。用户层面防住这件事的是 Tensor<D, K> 的 K（见 Tensor 与 Kind）。

GraphUnsupported 是一个没有变体的枚举。不支持图捕获的后端把它填进 GraphPrimitive，graph_stop_capture 的默认实现永远返回 Err，于是这个类型的值不可能出现，graph_replay 也就不可能被调用。用类型表达「这条路走不通」，不用运行期再查一次。

每个原语都要实现 TensorMetadata：dtype、shape、device，以及 can_mut。最后一个问的是「这个句柄是不是独占底层缓冲」，独占时后端可以原地写；不知道就必须答 false。Flex 上它就是 Arc::strong_count == 1，见 Flex。

Backend：op trait 的合集

Backend 本身没有算子方法，算子全在它的 supertrait 里：

pub trait Backend:
    BackendTypes
    + FloatTensorOps<Self>
    + BoolTensorOps<Self>
    + IntTensorOps<Self>
    + ModuleOps<Self>
    + ActivationOps<Self>
    + QTensorOps<Self>
    + TransactionOps<Self>
    + DistributedOps<Self>
    + Clone
    + Default
    + Sized
    + Send
    + Sync
    + core::fmt::Debug
    + 'static
{

所有 op 都是关联函数，没有 self：张量原语按值传进去、按值返回，设备从张量上取或作为参数传入。Backend 自己的方法是设备级的：name、seed、sync、flush、profile 一组、graph_* 一组、memory_* 一组、supports_dtype / dtype_usage、device_count。其中只有 name、seed、flush、dtype_usage、device_count 没有默认实现。

按 trait 数一下 0.22.1 里的方法（只数 trait 体内的 fn）：

合计 197 个必须实现、239 个有默认实现。量化（QTensorOps）几乎全靠默认：先反量化成浮点、调浮点 op，再按设备设置里的 QuantPropagation 决定量化回去还是留成浮点，后端只需提供量化、反量化和搬运这几个原语。ActivationOps 全是默认，一个后端什么都不写也有 relu、gelu、sigmoid。TransactionOps::tr_execute 的默认实现逐个 await 每个张量的 *_into_data，GPU 后端覆盖它，把多次读回合成一次。

默认实现：用别的 op 拼

默认实现只能用同一个 trait 体系里的其它方法，所以它们读起来像一小段用 op 写的程序。三个例子：

float_neg 是乘以 −1：

    fn float_neg(tensor: FloatTensor<B>) -> FloatTensor<B> {
        Self::float_mul_scalar(tensor, (-1f32).into())
    }

relu 是一次比较加一次按掩码填值，掩码的 bool dtype 从设备设置里取：

    fn relu(tensor: FloatTensor<B>) -> FloatTensor<B> {
        let bool_dtype = get_or_init_device_settings::<B>(&tensor.device()).bool_dtype;
        let mask = B::float_lower_equal_elem(tensor.clone(), 0f32.into(), bool_dtype);

        B::float_mask_fill(tensor, mask, 0f32.into())
    }

float_sort 干脆把数据读回主机，在 CPU 上排好再建回去：

    fn float_sort(tensor: FloatTensor<B>, dim: usize, descending: bool) -> FloatTensor<B> {
        let device = tensor.device();
        sort::<B, _, _, _>(
            tensor,
            dim,
            descending,
            device,
            |tensor| {
                let msg = "Failed to synchronously read tensor data. This operation is not supported until this backend has a GPU sorting implementation.";
                try_read_sync(B::float_into_data(tensor))
                    .expect(msg)
                    .expect(msg)
            },
            |data, device, _dtype| B::float_from_data(data, device),
        )
    }

哪些后端覆盖了它们（git grep "fn <名字>("，Dispatch 和 burn-router 是转发层，不计）：

几种取舍各有理由：

• Flex 几乎全覆盖。 CPU 上每个 op 是一次独立的遍历，neg 走默认就是「建一个 −1 标量、做一次乘法」，自己写一遍循环更省。relu 走默认要先生成一个布尔掩码张量再遍历一次，覆盖后一次遍历搞定。
• GPU 两层常用默认。 Cube 在开了 fusion feature 时是 Fusion<CubeBackend>（burn-cubecl/src/lib.rs），neg 变成一个 mul_scalar、relu 变成比较加 mask_fill，都会被记进融合流，和前后的逐元素 op 合成一个内核（见 Fusion 后端）。拆成通用小 op 反而方便融合。
• float_sort 在 GPU 上真的走默认。 burn-cubecl 和 burn-fusion 都没有覆盖它，所以 Metal 上的 x.sort(0) 会把张量读回主机、在 CPU 上排序、再上传。报错信息里那句「until this backend has a GPU sorting implementation」说的就是这个现状。在 wasm 上 try_read_sync 等不到结果，这个默认实现会直接 panic。
• Autodiff 几乎全覆盖，理由和速度无关。 默认实现由几个 op 拼成，在 Autodiff<B> 上走默认就会在图里记几个节点；覆盖之后只记一个节点、配一个专门的反向（relu 的反向是 B::relu_backward），内层算前向时再交给 B::relu，B 自己的快实现照样用得上。详见一个算子的反向。
• float_prod 的默认实现对负数不成立：负数的 log 是 NaN。树内的后端全都覆盖了它；自己写后端时如果不覆盖，这个默认值是错的。

AutodiffBackend

AutodiffBackend 在 Backend 之上加三类方法：backward / grad / grad_remove / grad_replace 管梯度；inner / from_inner（及 int_ / bool_ / q_ 版本）在「带图的张量」和「内层后端的张量」之间换；set_distributed_params 一组给分布式训练打标记。它要求 InnerBackend: Backend<Device = Self::Device>，即外层和内层用同一种设备类型。

Autodiff<B> 实现它，InnerBackend = B。Dispatch 也实现它，但 InnerBackend = Dispatch：在 Dispatch 层，「是否带图」不是类型的区别，而是 DispatchTensor 上的一个运行期标记，inner 只是把标记清掉。这些方法做什么、图怎么建，属于Autodiff 装饰器；本页只需要知道，0.22 的用户代码通过 Tensor::backward、Tensor::grad 间接用到它们。

DeviceOps 与设备设置

BackendTypes::Device 要实现 DeviceOps：

pub trait DeviceOps: Clone + Default + PartialEq + Send + Sync + core::fmt::Debug + Device {
    /// Returns the device id.
    fn id(&self) -> DeviceId {
        self.to_id()
    }

    /// Returns the default settings for this device.
    fn defaults(&self) -> DeviceSettings;
}

supertrait Device 来自 cubecl 的 cubecl-common（经 burn-std 重新导出），只有两个方法：from_id 和 to_id，DeviceId 是两个 u16（type_id、index_id）。burn 和 cubecl 用同一套设备编号，CPU 后端也一样：FlexDevice 是个零大小的结构体，to_id 固定返回 DeviceId::new(0, 0)。

defaults 返回 DeviceSettings：默认的 float / int / bool dtype 和量化配置。Flex 的是 F32、I32、Bool(Native)。这组设置按 (DeviceId, TypeId) 存在一个全局表里（get_or_init_device_settings），第一次用到某个设备时从 defaults() 初始化，set_default_dtypes 能换掉默认值，但只有一次机会，而且要赶在表项初始化之前：它先检查 dtype 在该设备上 supports_dtype，再往表项的 OnceLock 里写，写不进去就返回 DeviceError。实验 crate 里依次试了三次（NdArray 设备还没用过，Flex 设备已经建过一个张量）：

let fresh = Device::ndarray().as_dispatch().clone();
println!("before first use: {:?}", set_default_dtypes::<Dispatch>(&fresh, FloatDType::F16, IntDType::I32, BoolDType::Native));
println!("second call:      {:?}", set_default_dtypes::<Dispatch>(&fresh, FloatDType::F32, IntDType::I32, BoolDType::Native));
let used = Device::flex();
let _t = Tensor::<1>::from_floats([1.0, 2.0], &used);
println!("after first use:  {:?}", set_default_dtypes::<Dispatch>(used.as_dispatch(), FloatDType::F16, IntDType::I32, BoolDType::Native));

before first use: Err(UnsupportedDType { device: "NdArray(Cpu)", dtype: F16 })
second call:      Ok(())
after first use:  Err(AlreadyInitialized { device: "Flex(Cpu)" })

第一次因为 NdArray 不支持 F16 被拒，表项仍是空的，第二次设 F32 成功；Flex 设备建张量时已经按 defaults() 初始化了表项，再设就是 AlreadyInitialized。上面 relu 默认实现里的 get_or_init_device_settings::<B>(..).bool_dtype 就是查这张表。用户面的 Device::configure / Device::settings 见 Device。

dtype 支持：DTypeUsage

dtype_usage 返回一个 DTypeUsage 的集合，三个标志：Storage（能存、能转换）、Arithmetic（能做一般运算）、Accelerated（有硬件加速路径，比如矩阵乘的专用单元）。supports_dtype 的默认实现要求前两个都有。区分出 Accelerated 是因为有的类型只在矩阵乘这类专用路径上可用，不能拿来当默认 dtype。

Flex::dtype_usage 是一张写死的表：所有浮点、整数类型都是 Storage | Arithmetic，Bool(U32) 和大部分量化类型是空集。GPU 上的答案由 cubecl 按设备能力给出。

for (name, device) in [("flex", Device::flex()), ("metal", Device::metal(DeviceKind::DefaultDevice))] {
  let d = device.as_dispatch();
  println!("{name}: Dispatch::name = {}", Dispatch::name(d));
  for dt in [DType::F64, DType::F32, DType::F16, DType::BF16, DType::I64, DType::I32, DType::U8] {
    println!("  {dt:?}: usage={:?} supports={}", Dispatch::dtype_usage(d, dt), Dispatch::supports_dtype(d, dt));
  }
}

为什么 float_into_data 是 async

几乎所有 op 都是同步签名，读数据的 float_into_data 例外：

    fn float_into_data(
        tensor: FloatTensor<B>,
    ) -> impl Future<Output = Result<TensorData, ExecutionError>> + Send;

int_into_data、bool_into_data、q_into_data、tr_execute 也是同样的形状。原因在 GPU：op 只是把命令提交给设备，读数据要等设备算完、再把缓冲映射回主机，这是一次真正的等待。在原生平台上可以阻塞线程等；在浏览器里（WebGPU 跑在 wasm 上）主线程不允许阻塞，只能交出控制权、等回调。把签名定成 future，两种平台用同一份后端代码：原生上由调用方阻塞，wasm 上由调用方 .await。future 怎么被 poll、什么叫交出控制权，见 tokio 卷的 poll 与 wake。

签名写成返回位置的 impl Future + Send，实现方可以直接写 async fn。CPU 后端没有什么可等的，Flex::float_into_data 的 future 第一次 poll 就完成：

    async fn float_into_data(tensor: FloatTensor<Flex>) -> Result<TensorData, ExecutionError> {
        Ok(tensor.into_data())
    }

Result 来自另一处改动：GPU 上前面某个 op 的执行错误（比如设备丢失）往往要等到读数据时才浮出来，读数据是报告它的地方。

同步的 Tensor::into_data 最终调用 try_read_sync（cubecl-environment，经 burn-std 导出）：

pub fn try_read_sync<F: Future<Output = T>, T>(f: F) -> Option<T> {
    #[cfg(target_family = "wasm")]
    {
        use core::task::Poll;

        match embassy_futures::poll_once(f) {
            Poll::Ready(output) => Some(output),
            _ => None,
        }
    }
    #[cfg(not(target_family = "wasm"))]
    {
        Some(super::block_on(f))
    }
}

原生平台上 block_on 是 futures_lite::future::block_on，在当前线程上一直 poll 到完成。wasm 上只 poll 一次：Flex 的 future 当场完成，同步 API 照常可用；GPU 的 future 还没好，返回 None，Tensor::into_data 随即 panic，提示改用 into_data_async。这也是上面 float_sort 默认实现在 wasm 上会 panic 的原因。

动脑筋答案

cargo run --offline（debug），M2 Max、macOS 27、rustc 1.99.0：

flex: Dispatch::name = dispatch<flex>
  F64: usage=EnumSet(Storage | Arithmetic) supports=true
  F32: usage=EnumSet(Storage | Arithmetic) supports=true
  F16: usage=EnumSet(Storage | Arithmetic) supports=true
  BF16: usage=EnumSet(Storage | Arithmetic) supports=true
  I64: usage=EnumSet(Storage | Arithmetic) supports=true
  I32: usage=EnumSet(Storage | Arithmetic) supports=true
  U8: usage=EnumSet(Storage | Arithmetic) supports=true
metal: Dispatch::name = dispatch<fusion<cubecl<wgpu<msl>>>>
  F64: usage=EnumSet() supports=false
  F32: usage=EnumSet(Storage | Arithmetic | Accelerated) supports=true
  F16: usage=EnumSet(Storage | Arithmetic | Accelerated) supports=true
  BF16: usage=EnumSet(Storage | Arithmetic | Accelerated) supports=true
  I64: usage=EnumSet(Storage | Arithmetic) supports=true
  I32: usage=EnumSet(Storage | Arithmetic) supports=true
  U8: usage=EnumSet(Storage | Arithmetic) supports=true

• Metal 上 F64 是空集，连 Storage 都没有：Metal 着色语言没有 64 位浮点类型。同一份用 f64 的代码在 Flex 上能跑，换到 Metal 设备就得先改 dtype。
• F32、F16、BF16 在 Metal 上多一个 Accelerated，Flex 上一个也没有：Flex 不把任何路径标成「硬件加速」，gemm 再快也只是普通算术。
• name 把层次全打了出来：Dispatch 包着 Fusion，Fusion 包着 CubeBackend，运行时是 wgpu，内核编成 MSL。每一层的 name 都是「自己的名字 + 内层的名字」，Dispatch 层的写法见 Backend for Dispatch。

小结

• BackendTypes 定义设备和四种张量原语，Backend 是九个 op trait 的合集加设备级方法；op 都是不带 self 的关联函数，张量按值进出。
• 0.22.1 里共 197 个必须实现的方法、239 个默认实现。默认实现用别的 op 拼：neg 是乘 −1，relu 是比较加填值，sort 是读回主机排序。
• CPU 后端几乎全覆盖；GPU 的两层常留着默认，把它们拆成能融合的小 op；Autodiff 覆盖是为了一个 op 只记一个节点；Dispatch 必须逐个转发，否则内层的覆盖会被绕过。
• DeviceOps::defaults 给出每个设备的默认 dtype，存进全局设置表；dtype_usage 用三个标志说明一个 dtype 在设备上能做什么，Metal 上 F64 不可用。
• float_into_data 返回 future，原生平台阻塞等待，wasm 上只 poll 一次；Flex 的 future 当场完成。