计算图与算子：计算图基础、算子的定义与分类、计算图中的数据流

好，咱们今天聊聊算子融合的前置知识——计算图和算子。说实话，很多同学一上来就搞融合，结果连计算图怎么跑的都没搞明白，那肯定要踩坑。我刚开始接触深度学习框架时，也犯过这种错误，以为算子就是一个个黑盒子，拼起来就能用。后来才发现，不懂计算图，你连性能瓶颈在哪都找不到。

计算图基础：训练过程的“施工蓝图”

计算图是什么？说白了，它就是深度学习模型训练过程的“施工蓝图”。你把模型的前向传播、反向传播、梯度更新，全部画成一张有向无环图。每个节点是一个算子，每条边是张量数据。

我个人习惯把计算图分成两种：静态图和动态图。静态图是先搭好图，再执行；动态图是边执行边搭图。PyTorch 用的是动态图，TensorFlow 1.x 是静态图，TensorFlow 2.x 也转向了动态图。为什么？因为动态图调试起来太方便了，你可以在任意位置打断点，打印中间结果。

举个例子，一个简单的线性回归：

import torch

x = torch.randn(10, 5)
w = torch.randn(5, 1, requires_grad=True)
b = torch.randn(1, requires_grad=True)

# 前向传播
y_pred = x @ w + b
loss = (y_pred - target).pow(2).mean()

# 反向传播
loss.backward()

# 更新参数
with torch.no_grad():
    w -= lr * w.grad
    b -= lr * b.grad

这段代码背后，PyTorch 自动构建了一张计算图。你调用 backward() 时，它沿着图反向传播梯度。嗯，这里要注意：计算图在每次 backward() 后会被释放，除非你设置 retain_graph=True。我遇到过有人写循环训练时，忘记保留计算图，结果第二次迭代直接报错。

算子的定义与分类：从“原子操作”到“复合操作”

算子，就是计算图上的节点。每个算子完成一个特定的数学运算。比如加法、乘法、卷积、池化、激活函数，这些都是算子。

我习惯把算子分成三类：

• 元素级算子：对张量中的每个元素独立操作，比如 ReLU、加法、乘法。这类算子计算密度低，内存带宽往往是瓶颈。
• 规约算子：把多个元素合并成一个，比如求和、求均值、Softmax。这类算子需要跨线程通信，容易成为性能瓶颈。
• 变换算子：改变张量的形状或维度，比如转置、reshape、切片。这类算子不涉及计算，但会引发内存重排。

你想想看，为什么要把算子分类？因为不同类型的算子，融合策略完全不同。元素级算子可以随便融合，反正每个元素独立；规约算子融合时要小心，搞不好会改变数学语义。

算子类型	典型例子	计算特点	融合难度
元素级	ReLU, Add, Mul	计算密度低，带宽敏感	低
规约	Sum, Mean, Softmax	需要跨线程通信	中
变换	Transpose, Reshape	无计算，内存重排	高（需考虑内存布局）

计算图中的数据流：张量如何“流动”

数据流，就是张量在计算图中沿着边传递的过程。每个算子的输出，成为下一个算子的输入。这个过程看似简单，但里面藏着很多性能优化的机会。

我举个例子，假设你有这样一个计算图：

A = Conv2d(x)
B = ReLU(A)
C = BatchNorm(B)
D = Add(C, residual)

传统做法是：先算 Conv2d，把结果写回显存；再读出来算 ReLU，写回去；再读出来算 BatchNorm，写回去；最后读出来算 Add。你算算，这中间有多少次显存读写？四次写，四次读，总共八次内存访问。

如果做算子融合，把 Conv2d + ReLU + BatchNorm 合并成一个 kernel，中间结果直接留在寄存器或共享内存里，不写回显存。这样，显存访问次数从八次降到了两次（一次读输入，一次写输出）。

数据流还有一个重要概念：数据依赖。如果两个算子之间没有数据依赖，它们就可以并行执行。比如上面的 Conv2d 和 residual 分支，如果 residual 是直接从输入跳过来的，那 Conv2d 和 Add 之间就有依赖，但 Conv2d 和 residual 的预处理（如果有）可以并行。

我建议你在设计融合策略时，先画出数据流图，标出每个算子的输入输出，然后找出那些“计算密集但中间结果频繁读写”的路径。这些路径就是融合的黄金候选。

知识体系总览

下面这张图，是我自己总结的计算图与算子知识体系。你可以把它当作本章的“地图”，方便后续查阅。

这张图把计算图、算子分类、数据流三个核心概念串在了一起。你从左上角开始看：先理解计算图是静态还是动态，然后知道算子分三类，最后关注数据是怎么流动的。这三者结合起来，就是算子融合的理论基础。

好了，这一章的内容就到这里。计算图和算子是算子融合的基石，你把这些搞清楚了，后面讲具体的融合策略时，就会觉得顺理成章。