缘起
做了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-8；LLVM Tutorial |
| SSA（静态单赋值）形式 | FX Graph 和 LLVM IR 的核心表示 | 龙书 Ch9；SSA Book |
| 数据流分析 | 活跃变量分析、可达定义------理解 Inductor 调度 | 龙书 Ch9 |
| 循环优化（tiling/fusion/unrolling） | Inductor 的 kernel fusion 核心 | Muchnick Ch14-15 |
| 寄存器分配与指令调度 | 理解 LLVM 后端做了什么 | LLVM Developer Guide |
第二层：LLVM 体系（核心）
| 知识点 | 对应 PyTorch 环节 | 学习方式 |
|---|---|---|
| LLVM IR 语法与语义 | Triton 编译输出的中间表示 | llc, opt 工具实操 |
| LLVM Pass 框架（New/Legacy PM） | Triton 的优化 pipeline | 写一个自定义 Pass |
| LLVM 后端（SelectionDAG → MachineIR → MCInst） | PTX/AMDGPU 代码生成 | LLVM Backend Tutorial |
| NVPTX 后端 | Triton → PTX 的最终转换 | LLVM 源码 lib/Target/NVPTX/ |
| AMDGPU 后端 | ROCm 路径 .ll → .hsaco |
ROCm 文档 |
| MLIR（Multi-Level IR） | Triton 正在迁移到 MLIR Dialect | MLIR 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 Dialect（MLIR）：TTIR → TTGIR 的 lowering passes
- Auto-tuning ：PyTorch Inductor 如何选择 Triton kernel 的
num_warps,num_stages,BLOCK_M/N/K - 源码仓库 ：
github.com/triton-lang/triton（尤其python/triton/compiler/和lib/）
第四层：PyTorch 编译栈各组件
| 组件 | 需要掌握的核心机制 | 关键源文件 | 快速链接 |
|---|---|---|---|
| TorchDynamo | CPython PEP 523 帧求值钩子；符号执行；Graph Break 机制；Guard 生成 | symbolic_convert.py, eval_frame.py |
|
| FX Graph | Node DAG 表示；Proxy 追踪；Graph 变换（subgraph rewriting） |
graph.py, proxy.py | |
| AOT Autograd | make_fx 追踪；前向/反向图联合追踪；min-cut 分区算法 |
aot_autograd.py | 博客链接 |
| Decompositions | ATen 算子分解表；core_aten_decompositions |
decompositions.py | 博客链接 |
| Inductor Lowering | ATen → Inductor IR 映射 | lowering.py | 博客链接 |
| Inductor Scheduler | Kernel Fusion 策略；内存规划 | scheduler.py | 博客链接 |
| Triton Codegen | FX Node → Triton Python 代码生成 | triton.py | |
| C++ Codegen | FX Node → C++/OpenMP 向量化代码 | cpp.py | |
| FakeTensor | 元数据传播；符号形状（SymInt） | fake_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.dot |
| PTX ISA | 阅读 Triton 最终产物 |
| Warp Scheduling & Occupancy | 性能调优 |
| NVLink/NVSwitch 拓扑 | 理解 NCCL 通信优化 |
第六层：CPU 编译与向量化
| 知识点 | 对应 PyTorch 环节 |
|---|---|
| SIMD 指令集（SSE/AVX2/AVX512/NEON/SVE） | Inductor C++ codegen 输出、SLEEF 调用 |
| 自动向量化 | Clang/GCC 的 -O2 -march=native |
| OpenMP 并行化 | Inductor 的 #pragma omp parallel for |
| Cache-oblivious 算法 | 理解 tiling 的性能收益 |
推荐学习路径
阶段1（1-2月）: 编译原理 + LLVM IR 基础
→ 龙书前8章 + LLVM Kaleidoscope Tutorial
→ 能读写 LLVM IR，用 opt/llc 做实验
阶段2（1-2月）: CUDA + Triton
→ CUDA C Programming Guide
→ Triton Tutorial (triton-lang.org)
→ 手写 matmul/softmax Triton kernel
阶段3（1-2月）: PyTorch 编译栈
→ torch.compile 用户文档
→ 读 torch/_dynamo/CLAUDE.md（架构总览）
→ 用 TORCH_COMPILE_DEBUG=1 看完整编译产物
→ 单步调试 torch.compile 全流程
阶段4（1-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 是未来趋势