TorchInductor - Introduction

PyTorch 2.x通过TorchDynamo通过Python Bytecode的动态变换实现了图捕获功能，需要搭配一个Compiler Backend完成图编译。

Pytorch尝试集成了多个后端，并使用一个轻量级的autotuner来选择最优的后端图编译结果。这个解决方案存在2个问题：

• 这些后端的Execution Model和Pytorch差异较大，引入了许多转换步骤，导致性能损失较大。
• 大多数后端仅支持推理，一些设计决策使得支持训练非常困难。

因此Pytorch需要原生的Compiler Backend：TorchInductor。

整体设计

TorchInductor的整体设计思路是一个轻量级、易于扩展和实验的框架，用于将PyTorch的表示符号化映射到Compiler Backend：

• 完整表达Pytorch：torch.Tensor -> TensorBox，torch.Storage -> StorageBox 等。
• 通过symbolically strided tensor表达各种View：Reshape/Transpose/Slice等。
• 支持训练
• 支持多后端
• 支持高层级优化：如Memory Planning。

TorchInductor的初始设计支持2种不同的target：

• Triton：一种新型的编程语言，开发效率高于CUDA，同时性能可以媲美CUDA的库（如cuDNN），支持NVIDIA/AMD GPU等。
• C++/OpenMP：一种被广泛使用的编写高性能并行Kernel的API，支持CPU。

TorchInductor设计上优先考虑对Pytorch支持的完整性，包括：

• aliasing/mutation/views
• scatter (间接写)
• gather (间接读)
• pooling/windows/reductions
• masked/conditional execution（如padding）
• template epilogue fusions
• tiling
• horizontal/vertical fusions

TorchInductor使用SymPy库来支持动态形状（dynamic shapes）和步长（strides）：

使用SymPy符号化tensor的shape，并在整个程序中传播。

load和store直接通过SymPy的索引公式来表达。

通过guards来提升已编译好的子图的使用前提，在guards fail时，触发子图的重编译。

TorchInductor  IR

TorchInductor IR使用了一种define-by-run的loop-level IR。大部分IR是Python里的Callable，输入是SymPy Expression。基于这种IR做分析或代码生成的实现方式是改变ops.* 的实现，并运行IR。

例如对x.permute(1, 0) + x[2, :] 的IR：

def inner_fn(index: List[sympy.Expr]):i1, i0 = indextmp0 = ops.load("x", i1 + i0*size1)tmp1 = ops.load("x", 2*size1 + i0)return ops.add(tmp0, tmp1)torchinductor.ir.Pointwise(device=torch.device("cuda"),dtype=torch.float32,inner_fn=inner_fn,ranges=[size0, size1],
)

TODO：待补充

参考：

TorchInductor: a PyTorch-native Compiler with Define-by-Run IR and Symbolic Shapes - compiler - PyTorch Developer Mailing List