burn-cubecl 后端

主线位置：算子主线里 Fusion 之下、cubecl 之上的一层。Fusion<CubeBackend> 在 drain 流时，融合不了的操作调 CubeBackend::float_matmul 之类的方法（见 Fusion 装饰器），本页讲这些方法做什么：把 burn 的操作翻成一次或几次 cubecl 内核启动，matmul 交给 cubek 并经 autotune 选内核。启动之后的事归运行时与 ComputeClient 和 wgpu 与 Metal。

flowchart TB
  fus["Fusion 执行未融合的 op"] --> add["CubeBackend::float_add"]
  fus --> mm["CubeBackend::float_matmul"]
  add --> binop["launch_binop AddOp"]
  binop -->|"lhs 可原地写"| inl["写回 lhs 的显存"]
  binop -->|否则| fresh["申请新输出"]
  inl --> k["kernel_binop 启动"]
  fresh --> k
  mm --> tune["matmul_autotune"]
  tune --> cubek["cubek launch_ref"]
  k --> client["Client::launch"]
  cubek --> client

本页源码是 burn 0.22.1 的 burn-cubecl，matmul 部分读到 cubek 0.3.1 的入口为止。实验 crate rt-cube 绕过 Dispatch 和 Fusion，直接依赖 burn-cubecl（feature metal、tracing），在 CubeBackend 上调 FloatTensorOps 的方法，用 tracing 的 launch_inner span 数每一步启动了哪些内核（数法和一次 matmul 相同）。release 构建，M2 Max、macOS 27、rustc 1.99.0。

CubeBackend：一个空结构体

/// Tensor backend that compiles just-in-time for whichever runtime its device
/// names.
#[derive(new)]
pub struct CubeBackend;

impl BackendTypes for CubeBackend {
    type Device = CubeDevice;

    type FloatTensorPrimitive = CubeTensor;
    type IntTensorPrimitive = CubeTensor;
    type BoolTensorPrimitive = CubeTensor;
    type QuantizedTensorPrimitive = CubeTensor;

    type GraphPrimitive = CubeGraph;
}

CubeBackend 没有字段，也没有类型参数。四种原语都是同一个 CubeTensor，dtype 记在张量的值里；设备类型 CubeDevice 就是 cubecl::Device，CUDA、HIP、wgpu、CPU 运行时是这个枚举的不同变体。所以「后端是哪个运行时」完全由设备的值决定：Backend::name 是 format!("cubecl<{}>", client.name())，sync 是 block_on(client.sync())，都是先 device.client() 拿到这个设备的 Client。

0.22 之前它是 CubeBackend<R: CubeRuntime>，CubeTensor<R> 里存的是 ComputeClient<R> 和 R::Device，每个运行时一个单态化的后端；PR #5556 随 cubecl 0.11 的运行时擦除把 R 去掉了（这次改动在分派层的后果见 Dispatch）。Dispatch 用的不是它本身，而是 Cube：开了 fusion feature 时 Cube = Fusion<CubeBackend>，否则就是 CubeBackend。

CubeTensor：句柄加元数据

/// The basic tensor primitive struct.
pub struct CubeTensor {
    /// Compute client for the runtime this tensor's device names.
    pub client: Client,
    /// The buffer where the data are stored.
    pub handle: Handle,
    /// The metadata of the tensor.
    pub meta: Box<Metadata>,
    /// The device of the tensor.
    pub device: CubeDevice,
    /// The datatype of the tensor.
    pub dtype: DType,
    /// Runtime quantization parameters, if applicable
    pub qparams: Option<QParams>,
}

字段	是什么	谁用
`client`	这个设备的 `Client`，提交内核、分配、读回都经它	每个 kernel 的 `launch_unchecked(&client, ..)`
`handle`	cubecl 的 `Handle`：一块托管显存（`ManagedMemoryHandle`）加偏移和所在的流	绑定成内核参数；`can_mut` 看它的引用计数
`meta`	`Metadata`：`shape`、`strides`，以及存储分块（`tiling`，只有 cubek 的 matmul 会读分块存储的张量）	转置、reshape、广播只改它
`device`	`cubecl::Device`	跨设备检查、`to_device`
`dtype`	burn 的 `DType`，`dtype_to_storage_type` 转成 cubecl 的元素类型	选内核的编译期参数
`qparams`	量化张量的 scale 在同一块显存里的位置	量化路径

from_data 调 client.create_tensor(bytes, shape, elem_size)，由运行时给出句柄和步长，再拼成 CubeTensor。rt-cube 建一个 2×3 的 f32 张量，打印它的 Debug 和几个字段（buffer 是句柄里托管显存的 Debug，长行截断）：

name = cubecl<wgpu<msl>>
CubeTensor { shape: Shape { dims: [2, 3] }, device: Wgpu(WgpuDevice { kind: DefaultDevice, backend: Metal }), strides: Strides { dims: [3, 1] }, elem: f32, runtime: wgpu<msl>}
a        shape=Shape { dims: [2, 3] } strides=Strides { dims: [3, 1] } contiguous=true can_mut=true buffer=ManagedMemoryHandle { descriptor: ManagedMemoryDescriptor { id: ManagedMemoryId { value: 0 }, ... size=256B

句柄的 size 是 256 字节，6 个 f32 只占 24 字节：显存按池里的块分配，本页不展开（见运行时与 ComputeClient）。下面的实验只看 ManagedMemoryId：同一个 id 就是同一块显存。

一个逐元素操作：float_add

float_add 一行转给 numeric::add，再一行转给 launch_binop::<AddOp>。加法本身是一个 #[cube] trait 的实现：

#[cube]
impl<T: Numeric, N: Size> BinaryOp<T, N> for AddOp {
    fn execute(lhs: Vector<T, N>, rhs: Vector<T, N>) -> Vector<T, N> {
        lhs + rhs
    }
}

内核 kernel_binop 对所有二元运算通用，O::BinaryOp 在展开时换成具体的运算：每个线程按 ABSOLUTE_POS 读两个操作数的一个向量、算、写回，越界的线程 terminate!()。减法、乘法、取余、pow 都只是另一个 BinaryOp 实现（#[cube] 内核讲宏怎么展开这些）。

launch_binop 是 CPU 侧的准备，决定这次启动的编译期参数和网格：

pub(crate) fn launch_binop<O: BinaryOpFamily>(lhs: CubeTensor, rhs: CubeTensor) -> CubeTensor {
    let output_shape = broadcast_shape(&[&lhs, &rhs]);
    in_memory_order([lhs, rhs], output_shape, |[lhs, rhs], shape_out| {
        let vector_size_lhs = max_vector_size(&lhs);
        let vector_size_rhs = max_vector_size(&rhs);
        let vector_size = Ord::min(vector_size_lhs, vector_size_rhs);
        // ...
        let cube_dim = CubeDim::new(&lhs.client, working_units);
        let cube_count = calculate_cube_count_elemwise(&lhs.client, working_units, cube_dim);

        unsafe {
            if lhs.can_mut_broadcast(&rhs) {
                kernel_binop::launch_unchecked::<O>(
                    &client,
                    cube_count,
                    cube_dim,
                    address_type!(lhs, rhs),
                    vector_size,
                    lhs.clone().into_linear_view(),
                    rhs.into_linear_view_like(&lhs),
                    lhs.as_linear_view_alias(0),
                    dtype_to_storage_type(dtype),
                );

                lhs
            } else if rhs.can_mut_broadcast(&lhs) {
                // ... 同上，写回 rhs
                rhs
            } else {
                let output =
                    empty_device_dtype(lhs.client.clone(), lhs.device.clone(), shape_out, dtype);
                // ... 写进新的 output
                output
            }
        }
    })
}

广播形状和遍历顺序。 broadcast_shape 算输出形状；in_memory_order 必要时把操作数按显存顺序重排（下一节）。
向量宽度。 max_vector_size 看最后一维的形状和步长能一次读几个元素，取两边的较小值；它是内核的编译期参数，上一页 MSL 入口名里的 n_4 就是它。
网格。 工作单元数是元素数除以向量宽度，CubeDim::new 和 calculate_cube_count_elemwise 按设备的限制分成线程组。
输出放哪。 can_mut_broadcast 要求句柄可写（handle.can_mut()：这块显存只有这一个持有者）、不自重叠、形状不小于另一边。左边满足就把左边当输出（as_linear_view_alias(0)，编进 KernelId 是 inplace: Some(0)），否则看右边，都不行才申请新显存。

launch_unchecked 是 #[cube(launch_unchecked)] 生成的启动函数，越界由内核里的 is_in_bounds 自己检查。它把参数打包交给 client，这一步之后就是 cubecl 运行时的事了。

rt-cube 做两次加法：第一次 a 只有一个持有者，第二次把 c 克隆一份再加：

  let c = B::float_add(a, b.clone());
  // ...
  let d = B::float_add(c.clone(), b.clone());

--- c = a + b (a owned once)
c        shape=Shape { dims: [2, 3] } strides=Strides { dims: [3, 1] } contiguous=true can_mut=true buffer=ManagedMemoryHandle { descriptor: ManagedMemoryDescriptor { id: ManagedMemoryId { value: 0 }, ...
launches {"KernelBinop": 1}
--- d = c.clone() + b (c shared)
c        shape=Shape { dims: [2, 3] } strides=Strides { dims: [3, 1] } contiguous=true can_mut=true buffer=ManagedMemoryHandle { descriptor: ManagedMemoryDescriptor { id: ManagedMemoryId { value: 0 }, ...
d        shape=Shape { dims: [2, 3] } strides=Strides { dims: [3, 1] } contiguous=true can_mut=true buffer=ManagedMemoryHandle { descriptor: ManagedMemoryDescriptor { id: ManagedMemoryId { value: 6 }, ...
launches {"KernelBinop": 1}

c 的显存 id 还是 0：结果直接写进了 a 的显存。d 是新的 6 号，因为启动时 c 有两个持有者。打印时 c 的 can_mut 又是 true，那份克隆已经随这次调用用完丢掉了。

📝NOTE

没有蠢问题

问：原地写会不会把别人还在用的数据改掉？a 被写坏了怎么办？

答： float_add 按值拿走 lhs，调用者手里已经没有 a；还想留着 a 就得 clone，克隆共享同一个 Handle，引用计数大于 1，can_mut 就是 false，走新显存那一支，上面的 d 就是这种情况。所以「能不能原地写」在 CPU 侧按所有权决定，内核里不需要判断。burn 的 Tensor API 也是按值的（a + b 消耗两边），在 Fusion 之上同样成立，只是走的路径不同：融合内核有自己的原地写判断（生成融合内核）。

视图、into_contiguous 和显存顺序

改形状的操作大多只改 meta。swap_dims 是 tensor.meta.swap(dim1, dim2)；reshape 先问 reshape_action，能用新步长表达就只换 Metadata，表达不了才申请新显存、用 copy_into 搬一遍。需要行优先布局的地方调 into_contiguous：

/// Make a jit tensor contiguous.
pub fn into_contiguous(tensor: CubeTensor) -> CubeTensor {
    // A tiled buffer's strides order its tiles, not a matrix's rows: contiguous says nothing.
    let tensor = untile(tensor);
    if tensor.is_contiguous() {
        return tensor;
    }

    if tensor.qparams.is_some() {
        return into_contiguous_quantized(tensor, MemoryLayoutStrategy::Contiguous);
    }

    let (client, device, dtype) = (tensor.client.clone(), tensor.device.clone(), tensor.dtype);

    let output = cubecl::std::tensor::into_contiguous(
        &client,
        tensor.binding(),
        dtype_to_storage_type(dtype),
    );
    // ...
}

已经连续就原样返回（不启动内核）；分块存储的先用 cubek 的 untile 还原成行；量化张量另走一条；普通张量交给 cubecl 标准库的 into_contiguous，它启动一个拷贝内核。读回数据用的是另一个版本 into_contiguous_aligned：先问设备「这个步长能不能直接读」（wgpu 上就是步长是否等于连续步长），不能才拷贝，拷贝时允许行末留对齐用的空隙。into_data 第一行就是它，所以读一个转置过的张量会多一次拷贝。

rt-cube 接着把 d 转置成 t，再分别 into_contiguous、做加法、读回：

--- t = swap_dims(d, 0, 1)
t        shape=Shape { dims: [3, 2] } strides=Strides { dims: [1, 3] } contiguous=false can_mut=false buffer=ManagedMemoryHandle { descriptor: ManagedMemoryDescriptor { id: ManagedMemoryId { value: 6 }, ...
launches {}
--- into_contiguous(t)
tc       shape=Shape { dims: [3, 2] } strides=Strides { dims: [2, 1] } contiguous=true can_mut=true buffer=ManagedMemoryHandle { descriptor: ManagedMemoryDescriptor { id: ManagedMemoryId { value: 9 }, ...
launches {"CopyKernel": 1}
...
--- into_data(t)
t data [21.0, 84.0, 42.0, 105.0, 63.0, 126.0]
launches {"CopyKernel": 1}

转置没有启动内核，t 和 d 是同一块 6 号显存，步长从 [3, 1] 变成 [1, 3]；d 还活着，所以 t 不可写。
into_contiguous 启动一个 CopyKernel，结果在新的 9 号显存里，步长 [2, 1]。
into_data(t) 又启动了一个 CopyKernel：into_contiguous_aligned 判断 wgpu 读不了这种步长。

❗IMPORTANT

动脑筋：两个转置过的张量相加，结果是什么步长？

t 和 t2 都是 [2, 3] 张量 swap_dims(0, 1) 得到的 [3, 2] 视图，步长都是 [1, 3]。B::float_add(t.clone(), t2) 的结果 s：步长是 [2, 1]（新建一个行优先的输出），还是 [1, 3]？启动了几个内核？答案在本页末尾。

这和 in_memory_order 有关。in_memory_order 让和输出同形状的操作数按字节数投票，选它们在显存里的维度顺序；票数最多的顺序不是逻辑顺序时，先把操作数按那个顺序「摆正」（只改元数据），内核按显存顺序线性地走、能向量化，最后把结果的元数据换回逻辑顺序。平票或者已经是逻辑顺序时什么都不做，和以前完全一样。

📝NOTE

冷知识：逐元素内核以前只按逻辑顺序走

PR #5625（「perf(cubecl): walk elementwise kernels in operands' memory order」，首次进入 v0.22.0-pre.4）之前，launch_binop 直接按逻辑维度顺序遍历，输出也按逻辑顺序连续分配。提交说明讲了为什么要改：卷积把 channels-last 的张量以 NCHW 的形状交给下游，没被融合掉的逐元素操作（反向里的 grad * x、乘标量、激活）就按步长跳着读它，而且逻辑上的最后一维不是连续的那一维，向量化也用不上；融合规划器开始保留 channels-last 布局以后，「every unfused consumer paid this on every activation」。现在由 in_memory_order 按操作数的字节数投票选遍历顺序，「A permute is a change of metadata, so no data moves」。同一个 PR 先把 dim_order、nested_dim_order 这些只看步长的函数从 burn-cubecl-fusion 搬到了 burn-std，因为 burn-cubecl 不能依赖融合 crate。

matmul：带着步长交给 cubek

float_matmul 调 kernel::matmul::matmul，策略取 MatmulStrategy 的默认值：开了 autotune feature 就是 Autotune，否则是 Cube（cubek 的默认策略）。matmul 做三件事：

按 lhs、rhs 的形状建输出（init_matmul_output）；两级量化的输入先反量化，因为「No quantized matmul kernel applies a per-tensor scale」。
MatmulTransformPolicy 检查「rhs 在批维上广播」的情况：[.., b, m, k] @ [.., 1, k, n] 改写成一个 [.., 1, b*m, k] 的 matmul，只改 lhs 和输出的元数据，免得 rhs 被读 b 遍。
Autotune 走 matmul_autotune（选候选的过程见 autotune），每个候选最后都调 launch_matmul：

    let mut dtypes = MatmulElems::from_globals(&MatmulGlobalElems {
        lhs: dtype_to_storage_type(lhs_dtype),
        rhs: dtype_to_storage_type(rhs_dtype),
        out: dtype_to_storage_type(out_dtype),
    });

    cubek::matmul::launch::launch_ref(
        &strategy,
        client,
        lhs_handle,
        rhs_handle,
        out.clone().binding(),
        &mut dtypes,
    )?;

lhs_handle / rhs_handle 是 InputBinding::new(tensor.binding(), ..)：句柄、形状、步长、分块原样交出去，量化张量再附上 scale 的句柄。burn 这边没有 into_contiguous。cubek 的 launch_ref 只有一行 strategy.launch_ref(..)，由策略决定内核族和分块；步长怎么读是内核和调优键的事（调优键里的 matrix_layout_rhs 会记下「转置」，见 autotune）。MatmulElems 是可变引用，选内核时 cubek 可能把寄存器或共享内存里的元素类型换掉（比如 f32 算成 tf32），全局的输入输出类型不变。

rt-cube 用上面的 d（[2, 3]）乘它的转置视图 t（[3, 2]，步长 [1, 3]），这个形状的 autotune 结果已经在缓存里：

--- matmul(d, t)  ([2,3] x [3,2], rhs transposed view)
m        shape=Shape { dims: [2, 2] } strides=Strides { dims: [2, 1] } contiguous=true can_mut=true buffer=ManagedMemoryHandle { descriptor: ManagedMemoryDescriptor { id: ManagedMemoryId { value: 14 }, ...
launches {"MatmulEntry": 1}
...
m data [6174.0, 14112.0, 14112.0, 33957.0]

只有一个 MatmulEntry，没有 CopyKernel：转置的 rhs 直接按步长读。结果 d · dᵀ，第一个元素 21² + 42² + 63² = 6174。

动脑筋答案

步长是 [1, 3]，只启动了一个内核：

--- t + t2 (both transposed)
t2       shape=Shape { dims: [3, 2] } strides=Strides { dims: [1, 3] } contiguous=false can_mut=true buffer=ManagedMemoryHandle { descriptor: ManagedMemoryDescriptor { id: ManagedMemoryId { value: 12 }, ...
s        shape=Shape { dims: [3, 2] } strides=Strides { dims: [1, 3] } contiguous=false can_mut=true buffer=ManagedMemoryHandle { descriptor: ManagedMemoryDescriptor { id: ManagedMemoryId { value: 12 }, ...
launches {"KernelBinop": 1}

两个操作数都和输出同形状，显存顺序都是「第 0 维最快」，投票一致，in_memory_order 把它们摆成 [2, 3]、步长 [3, 1] 交给内核，内核线性地走。
结果写进了 t2 的显存（12 号）：t2 只有一个持有者，t 被克隆过不可写，于是 rhs.can_mut_broadcast(&lhs) 那一支成立，原地写回右操作数。
结果换回逻辑形状 [3, 2]，步长随之是 [1, 3]，不是连续的。之后要读回或做真正的 reshape 时才付拷贝的钱；模块注释也提醒了这一点：「The extra copy can outweigh the elementwise speedup」。

小结

CubeBackend 是空结构体，四种原语都是 CubeTensor，设备是 cubecl::Device；哪个运行时由设备的值决定，name / sync 都先 device.client()。Dispatch 用的是 Cube = Fusion<CubeBackend>。
CubeTensor = client + handle（托管显存）+ meta（形状、步长、分块）+ device + dtype + qparams。
逐元素操作 = 一个 #[cube] 的 BinaryOp 实现 + 通用的 kernel_binop + CPU 侧的 launch_binop（广播、显存顺序、向量宽度、网格、原地写）。原地写由句柄的引用计数决定。
转置、多数 reshape 只改元数据；into_contiguous 已连续就不动，否则启动拷贝内核；into_data 用 into_contiguous_aligned，读转置张量会多一次拷贝。
matmul 不先变连续，带着步长和句柄交给 cubek 的 launch_ref；内核由 autotune 在 cubek 的策略里挑。