算子融合 编译器：Conv-BN-ReLU 合一、纵向/横向融合与融合边界
更新时间：2026-09-25。神经推理算子拆解 把一张网络拆成了 conv、matmul、batch norm、激活函数等几十上百个算子，本篇回答紧随其后的工程问题：这些算子逐个执行是不是最优？答案通常是否定的。Conv-BN-ReLU 在模型定义里是三层，在 GPU 上 naive 实现却是三次 kernel launch、三段显存写了又读的往返；而把它们融合成一个 kernel，中间结果可以根本不离开寄存器。这是推理性能优化中投入产出比最高的一类改写，也是 TVM、XLA、MLIR 等深度学习编译器的核心能力。配合阅读的还有 推理 Runtime 与调度 和 cuBLAS/cuDNN 高性能库：前者讲算子在运行时如何被编排，后者讲库里那些手工优化 kernel 为什么仍然难被自动生成超越。
本文要回答的问题
- Conv、BN、ReLU 分开跑三个 kernel 到底慢在哪？中间张量写回显存的代价怎样量化？
- 什么是纵向融合、什么是横向融合？BN 的参数为什么可以直接折叠进 Conv？
- TVM、XLA、MLIR 分别用什么机制自动发现并执行融合？polyhedral tiling 在这里扮演什么角色？
- 融合的边界在哪：为什么 reduction 算子两侧往往不能融、shape 不一致为什么让融合失败？
- 融合后寄存器和共享内存压力变大怎么办？典型带宽数字和真实提速大概是多少？
一、三个 kernel 的账本：算力闲着，显存在跑
先把 Conv-BN-ReLU 分开执行的完整路径走一遍。Conv 是计算密集的大 kernel，输出特征图 t₁ 写到显存；batch normalization（以下简称 BN）是逐通道仿射变换，从显存读出 t₁、算出 t₂ 再写回；ReLU 逐元素取 max(x, 0)，又把 t₂ 读出、把最终结果 y 写回。三个算子之间没有任何数据依赖上的障碍阻止合并，但在「一个算子对应一个 kernel」的朴素实现里，它们各自是独立的调度单位。
算两笔账。第一笔是访存账。BN 和 ReLU 每个元素只做一两次简单运算，却要把整个特征图读一次、写一次，算术强度（每搬运一字节换来的 FLOP）不足 0.5，稳稳落在 roofline 的带宽斜线区——乘法器再快也无济于事。以一张 256 通道、56×56 的 FP32 特征图为例，体积约 3.2 MB；在约 1.5 TB/s 带宽的当代数据中心 GPU 上，仅「写回再读出」一次往返就要约 4 微秒，而这期间核上的张量计算单元基本闲置。第二笔是启动账：一次 kernel launch 通常有几微秒到十几微秒的固定开销（驱动入队、grid 配置、warp 预热），三个 kernel 就是三次串行付出。对大 batch 下的卷积主体，这些零头可以忽略；但在 batch=1、小通道数、分割模型的后处理链以及手机 NPU 上，launch 开销和逐元素访存可以吃掉网络总时延的三分之一以上。
问题的要害是：t₁、t₂ 这两个中间张量除了喂给紧接着的下一个算子，别无他用。它们被写进显存、又立刻被原样读出，属于纯粹的无效搬运。融合要消灭的正是这两次往返。
还要注意中间张量带来的不只是时延，还有容量与分配开销。逐算子执行时，运行时必须为 t₁、t₂ 预留与输入特征图同量级的显存空间；浅层网络的特征图往往是全网络最大的（通道虽少但空间尺寸大），一段未融合的后处理链在 batch 稍大时可能额外占用几十 MB，移动端和边缘 NPU 的显存预算因此先于算力被击穿。显存分配本身也不是免费的：朴素实现每层临时 malloc/free，runtime 内部的分配器锁和块管理会再叠加微秒级延迟，生产系统普遍靠内存池缓解，而融合直接让这些临时分配从计划里消失——三个算子只剩输入、输出两块张量。换句话说，融合同时改善了时延、显存峰值和分配行为三件事，只是后两件不像带宽数字那样直观。
二、融合后：一个 kernel，中间结果留寄存器
融合版 Conv-BN-ReLU 只有一个 kernel：每个线程块算出自己负责的输出元素后，Conv 的部分和在寄存器里完成累加，紧接着在同一线程内做 BN 的缩放偏移和 ReLU 截断，最终只把 y 写一次显存。中间值 t₁、t₂ 的生命周期收缩到几条寄存器指令之间，物理上从不离开计算核心。
这张图也点明了融合收益的来源结构：省的不是计算，是搬运和调度。Conv 的乘法次数一个没变，但网络中大量逐元素算子的访存量被整段砍掉，launch 次数随链长线性减少。一个经验性的量级：纯逐元素链（add-mul-sigmoid 之类）融合后，这些算子自身的耗时常常下降到原来的 1/3 到 1/10，因为它们本来就是 memory-bound，融合把算术强度从「每元素读写 8 字节做 1 次运算」抬到「每元素只读输入、只写输出」；包含卷积主体的整网推理，端到端提速常见 1.2~1.5 倍，后处理占比高的分割、生成模型可以更高。顺带一提，BN 还可以更彻底地在编译期消掉：推理时 BN 的均值方差都是常量，仿射变换 y = γ(x−μ)/σ + β 可以合并进 Conv 的权重和偏置（W′ = W·γ/σ，b′ = γ(b−μ)/σ + β），连运行时的逐元素计算都不剩下，这叫 BN folding，属于常量折叠而非 kernel 融合，但两者通常由编译器的同一条 pass 流水线完成。
三、纵向融合与横向融合：融合类型全景
算子融合按拓扑关系分为两大类。纵向融合（vertical fusion） 融合生产者与消费者构成的链：上游算子的输出直接被下游消费，中间张量不必落盘，Conv-BN-ReLU、add-mul-gelu 都是典型。横向融合（horizontal fusion） 融合共享同一输入、彼此无依赖的多个分支：同一张量需要同时喂给 sigmoid 和 tanh（门控单元常见），或多个同形输出后处理，分开跑意味着输入被读多遍，融合成一个 kernel 后输入只读一次、多个输出并在同一次遍历中算出。横向融合在 NLP 模型中价值突出，例如 LSTM/GRU 的多个门、注意力中并行的多个投影后处理。
| 融合类型 | 典型例子 | 主要收益 | 限制条件 |
|---|---|---|---|
| 纵向：逐元素链 | Conv-BN-ReLU、Add-Mul-GELU | 消灭中间张量往返，省 launch；链上提速 3~10 倍 | 下游需逐元素（injective），不能是大 reduction |
| 纵向：BN folding | BN 常量参数折叠进 Conv 权重 | 运行时整个 BN 消失 | 仅推理或 eval 模式；训练中均值方差在变化 |
| 纵向：reduce + 逐元素 | Softmax 内部 exp/sum/div、ReduceSum+cast | 省 reduce 输出的一次往返 | 融合点必须在 reduction 边界内侧，不能跨两次 reduce |
| 横向：同输入多输出 | 多个门控激活并行、多分支后处理 | 共享输入只读一遍，合并 launch | 各分支输出 shape 需兼容同一调度；寄存器压力上升 |
| 横向：GEMM epilogue 融合 | matmul 后接 bias+激活+残差相加（cuDNN/cuBLAS fused epilogue） | 累加器出寄存器前完成后处理 | epilogue 逻辑须能映射到固定模板 |
| 跨层流水融合 | 层间首尾相接，片上 SRAM 直传（NPU 常见） | 整层输出不回 HBM | 各层吞吐需均衡，缓冲深，shape 要对齐 tile |
判断一个融合值不值得，除了看省掉多少字节，还要看融合后 kernel 的资源占用：融合拉长了每个线程的活跃区间、增加了同时存活的中间值，寄存器用量上升可能迫使编译器降低 occupancy（每个 SM 上可驻留的 warp 数），反而损害 Conv 主体的延迟掩盖。因此编译器在 fuse 之后通常还要做寄存器压力检查，压力过大时放弃部分融合或重排计算顺序。
四、编译器怎样自动做：TVM、XLA、MLIR 三条路线
手工融合每个网络显然不可持续，编译器的任务是从计算图中自动识别可融合的算子簇并生成合并后的 kernel。三家代表性系统的做法各有侧重。
TVM 的 Relay/Relax 层先把算子按融合性质分类：injective（逐元素、一对一映射，如 ReLU、scale）、reduction（如 sum、softmax 的归约阶段）、complex/out-dynamic（如 conv，视为不透明锚点）。融合规则简化说就是：injective 链可以一路吞并，reduction 可以把尾随的 injective 拉进来（epilogue 融合），但两个 reduction 之间不跨、opaque 算子作为簇的边界。划分完的每个簇 lowering 到 TensorIR，再由调度原语（tile、reorder、vectorize、tensor core 绑定）生成代码。XLA（TensorFlow/JAX 的编译器）走的是 HLO 图上的 cluster fusion：以 producer-consumer 边构造融合候选，用代价模型评估「融合后省的访存 vs 多占的寄存器/可能的重复计算」，并内置对 dot、convolution 周边 epilogue 的专门融合规则；JAX 用户感受到的「加一句 jit 就变快」，大部分收益正来自这里。MLIR 不直接规定融合策略，而是提供分层 IR：linalg dialect 把算子表示成带索引语义的命名/泛型 op，融合表现为 linalg tiling + fusion on tensors（把生产者的 tile 嵌进消费者的循环嵌套），再逐步 lowering 到 affine、scf、gpu dialect；多数学术界和芯片厂商的自研编译器（包括多家 NPU 工具链）都在这条路径上定制自己的融合 pass。
三家路线的差异背后是两种融合哲学的取舍。TVM 的分类规则是白盒的：什么算子和什么算子能融，规则写在框架里，性能可预期、出问题容易定位，但新算子要想被融合必须接入分类体系，否则被当成 opaque 永远孤立。XLA 更倚赖代价模型的自动裁决：同一个候选，shape 不同、带宽参数不同，决策可以不同，覆盖面广但偶尔做出让用户意外的选择（这也是 XLA 提供 fusion 调试 dump 的原因——大模型调优时常要逐簇检查它到底融了什么）。MLIR 则把策略让渡给上层：同样的 linalg IR，面向 GPU 的 pipeline 偏向向量化与 shared memory 搬运，面向 NPU 的 pipeline 偏向固定大小的 SRAM tile，融合判定与目标硬件的内存层级绑在一起。理解这套光谱，读各家厂商「我们的编译器支持算子融合」的宣传时才会追问正确的问题：融合规则覆盖到哪些算子、代价模型看哪些量、生成的代码落在哪一级存储。
图中的 polyhedral tiling 值得单独解释。传统融合只回答「哪几个算子合并」，合并后的嵌套循环仍可能很糟：两层算子各写一个全尺寸临时张量，即使在同一 kernel 里，中间数组仍可能被编译器放进全局内存。多面体模型把循环嵌套建模成整数空间中的多面体，用统一的依赖分析做循环交换、分块和嵌套融合——把 Conv 输出的一个 tile 算出来，立刻在同一内层循环位置上喂给 BN 和 ReLU 的对应 tile，中间数组退化为标量或寄存器向量。代价是编译时间显著增加，且对动态 shape、复杂控制流的支持需要保守处理，这也是实际编译器普遍采用「模式匹配处理高频结构 + polyhedral 处理规则算子」混合路线的原因。
五、融合的边界：什么融不了，以及为什么
融合不是免费的万能药，编译器在以下几类边界前必须止步或保守，理解这些边界比记住融合规则更重要。
第一，不能跨 reduction 融合两个不兼容的归约。逐元素算子之间的依赖是点对点的，融合只是把两个映射串起来；而 reduction（sum、mean、全局 norm）要求先看到整个归约域才能输出一个元素，它两侧的逐元素操作可以融进它的 prologue/epilogue，但两个不同归约域的 reduction 无法点对点对齐——Mean 之后的算子依赖全通道统计量，不能在部分和阶段执行。Softmax 是个精巧的反例：它形式上含一次全局 reduce，但编译器可以把整个 exp-max-sum-div 融成一个 kernel（甚至在线算法单趟完成），融的是 reduction 边界内侧的全部逐元素计算，而不是跨过去。
第二，shape 与 broadcast 约束。纵向融合要求逐元素位置可对齐；若上游输出与下游输入之间存在 reshape、transpose、非平凡 broadcast，点对点融合会引入 gathers/scatters，省掉的带宽被寻址开销吃掉。此时要么先做 layout 归一化（把 transpose 吸收进相邻算子的索引计算），要么放弃融合。动态 shape 更是老一代编译器直接放弃融合的常见原因，当代方案靠 shape constraint 推断和多版本特化代码部分缓解。
第三，重复计算风险。横向融合共享输入是纯赚，但当一个生产者有多个消费者、且消费者只各用其一小部分时，激进的纵向融合可能让同一段上游计算在多个 kernel 里重复执行。XLA 代价模型的核心检查之一就是「融合节省的访存是否超过重复计算的 FLOP」，不划算则允许公共上游保留独立 kernel、由下游共享其输出。
第四，硬件资源上限。融合 kernel 变长后寄存器/片上缓冲压力增大，可能触发 register spill（寄存器溢出到本地内存，反而更慢）或 occupancy 下降；NPU 上还要受 tile 尺寸整除约束、scratchpad 容量限制。工程上常见的反直觉现象是：关掉某几个融合后整网反而变快——这通常就是资源压力超界的信号，应当回看每个融合 kernel 的寄存器用量和 SM 占用，而不是把融合数当成越多越好的指标。
六、融合与库、Runtime 的分工：自动生成的边界在哪
现实推理栈不是编译器一家独大。cuBLAS/cuDNN 这类手写库在 GEMM 和卷积本体上仍有优势：tensor core 的 WMMA 指令排布、swizzle 访存、split-K、数十种 shape 特化模板是多年手工打磨的结果，编译器对这些 opaque 锚点通常直接调用库实现，只融合其 epilogue（cuDNN 的 convolution fused with bias/activation 本质就是预置的纵向融合模板）。而 推理 Runtime 负责在图执行层面编排：哪些子图交给编译器融合编译、哪些直接调库、多流并行下融合簇之间怎样并发，以及动态 shape 触发时的重编译缓存。三者分工可以概括为：库守住计算主体的性能下限，编译器在逐元素和后处理链上榨取系统性收益，runtime 决定它们何时上场。
给部署者的实操判据收个尾。第一步，用 profiler 按 kernel 看时间占比：若大量小于十几微秒的逐元素 kernel 首尾相接、显存读写流量明显大于模型张量体积，就是融合缺失的典型信号。第二步，确认融合真的发生——读编译后代码或 kernel 名，而不是假设开了 jit 就万事大吉，个别编译器在动态 shape、自定义算子前会静默回退。第三步，对融合后的大 kernel 核对寄存器用量和 occupancy，必要时调整 tile 或拆分。记住融合改变的是 roofline 上的工作点：它把 memory-bound 的逐元素操作从带宽斜线上抬走，却动不了计算平顶上的 Conv/GEMM 本体——收益大小，永远由被省掉的那部分访存和 launch 在总账单里的占比决定。
这一篇在知识体系里的位置
- 上游基础：神经推理算子拆解——先有算子清单，才谈得上哪些算子可以合并；
- 运行时配套：推理 Runtime 与调度——融合簇在执行期如何被编排、缓存与并发；
- 手写库对照：cuBLAS/cuDNN 高性能库——编译器选择调用而非生成的那些 kernel；
- 硬件侧前提：AI 加速器脉动阵列 讲清了融合所省的访存在内存层级中为何昂贵，tile 与双缓冲也发生在融合生成的循环嵌套里。
总结
算子融合 编译器优化的本质是「省搬运、省调度而非省计算」：Conv-BN-ReLU 分开跑是三个 kernel、两次中间张量写回显存再读出，BN、ReLU 等逐元素算子因此被死死按在 roofline 的带宽斜线区，还要付多轮 launch 开销；融合成单个 kernel 后，中间结果在寄存器内从 conv 部分和一路流过 scale、bias 与截断，BN 参数更可在编译期折叠进权重；纵向融合吞并生产者-消费者链、横向融合合并同输入的多输出分支，TVM 靠 injective/reduction/opaque 分簇、XLA 靠代价模型驱动的 cluster fusion、MLIR 靠 linalg tiling-fusion 加 polyhedral 调度自动完成这一切；但融合不能跨不兼容的 reduction、怕 shape 不匹配和动态 shape、要防止重复计算与寄存器溢出，并把 GEMM/卷积本体留给 cuBLAS、cuDNN 这类手写库——典型收益是逐元素链自身提速 3~10 倍、整网推理 1.2~1.5 倍，而收益上限永远等于被消灭的那部分显存流量与启动开销在总时延账单里的占比。