Rust in Everything
基于 burn 0.22.1

融合搜索

主线位置：算子主线第 6 步的前一半。op 已经进了某条流的队列（见 Fusion 装饰器），每来一条，服务端就问一次：现在该执行什么、按什么方案执行。本页讲这个决定怎么做出来、做出来之后怎么缓存；方案里的融合内核怎么生成见生成融合内核。

flowchart TB
  newop["新 op 进队列"] --> rel["换成相对形式"]
  rel --> pol{"Policy 查计划缓存"}
  pol -->|命中| exec["按缓存的计划执行"]
  pol -->|可能命中| defer["Defer 再等等"]
  pol -->|没有| explore["Explorer 探索"]
  explore --> blocks["StreamOptimizer 分块"]
  blocks --> fusers["每块各融合器竞争打分"]
  fusers --> plan["新计划 存进缓存"]
  plan --> exec

每条流一个 Processor

MultiStream 给每条流配一个 Processor，里面两样东西：Policy（查缓存）和 Explorer（找新方案）。它在两种模式下运行（ExecutionMode）：

  • Lazy：每登记一个 op 跑一次。可以「再等等」，等更多 op 来了再决定。
  • Sync：drain 时跑。队列里的东西必须全部执行完，不能再等。

process 是一个循环，每轮问 Policy 要一个 Action：

        loop {
            if segment.operations().is_empty() {
                break;
            }

            let action = self.policy.action(store, segment.operations(), mode);

            match action {
                Action::Explore => {
                    self.explore(&mut segment, store, mode);

                    if self.explorer.is_up_to_date() {
                        break;
                    }
                }
                Action::Defer => {
                    match mode {
                        ExecutionMode::Lazy => break,
                        ExecutionMode::Sync => panic!("Can't defer while sync"),
                    };
                }
                Action::Execute(id) => {
                    // ...
                    segment.execute(id, store);
                    self.reset(store, segment.operations());
                }
            };
        }

三种动作：

动作什么时候结果
Execute(id)缓存里有一个计划，它的 op 序列和队列开头对上了，触发条件也满足了按计划执行这一段，队列里剩下的重新开始匹配
Defer只在 Lazy 下：有候选计划正在逐条比对，还没比完什么也不做，等下一个 op
Explore没有候选让 Explorer 找方案；Lazy 下融合器还开着就继续等，都关了或者是 Sync 就定下方案、存进缓存、执行

相对形式：让不同的调用看起来一样

缓存要能命中，第二次的 op 序列得和第一次「一样」。但每次调用的张量 id 都是新发的，标量也可能不同。所以 OperationQueue::add 每收一条 IR 存两份：global 是原样，relative 是换算过的，搜索和缓存只看后者。换算由 RelativeOps 做，规则三条：

  • 张量 id 按在这一段里第一次出现的顺序重新编号，0、1、2……
  • 形状 不存具体数，存「编号」：impl RelativeOps for TensorIr 给每个没见过的维度值发一个号，见过的复用。维度值 1 固定是 0 号（「global 1 is always shape id 0」），因为 1 意味着广播，必须和别的值分开。
  • 标量 一律换成占位符 ScalarIr::UInt(序号)（impl RelativeOps for ScalarIr），真值存在旁边，执行时再绑回去。切片范围也一样。

于是 [64, 64] 和 [32, 32] 的相对形状都是 [1, 1]（第一个非 1 的维度值得到 1 号，同一个值第二次出现复用），x * 2.0 和 x * 3.0 的相对形式都是 MulScalar(t0, 标量#0)。代价是缓存认的是这种「编号结构」，不是具体的数。

❗IMPORTANT
动脑筋：形状换成 64×32，还能命中同一个计划吗？

下面「看日志」一节的程序先在 64×64 上跑 (a * b + c * 2.0).exp()，找到计划 #2，换标量、换成 32×32 都命中了它。最后再跑一次，a、b、c 换成 64×32。日志里会出现 cache hit: execute plan #2，还是一次新的 exploration completed？答案在本页末尾。

搜索：分块，让融合器竞争

Explorer 把队列里的 op 一条条交给 StreamOptimizer。它维护若干个 Block：

  • 新 op 和某个块共用张量，就交给那个块；和谁都不共用，就开一个新块；同时连着几个合不到一起的块，就单开一个依赖它们的块，块集合成了一个有向无环图。
  • 块数有上限 max_blocks，默认 5（burn.toml 的 [fusion.beam_search]，BeamSearchConfig）。满了还要开新块，就先试着合并现有的块；合不下去，搜索停止，这一段到此为止。注释说上限不能太高：「Too high and it may break the fusion cache always retriggering explorations」，而且块之间的依赖分析随块数平方增长。

每个块手里有一套全新克隆的融合器，也就是 FusionRuntime::fusers 给的那几个（cubecl 后端是五个，见生成融合内核）。块里每进一个 op，就喂给它的每个融合器。融合器的接口是 OperationFuser：

pub trait OperationFuser<O>: Send {
    /// Register a new [tensor operation](OperationIr).
    fn fuse(&mut self, operation: &OperationIr);
    /// Finish the optimization and create a fusion operation.
    fn finish(&mut self) -> O;
    /// Reset the state.
    fn reset(&mut self);
    /// Return the builder [status](FuserStatus).
    fn status(&self) -> FuserStatus;
    /// Return the builder [properties](FuserProperties).
    fn properties(&self) -> FuserProperties;
    /// The number of operation fused.
    fn len(&self) -> usize;
    // ...
    /// Clone the optimization builder.
    fn clone_dyn(&self) -> Box<dyn OperationFuser<O>>;
}

融合器自己决定收不收：收不下就把 FuserStatus 置为 Closed，之后再喂什么都忽略。它随时报告 FuserProperties：ready（已经收到的 op 能不能生成一个可执行的方案）和 score（越高越好）。一个块里所有融合器都关了，这个块就不再「优化中」；所有块都不在优化中，Explorer 才在 Lazy 下定下方案。

定方案时，BlocksOptimizer 先尝试合并块，再按依赖的拓扑序让每个块各自出方案。块的方案就是它的融合器里 ready 且分数最高的那个（find_best_optimization_index）：

    for (i, optimization) in optimizations.iter().enumerate() {
        let properties = optimization.properties();

        // A score of zero is worse than fusing.
        if properties.ready && properties.score > best_score {
            best_index = BestOptimization::Found {
                index: i,
                score: properties.score,
            };
            best_score = properties.score;
        }
    }

分数为 0 不算数，块里的 op 就逐个执行（ExecutionStrategy::Operations），每个 op 调登记时附的那个 Operation。逐元素融合器的分数是「省下的读写次数 × 100 + 省下的内核启动次数 × 10」（Scoring::evaluate）。(a * b + c * 2.0).exp() 的四个 op 不融合时要读 6 次、写 4 次，融合后读 3 次（a、b、c）、写 1 次，省下 6 次读写，加上少启动 3 个内核：6 × 100 + 3 × 10 = 630，和下面日志里的 score=630 一致。单个 op 融合不省任何东西，分数是 0，所以一个孤零零的 x + 1.0 走的是 CubeBackend 的普通内核。

找到的方案是一个 ExecutionStrategy：

pub enum ExecutionStrategy<O> {
    /// An optimization was found, and therefore should be executed.
    Optimization {
        opt: O,
        ordering: Arc<Vec<usize>>,
        score: u64,
    },
    /// No optimization was found, each operation should be executed individually.
    Operations { ordering: Arc<Vec<usize>> },
    /// A composition of multiple execution strategies.
    Composed(Vec<Box<Self>>),
}

ordering 是这些 op 在段里的下标：块不一定是连续的一段，执行顺序也不一定是登记顺序，只要不违反「先产生后读、读完才释放」。几个块的方案拼成 Composed。中间有 op 谁都没认领（「洞」）时，再对这些 op 单独搜一次补上。

OperationFuser 在大纲和一些文档注释里还叫 builder，那是旧名：#4044（首次出现在 v0.20.0-pre.3）的提交信息里有一条「Builder => Fuser」，diff 把 pub trait OptimizationBuilder<O> 改成了 pub trait OperationFuser<O>，OptimizationStatus 改成 FuserStatus；trait 方法的文档注释没跟着改。

📝NOTE
冷知识：一个合不起来的头，曾经拖着整个块不融合

Block::optimize 里没找到方案的那个分支，来自 #5622（「perf(fusion): settle only the unfusable head of a block」，首次出现在 v0.22.0-pre.4）。提交信息说，之前没有融合器出方案时，Block::optimize「returned the whole block as unfused operations」，块是一条依赖链，头之后本来可以融合的 op 也一个个单独执行；常见的头是「a lone elementwise op followed by a view of its own output, a convolution backward, or a reduction」，正是反向传播的形状，一个卷积网络的训练步里「3,588 one-op segments per step」。改成只把头（到最后一个融合器关闭的那个 op 为止，加上后面没有新融合器愿意接的 op）定成不融合，其余留给下一轮搜索。作者在 L4（CUDA）上测的训练步从 366.1 ms 降到 320.0 ms。现在代码里的 operations_settled_once_every_builder_closed 和 operations_settled_while_no_builder_would_start 就是这两段「头」的算法。

不相干的计算也会被合并

块的合并规则（Block::merge）有一句：合并后「can still be improved, keep it」。只要融合器还开着，两个毫不相干的块也会合成一个。Fusion 装饰器的计时里，5 组互不相干的 8-op 链在一次 sync 前登记，BURN_FUSION_LOG=full 的记录（release，256×256；中间的 accepted 行省略）：

[stream] op 0 NumericFloat::MulScalar → new block (total: 1)
...
[stream] op 8 NumericFloat::MulScalar → new block (total: 2)
...
[stream] op 32 NumericFloat::MulScalar → new block (total: 5)
...
[stream] op 39 NumericFloat::Abs → accepted in 1/5 block(s)
[explorer] explore (sync): 0 deferred of 40 queued
selected single strategy
[plan] exploration completed: sync, 40/40 ops optimized
...
fusion trace (40 ops, 1 block)
...
0    ▸ fused ElementWise (score=7390, 40 ops)    NumericFloat::MulScalar  t1:f32[256,256]   t24:f32[256,256]

登记时 5 条链各占一块，正好碰到 max_blocks；sync 时合并成一个 40-op 的块，生成一个读 5 个输入、写 5 个输出的内核。所以那次计时里每轮只有 1 次提交。

执行计划和缓存

搜索的结果存成 ExecutionPlan，放进设备级的 ExecutionPlanStore：

pub struct ExecutionPlan<O> {
    /// The operations on which the exploration is related to.
    pub operations: Vec<OperationIr>,
    /// The criteria that signal when this plan should be executed. Only one trigger is necessary.
    pub triggers: Vec<ExecutionTrigger>,
    /// The optimization that should be used when executing this plan.
    pub optimization: BlockOptimization<O>,
}

operations 是相对形式的 op 序列，triggers 说「什么情况下用这个计划」，ExecutionTrigger 三种：

触发来自意思
OnOperations(next_ops)Lazy 下探索完成，后面已经有 op 在排队序列对上、并且后面跟着的正是这几条 op 时执行
OnSyncSync 下探索完成序列对上、并且正在 drain 时执行
AlwaysLazy 下探索完成，后面没有 op序列一对上就执行

同一个计划可以攒多个触发条件。查找靠 ExecutionPlanIndex：按计划第一条相对 op 的哈希建索引（OperationIr 不能实现 Eq，因为里面有浮点，所以先哈希再用 PartialEq 处理冲突；标量已经换成占位符，比较不受具体值影响）。Policy::update 在段的第一条 op 到来时取出所有以它开头的计划作候选，之后每来一条就逐条比对，比对失败的剔除，比完整个序列的转成「可用」，再等触发条件。Lazy 下还有候选在比对，action_lazy 就返回 Defer。

还有两道保险。命中的计划如果用到这条流没分配过的相对形状编号，或者用到超出队列长度的下标（缓存是跨流共用的），OperationQueue::execute 就按登记顺序逐个执行这些 op，不冒险。另一个是给动态图的：每次形状都不一样的程序永远命中不了缓存，每段都要重新探索，缓存也越长越大；max_explorations（默认不限，环境变量 BURN_FUSION_MAX_EXPLORATIONS）设了之后，探索次数到顶的流对没命中的段直接不融合执行，也不再往缓存里加。

计划缓存只在进程内存里，没有落盘；Optimization::to_state / from_state 定义了可序列化的状态，但 burn-fusion 里没有调用者。这和 cubecl 会落盘的 autotune 缓存不同（见 autotune）。

看日志：BURN_FUSION_LOG

burn-std 的配置读环境变量 BURN_FUSION_LOG（override_from_env），也可以写在 burn.toml 的 [fusion.logger] 里。四档（FusionLogLevel）：disabled；basic 打每次选了单一还是组合策略；medium 加上块的合并拆分和缓存命中；full 打每次登记、拒绝和打分。设了就默认写到 stderr。这是看融合决定最直接的办法，比数内核更细。

fusion-lab reuse 模式：同一个 (a * b + c * s).exp()，先在 64×64、s = 2.0 上跑三次，再换标量 s = 3.0，再换形状 32×32，每次都 into_data。BURN_FUSION_LOG=full，debug 构建，摘出关键的几段（Random 的计划命中和表格的分隔线省略，标 ...）：

--- iteration 0: (64x64) (a * b + c * 2).exp()
[explorer] explore (lazy): 1 deferred of 1 queued
[stream] op 0 NumericFloat::Mul → new block (total: 1)
[explorer] explore (lazy): 1 deferred of 2 queued
[stream] op 1 NumericFloat::MulScalar → new block (total: 2)
[explorer] explore (lazy): 1 deferred of 3 queued
[stream] op 2 NumericFloat::Add → accepted in 1/1 block(s)
[explorer] explore (lazy): 1 deferred of 4 queued
[stream] op 3 Float::Exp → accepted in 1/1 block(s)
[explorer] explore (sync): 0 deferred of 4 queued
selected single strategy
[plan] exploration completed: sync, 4/4 ops optimized
...
fusion trace (4 ops, 1 block)
idx  block                                     op                       inputs                        outputs      
...
0    ▸ fused ElementWise (score=630, 4 ops)    NumericFloat::Mul        t1:f32[64,64], t2:f32[64,64]  t4:f32[64,64]
1                                              NumericFloat::MulScalar  t3:f32[64,64]                 t5:f32[64,64]
2                                              NumericFloat::Add        t4:f32[64,64], t5:f32[64,64]  t6:f32[64,64]
3                                              Float::Exp               t6:f32[64,64]                 t7:f32[64,64]
    1 launch(es)
       1 burn_cubecl_fusion::optim::elemwise::optimization::elemwise_fuse::ElemwiseFuse
...
--- iteration 1: (64x64) (a * b + c * 2).exp()
[plan] cache hit: execute plan #2 (sync, segment has 4 ops)
...
--- other scalar: (64x64) (a * b + c * 3).exp()
[plan] cache hit: execute plan #2 (sync, segment has 4 ops)
...
--- other shape: (32x32) (a * b + c * 2).exp()
[plan] cache hit: execute plan #2 (sync, segment has 4 ops)
...
    1 launch(es)
       1 burn_cubecl_fusion::optim::elemwise::optimization::elemwise_fuse::ElemwiseFuse

第一次：Mul 和 MulScalar 不共用张量，各开一块；Add 同时读两块的输出，两块被合并后收下它（日志里是「accepted in 1/1」，这时已经只剩一块）；Exp 也进了这一块。四个 op 都进来时融合器仍然开着，Lazy 下一直在等；读数据触发 Sync，这时才定下方案，存成计划 #2（#0、#1 是前面的 Init 和 Random），触发条件 OnSync。之后三次直接 cache hit，没有任何 explore 行，包括标量和形状都换了的两次。

动脑筋答案

新的探索。日志（debug 构建，表格分隔线省略）：

--- non-square: (64x32) (a * b + c * 2).exp()
[explorer] explore (lazy): 1 deferred of 1 queued
[stream] op 0 NumericFloat::Mul → new block (total: 1)
[explorer] explore (lazy): 1 deferred of 2 queued
[stream] op 1 NumericFloat::MulScalar → new block (total: 2)
[explorer] explore (lazy): 1 deferred of 3 queued
[stream] op 2 NumericFloat::Add → accepted in 1/1 block(s)
[explorer] explore (lazy): 1 deferred of 4 queued
[stream] op 3 Float::Exp → accepted in 1/1 block(s)
[explorer] explore (sync): 0 deferred of 4 queued
selected single strategy
[plan] exploration completed: sync, 4/4 ops optimized
...
0    ▸ fused ElementWise (score=630, 4 ops)    NumericFloat::Mul        t36:f32[64,32], t37:f32[64,32]  t39:f32[64,32]
  • 64×64 和 32×32 的相对形状都是 [1, 1]：两个维度值相等，共用一个编号。64×32 的两个维度值不同，相对形状是 [1, 2]，第一条 Mul 的相对形式就变了，Policy 按它的哈希找不到任何候选计划，只能重新探索。
  • 不光表达式本身，紧接着前面建 a、b、c 的 Random 也没命中原来的计划 #1，日志里多了一次 exploration completed: lazy, 1/1 ops optimized，原因相同。
  • 新方案和旧的一模一样（同样 4 个 op 融成一个，score=630），只是缓存的键不同。形状每次都变的程序会一直探索下去，这正是 max_explorations 要管的情况。

小结

  • 每条流一个 Processor：Policy 查执行计划缓存，Explorer 找新方案；Lazy（登记时）可以再等，Sync（drain 时）必须定下来。
  • 队列同时存原始 IR 和相对形式：张量 id 按出现顺序编号，维度值换成编号（1 固定为 0 号），标量换成占位符；缓存只看相对形式，所以换标量、换成同结构的形状都能命中。
  • StreamOptimizer 按共享张量把 op 分进至多 max_blocks（默认 5）个块，每块里五个融合器各自决定收不收、报告 ready 和分数；块的方案取分数最高的那个，分数 0 就逐个执行。
  • 逐元素融合的分数是省下的读写 × 100 加省下的启动 × 10；单个 op 分数为 0，不融合。只要融合器还开着，不相干的块也会合并成一个内核。
  • 方案存成 ExecutionPlan（相对 op 序列 + 触发条件 + 策略），按第一条 op 的哈希索引，只在进程内存里；BURN_FUSION_LOG=full 能看到每一步的决定。

💬 关联讨论

发起讨论
加载中...