【编译系列】Torch.compile()流程解析——1. torch.compile介绍

摘要：

本系列主要分享训练编译的相关知识，第一章主要解析torch.compile的工作流程以及各个组件的作用和源码调用关系，配合代码示例进行解析。主要分为两大块：前端静态图捕获和后端编译。
a. 前端对应TorchDynamo，主要介绍python编译执行的基本流程、PEP 523、TorchDynamo模拟执行字节码、FX Graph的构建过程；
b. 后端编译分析inductor函数实现，里面包含AOTAutograd、PrimTorch、TorchInductor等，其中AOTAutograd涉及__torch_dispatch__如何从FX Graph一步步构建前向反向joint graph，PrimTorch涉及算子decompose，TorchInductor主要涉及图优化、Triton内核代码生成；
本文主要基于torch2.4.1+cu118的源码进行解析。

背景：

为了改善模型训练推理性能，Pytorch推出torch.compile，并从torch 1.xx更为torch 2.xx，改变并增强了Pytorch在编译器级别的运行方式，同时开始将pytorch的部分代码从C++迁移回python，主要有以下特点：
a. 将动态的 PyTorch 模型转换为静态的、优化过的执行图，从而减少运行时的开销，提高计算效率。
b. 用户无需手动修改模型代码或选择特定的优化策略，torch.compile可以自动识别和应用最佳的优化路径，使得模型加速更加便捷。
c. 支持广泛的 PyTorch API 和操作，确保大部分现有模型在使用 torch.compile后无需进行大的调整即可获益于性能提升。

torch.compile主要包含四个基础组件：

• TorchDynamo：从python bytecode中解析构建计算图，是一个动态的、Python级别的编译器，旨在捕捉 PyTorch 模型的动态执行路径，将其转换为优化代码，实现bytecode-to-bytecode编译。
• AOTAutograd：是 PyTorch 引入的一种自动求导机制，旨在在模型执行之前预先生成梯度计算的代码。这种方法通过静态分析模型的前向计算图，提前生成反向传播所需的梯度计算逻辑，从而减少运行时的开销，提升训练效率。
• PrimTorch：将2000+ Pytorch op规范化为约250个原始op封闭集合，是 PyTorch 的一个中间表示（Intermediate Representation, IR）层，负责将高层的 PyTorch 操作转换为更低层次的、适合进一步优化和编译的基础操作。它通过简化和标准化操作，提高编译器和后端优化器处理计算图的效率。
• TorchInductor：深度学习编译器，可为多种加速器和后端生成代码，生成OpenAI Triton(Nvidia/AMD GPU)和OpenMP/C++(CPU)代码，它利用多种优化技术，包括内存优化、并行化和低层次的代码生成，以最大化计算性能。
  以上四个组件都是用Python编写的，支持动态形状（即能够发送不同大小的张量而无需重新编译），实现灵活并减低开发人员和供应商开发门槛的效果。
  

torch.compile使用

• torch.compile使用示例：

import torch
import torchvision.models as models

model = models.resnet18().cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
compiled_model = torch.compile(model, backend="inductor")    # 通过backend参数指定后端，默认为inductor
# compiled_model = torch._dynamo.optimize("inductor")(fn)    # 也可以通过torch._dynamo.optmize函数进行编译

x = torch.randn(16, 3, 224, 224).cuda()
optimizer.zero_grad()
out = compiled_model(x)
out.sum().backward(