TVM计算图优化三大技术

可视化演示

动态演示TVM计算图优化三大技术的核心机制，通过输入/处理/存储/输出四种模式展示关键概念与交互流程。

一、计算图概念：一份可以改写的菜谱

想象你是后厨总管，手头有一份“结构化菜谱”：每个步骤画成一个方框，食材从上一站流向下一站，箭头标明顺序。只有菜谱写成这种规范格式，你才能分析哪一步可以提前、哪几步能合并、哪一步纯属多余。深度学习模型也一样——只有先把模型翻译成一张计算图，编译器才能对它动手脚。

计算图是模型运算的可视化结构：节点是算子（卷积、加法、ReLU），边是张量流动（数据从上游传到下游）。TVM 用 Relay 中间表示（IR）承载这张图，图优化 Pass 在 IR 上做分析和改写。训练框架负责“写出菜谱”，TVM 负责“改造菜谱”。

计算图是一张有向无环图（DAG），数学上记作：

$$\displaystyle G = (V, E), \quad v \in V \text{ 表示算子，} e = (u, v) \in E \text{ 表示张量从 } u \text{ 流向 } v$$

1. 节点集合：$G=(V,E)$ 中的 $V$ 是所有算子的集合，如卷积、加法、ReLU，相当于菜谱中一个个具体动作。

$$V = \{\text{conv}, \text{add}, \text{relu}, \dots\}$$

2. 有向边集合：$e=(u,v)$ 表示张量从算子 $u$ 输出后作为输入流入 $v$，相当于食材盆从上一站端到下一站。

$$E = \{(u,v) \mid u \text{ 的输出是 } v \text{ 的输入}\}$$

3. 无环约束：图是 DAG，不存在循环依赖，保证每道菜都能按序完成，不会出现“回锅炖自己”的死循环。

$$G \text{ 中不存在路径 } v_1 \to v_2 \to \cdots \to v_1$$

菜谱结构化之后，后厨总管才能火眼金睛地找到改造点。同样的道理，Relay IR 就是那张“可以改写的菜谱”，接下来的三大技术都在这张图上施展。

二、常量折叠：提前把酱汁调好

宫保鸡丁的酱汁是糖、醋、酱油按固定比例调的。如果每来一桌客人，厨师都现场称量、现场搅拌，高峰期必然手忙脚乱。聪明的后厨会提前把酱汁批量调好、装瓶备用，客人点菜时直接淋上去。常量折叠就是这瓶提前备好的酱汁。

在计算图中，如果某个节点的所有输入都是常量——编译期就能确定数值的节点——那么它的输出也一定可以在编译期算出来。TVM 的做法是：遍历计算图，找到这样的节点后，直接在编译阶段执行它，把整个子图替换成一个常量节点。比如 2 × 3 在编译时直接变成 6，推理时少执行一个乘法算子。

这个策略可以严格写成：

$$\displaystyle v_{\text{fold}} = \begin{cases} \text{eval}(v), & \forall u \in \text{inputs}(v):\ u \in \text{Const} \\ v, & \text{otherwise} \end{cases}$$

1. 输入全为常量：$\forall u \in \text{inputs}(v): u \in \text{Const}$ 表示节点 $v$ 的每个输入 $u$ 都是常量节点，就像酱汁配方里每样原料都是已知的。

$$\forall u \in \text{inputs}(v):\ u \in \text{Const}$$

2. 编译期求值：$\text{eval}(v)$ 表示在编译阶段真的执行一遍这个算子，像后厨提前把酱汁熬好、装瓶。

$$\text{result} = \text{eval}(v)$$

3. 替换或保留：条件满足就替换成常量节点，不满足就保留原样，绝不动任何没把握的步骤。

$$v_{\text{fold}} = \text{eval}(v) \ \text{或} \ v$$

别忘了边界情况：依赖运行时数据的算子（如 dropout、受 batch 影响的 BN）不能折叠，含随机性的算子也不该折叠，而且折叠本身要花编译时间——备菜也得算人力成本。提前备好酱汁，高峰期上菜自然快；常量折叠让推理时的算子更少、启动更快。

三、算子融合：一锅炖的带宽革命

后厨真正的瓶颈往往不是切菜，而是搬运：食材从冰箱取、洗完放回台面、切完再端到灶台，每一步都让食材在冷藏室和工作台之间来回折腾。推理引擎也一样——现代硬件算力增长飞快，真正的瓶颈是内存带宽：中间张量写回全局内存、再读出来，每一趟都是真金白银的时间。

以 conv → bn → relu 为例：卷积的输出明明马上要被 BN 消费，却必须先完整写回内存，BN 再从头读一遍，ReLU 又重复一遍。算子融合把这三个算子合并成单个内核，让中间结果留在寄存器或缓存里，根本不出芯片。一次融合就省掉两次中间张量的全局内存访问：

$$\displaystyle \text{Saving} = 2 \times \sum_{k=1}^{m} |T_k| \ \text{字节}$$

1. 中间张量求和：$\sum_{k=1}^{m} |T_k|$ 是把 $m$ 个中间结果的大小累加，每个被融合的中间张量都不再落地。

$$\sum_{k=1}^{m} |T_k|$$

2. 写读双省：乘以 $2$ 是因为每个中间张量原本要“写回内存一次、读回一次”，共计两次访存。

$$2 \times \sum_{k=1}^{m} |T_k|$$

3. 带宽即瓶颈：省内存访问比省运算更划算，因为内存带宽的增长远远慢于算力增长。

$$\text{Time} \approx \frac{\text{访存字节数}}{\text{内存带宽}}$$

融合不是想融就融，TVM 通过模式匹配找到 conv + bn + relu 这种固定套路，再检查融合合法性：中间节点不能被其他节点旁路消费、数据类型要一致。一锅炖省下的是搬食材的功夫，这也是图优化里收益最直观的一招。

四、死代码消除：划掉没人吃的装饰菜

菜谱上写着“雕一朵萝卜花放在盘边”，可这道装饰菜既没人吃、也不影响任何主菜的味道，纯粹浪费雕花师傅的时间。死代码消除就是那个划掉萝卜花的狠角色。

计算图里也会长出“萝卜花”：某些节点的输出没有任何下游节点消费，或者根本不影响最终输出。它们从哪来？多半是前两步优化的“残留物”——常量折叠替换了旧子图、算子融合合并了旧节点，旧的节点没人引用了，就成了死代码。死代码消除把这些节点从图里永久删掉，让图更干净、启动更快。

具体算法是经典的标记-清除（mark-and-sweep）：从输出节点出发，沿反向边遍历，标记所有“被真正需要”的节点，未标记的就是死代码：

$$\displaystyle \text{Reach} = \{ u \mid u \text{ 能到达某个输出节点} \}, \quad V_{\text{dead}} = V \setminus \text{Reach}$$

1. 反向可达集：$\text{Reach}$ 是从输出节点沿边反向能走到的所有节点，它们是“被需要”的步骤。

$$\text{Reach} = \{ u \mid u \Rightarrow \text{output} \}$$

2. 补集即死节点：$V \setminus \text{Reach}$ 是可达集的补集，即所有不被输出需要的节点。

$$V_{\text{dead}} = V \setminus \text{Reach}$$

3. 安全删除：删掉死节点不会改变任何输出结果，就像划掉装饰步骤不影响菜的味道。

$$\text{output}(G) = \text{output}(G \setminus V_{\text{dead}})$$

死代码消除是“保洁员”：它自己很少制造收益，但能把前两步优化留下的垃圾扫干净。后厨定期清理闲置设备和多余步骤，空间大了、人也不堵了。

五、优化流水线：后厨的标准化改造流程

三种优化不是随随便便乱做，后厨改造必须按流程：先备菜，再并锅，最后打扫。如果先打扫卫生再备菜，融合之后又会产生新垃圾，等于白扫。TVM 里每个优化都是一个 Pass（一遍图变换），Pass 管理器按顺序调度它们。

TVM Relay 中，每个 Pass 接收一张计算图、产出一张新图。三大 Pass 的标准顺序是：先 FoldConstant 做常量折叠，再 FuseOps 做算子融合，最后 EliminateDeadCode 清理死代码。写成数学形式：

$$\displaystyle \text{PassSeq} = \text{EliminateDead} \circ \text{FuseOps} \circ \text{FoldConst}$$

1. 复合顺序：$\circ$ 表示函数复合，右边的 $\text{FoldConst}$ 最先执行，符合“先备菜”的直觉。

$$f \circ g (x) = f(g(x))$$

2. 先折叠：$\text{FoldConst}$ 把常量子图压成数值，让图变小，为后续匹配提供更干净的输入。

$$G_1 = \text{FoldConst}(G_0)$$

3. 再融合：$\text{FuseOps}$ 在折叠后的图上做模式匹配并合并且，产生新的中间结构。

$$G_2 = \text{FuseOps}(G_1)$$

4. 最后清扫：$\text{EliminateDead}$ 删除 $G_2$ 中残留的旧节点，输出最终优化图。

$$G_3 = \text{EliminateDead}(G_2)$$

为什么顺序如此重要？如果先做死代码消除，融合后新产生的废弃节点就没人管了。标准化流程保证每个模型都经历同样的严格改造，这正是 TVM Pass 管理器的价值。

六、性能评估：改造效果怎么算账

改造完的后厨到底行不行？不能靠感觉，要算账：翻台率（推理延迟）、食材损耗（内存占用）、开档准备时间（启动开销）。图优化的效果同样要量化，用优化前后对比说话。

最常用的两个指标是推理延迟加速比和内存节省率。前者看速度提升了多少倍，后者看内存省了多少百分比：

$$\displaystyle \text{Speedup} = \frac{L_{\text{base}}}{L_{\text{opt}}}, \quad \text{MemSave} = \frac{M_{\text{base}} - M_{\text{opt}}}{M_{\text{base}}} \times 100\%$$

1. 延迟加速比：$L_{\text{base}}$ 是优化前单张图的推理延迟，$L_{\text{opt}}$ 是优化后的延迟，比值衡量提速倍数。

$$\text{Speedup} = \frac{L_{\text{base}}}{L_{\text{opt}}}$$

2. 内存节省率：$\frac{M_{\text{base}} - M_{\text{opt}}}{M_{\text{base}}}$ 是内存下降比例，乘以 $100\%$ 转成百分比。

$$\frac{M_{\text{base}} - M_{\text{opt}}}{M_{\text{base}}} \times 100\%$$

3. 端到端权衡：融合可能让单个内核变复杂、编译变慢，所以要综合看“编译一次 + 推理 N 次”的总账。

$$\text{cost}_{\text{total}} = t_{\text{compile}} + N \times L_{\text{opt}}$$

后厨用数据说话：不看单个动作快不快，只看翻台率和损耗。评测时务必和基线（优化前的模型）做 A/B 对比，排除硬件波动的影响。

七、工程实践：在 TVM 里跑起来

理论学完，该回后厨实操了。打开 TVM 的“后厨手册”，把三大 Pass 装进编译流程。TVM 的 Relay 前端提供了开箱即用的接口：relay.transform.FoldConstant()、relay.transform.FuseOps() 和 relay.transform.EliminateDeadCode()，它们都挂在 Pass 管理器下，一行代码就能调度。

实际部署中，这三个 Pass 几乎总是一起出现：移动端 App 内存带宽小，融合收益立竿见影；云端推理集群显存昂贵，常量折叠和死代码消除能压低峰值显存；自动驾驶对延迟极度敏感，任何多余算子都是安全隐患。调试技巧很简单：每个 Pass 跑完后打印一次 mod 里的计算图，肉眼对比图的变化，再配合 profiler 验证收益。

工程上的取舍可以写成这样：

$$\displaystyle \text{cost}_{\text{end2end}} = t_{\text{compile}} + \sum_{i=1}^{N} t_{\text{infer}}^{(i)}$$

1. 端到端成本：$t_{\text{compile}}$ 是编译阶段的耗时，只在部署前发生一次。

$$t_{\text{compile}}$$

2. 推理总时长：$\sum_{i=1}^{N} t_{\text{infer}}^{(i)}$ 是线上服务 $N$ 次推理的总时长，这是优化的核心目标。

$$\sum_{i=1}^{N} t_{\text{infer}}^{(i)}$$

3. 工程取舍：用一次性编译开销换取长期推理收益，opt_level 调得越高，Pass 越多、编译越久，但线上推理越快。

$$\min_{\text{Pass 配置}} \left( t_{\text{compile}} + \sum_{i} t_{\text{infer}}^{(i)} \right)$$

从菜谱分析到后厨落地，三大技术就这样变成了生产环境中的真实提速。记住了这三板斧，你再看 TVM 的编译日志，就能读出每一行优化背后的用意。

AI 应用场景

Python 演示

class Node:    def __init__(self, name, op, inputs, const=None):        self.name = name        self.op = op        self.inputs = list(inputs)        self.const = const        self.users = []        self.alive = True        self.fused = None def build_links(nodes):    for n in nodes:        for u in n.inputs:            u.users.append(n) def topo_sort(nodes):    seen = set()    order = []     def visit(n):        if id(n) in seen:            return        seen.add(id(n))        for u in n.inputs:            visit(u)        order.append(n)     for n in nodes:        visit(n)    return order def fold(op, vals):    if op == 'mul':        return vals[0] * vals[1]    if op == 'add':        return vals[0] + vals[1]    raise ValueError('unknown op: ' + op) def constant_folding(nodes):    for n in topo_sort(nodes):        if n.const is not None:            continue        if all(u.const is not None for u in n.inputs):            n.const = fold(n.op, [u.const for u in n.inputs]) def operator_fusion(nodes):    found = []    for n in nodes:        if n.op == 'relu' and n.inputs and n.inputs[0].op == 'bn':            bn = n.inputs[0]            if bn.inputs and bn.inputs[0].op == 'conv':                conv = bn.inputs[0]                fused = Node(n.name + '_fused', 'fused', conv.inputs[:])                fused.fused = [conv, bn, n]                found.append(fused)    return found def dead_code_elimination(nodes):    marked = set()    stack = [n for n in nodes if n.name == 'out']    while stack:        n = stack.pop()        if id(n) in marked:            continue        marked.add(id(n))        stack.extend(n.inputs)    for n in nodes:        if id(n) not in marked:            n.alive = False def dump(nodes):    for n in nodes:        if not n.alive:            continue        ins = []        for u in n.inputs:            tag = u.name            if u.const is not None:                tag += '=' + str(u.const)            ins.append(tag)        print(n.name, '(' + n.op + ')', '<-', ', '.join(ins) if ins else '-') def main():    c1 = Node('c1', 'const', [], const=2)    c2 = Node('c2', 'const', [], const=3)    mul = Node('mul', 'mul', [c1, c2])    data = Node('data', 'input', [])    conv = Node('conv', 'conv', [data])    bn = Node('bn', 'bn', [conv])    relu = Node('relu', 'relu', [bn])    add = Node('add', 'add', [relu, mul])    out = Node('out', 'out', [add])    dead = Node('dead', 'id', [data])    nodes = [data, conv, bn, relu, mul, add, out, c1, c2, dead]    build_links(nodes)    print('== 优化前 ==')    dump(nodes)    constant_folding(nodes)    print('== 常量折叠后 (mul 变成常量 6) ==')    dump(nodes)    fused_list = operator_fusion(nodes)    for f in fused_list:        print('融合:', '+'.join(x.name for x in f.fused), '->', f.name)    dead_code_elimination(nodes)    print('== 死代码消除后 (dead 被移除) ==')    dump(nodes) if __name__ == '__main__':    main()

1. if all(u.const is not None for u in n.inputs)检查一个节点的所有输入是否都是常量，这是触发折叠的开关。
2. n.const = fold(n.op, [u.const for u in n.inputs])在编译期执行算子，把结果写回节点，子图悄悄变成了常量值。
3. if n.op == 'relu' and n.inputs and n.inputs[0].op == 'bn'从链尾往前做模式匹配，识别 conv→bn→relu 融合套路。
4. fused = Node(n.name + '_fused', 'fused', conv.inputs[:])构造融合节点，它的输入直接是 conv 的输入，中间算子被“吞并”。
5. stack.extend(n.inputs)从输出节点出发，把上游节点不断压栈，这就是反向可达标记的 BFS/DFS。
6. if id(n) not in marked: n.alive = False遍历结束后，没有被标记的节点就是死代码，直接标记为不可用。
7. mul = Node('mul', 'mul', [c1, c2])构造 c1×c2 常量子图，专门用来演示折叠。
8. dead = Node('dead', 'id', [data])构造一个只被“喂”却没人“吃”的节点，专门等着被 DCE 清除。

常见误区

面试题