算子融合：融合规则

可视化演示

展示融合规则三维度：模式分类（垂直融合生产者消费者串行合并消除中间访存/水平融合共享输入独立算子合并减少kernel启动）、合法性判断（单消费者安全融合消除访存/多消费者部分融合未融合方仍需VRAM读取不改变计算顺序不破坏依赖）、代价模型DP搜索（P0未融合100%开销→P3最优融合25%开销动态规划枚举所有融合分组选最小代价），三模式切换理解编译器融合决策机制。

融合动机：把“端盘子”省下来

你后厨里有一锅汤已经炖好，下一步要调味、勾芡。按老规矩，每次工序切换都得把半成品端回冰箱，下道菜再端出来。结果你发现：一天下来，腿比锅铲还累。AI 芯片也一样，一个个算子各干各的活，每完成一个，中间结果就要写回显存；下一个算子再把它读出来。这趟“端盘子”的开销，就是访存开销。

在 AI 编译器的眼里，GPU 上执行一个算子就是启动一次内核，也就是“生一次火”。生火要时间，热锅要时间，把食材从冰箱搬进搬出也要时间。当算子本身算得飞快、而数据搬运占了大头时，这类算子称为访存密集算子；只有像矩阵乘法那样长时间“炖煮”的算子，才是计算密集算子。现实中的模型往往由大量访存密集的小算子组成，比如激活函数、张量加法、归一化，它们的计算量不大，却把时间都耗在了端盘子上。

算子融合的思路很简单，却极为有效：在编译阶段把多个算子揉成一个内核。中间结果不再写回显存，而是直接留在寄存器或缓存里，递给下一个算子。假设一个逐元素算子链有 N 个环节、每个环节处理 S 个元素，融合前后的访存开销差距一目了然：

这就是融合的第一性原理：把端盘子的力气省下来，让厨师专心做菜。明白了为什么省、省在哪里，你才能真正理解后面所有规则，都是为了守住这份“省”。

融合规则：什么能一锅出

“什么都往一锅炖”是新手才犯的错。后厨里的两样菜要同锅，得先问三个问题：下料顺序能不能颠倒？这道半成品是不是只有一个人要用？火候时间合不合拍？算子融合的判断逻辑，和这几乎一模一样。

先看数据依赖。算子构成的执行图里，每条边代表张量从生产者流向消费者。如果算子 A 算出张量 T，算子 B 要读 T 才能继续，那么 A 必须严格排在 B 前面，这叫先写后读依赖，编译术语叫 RAW（Read After Write）。融合必须保持这条边原样存在，绝不能因为合并把先后顺序打乱，否则就像红烧肉还没炖烂就出锅，后面的菜全毁了。

再看“这道菜是不是只有一个人要用”。如果张量 T 被两个下游算子同时消费，比如一个分支去做分类，另一个分支去做重构，那就绝不能把 B 与 A 融合后“顺手”把 T 减少成一次计算——因为融合成单内核后，T 可能不再落到显存里，另一个消费者就拿不到它了。强行融合只能靠重复计算来弥补，往往得不偿失。把这两个条件写成一个判定式：

回到厨房：高汤罐子炖一次，好几个菜都要舀一勺，那就不能为了省一道工序把汤并进某一个菜里。融合规则的底线，是任何消费者都还能拿到属于它的那份数据。

资源约束：灶台不够大

你终于找到了几道可以同锅的菜，一股脑全倒进锅里，结果发现锅装不下，灶台也摆不下碗。硬塞的下场是：本来想省下端盘子的功夫，反而因为手忙脚乱打翻了两个碗。硬件对融合的限制，就是这口锅的尺寸。

GPU 里最宝贵的“工作台”是寄存器。融合后，中间张量不再返回显存，而是以临时值的形式保存在寄存器里，等待下一个算子取用。可每个线程的寄存器数量是有限的，比如一块 GPU 给每个线程最多 255 个寄存器。融合链越长，中间值越多，需要的寄存器就越多。一旦超出上限，编译器只能把部分临时值“ spill 溢出”到本地内存——相当于把碗从灶台挪回冰箱。这一挪，恰恰又产生了访存，之前的融合收益灰飞烟灭。

除了寄存器，还有共享内存这道“操作台”。归约、矩阵分块等操作都要在共享内存里暂存数据，容量一般只有几十到一百多 KB。因此，融合决策本质上是一道预算题：

好菜谱不是把所有菜都塞进一口锅，而是在工作台摆得下的前提下让流程最短。融合规则里最常被忽略的，正是这个“摆不摆得下”的问题。

融合决策：算细账还是凭经验

面对一长串候选融合方案，你怎么定夺？老厨师凭手感：看到“切丝-焯水-爆炒”三连，闭着眼就知道合在一起划算，这是启发式。精算师父拿账本：把每道工序的耗时、搬运成本列成表，选总成本最低的方案，这是代价模型。还有一位痴迷的学徒，把所有配方排列组合挨个试一遍，这是自动搜索。

真正的 AI 编译器里，三种办法同时存在。最基础的是基于规则的启发式，比如“相邻逐元素算子必融”“归约后接逐元素算子必融”，规则简单、编译快，但容易看走眼。更进一步的是代价模型：编译器为每种融合候选估算访存次数、内核启动次数、计算量，再乘上权重系数求和，得到预估代价 J。最后，像 TVM 的 Ansor 这样的自动调度器，会在搜索空间里采样若干融合方案，直接在真实硬件上跑一遍，用实测数据反馈给搜索算法。

代价模型的核心是一个加权求和：

定菜谱不能只看“看起来顺眼”，要有成本账撑腰。融合决策的进阶之路，就是从“凭经验”到“算细账”，再到“让机器自己试菜”。

融合收益：省了几趟路

花这么大力气去融合，到底图什么？答案就是图“少跑几趟”。你以前做一道三菜一汤，要往返冰箱七八次；现在配菜师把所有食材按顺序摆上灶台，一趟流水线走完，你只需要站在锅前翻勺。

从数据流的角度看，融合收益最直观地体现在“算术强度”上。算术强度（Arithmetic Intensity）定义为单位字节访存对应的浮点运算次数。它越高，说明每次从显存取来的数据都被充分“压榨”了。对 N 个逐元素算子链，每个算子做 F 次浮点运算，输出张量大小为 S；融合前每算子各读各的输入、写各的输出，访存总量约为 2NS；融合后整体只读一次输入、写一次输出，访存降为 2S。于是融合后的算术强度变成：

当然，算术强度提升不等于总时间也省 N 倍，还要看硬件是否“喂得饱”计算单元，但方向很明确：少搬数据，就是省钱。这份收益账，是验证一切融合方案好坏的试金石。

融合模式：经典组合拳

后厨里流传着几套固定搭配：青菜要先焯水再爆炒，牛肉要切好腌制再下锅，汤要大火烧开再转小火。这些搭配是几代厨师验证过的“模式”（Pattern）。编译器里也有这样一组经典的融合模式，见到就能直接套用。

第一种是逐元素链融合。加法、乘法、激活函数、缩放这类算子，输入输出形状一一对应，每个输出元素只依赖输入的同位置元素。它们像“切丝-焯水-爆炒”一样环环相扣，中间结果天然适合留在寄存器里，等号右边的一串复合函数可以整体压成单个内核：

第二种是归约加逐元素融合。比如 Softmax 要“先求总和再逐元素除以总和”，归一化则是“先算均值方差再逐元素缩放”。归约结果很小，刚好能塞进共享内存，紧随其后的逐元素操作就可以直接使用这个局部量，不必再访存一次。第三种是矩阵乘法加逐元素融合，典型的如卷积后接批归一化、全连接后接激活；QKV 投影三个矩阵乘法拼接成一个大矩阵乘法，也属于这类“组合拳”的延伸。

把经典搭配做成模板，编译器遇到相似结构时就能秒认秒融。这也是 TVM、XLA、oneDNN 等框架里融合实现的核心套路。

工程实现：流程化改造

经验再丰富，也要沉淀成标准流程，厨房才能规模化。融合规则的工程实现，就是把“什么能融、怎么融”固化成编译器里的一个个 Pass（编译优化阶段），像中央厨房的标准化作业手册。

首先，每一道算子都要在“算子注册表”里登记自己的属性：是逐元素型还是归约型？输出形状和输入形状是否一一对应？能否原地写回？这张注册表相当于食材档案——哪些食材能生吃、哪些必须煮熟，系统心里有数。然后，编译器的图遍历器从计算图入口出发，按拓扑序扫描节点，寻找与注册表里融合模式匹配的相邻算子，比如发现“逐元素加 + 逐元素乘”这对组合，就标记为候选。最后，改写器把这对节点替换成一个新的融合算子节点，并复用已有的输出缓冲区，避免额外开辟显存。

这个流程可以表示为一次图重写：

流程化之后，融合不再是某个高手的灵光一现，而是每次编译都会自动执行的固定步骤。老厨师的经验由此变成了后厨的标准作业手册。

实践陷阱：乱炖的代价

融合不是万灵药，乱炖一锅的代价，往往比分开炒还大。后厨新手最容易踩的坑，有三个。

第一个坑是数值精度变化。浮点数在计算机里不是精确的，加法不满足结合律：同一组数，先加后两个还是先加前两个，结果可能有细微差别。融合改变了中间结果的舍入时机——原来中间值会写回显存“四舍五入”一次，融合后它一直在高精度寄存器里，看似更精确，但最终结果和未融合版本并不完全一致。对精度敏感的模型，这一步差异可能被放大：

第二个坑是调试困难。融合后中间张量不再落盘，你想看看“这一步算出来到底是什么”都没机会，只能靠数学校验。第三个坑是性能回退：前面说过，寄存器溢出、并行度下降都可能让融合后反而更慢。成熟编译器通常提供融合开关、保留中间张量的调试模式，以及“融合后实测变慢就回退”的保险机制。

好菜谱是反复试吃调整出来的。融合规则也一样，必须配合数值校验和性能基准测试，在实践中不断修正，才能从“纸上谈兵”变成“后厨真经”。

AI 应用场景

Python 演示

# 一个简单的算子融合规划器: 判断 DAG 中哪些相邻算子能融合,# 并估算融合前后的显存访存字节数.from dataclasses import dataclass, field @dataclassclass OpNode:    name: str    op_type: str          # elementwise / reduce / matmul    size: int             # 输出张量元素个数    inputs: list = field(default_factory=list)    outputs: list = field(default_factory=list) class FusionPlanner:    def __init__(self, nodes):        self.nodes = nodes        self.fused_groups = []     def can_fuse(self, a, b):        # 逐元素链必然可融; 归约后接逐元素也可融        if a.op_type == "elementwise" and b.op_type == "elementwise":            return True        if a.op_type == "reduce" and b.op_type == "elementwise":            return True        return False     def plan(self):        visited = set()        for node in self.nodes:            if node.name in visited:                continue            group = [node]            cur = node            while cur.outputs:                nxt = cur.outputs[0]                if len(nxt.inputs) == 1 and self.can_fuse(cur, nxt):                    group.append(nxt)                    visited.add(nxt.name)                    cur = nxt                else:                    break            self.fused_groups.append(group)        return self.fused_groups def memory_bytes(size):    return size * 4  # 每个元素按 4 字节 float32 计 def report(planner):    total_before = 0    total_after = 0    for group in planner.fused_groups:        for node in group:            total_before += memory_bytes(node.size) * 2  # 写一次 + 读一次        if len(group) > 1:            total_after += memory_bytes(group[0].size) + memory_bytes(group[-1].size)        else:            total_after += memory_bytes(group[0].size) * 2    print("Before fusion:", total_before, "bytes")    print("After fusion:", total_after, "bytes")    print("Saved:", total_before - total_after, "bytes") if __name__ == "__main__":    # 构造一个小图: add -> relu -> sum -> normalize    a = OpNode("add", "elementwise", 1024)    b = OpNode("relu", "elementwise", 1024)    c = OpNode("sum", "reduce", 256)    d = OpNode("norm", "elementwise", 256)    a.outputs.append(b)    b.inputs.append(a)    b.outputs.append(c)    c.inputs.append(b)    c.outputs.append(d)    d.inputs.append(c)    planner = FusionPlanner([a, b, c, d])    groups = planner.plan()    print("Fused groups:", [[n.name for n in g] for g in groups])    report(planner)

1. @dataclass class OpNode定义算子节点，记录名字、类型和张量大小，inputs/outputs 表示图里的数据依赖边。
2. def can_fuse(self, a, b)这里实现了文中的融合规则：逐元素接逐元素、归约接逐元素都合法。
3. while cur.outputs从当前算子沿输出边向下走，尝试把整条链尽可能长地融在一起。
4. if len(nxt.inputs) == 1这是“单一消费者”约束的变体：下一个算子只有一个输入时，融合后的中间值不会破坏其他分支。
5. memory_bytes(node.size) * 2每个中间张量按“写一次 + 读一次”估算访存，对应文中 M_before = 2NS。
6. group[0].size + group[-1].size融合后整组只读首、写尾，对应 M_after = 2S。

常见误区

面试题