4、TVM算子融合Pass：FuseOps、规则与图结构分析
算子融合，说白了就是把多个小算子合并成一个大算子。为什么要这么做？因为每次调用一个算子，都有启动开销、内存读写开销。你想想看，如果能把几个连续的操作合并成一个，中间结果就不用写回内存了，直接在寄存器或者共享内存里传递，速度能快不少。
TVM里的FuseOps Pass，就是干这个活的。我刚开始接触TVM时，觉得这玩意儿挺玄乎的，后来自己动手改过几次融合规则，才真正摸透了它的脾气。
4.1 FuseOps Pass 干了什么
FuseOps Pass 是TVM Relay前端的一个图优化Pass。它的目标很明确：把可以融合的算子节点合并成一个“融合组”，然后交给后端代码生成器，生成一个大的kernel函数。
我个人习惯把融合分成两类：
- 横向融合：相同类型的算子，比如两个连续的Conv2D，可以合并成一个更大的Conv2D（如果支持的话）。
- 纵向融合：不同类型的算子，比如Conv2D + ReLU + BiasAdd，这是最常见的融合模式。
TVM的FuseOps Pass主要做的是纵向融合。它沿着计算图从输出往输入遍历，判断每个算子能不能和它的消费者融合在一起。
核心思想：如果一个算子的输出只被一个消费者使用，并且这个消费者支持融合，那就可以把它们合并。
4.2 融合规则的定义
融合规则不是随便定的。TVM里有一套算子分类体系，每个算子被标记为不同的“融合类型”。我记得主要有这么几类：
| 融合类型 | 说明 | 常见算子 |
|---|---|---|
| kElemWise | 逐元素操作，可以任意融合 | ReLU, Add, Mul, Clip |
| kBroadcast | 广播操作，可以融合 | BroadcastTo, ExpandDims |
| kInjective | 单射操作，输入输出一一对应 | Reshape, Transpose, Pad |
| kCommReduce | 规约操作，融合有限制 | Sum, Max, Min |
| kOutEWiseFusable | 可以和逐元素操作融合 | Conv2D, Dense |
规则其实很简单：一个算子能不能和它的父节点融合，取决于父节点的类型。比如，Conv2D（kOutEWiseFusable）后面跟一个ReLU（kElemWise），那肯定能融合。但如果Conv2D后面跟一个Sum（kCommReduce），那就要看情况了。
我曾经踩过一个坑：把Reshape和Conv2D强行融合，结果生成的代码性能反而下降了。为什么？因为Reshape改变了数据的布局，融合后Conv2D的访存模式变得不连续了。嗯，这里要注意：融合不是万能的，有时候不融合反而更快。
4.3 融合后的图结构分析
融合之后，计算图会变成什么样？我画了一张图，帮你直观理解：
融合后的图，结构变得非常简洁。原来需要4次kernel launch，现在只需要1次。中间结果（Conv2D的输出、BiasAdd的输出）都不需要写回全局内存了。
但这里有个细节：融合后的节点在TVM内部其实是一个“复合节点”，它包含了一个子图。代码生成时，TVM会把这个子图展开，生成一个大的循环嵌套。
小技巧：如果你想查看融合后的图结构，可以用 relay.analysis.list_fuse_ops(mod) 来打印融合组的信息。我调试时经常用这个命令。
4.4 融合的代价与收益
融合不是免费的午餐。我总结了一个简单的收益模型：
- 收益：减少kernel launch次数、减少全局内存读写、提高计算密度
- 代价：单个kernel变复杂、寄存器压力增大、可能降低并行度
举个例子，如果你把一个大矩阵乘法和一个逐元素操作融合，矩阵乘法的计算量远大于逐元素操作，那融合的收益就很明显。但如果你把两个都很小的算子融合，收益可能微乎其微。
注意：我曾经遇到过一种情况——把多个Reshape和Transpose融合后，生成的代码因为寄存器溢出，性能反而下降了20%。所以，融合前最好先做一次性能预估，或者用TVM的AutoTVM自动搜索最优配置。
4.5 实战：如何自定义融合规则
如果你想在TVM里添加自己的融合规则，可以这样做：
from tvm import relay
from tvm.relay.op import op as _op
# 注册一个新的融合规则
@_op.register_fuse("nn.conv2d")
def fuse_conv2d(expr):
# 检查输入是否支持融合
if isinstance(expr, relay.Call):
# 获取输入节点
input_node = expr.args[0]
# 判断是否可以融合
if input_node.op.name == "nn.relu":
return True
return False
当然，这只是一个简化示例。实际项目中，你还需要考虑数据类型、形状、布局等因素。我建议你从简单的规则开始，逐步增加复杂度。
好了，关于TVM的FuseOps Pass，核心内容就这些。记住一句话：融合是手段，不是目的。最终目标永远是让模型跑得更快。