Rust in Everything
基于 burn 0.22.1

生成融合内核

主线位置：算子主线第 6 步的后一半到第 7 步。融合搜索选定了「这几个 op 用逐元素融合器的方案」，执行时调用这个方案的 execute，它生成（或从缓存取出）一个 cubecl 内核并提交。本页讲 burn-cubecl-fusion（16 660 行）怎么把一串 OperationIr 变成一个内核，以及 matmul、归约怎么带上前后的逐元素 op。#[cube] 宏怎么展开、内核怎么编成 MSL、编译缓存和 autotune 属于 cubecl 0.11.3，见 #[cube] 内核、wgpu 与 Metal、autotune。

flowchart TB
  ops["一块 OperationIr"] --> tf["TraceOperationFuser fuse"]
  tf --> trace["FuseTrace 输入 输出 FuseOp 列表"]
  trace --> launch["FuseTraceLauncher 绑定句柄"]
  launch --> vec["选向量宽度 选参考布局"]
  vec --> cfg["FuseBlockConfig comptime"]
  cfg --> kern["elemwise_fuse 内核"]
  kern --> msl["cubecl 编成 MSL 提交"]

五个内置融合器

burn-cubecl 的注册表（fusion/registry.rs）给每条流发这五个融合器，顺序就是 BUILTIN_NAMES：

名字融合器收什么
ElementWiseElementWiseFuser一串逐元素 op（算术、比较、exp 之类、mask_fill、gather、cat……）
MatmulMatmulFuser一个 Matmul 开头，后面接的逐元素 op
ReduceReduceFuser一个按维度的归约（SumDim、MeanDim、ArgMax……），前面和后面的逐元素 op
ReduceBroadcastedReduceBroadcastedFuser以 ReduceFuser 为单元，把几个归约块和它们周围的逐元素 op 串成一个方案（只读了源码）
NHWCRelayoutNHWCRelayoutFuser为模块 op 把张量重排成通道在最后（NHWC）的布局（只读了源码）

五个都实现 OperationFuser，finish 产出包在 CubeOptimization 里的 FusedOperation。用户也能用 register 加自己的融合器，必须在第一个张量 op 之前，否则报 ServiceRunning（「register or remove fusion optimizations at the start of your program」）。

前三个共用同一个引擎：TraceOperationFuser。它的 fuse 是一个 match：BaseFloat / BaseInt / BaseBool / Float / Int / NumericFloat / NumericInt 交给对应的 fuse_* 方法，认得就收；Drop 记下「这个张量用完了」；其余顶层变体（Init、Module、Custom、Activation、Bool、Distributed）一律关闭，日志里的原因是 unsupported op variant。

逐元素：从 IR 到 FuseTrace

ElementWiseFuser 只是给 TraceOperationFuser 配了一组设置（new）：

            fuser: TraceOperationFuser::new(
                max_bindings,
                FuseSettings {
                    broadcast: true,
                    output_shape_updates: true,
                    inplace: true,
                    vectorization: VectorizationSetting::Activated,
                    ref_layout: RefLayoutSetting::Any,
                    // The elementwise runner reads and writes every operand through
                    // the generic fused paths, so it is free to iterate in whatever
                    // order its inputs already sit in.
                    choose_output_layout: true,
                },
            ),

收 op 的过程是在记一份 FuseTrace：哪些张量是输入（从全局内存读）、哪些是输出（要写回全局内存）、中间是一串 FuseOp（Add、Mul、Exp、Assign、ConditionalAssign、Gather、Cat……）。只在块内产生、块内用完的张量（Drop 过的，或者以 ReadWrite 被下一个 op 读掉的）不算输出，变成 BlockLocal 局部变量，永远不进显存。

每收一个 op 前，包装层 TryTraceFuser 先在一份克隆上试加，过两道限制才真正收下（fuse）：

    fn fuse(&mut self, add_ops: impl FnOnce(&mut TraceFuser) -> Option<()>) -> bool {
        if self.fuser.num_ops_fused() > self.max_ops {
            return false;
        }

        // Optimization to avoid cloning the fuser when adding the first operation
        if !self.added_ops {
            self.added_ops = true;

            if add_ops(&mut self.fuser).is_none() {
                return false;
            }
            return self.fuser.estimate_bindings() <= self.max_bindings;
        }

        let mut cloned = self.fuser.clone();
        if add_ops(&mut cloned).is_none() {
            return false;
        }

        if cloned.estimate_bindings() > self.max_bindings {
            return false;
        }

        self.fuser = cloned;
        true
    }
  • max_ops 在构造时写死为 64（注释：「A good default, avoid errors with for loops over only memory bound operations」）。检查用的是 >，所以第 66 个 op 才被拒，一个内核最多 65 个 op；Fusion 装饰器里 100 个 op 的链被切成 65 + 35，就是这里。
  • max_bindings 来自设备属性。wgpu 运行时取「每个着色器阶段的存储缓冲区上限减一」（cubecl-wgpu/src/runtime.rs，注释说 wgpu 自己要占一个），这台机器的融合日志里是 est_bindings=1/28。估算是每个输入输出一个绑定、再加一个元数据。
📝NOTE
冷知识：第一个 op 曾经不查绑定数

上面代码里 return self.fuser.estimate_bindings() <= self.max_bindings; 这一行来自 #6053（「fix(fusion): check binding limit for the first operation」，首次发布于 v0.22.1，0.22.1 里刚 cherry-pick 进来）。改之前这里是注释「Always allow the first operation to be added」加 return true;：第一个 op 无条件收下。提交里新增的测试写明了漏洞：「A single op can exceed max_bindings on its own (here a cat with many inputs), so the binding check must also apply to the first op of a kernel」，对应 issue #6049。一个输入很多的 cat 作为块的第一个 op，生成的内核绑定数超过设备上限。

内核：comptime 展开的一串局部变量

方案执行时，ElemwiseRunner::run 按参考形状的元素数除以向量宽度算出要多少个工作单元，启动同一个 #[cube] 内核 elemwise_fuse：

#[cube(launch_unchecked, address_type = "dynamic")]
fn elemwise_fuse(
    inputs: &GlobalArgs,
    outputs: &mut GlobalArgs,
    #[comptime] config: FuseBlockConfig,
) {
    // We write no values for this fusion.
    let values = Registry::<FuseArg, Vector<f32, DynSize>>::new();
    let args = comptime![Vec::<FuseArg>::new()];
    let pos = ABSOLUTE_POS;

    multi_block_variables_init(&config, &mut outputs.variables);

    let mut locals = init_locals(inputs, outputs, &config);
    let length = ref_len(inputs, &*outputs, &locals, &config);

    if pos < length {
        fuse_on_write::<f32, DynSize>(inputs, outputs, &mut locals, pos, values, args, &config)
    }
}

所有逐元素融合共用这一个 Rust 函数。区别全在 #[comptime] config: FuseBlockConfig 里：op 列表、向量宽度、参考布局都是编译期常量。fuse_on_write 调 fuse，后者对 config.ops 做一个 #[unroll] 循环，每个 FuseOp 展开成对应的运算：

    #[unroll]
    for index in 0..config.ops.len() {
        let op = config.ops[index].clone();
        let define!(E) = op.cmp_storage_ty();
        let size!(N) = config.width;

        match op {
            FuseOp::Add(op) => add::<E, N>(inputs, outputs, locals, pos, op, config),
            FuseOp::Div(op) => div::<E, N>(inputs, outputs, locals, pos, op, config),
            FuseOp::Sub(op) => sub::<E, N>(inputs, outputs, locals, pos, op, config),
            FuseOp::Mul(op) => mul::<E, N>(inputs, outputs, locals, pos, op, config),
            // ...
            FuseOp::Exp(op) => exp::<E, N>(inputs, outputs, locals, pos, op, config),
            // ...

循环和 match 都在 cubecl 展开内核的时候执行完，生成的代码里只剩这一串运算本身。所以不同的 op 序列是不同的内核，同一个序列换形状还是同一个内核（形状、步长是运行期参数）。

实验：(a * b + c).exp() 生成的 MSL

CUBECL_DEBUG_LOG=stdout 让 cubecl 打印每次编译（做法见一次 matmul）。跑 Fusion 装饰器里的 ir 模式（Metal，2×2，a、b、c 都被 move 进表达式），编译日志先打出 KernelId 里的 comptime 配置（摘出 ops 和输出参数，结构体的收尾行省略，标 ...）：

name: burn_cubecl_fusion::optim::elemwise::optimization::elemwise_fuse::ElemwiseFuse
...
            ops: [
                Assign(
                    UnaryFuseArgs {
                        input: Input(
                            0,
                            F32,
                            SameAsRef,
                        ),
                        out: BlockLocal {
                            pos: 0,
                            ty: F32,
...
                Mul(
                    BinaryFuseArgs {
                        lhs: BlockLocal {
                            pos: 0,
                            ty: F32,
                        ...
                        out: BlockLocal {
                            pos: 2,
                            ty: F32,
...
                Exp(
                    UnaryFuseArgs {
                        input: BlockLocal {
                            pos: 4,
                            ty: F32,
                        ...
                Assign(
                    UnaryFuseArgs {
                        input: BlockLocal {
                            pos: 5,
                            ty: F32,
                        ...
                        out: Output(
                            0,
                            F32,
                            IsRef,
                        ),
...
            width: 2,
            vector_axis: 1,
...
        outputs: GlobalArgs {
            tensors: Sequence [
                GlobalTensorCompilationArg {
                    tensor: TensorCompilationArg {
                        ...
                        buffer: BufferCompilationArg {
                            inplace: Some(
                                0,
                            ),

六个 FuseOp：三个 Assign 把输入读进局部变量 0、1、3，Mul、Add、Exp 在局部变量之间算，最后一个 Assign 写到输出 0。输出的缓冲区 inplace: Some(0)：复用输入 0（也就是 a）的显存。接着是生成的 MSL（类型定义和下标计算省略，标 ...）：

struct alignas(8) float_2 { float i_0; float i_1; };
...
[[max_total_threads_per_threadgroup(64)]] [[kernel]] void elemwise_fuse_cecf0de9(device float_2 * const v1 [[buffer(0)]], device float_2 const* const v7 [[buffer(1)]], device float_2 const* const v13 [[buffer(2)]], constant info_st const* const v103, uint3 const threadgroups_per_grid_v177 [[threadgroups_per_grid]], uint3 const threadgroup_position_in_grid_v178 [[threadgroup_position_in_grid]], uint32_t const thread_index_in_threadgroup_v179 [[thread_index_in_threadgroup]]) {
...
uint32_t const v170 = v167 + v169;
// Init locals begin
...
// Init locals end
bool const v76 = v170 < v112;
if(v76) {
// Fuse on write begin
device float_2 * const v81 = &v1[v170];
float_2 const v82 = *v81;
device float_2 const* const v87 = &v7[v170];
float_2 const v88 = *v87;
device float_2 const* const v93 = &v13[v170];
float_2 const v94 = *v93;
float const v129 = v82.i_0;
float const v130 = v88.i_0;
float const v131 = v129 * v130;
float const v132 = v82.i_1;
float const v133 = v88.i_1;
float const v134 = v132 * v133;
float const v137 = v94.i_0;
float const v138 = v131 + v137;
float const v140 = v94.i_1;
float const v141 = v134 + v140;
float const v144 = exp(v138);float const v146 = exp(v141);float_2 const v147 = float_2{v144, v146};
*v81 = v147;
// Fuse on write end
}
return;
}

对着读：

  • // Fuse on write begin / end 是 fuse_on_write 里 comment! 宏留下的注释，中间就是 config.ops 展开的结果：三次读（v82、v88、v94）、一次乘、一次加、一次 exp、一次写。a * b 的结果 v131、v134 和加完的 v138、v141 只是寄存器里的局部变量，不进显存。不融合时这两个中间张量各要写一遍、读一遍显存。
  • 三个缓冲区里只有 v1（buffer(0)，即 a）是可写的 device float_2 *，结果 *v81 = v147 写回读 a 的同一个位置：就地改写。条件是 a 在 IR 里是 ReadWrite（最后一次使用，见 burn-ir），融合设置里 inplace: true。把 a 写成 a.clone()，a 变成 ReadOnly，输出就会另分配一个缓冲区：下面计时用的 bench expr 模式里 a 是 clone 进去的，CUBECL_DEBUG_LOG 打出的内核签名第一个参数是只读的 device float_4 const*。
  • float_2 和 exp 被逐分量展开：cubecl 的 MSL 后端把向量写成结构体，每个分量单独算（编译细节见 wgpu 与 Metal）。宽度为什么是 2，下一节。

向量化：按末维选宽度

每个输入的向量宽度由 vectorization_input 决定：

    let axis = axis.get(desc.id, || handle.strides.len() - 1);
    let shape_axis = desc.shape[axis];

    if shape_axis == 1 {
        return Vect::Broadcasted;
    }

    // Last dimension strides should be 1, otherwise vecX won't be contiguous.
    if handle.strides[axis] != 1 {
        return Vect::Aligned(1);
    }
    // ...
    let inner = |s: VectorSize| {
        // The last dimension should be a multiple of the vector size or broadcated.
        if shape_axis.is_multiple_of(s) && next_stride % s == 0 {
            return Some(Vect::Aligned(s));
        }
        None
    };

候选宽度来自 cubecl 客户端的 io_optimized_vector_sizes(元素字节数)，从大到小试，取第一个既能整除向量轴（默认是最后一维）长度、又能整除其他维最小步长的；末维步长不是 1（转置过、切片过）就只能是 1；末维是 1 的输入算广播，不参与投票。整个块的宽度再从各输入、输出的结果里协调出来。

❗IMPORTANT
动脑筋：[4, 6] 的 (a * b + c).exp()，MSL 里是 float、float_2 还是 float_4？

同一个表达式，a、b、c 都是 Tensor::<2>::random([4, 6], ..)，Metal，打开 CUBECL_DEBUG_LOG=stdout。生成的融合内核的缓冲区元素类型是哪个？对照：[4, 3] 和 [64, 64]。答案在本页末尾。

matmul 和归约：把逐元素 op 收进 cubek 内核

MatmulFuser 的 fuse 只接受一种开头：第一个 op 必须是 Float(_, Matmul(..))，否则立刻关闭。之后的每个 op 交给两份 trace：fuser 用来生成「matmul 内核读写时顺带做这些逐元素运算」的版本，fuser_fallback 用来生成「matmul 单独跑，后面的逐元素 op 单独融合成一个内核」的版本。两份都收得下才收。执行时 fused_matmul_autotune 在两者之间挑，第一个候选就叫 fused_matmul_fallback（注释「First entry should always work, since it is considered the fallback」），其余是各种 cubek matmul 选择器的融合版本。

ReduceFuser 认的是按维度的归约：SumDim、MeanDim、ProdDim、ArgMax、ArgMin、MinDim、MaxDim、MaxAbsDim、AnyDim、AllDim。归约之前的逐元素 op 在读输入时做（fuse_on_read），之后的在写输出时做（fuse_on_write），内核是 reduce_kernel_fused。

mm 模式（Metal，64×64，debug；第一次运行时这两个形状的 autotune 结果还不在缓存里，第二次运行已在）：

let d = (a.clone().matmul(b.clone()) + c.clone()).exp();
let _ = d.into_data();
let e = (a * b).sum_dim(1).exp();
let _ = e.into_data();

第一次（内核名截到 110 个字符）：

--- (a.matmul(b) + c).exp()
  [block] 3 op(s): Float(F32, Matmul), NumericFloat(F32, Add), Float(F32, Exp)
  [launch] [DSD-4-16384] cubek_matmul::multi_level::components::batch::partitioned_matmul::matmul::matmul_entry::MatmulEntry<burn_cubec...
  [launch] [DSD-4-16384] cubek_matmul::multi_level::components::batch::partitioned_matmul::matmul::matmul_entry::MatmulEntry<burn_cubec...
  [launch] [DSD-4-16384] cubek_matmul::multi_level::components::batch::partitioned_matmul::matmul::matmul_entry::MatmulEntry<burn_cubec...
  [launch] [DSD-4-16384] cubek_matmul::multi_level::components::batch::partitioned_matmul::matmul::matmul_entry::MatmulEntry<burn_cubec...
  [launch] [DSD-4-16384] cubek_matmul::multi_level::components::batch::partitioned_matmul::matmul::matmul_entry::MatmulEntry<burn_cubec...
  [launch] [DSU-100-16384] cubek_matmul::multi_level::components::batch::partitioned_matmul::matmul::matmul_entry::MatmulEntry<burn_cubec...
--- (a * b).sum_dim(1).exp()
  [block] 3 op(s): NumericFloat(F32, Mul), NumericFloat(F32, SumDim), Float(F32, Exp)
  [launch] [DSD-4-16384] burn_cubecl_fusion::optim::reduce::optimization::reduce_kernel_fused::ReduceKernelFused
  [launch] [DSD-4-16384] burn_cubecl_fusion::optim::reduce::optimization::reduce_kernel_fused::ReduceKernelFused
  [launch] [DSD-4-16384] burn_cubecl_fusion::optim::reduce::optimization::reduce_kernel_fused::ReduceKernelFused
  [launch] [DSU-100-16384] burn_cubecl_fusion::optim::reduce::optimization::reduce_kernel_fused::ReduceKernelFused
--- done

第二次（内核名截到 400 个字符）：

--- (a.matmul(b) + c).exp()
  [block] 3 op(s): Float(F32, Matmul), NumericFloat(F32, Add), Float(F32, Exp)
  [launch] [DSU-100-16384] cubek_matmul::multi_level::components::batch::partitioned_matmul::matmul::matmul_entry::MatmulEntry<burn_cubecl_fusion::optim::matmul::args::FusedMatmulArgs, cubecl_core::frontend::element::typemap::DynamicScalar<cubek_matmul::multi_level::args::__Lhs>, cubecl_core::frontend::element::typemap::DynamicSize<cubek_matmul::multi_level::args::__LhsSize>, cubecl_core::frontend::element::typemap::Dynamic...
--- (a * b).sum_dim(1).exp()
  [block] 3 op(s): NumericFloat(F32, Mul), NumericFloat(F32, SumDim), Float(F32, Exp)
  [launch] [DSU-100-16384] burn_cubecl_fusion::optim::reduce::optimization::reduce_kernel_fused::ReduceKernelFused
--- done
  • 两个表达式各是一个块、各提交一个内核。matmul 用的是 cubek 的 MatmulEntry，类型参数第一个是 burn_cubecl_fusion::optim::matmul::args::FusedMatmulArgs：cubek 的 matmul 内核对「怎么读输入、怎么写输出」是泛型的，burn 换进一个会顺带算 + c 和 exp 的参数类型。cubek 内核怎么写见 burn-cubecl 后端。
  • 第一次在 DSD 线程上多出来的几次提交是 autotune 在试候选，第二次结果已经在缓存里，只提交一次。
  • 不是所有归约都能融合：Fusion 装饰器那个 lazy 模式的 y.sum()（对全部元素求和，IR 是 NumericFloat(F32, Sum)，不带维度）不在 ReduceFuser 的名单里，单独执行，提交了两个 ReduceKernel，后面的 + 1.0 也只能单独一个内核。

实验：(a * b + c).exp() 融合前后

和 Fusion 装饰器里一样，fusion-lab 和 fusion-nofuse 用同一份代码，只差 burn 的默认 feature。先数内核（expr-count，1024×1024，debug，每次 sync）：

dispatch<fusion<cubecl<wgpu<msl>>>>
run 0: 31.018125ms, 1 launch(es)
       1 burn_cubecl_fusion::optim::elemwise::optimization::elemwise_fuse::ElemwiseFuse
run 1: 1.357959ms, 1 launch(es)
       1 burn_cubecl_fusion::optim::elemwise::optimization::elemwise_fuse::ElemwiseFuse
dispatch<cubecl<wgpu<msl>>>
run 0: 52.474959ms, 3 launch(es)
       1 burn_cubecl::kernel::binary::kernel_binop::KernelBinop<burn_cubecl::kernel::binary::MulOp>
       1 burn_cubecl::kernel::binary::kernel_binop::KernelBinop<burn_cubecl::kernel::binary::AddOp>
       1 burn_cubecl::kernel::unary_float::unary_float::UnaryFloat<burn_cubecl::kernel::unary_float::unary_ba...
run 1: 1.2875ms, 3 launch(es)
       ...

再计时（bench expr，release，4096×4096 f32；每轮 5 组不相干的输入各算一遍、保留结果、sync 一次，取每次的平均；方法和原因见 Fusion 装饰器）：

进程里第一次（含编译）稳定后每次（7 轮中位数）每轮 5 次的内核数
开 fusion19.5 ms883 µs1
不开 fusion39.7 ms1.61 ms15
  • 冷：第一次的时间主要是编译。开 fusion 编一个内核，不开要编三个（KernelBinop<MulOp>、KernelBinop<AddOp>、UnaryFloat），所以第一次反而快。这几个逐元素内核不经 autotune：fusion-nofuse 跑完所有实验，crate 目录下连 target/environment/ 都没有生成。上面 debug 构建的单次运行里，第二次 1.36 ms 对 1.29 ms 几乎一样：1024² 上一次 sync 的往返就占了大头，看不出差别。
  • 热：稳定后约 1.8 倍。按显存流量算：融合后每次读 a、b、c 三个 64 MiB、写一个，共 256 MiB；不融合时 mul 读二写一、add 读二写一、exp 读一写一，共 512 MiB。流量减半，时间也接近减半。
  • 开 fusion 的第 0 轮要再编一个「5 个输出」的内核（5 组输入在一次 sync 前登记，被合进同一个块，见融合搜索），那一轮是 7.76 ms，不计入中位数。

数字只对这台机器（M2 Max、macOS 27、rustc 1.99.0）有效；CUDA、ROCm、Vulkan 上走的是同一套 burn-cubecl-fusion 代码，没有实测。

动脑筋答案

float_2。vec 模式在一个进程里依次跑 [4, 3]、[4, 6]、[64, 64]，CUBECL_DEBUG_LOG=stdout，摘出三个融合内核的签名开头（参数列表后半省略，标 ...）：

--- shape [4, 3]
[[max_total_threads_per_threadgroup(64)]] [[kernel]] void elemwise_fuse_c7c89564(device float * const v1 [[buffer(0)]], device float const* const v7 [[buffer(1)]], device float const* const v13 [[...
--- shape [4, 6]
[[max_total_threads_per_threadgroup(64)]] [[kernel]] void elemwise_fuse_cecf0de9(device float_2 * const v1 [[buffer(0)]], device float_2 const* const v7 [[buffer(1)]], device float_2 const* const ...
--- shape [64, 64]
[[max_total_threads_per_threadgroup(256)]] [[kernel]] void elemwise_fuse_3779af24(device float_4 * const v1 [[buffer(0)]], device float_4 const* const v7 [[buffer(1)]], device float_4 const* cons...
  • 末维 3 不能被 4 或 2 整除，只能逐个元素：float。末维 6 能被 2 整除、不能被 4 整除：float_2。末维 64：float_4，每个线程一次处理 4 个 float。
  • [4, 6] 的内核名后缀 cecf0de9 和前面 2×2 实验里的一模一样：两者的 op 序列、宽度、布局相同，comptime 配置相同，就是同一个内核，形状只是运行期参数。
  • [64, 64] 的线程组上限从 64 变成 256：工作单元多了，cubecl 按工作量选了更大的 CubeDim（ElemwiseRunner::run 里的 CubeDim::new(client, working_units)）。

小结

  • cubecl 后端有五个内置融合器；逐元素、matmul、归约三个共用 TraceOperationFuser，它只收 Base*、Float、Int、Numeric* 几类 IR，Init、Activation、Module 之类一来就关闭。
  • 收下的 op 记成 FuseTrace：输入、输出和一串 FuseOp；块内产生又用完的张量是局部变量，不进显存。一个内核最多 65 个 op、绑定数不超过设备上限（Metal 上 28）。
  • 所有逐元素融合都是同一个 #[cube] 函数 elemwise_fuse，op 列表、宽度、布局是 comptime 配置，#[unroll] 把它展开成一串寄存器运算；实测的 MSL 里三次读、三次运算、一次写，结果就地写回 ReadWrite 的输入。
  • 向量宽度按向量轴长度和步长从大到小试：末维 3、6、64 分别得到 float、float_2、float_4。
  • matmul 融合换掉 cubek 内核的读写参数类型（FusedMatmulArgs），并和「不融合的 matmul + 融合的尾巴」一起参加 autotune；按维度的归约在读时和写时各带上逐元素 op，全量 sum 不融合。
  • (a * b + c).exp() 在 4096² 上：3 个内核变 1 个，稳定后 1.61 ms 变 883 µs，第一次运行少编译两个内核。

💬 关联讨论

发起讨论
加载中...