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