AI算子融合原理与编译器优化实战 - 第2节
一、计算图：AI模型的骨架
做AI系统优化这些年，我越来越觉得，理解计算图就像理解汽车的底盘。你光看引擎参数没用，得知道动力怎么传递到轮子上。计算图就是干这个的——它把模型里的每个操作都画成一个节点，数据流就是连接节点的边。
说白了，一个神经网络就是一张有向无环图（DAG）。每个节点代表一个算子，比如卷积、ReLU、矩阵乘法。边代表张量（Tensor）的流向。我刚开始接触这个概念时，总觉得它太抽象。直到有一次调试性能瓶颈，发现两个相邻的算子之间反复读写显存，才恍然大悟——哦，原来计算图能帮我看到这些隐藏的开销。
核心要点：算子就是计算图的节点。每个节点完成一个特定的数学运算，节点之间的连线就是数据依赖关系。
二、静态图 vs 动态图：两种执行哲学
这个问题，我面试过不少候选人，十有八九会背定义。但真问到「为什么静态图能优化得更狠」，很多人就卡壳了。咱们今天把它彻底讲透。
2.1 静态图：先规划，再执行
静态图的特点是：你先搭好整个计算图，然后交给框架去执行。TensorFlow 1.x 就是典型代表。你想想看，这就像写代码之前先画好流程图——所有分支、循环、条件都得提前定死。
这样做的好处很明显：
- 全局可见性：框架能看到整张图，可以做跨节点的优化
- 内存预分配：知道每个张量的大小，提前分配好内存，避免动态申请
- 算子融合：把多个小算子合并成一个大算子，减少内核启动开销
我记得有一次优化一个BERT模型，用静态图把连续的三个小算子（BiasAdd + ReLU + Dropout）融合成一个算子，推理速度直接提升了15%。这就是静态图的威力——它能看到全局，敢做激进优化。
个人经验：静态图调试起来比较痛苦。我曾经花了一整天，就为了找出图里一个形状不匹配的bug。因为错误信息是在图执行时才抛出来的，你很难定位到具体是哪行代码导致的。
2.2 动态图：边执行边构建
动态图正好相反。PyTorch 是它的代言人。你写一行代码，它就执行一行，同时构建计算图。这就像即兴演讲——说到哪算哪，灵活得很。
动态图的优势：
- 调试友好：可以用Python的pdb断点，逐行检查中间结果
- 控制流自然：if-else、for循环直接用Python原生的，不用搞什么tf.cond
- 开发效率高：改模型结构就像改普通Python代码一样
但代价呢？性能。动态图每次执行都要重新构建图，没法做全局优化。你想想看，每次前向传播都得重新解析一遍算子依赖关系，这开销可不小。
避坑指南：我曾经在一个项目里，用PyTorch动态图跑一个很深的ResNet。训练时发现GPU利用率只有30%左右。后来一分析，发现大量时间花在了Python层面的图构建和算子调度上。换成静态图后，利用率直接飙到85%。
三、一张图看懂区别
下面这张SVG图，是我自己画的。它把静态图和动态图的执行流程做了对比。你看完应该能秒懂。
四、算子作为节点：为什么这个视角重要？
把算子看作计算图的节点，这个视角有什么用？我举两个实际例子。
4.1 算子融合：减少内核启动开销
GPU上执行一个算子，需要启动一个CUDA内核。启动内核本身有开销——大概几微秒到几十微秒。如果你的模型里有几百个很小的算子，光启动开销就能吃掉不少性能。
算子融合就是把多个小算子合并成一个。比如：
# 融合前：3次内核启动
y1 = torch.add(x, bias) # 内核启动1
y2 = torch.relu(y1) # 内核启动2
y3 = torch.dropout(y2, p=0.5) # 内核启动3
# 融合后：1次内核启动
y3 = fused_add_relu_dropout(x, bias, p=0.5) # 内核启动1你看，三次启动变成一次。我实测过，对于小算子密集的网络，融合后能省掉30%-50%的内核启动时间。
4.2 内存规划：减少显存碎片
静态图因为知道所有张量的生命周期，可以提前规划好内存复用策略。比如两个张量，一个在时间点T1创建、T2销毁，另一个在T3创建、T4销毁——如果T2 < T3，它们就可以共用同一块显存。
动态图做不到这一点。它只能走一步看一步，每次分配新内存。时间长了，显存碎片化严重。我见过一个极端案例，动态图跑一个模型，显存占用比静态图多了40%。
一句话总结：静态图用「规划」换「性能」，动态图用「性能」换「灵活」。没有绝对的好坏，看你的场景。
五、实际项目中的选择建议
说了这么多，到底该用哪个？我个人的经验是这样的：
| 场景 | 推荐模式 | 原因 |
|---|---|---|
| 研究探索、快速原型 | 动态图 | 调试方便，改结构快 |
| 生产部署、推理优化 | 静态图 | 性能好，可做算子融合 |
| 大规模分布式训练 | 静态图 | 通信调度可提前规划 |
| 模型结构频繁变化 | 动态图 | 静态图改结构成本高 |
不过现在框架也在互相学习。PyTorch 2.0 引入了 torch.compile，可以把动态图编译成静态图。TensorFlow 2.x 也支持了 eager mode。说白了，两者在趋同。但理解它们的本质区别，对你做性能优化绝对有帮助。
一个小技巧：如果你用PyTorch做推理优化，可以试试 torch.jit.script 或 torch.jit.trace。它们能把动态图「固化」成静态图，享受算子融合的好处。我有个项目就是这么干的，推理延迟从12ms降到了8ms。
好了，这一章就到这里。计算图是算子融合的基础，理解它，你才能明白后面为什么要做融合、怎么做融合。下一章我们会深入算子的具体实现，看看GPU上到底发生了什么。