【infra之路】编译优化MLIR做pass的详细过程

本质上不管是静态图还是动态图都是由3个部分组成的Tensor、算子（Ops）、依赖关系Edge。这是我们在图捕获得到静态/动态图之后需要这个图Lowering到相应的中间层IR做相关的优化。

我们假设你有一个简单的 PyTorch 模型，比如一个包含 Linear -> ReLU -> Linear 的简单全连接网络。下面我为你详细拆解它进入 MLIR 后，经历各种 Pass 的"奇幻漂流"。

整个过程可以分为 5 个主要阶段。

---

阶段一：前端接入与高层 IR 生成 (Frontend & High-Level IR)

PyTorch 本身是一个动态图框架，MLIR 需要静态结构。所以第一步是捕获计算图。

1. TorchScript / FX Graph 捕获：
   PyTorch 先将模型转换为 TorchScript 或者 FX Graph（一种静态的 Python 级别的计算图）。
2. Torch Dialect 生成 (Torch-MLIR)：
   通过 torch-mlir 项目，将 PyTorch 的图转换为 MLIR 的 torch dialect 。
   此时 IR 看起来非常像 PyTorch 代码：
   • torch.operator 表示具体的算子（如 torch.aten.linear, torch.aten.relu）。
   • torch.tensor 表示张量类型（带有 shape 和 dtype）。
   • 这时的 IR 是最高层的，完全保留了 PyTorch 的语义。

阶段二：高层图优化与 Dialect 转换 (Graph Optimization)

现在 IR 在 MLIR 手里了。第一步是"去 PyTorch 化"，将其转换为更通用的 MLIR 方言。

1. Torch -> Linalg / TOSA 转换：
   调用 convert-torch-to-linalg 或 convert-torch-to-tosa Pass。
   • torch.aten.linear 被分解为 linalg.matmul + linalg.add（加上 bias）。
   • torch.aten.relu 被转换为 linalg.generic（一个逐元素的最大值操作）或者 arith.maxf。
   • torch.tensor 变成了 tensor dialect。
2. 常量折叠与死代码消除 (Canonicalization)：
   此时会运行 canonicalize Pass。比如如果权重是常量，可能会在此时进行计算；如果某个算子对最终结果没影响，会被删掉。

此时的 IR 状态： 混合了 tensor（数据）和 linalg（计算）方言。这是 MLIR 最适合做高层图优化的层级。

阶段三：中层循环优化与 Bufferization (Loop Optimization & Memory Planning)

这是 MLIR 区别于传统编译器最精彩的部分。linalg 只是说"做个矩阵乘法"，但没说怎么做。

1. Tiling (循环分块)：
   调用 linalg-tile Pass。考虑到 Cache 大小，编译器决定把大矩阵切成小块（比如 32x32）。
   • linalg.matmul 被转化为 scf.for 循环嵌套，内部包含小的 linalg.matmul。
2. Fusion (算子融合)：
   这是 AI 编译器的关键。linalg.matmul 后面跟着 linalg.generic (ReLU)。
   通过 linalg-fuse Pass，编译器把 ReLU 的计算融合进了矩阵乘法的循环里。这样就不需要把矩阵乘法的结果写回内存再读出来做 ReLU，大大节省了带宽。
3. Bufferization (内存分配)：
   之前我们用的是 tensor（抽象值，不可变）。现在必须变成 memref（内存缓冲区，可变）。
   调用 iree-bufferize 或 one-shot-bufferize Pass。
   • tensor 类型变成 memref 类型。
   • 编译器在此刻进行内存规划 ：分析哪些内存可以复用，插入 memref.alloc 和 memref.dealloc，或者尽量提升到栈上分配。
4. Lowering to SCF / Affine：
   linalg 结构被完全打散，变成 scf.for / affine.for 循环和 memref.load / memref.store 操作。

此时的 IR 状态： memref（内存） + scf/affine（循环） + arith（标量计算）。

阶段四：底层硬件相关优化 (Low-Level Optimization)

现在代码已经很像 C 语言了，接下来要针对硬件（CPU/GPU）做优化。

1. Vectorization (向量化)：
   调用 scf-vectorize 或 linalg-vectorize Pass。
   将 scf.for 循环中的标量操作（一次算一个 float）转换为 vector dialect 操作（一次算 8 个 float，利用 AVX2/AVX512 指令）。
   • arith.addf -> vector.add (SIMD 指令)。
2. 循环展开与不变量提取：
   运行 affine-loop-unroll 等 Pass，进一步优化指令流水线。
3. Memref 降级为裸指针：
   调用 memref-to-llvm Pass。memref 的多维结构被展平，变成 llvm.ptr（裸指针）和 GEP 计算。

阶段五：LLVM IR 生成与后端编译 (LLVM Lowering)

这是最后一步，把所有剩余的 MLIR 方言全部翻译成 LLVM IR。

1. 转换为 LLVM Dialect：
   运行 convert-scf-to-cf（结构化控制流变跳转），convert-arith-to-llvm，convert-func-to-llvm 等 Pass。
   此时，所有的 scf.for 变成了 cf.br（基本块跳转），所有的 arith.addf 变成了 llvm.add。
2. 导出 LLVM IR：
   调用 mlir-translate --mlir-to-llvmir。
   生成标准的 .ll 文件。
3. LLVM 后端：
   LLVM 接手，进行指令调度、寄存器分配，最终生成 .o 或 .s 汇编文件，链接成可执行程序。

---

总结：一个 Pass Pipeline 的伪代码

在实际工程（如 IREE 或 Torch-MLIR）中，这一长串过程是被封装在一个 Pass Pipeline 里的。看起来大概是这样：

// 1. 接入与高层转换
pass pipeline:
  convert-torch-to-linalg
  canonicalize
  inline

// 2. 中层优化 (核心魔法发生的地方)
  linalg-tile (tile sizes: 32, 32, 32)
  linalg-fuse
  one-shot-bufferize
  convert-linalg-to-loops

// 3. 底层与硬件优化
  scf-vectorize
  affine-loop-unroll
  convert-vector-to-llvm
  convert-memref-to-llvm

// 4. 收尾
  convert-scf-to-cf
  convert-arith-to-llvm
  convert-func-to-llvm
  reconcile-unrealized-casts

在这个过程中，你可以清晰地看到：

1. Dialect 是逐渐混合又逐渐统一的： 一开始全是 torch，然后变成 tensor+linalg，再加入 scf，再加入 memref，最后全部统一为 llvm。
2. 不是所有 Dialect 都走完全程： 比如 torch dialect 在阶段二就消失了，linalg 在阶段三消失了。它们只在最适合自己的层级存在，做完优化就"功成身退"。
3. 优化是分层进行的：
   • 在 linalg 层面做算子融合（图优化）。
   • 在 scf 层面做循环分块（访存优化）。
   • 在 vector 层面做SIMD 指令生成（指令级优化）。

这就是 MLIR 设计哲学的魅力：让最适合的 Dialect 在最合适的阶段做最擅长的事。