AI 编译器：TVM、MLIR 与 torch.compile 的算子编译优化

手写 Kernel 的困境：
□ 新算子（FlashAttention/量化新算法）→ 手写很慢
□ 硬件多样性（NVIDIA/AMD/Intel/ARM/NPU）→ 每硬件一份
□ 手工优化（内存布局/向量化/分块）→ 靠专家经验
□ 演进慢：模型迭代速度 >> 算子库更新速度

AI 编译器解决的问题：
□ 自动生成：从计算图自动生成高效 Kernel
□ 可移植：一套描述 → 多硬件后端
□ 可调优：自动搜索最优调度（分块/向量化/展开）
□ 融合：跨算子融合减少内存访问（见融合篇）

传统 vs 编译器：
□ 传统：人写 Kernel（准但慢，覆盖有限）
□ 编译器：机器生成（快覆盖，质量接近手写）

前端（Frontend）：模型 → 计算图 IR
□ 从 PyTorch/ONNX/TensorFlow 导入模型
□ 转换成「框架无关」的计算图 IR
□ 完成第一层图优化（常量折叠、死代码消除）

中端（Optimization）：图 IR → 优化后的图/张量 IR
□ 算子融合、图重写、内存规划
□ 张量级优化（布局转换、代数化简）

后端（Backend）：优化 IR → 目标代码
□ 调度选择（分块/向量化/并行）
□ 代码生成（生成 CUDA/C++/OpenCL/汇编）
□ 编译成可执行 Kernel + 加载

关键：三层解耦 → 前端支持多框架，后端支持多硬件

数据流示例：
PyTorch 模型 → (前端) 计算图 → (中端) 融合优化 → (后端) CUDA Kernel
                ONNX 模型  →  计算图   →  优化图   →  目标代码

IR 层级：
□ 高层 IR（计算图）：算子级别的 DAG（add/matmul/softmax）
  → 做图优化（融合、常量折叠）
□ 张量 IR：逐算子展开成「循环 + 内存操作」
  → 表达分块、数据布局
□ 调度 IR / 代码生成 IR：循环变换 + 向量化 + 并行
  → 决定最终怎么跑

层级之间的「降低」（lowering）：
高层算子 → 展开成低层循环 → 调度变换 → 目标代码
每个降低步骤都有对应的「可优化点」

例：softmax 的降低
高层：Softmax(x)                      # 一个算子
张量级：exp(x) → sum → div             # 拆成基础操作
调度级：分块计算、在线 softmax、向量化  # 决定怎么跑

主要优化：
□ 算子融合（Fusion）：多个算子合成一个
  例：conv + bn + relu → 一个融合 Kernel
  → 减少「中间结果的显存读写」（带宽收益，见融合篇）
□ 常量折叠：常数子图 → 编译期算好
□ 死代码消除：无用分支/节点删除
□ 内存规划：张量内存复用（in-place / 别名分析）
□ 布局转换：NCHW → NHWC 等（匹配硬件偏好）

融合的收益来源：
中间张量不进显存（留在寄存器/片上内存）
→ 对带宽受限的推理，这是最大的单点优化

融合示例：
x → conv → bn → relu → y
融合成：kernel_conv_bn_relu(x) → y
→ 省去 conv 输出、bn 输出的两次显存往返

核心调度原语：
□ 分块（Tiling/Split）：把大循环切成小块
  → 匹配片上内存/缓存大小，减少重复读
□ 向量化（Vectorize）：相邻元素并行处理（SIMD）
  → 利用 GPU 的向量单元 / CPU 的 SIMD
□ 并行化（Parallelize）：循环维度映射到线程
  → GPU：线程块/线程分配
□ 展开（Unroll）：展开循环体减少分支开销
□ 重排序（Reorder）：循环嵌套顺序（影响缓存局部性）

调度的目标：
□ 最大化「数据复用」：片上/寄存器反复用
□ 最大化「带宽利用」：合并访问、充分加载
□ 最小化「同步/分支」开销

调度例（矩阵乘）：
分块：把 M×N×K 切成 tile（如 64×64×64）
向量化：内层乘加用向量指令
并行：每个线程块算一个 tile
重排序：让 tile 的加载与计算重叠（流水线）

TVM 架构：
□ Relay（图 IR）+ TensorIR（张量 IR）
□ TVMScript：用 Python 表达「算子 + 调度」（可读）
□ 自动调优：让机器搜索最优调度

自动调优两条路：
□ AutoTVM：模板 + 搜索（预定义调度模板，搜索参数）
□ Ansor：程序生成 + 搜索（自动生成调度候选）
  → Ansor 更自动化，覆盖更广

搜索策略：
□ 随机/遗传/贝叶斯搜索
□ 在真实硬件上测量（硬件在环）
□ 调优成本高（GPU 上几小时）→ 调优一次、保存 artifact

# TVMScript 调度示例（概念）
@T.prim_func
def matmul(A: T.Tensor, B: T.Tensor, C: T.Tensor):
    ...
    T.parallel(M)          # 并行
    T.vectorize(K)         # 向量化
# 或用 Ansor 自动搜索这些调度参数

MLIR 的核心思想：
□ 多级 IR：不是一层 IR，而是「一系列」dialect（方言）
□ Dialect：每种 IR 风格是一个 dialect
  （tensor/arith/memref/gpu/llvm 等）
□ 渐进降低：从高层 dialect 逐级降到低层
  → 每级只做「该级能做的优化」

MLIR 的价值：
□ 可组合：不同框架/后端共享 IR 基础设施
□ 可移植：一套基础设施 → 多硬件（GPU/CPU/TPU/FPGA）
□ 可定制：新的硬件后端 = 新增 dialect + 降低路径

在 AI 推理的角色：
□ ONNX-MLIR、torch-mlir 等用 MLIR 做中间层
□ 与 TVM 的关系：MLIR 偏「基础设施」，TVM 偏「完整编译栈」

MLIR 降低路径（示意）：
onnx dialect → linalg dialect → affine dialect → gpu dialect → llvm
（框架算子）   （张量/循环）   （仿射循环）     （GPU 调度）   （目标码）

torch.compile 做什么：
□ 把 PyTorch 模型「图化」→ 图优化 + 算子编译
□ 动态图 → 捕获成图 → 编译优化 → 生成 Triton Kernel
□ 一行代码：model = torch.compile(model)

底层技术：
□ TorchDynamo：动态图捕获（Python 层字节码分析）
□ TorchInductor：图 → Triton/C++ 代码生成
□ Triton：GPU 上的「类 Python」编程语言（写 Kernel 简单）
□ 与 CUDAGraph 结合：进一步减少启动开销

收益：
□ 训练：自动融合 + 内核优化 → 提速
□ 推理：同模型代码「编译后」延迟/吞吐提升
□ 开发友好：不用换框架、不用手写 Kernel

import torch
model = torch.compile(model, mode="reduce-overhead")
out = model(x)   # 编译 + 执行（首次编译慢，之后走缓存）

结合形态：
□ 推理引擎（TensorRT/vLLM/Triton）内部就用编译技术
  → TensorRT 的图优化 + Kernel 选择 = 编译器思想
  → vLLM 的融合算子（FlashAttention 等）是手工/编译混合
□ torch.compile 产物 → 导出 ONNX/TensorRT → 部署
□ 编译栈直出：TVM 编译整个模型 → 部署 artifact

生产流程（推荐）：
1. 训练后 → 导出（ONNX/IR）
2. 编译器（TVM/torch.compile）→ 优化 + 生成 Kernel
3. 离线调优（搜索调度/校准量化）→ 保存 artifact
4. 部署 artifact → 推理引擎加载执行

关键：离线编译、在线执行
□ 编译/调优是离线任务（花时间）
□ 推理是执行编译产物（快）

选型路径：
PyTorch 生态内 → torch.compile + Triton
跨框架/跨硬件 → ONNX + TVM/ONNX-MLIR
极致性能 + NVIDIA → TensorRT（内部含编译器栈）