PyTorch 学习笔记（14）：PyTorch/LLVM 编译栈

缘起做了7年GPU编译器记录一下 Pytorch 端到端的技术栈, 记录一下查漏补缺.完整编译链路Python 用户代码 │ ▼ PEP 523 字节码拦截 TorchDynamo (torch._dynamo) ← 需要掌握CPython 内部机制 │ ▼ FX Graph (torch 级算子) AOT Autograd (torch._functorch) ← 需要掌握自动微分理论 │ ▼ FX Graph (ATen 级算子, decomposed) Inductor (torch._inductor) ← 需要掌握编译器中端优化 │ ├──→ Triton codegen (GPU) ← 需要掌握Triton 语言 GPU 架构 │ │ │ ▼ Triton IR → LLVM IR → PTX ← 需要掌握LLVM IR GPU ISA │ └──→ C/OpenMP codegen (CPU) ← 需要掌握向量化 OpenMP │ ▼ GCC/Clang 编译 → .so分层学习路线图第一层编译器基础理论必修知识点为什么需要推荐资源编译原理前端/中端/后端理解 Pass 管线、IR 设计龙书 Ch1-8LLVM TutorialSSA静态单赋值形式FX Graph 和 LLVM IR 的核心表示龙书 Ch9SSA Book数据流分析活跃变量分析、可达定义——理解 Inductor 调度龙书 Ch9循环优化tiling/fusion/unrollingInductor 的 kernel fusion 核心Muchnick Ch14-15寄存器分配与指令调度理解 LLVM 后端做了什么LLVM Developer Guide第二层LLVM 体系核心知识点对应 PyTorch 环节学习方式LLVM IR 语法与语义Triton 编译输出的中间表示llc,opt工具实操LLVM Pass 框架New/Legacy PMTriton 的优化 pipeline写一个自定义 PassLLVM 后端SelectionDAG → MachineIR → MCInstPTX/AMDGPU 代码生成LLVM Backend TutorialNVPTX 后端Triton → PTX 的最终转换LLVM 源码lib/Target/NVPTX/AMDGPU 后端ROCm 路径.ll → .hsacoROCm 文档MLIRMulti-Level IRTriton 正在迁移到 MLIR DialectMLIR Tutorial Triton 源码第三层Triton 编译器关键桥梁Triton 是连接 PyTorch 和 LLVM 的桥梁必须深入掌握Triton Python DSL │ ▼ AST → Triton IR (TTIR) ← Triton 前端 │ ▼ Triton GPU IR (TTGIR) ← GPU 特化shared memory, warp 调度 │ ▼ LLVM IR ← 通过 MLIR 的 LLVM Dialect lowering │ ▼ PTX / AMDGPU Assembly ← LLVM 后端需要学习的 Triton 知识Triton 编程模型tl.load,tl.store,tl.dot, block-level 并行Triton DialectMLIRTTIR → TTGIR 的 lowering passesAuto-tuningPyTorch Inductor 如何选择 Triton kernel 的num_warps,num_stages,BLOCK_M/N/K源码仓库github.com/triton-lang/triton尤其python/triton/compiler/和lib/第四层PyTorch 编译栈各组件组件需要掌握的核心机制关键源文件快速链接TorchDynamoCPython PEP 523 帧求值钩子符号执行Graph Break 机制Guard 生成symbolic_convert.py,eval_frame.pyFX GraphNode DAG 表示Proxy追踪Graph 变换subgraph rewritinggraph.py, proxy.pyAOT Autogradmake_fx追踪前向/反向图联合追踪min-cut 分区算法aot_autograd.py博客链接DecompositionsATen 算子分解表core_aten_decompositionsdecompositions.py博客链接Inductor LoweringATen → Inductor IR 映射lowering.py博客链接Inductor SchedulerKernel Fusion 策略内存规划scheduler.py博客链接Triton CodegenFX Node → Triton Python 代码生成triton.pyC CodegenFX Node → C/OpenMP 向量化代码cpp.pyFakeTensor元数据传播符号形状SymIntfake_tensor.py第五层GPU 架构与并行计算知识点重要性CUDA 编程模型Thread/Warp/Block/Grid理解 Triton 的 block-level 抽象GPU 内存层次Registers → Shared Memory → L2 → HBM理解 tiling 和数据复用策略Tensor Core / Matrix Core理解 CUTLASS 和 Triton 的tl.dotPTX ISA阅读 Triton 最终产物Warp Scheduling Occupancy性能调优NVLink/NVSwitch 拓扑理解 NCCL 通信优化第六层CPU 编译与向量化知识点对应 PyTorch 环节SIMD 指令集SSE/AVX2/AVX512/NEON/SVEInductor C codegen 输出、SLEEF 调用自动向量化Clang/GCC 的-O2 -marchnativeOpenMP 并行化Inductor 的#pragma omp parallel forCache-oblivious 算法理解 tiling 的性能收益推荐学习路径阶段11-2月: 编译原理 LLVM IR 基础 → 龙书前8章 LLVM Kaleidoscope Tutorial → 能读写 LLVM IR用 opt/llc 做实验 阶段21-2月: CUDA Triton → CUDA C Programming Guide → Triton Tutorial (triton-lang.org) → 手写 matmul/softmax Triton kernel 阶段31-2月: PyTorch 编译栈 → torch.compile 用户文档 → 读 torch/_dynamo/CLAUDE.md架构总览 → 用 TORCH_COMPILE_DEBUG1 看完整编译产物 → 单步调试 torch.compile 全流程 阶段41-2月: MLIR Triton 编译器内部 → MLIR Tutorial Triton 源码 → 理解 TTIR → TTGIR → LLVM Dialect 管线 → 尝试添加一个 Triton optimization pass 阶段5持续: 深入专项 → Flash Attention 论文 源码 → Inductor 的 Fusion 策略源码 → AOTInductor 部署流程 → 给 PyTorch/Triton 提 PR补充容易忽略但重要的知识CPython 内部机制——Dynamo 直接操作字节码需要理解dis模块、frame 对象、co_code符号形状推导Symbolic Shapes——SymInt/SymFloat和torch.fx.experimental.symbolic_shapes动态 shape 编译的核心难点Polyhedral 编译理论——理解 Halide/TVM 的调度语言思想有助于理解 Inductor 的 tiling 决策MLIR 正在成为新标准——Triton 已基于 MLIR 重写StableHLO/TOSA/Linalg 等 Dialect 是未来趋势