AI 编译器工作机理：前端融合、后端 Kernel 与自动调优

PyTorch 写出来的模型，为什么不能直接在 GPU 上跑？中间隔着 AI 编译器——它把"框架描述的计算图"翻译成"硬件上高效的算子"。本文把 AI 编译器的三层结构、算子融合与自动调优讲透，衔接本站 MLIR 架构 与 Triton 实战。

1. 为什么需要 AI 编译器

框架与硬件之间存在语义鸿沟：

PyTorch（高层计算图）        GPU 指令（底层 Kernel）
      │                            ▲
      └──────  AI 编译器  ─────────┘
        图优化 → IR → Kernel 生成/调优

直接逐算子翻译的结果：每个算子单独启动、中间张量反复落显存、带宽利用率惨不忍睹（见 算子运行时账）。编译器存在的意义就是在翻译时做全局优化。

2. 三层结构

层	干什么	典型手段
前端	计算图优化	算子融合、常量折叠、死代码消除
中端	统一 IR 表示	MLIR、TOSA、ONNX
后端	Kernel 生成与调度	代码生成、AutoTuning、内存规划

判断：前端优化与硬件无关（什么卡都受益），后端优化与硬件强绑定（每家卡的后端不同）——这就是 MLIR 分层设计的核心动机：前端写一次，后端各自实现。

3. 算子融合：前端最重要的优化

融合的收益已经在 算子账 里算过（省 2 写 2 读）。按类型分：

融合类型	例子	收益
相邻算子融合	Conv+BN、GELU+Add	少一次中间张量落盘
跨层结构融合	Attention 全融合（FlashAttention）	显存 IO 复杂度降一档
计算-通信融合	梯度切分与通信重叠	通信时间藏进计算

判断：融合的本质是"把 IO 账算进算子边界"——带宽受限的算子（大多数）融合一个中间张量省 30-50%。

4. 后端自动调优：三条路线

Kernel 参数（block 大小、循环展开、共享内存划分）组合爆炸，手调不可行：

路线	原理	代表
手工模板 + 启发式	预置模板，规则选参	cuDNN/cuBLAS 启发式
自动搜索	编译期枚举 + 实测打分	TVM AutoTVM、Triton autotune
机器学习驱动	学习 kernel 性能模型	XLA、部分研究系统

判断：生产环境"模板 + 启发式"最稳（cuDNN 的算法选择器），研究/新卡"自动搜索"更通用。Triton 的 @triton.autotune 是实践入口（见 Triton 学习路径）。

5. PyTorch2.0 编译栈：torch.compile 拆解

torch.compile(model)
   ├── FX 图捕获（Python 层 → 计算图）
   ├── TorchInductor（IR 优化 + 代码生成）
   └── 生成 Triton / C++ Kernel

• 图捕获解决"Python 动态语言难编译"的问题——先固化成图；
• TorchInductor 默认生成 Triton kernel，兼顾性能与可读性；
• 收益：训练端到端提速 10-40%（编译 + 算子融合 + 内核优化），与 PyTorch 编译栈 一文呼应。

6. 交付清单

1. 训练/推理性能上不去，先看算子是否"每步落显存"——启用 torch.compile 或手写融合算子；
2. 新卡/新算子优先试 Triton autotune，生产固化模板参数；
3. 跨硬件部署用 MLIR/TOSA 中间层，避免每个硬件各写一套前端；
4. 每次编译优化用 profiler 前后对比（见 验收基线）。

---

本文为本站整理思考后撰写，用于知识库学习交流；具体版本号与实测数据以现场为准。