在实际框架中集成 MLIR（以 IREE 为例）

前面十六章都在讲 MLIR 本身的机制和用法。这一篇换个视角——看一个真实的项目如何将 MLIR 作为编译基础设施来用。

IREE（Intermediate Representation Execution Environment）是 Google 维护的开源端到端 ML 编译器与运行时。它的编译流程完全基于 MLIR，是目前 MLIR 在 AI 编译器领域最全面的应用实例。

一、IREE 的定位

IREE 解决的问题很具体：将 TensorFlow / PyTorch / JAX 训练好的模型，编译到移动端、边缘设备、嵌入式系统上高效运行。

这是一个已有产品覆盖的领域——TensorFlow Lite、ONNX Runtime、Apple CoreML、NVIDIA TensorRT 都做端侧推理。IREE 的差异化在于：

1. 基于 MLIR 的编译管线：所有优化 Pass 是 MLIR Pass，方言是 MLIR Dialect——没有自研的”另一种 IR”。
2. 统一的 HAL 运行时：通过硬件抽象层（HAL），同一套编译管线生成的代码可以在 CPU（LLVM）、GPU（Vulkan/SPIR-V）、Metal（Apple GPU）上运行。
3. 从训练框架到设备代码的完整编译：不像 TVM 那样依赖用户手动调度，也不像 TensorRT 那样是黑箱优化。

二、IREE 的编译流程

输入模型 (.mlir / .tflite / .onnx)
   │
   ▼
┌─────────────────────────────────────────────────┐
│ 前端导入                                         │
│ PyTorch → Torch-MLIR → MHLO 方言                 │
│ TF → MHLO / StableHLO 方言                       │
│ JAX → StableHLO 方言                             │
└─────────────────────────────────────────────────┘
   │
   ▼
┌─────────────────────────────────────────────────┐
│ IREE 编译器 (iree-compile)                      │
│                                                  │
│ 1. 输入方言 (mhlo/stablehlo/torch) → linalg       │
│ 2. Linalg 优化 (tiling, fusion, vectorization)    │
│ 3. Bufferization (tensor → memref)               │
│ 4. 降阶到目标方言 (llvm / spirv / vmvx)           │
│ 5. 生成 .vmfb (IREE FlatBuffer)                  │
└─────────────────────────────────────────────────┘
   │
   ▼
.vmfb 文件 ──→ IREE 运行时 (iree-run-module)
                  │
                  ├── CPU 后端 (llvm-cpu)
                  ├── GPU 后端 (vulkan-spirv)
                  ├── Metal 后端 (metal-spirv)
                  └── VMVX 后端 (微型解释器)

三、编译命令实战

以下命令展示 IREE 的典型调用方式（标志名与后端选项随 IREE 版本变化，以 IREE 文档 为准）。本篇未附特定版本的实测输出——读者应在本地 iree-compile --help 核对后再运行。

3.1 从 MLIR 输入编译

# 编译一个简单的 MLIR 程序到 VMVX（IREE 的轻量虚拟机解释器）
iree-compile \
  --iree-hal-target-backends=vmvx \
  input.mlir \
  -o output.vmfb

# 编译到 CPU（通过 LLVM JIT）
iree-compile \
  --iree-hal-target-backends=llvm-cpu \
  --iree-llvmcpu-target-triple=x86_64-linux-gnu \
  input.mlir \
  -o output.vmfb

# 编译到 Vulkan GPU
iree-compile \
  --iree-hal-target-backends=vulkan-spirv \
  --iree-vulkan-target-triple=rdna3-unknown-linux \
  input.mlir \
  -o output.vmfb

3.2 运行编译产物

# 用 IREE 的 VMVX 后端运行
iree-run-module \
  --module=output.vmfb \
  --function=main \
  --input=@input_tensor.npy

# 用 trace 模式观察执行过程
iree-run-module \
  --module=output.vmfb \
  --function=main \
  --trace_execution

四、HAL：硬件抽象层

HAL 是 IREE 运行时设计的核心。它抽象了三个概念：

1. 设备（Device）：CPU、GPU 的物理设备表示。HAL 管理和发现设备、分配内存、提交任务。
2. 命令缓冲区（Command Buffer）：记录在指定设备上执行的操作序列（dispatch、copy、barrier）。编译管线生成的就是填满 HAL command buffer 的代码。
3. 内存缓冲区（Buffer）：设备内存的抽象。HAL 管理不同设备间的内存传输和同步。

// HAL 使用示例（概念层面，非可直接编译的完整摘录）
iree_hal_device_t *device;
iree_hal_create_device(..., &device);

iree_hal_buffer_t *input, *output;
iree_hal_allocator_allocate_buffer(..., &input);

iree_hal_command_buffer_t *cmd;
iree_hal_create_command_buffer(..., &cmd);

// 记录 dispatch
iree_hal_command_buffer_dispatch(cmd, executable, entry_point, ...);
iree_hal_command_buffer_execution_barrier(cmd, ...);

// 提交到设备执行
iree_hal_device_queue_execute(device, cmd, ...);

HAL 的设计来自 Vulkan 的命令缓冲模型——Vulkan 的 vkCmdDispatch、vkCmdPipelineBarrier 等概念在 HAL 中有直接对应。这意味着 IREE 对 Vulkan 的映射是最自然的，对 CPU 的映射则需要将 HAL command buffer 翻译为 LLVM JIT 的函数调用。

五、IREE 的 MLIR 方言栈

IREE 在标准 MLIR 方言上构建了三个自定义方言：

5.1 Flow 方言

Flow 方言将标准 MLIR 方言（mhlo、linalg）的 IR 转换为 IREE 的专用表示——执行 dispatch region 的分组、输入/输出 tensor 的标注、内存布局的初步决策。

// flow.dispatch.region — 将可独立执行的 Op 分组
flow.dispatch.region @matmul_region {
  %0 = linalg.matmul ins(...) outs(...) -> ...
  flow.return %0
}

5.2 HAL 方言

HAL 方言将 dispatch regions 映射到 HAL API 调用——分配设备 buffer、填 command buffer、调度执行：

// hal.command_buffer.dispatch — 记录 dispatch 到 command buffer
hal.command_buffer.dispatch @exe::@matmul_entry[%grid_x, %grid_y, %grid_z]
    (%arg0, %arg1, %arg2 : !hal.buffer, !hal.buffer, !hal.buffer)

5.3 VM 方言

VM 方言是 IREE 的轻量虚拟机字节码——一个低级、高度定制的 IR，用于在运行时执行编译后的计算：

// vm.call @exe::@matmul_entry(%buf0, %buf1, %buf2)
vm.call @exe::@matmul_entry(%buf0, %buf1, %buf2)

这三个方言的降阶关系：

linalg / mhlo
    │
    ▼ (flow 方言——分组和标注)
flow.dispatch.region
    │
    ▼ (hal 方言——设备抽象)
hal.command_buffer.dispatch
    │
    ▼ (vm 方言——虚拟机字节码)
vm.call
    │
    ▼ (序列化为 FlatBuffer)
.vmfb 文件

六、IREE 的 Linalg Tiling 策略

IREE 的 GPU tiling 是通过 Linalg 和 Transform 方言驱动的——用户或 auto-scheduler 可以指定每个 linalg.generic / linalg.matmul 的分块方式和内存提升策略。

IREE 提供了一组默认的 tiling 配置（针对不同硬件预设），根据目标后端自动选择：

# 指定 tile 大小（用于 CPU）
iree-compile --iree-llvmcpu-target-arch=x86_64-avx512f \
  --iree-flow-dispatch-tile-sizes=f32=16,16,32

# 指定 GPU 的 workgroup 大小
iree-compile --iree-vulkan-target-triple=rdna3 \
  --iree-flow-dispatch-tile-sizes=f32=64,64,16

七、部署链路

编译产物 .vmfb 可以部署到多种平台：

部署流程：

# 1. 编译到目标后端
iree-compile --iree-hal-target-backends=vulkan-spirv \
  model.mlir -o model.vmfb

# 2. 集成到移动应用
# Android: 通过 IREE 的 Java/Kotlin API
# iOS: 通过 IREE 的 Swift API 或 C API

# 3. 运行时调用
iree_run_module(model.vmfb, inputs, &outputs)

八、IREE 与其他方案的对比

IREE 的最大优势：MLIR 的方言可组合性意味着新的硬件后端和新的优化 Pass 可以以模块化方式添加，而不需要修改整个编译管线。一个新的 NPU 芯片只需要实现对应的 HAL 后端和设备内存管理——Linalg 优化管线和 HAL 抽象层都能共享。

九、学习 IREE 的推荐路径

1. 从 iree-compile 编译一个简单的 linalg.matmul 开始，逐步添加 tiling 和优化 flags。
2. 用 --mlir-print-ir-after-all 观察每个 Pass 后的 IR 变化。
3. 阅读 IREE 的 flow 方言源码，理解 dispatch region 分组算法。
4. 阅读 IREE 的 hal 方言源码，理解设备抽象和 command buffer 模型。
5. 最后看 VM 方言——它是 IREE 运行时的最终执行层。

参考资料

社区项目（B 级）

• IREE — https://github.com/iree-org/iree
• IREE 文档 — https://iree.dev/
• IREE Compiler Overview — https://iree.dev/developers/design-docs/compiler-overview/

论文（A 级）

• Lattner, C. et al. MLIR: A Compiler Infrastructure for the End of Moore’s Law. arXiv:2002.11054, 2020.（IREE 是 MLIR 论文中提及的核心用例之一）