torch反向计算图提取

[fangfangssj]

流程

反向计算图的捕捉，先使用Dynamo将model.py中Python代码转换成一个前向的FX Graph，再将这个Graph传给AOT Autograd来自动生成反向图，使用Autograd来获取每个算子对应的反向操作，AOT Autograd引擎输出的是一个联合优化的计算图Joint Graph (Fwd + Bwd)，最后保存为backward计算图。PR已经实现 #490
model.py (Python)
⬇
[TorchDynamo] <-- 1. 捕捉 Python 逻辑
⬇
FX Graph (Forward Only)
⬇
[AOT Autograd] <-- 2. 这里的 Autograd 引擎介入
┃
┣━━ 运行 Forward (FakeTensor) ━━┓
┃ ⬇
┃ [PyTorch C++ Autograd]
┃ ⬇
┗━━ 记录并生成 Joint Graph (Fwd + Bwd)
⬇
导出存为backward计算图

待讨论的地方：

#单独的一个反向计算图
    def forward(self, primals_1, primals_2, primals_4, primals_5, primals_6, convolution, getitem_1, getitem_2, tangents_1):
        native_batch_norm_backward = torch.ops.aten.native_batch_norm_backward.default(tangents_1, convolution, primals_4, primals_1, primals_2, getitem_1, getitem_2, False, 0.001, [True, True, True]);  tangents_1 = convolution = primals_4 = primals_1 = primals_2 = getitem_1 = getitem_2 = None
        getitem_3 = native_batch_norm_backward[0]
        getitem_4 = native_batch_norm_backward[1]
        getitem_5 = native_batch_norm_backward[2];  native_batch_norm_backward = None
        convolution_backward = torch.ops.aten.convolution_backward.default(getitem_3, primals_6, primals_5, [0], [1, 1], [1, 1], [1, 1], False, [0, 0], 1, [True, True, False]);  getitem_3 = primals_6 = primals_5 = None
        getitem_6 = convolution_backward[0]
        getitem_7 = convolution_backward[1];  convolution_backward = None
        return (getitem_5, getitem_4, getitem_7, getitem_6)

1. 保存的计算图是单独的backward，还是Joint Graph (Fwd + Bwd)，后者更方便进行联合优化
2. backward中都是torch内部的一些反向算子（native_batch_norm_backward，convolution_backward等），作为计算图样本的意义可能不大，需要考虑将计算图使用torch._decomp来处理嘛，分成更细粒度的算子，方便编译器融合，（例如将native_batch_norm_backward拆分为aten.sub，aten.mul，aten.sum，aten.rsqrt），基础算子是一个封闭的集合，更方便ai4c来学习融合。同时在国产新硬件上，基本上会支持torch的所有的基础算子，而某些复杂的算子则不会提供支持。

[fangfangssj]

传统标准流水线如下：

1. Input: torch.ops.aten.native_batch_norm (High Level)
2. Decomposition Pass:

• 查表：发现是 BN。
• 动作：替换为 (input - mean) * inv_std + beta。
• 结果：图里全是 sub, mul, add。

3. Optimization Pass (Pattern Match):

• 扫描：发现 mul 后面接了个 add，再接了个 relu。
• 决策：这是一个 Vertical Fusion 组。

4. Codegen:

• 生成单个 GPU Kernel，在一个线程内完成这三个操作。

ai4c是训练一个agent来代替第三部分和第四部分嘛