MLIR 编译器基础设施：多级中间表示与 AI/ML 系统编译的全栈工程

引言：为什么 MLIR 正在统一编译器世界

2017 年 Google 发布 IREE（Intermediate Representation Execution Environment），2019 年 LLVM 社区正式接纳 MLIR 为官方子项目，到 2024 年 MLIR 已成为 TensorFlow/XLA、PyTorch inductor、oneAPI、Triton 等 AI 基础设施的核心编译层。这不是偶然。

传统编译器采用单一 LLVM IR 表示程序，从高层语义到硬件指令的 lowering 过程中丢失了大量领域知识。当面对 AI 工作负载（张量卷积、注意力分块、算子融合）时，编译器无法在高层表达这些结构，只能在已经降级为基本块和内存操作的 LLVM IR 上苦苦重建高级语义——这就是所谓的"语义下降悬崖"（Semantic Cliff）。

MLIR（Multi-Level Intermediate Representation）通过可组合方言（Dialects）体系解决了这个问题：每个层级保留自己的语义信息，通过一系列 lowering pass 逐步降级，让每一层优化都能在最适合的抽象级别上工作。

本文将从 MLIR 的 IR 设计哲学切入，深入方言体系、Pass 基础设施、转换流水线，并以实际代码示例展示如何构建自定义 MLIR Pass 实现算子融合优化。

一、MLIR 设计哲学：方言、操作与类型系统

1.1 核心抽象：Operation

MLIR 中一切皆操作（Operation）。不像 LLVM IR 中只有指令、基本块、函数三种结构体，MLIR 通过 Operation 统一抽象，每个操作包含操作名（命名空间 + 操作名，如 arith.addi）、操作数（SSA 值引用列表）、结果（SSA 值产出列表）、属性（编译期常量如张量维度、广播标志）、区域（用于嵌套结构化控制流，如 scf.for、linalg.generic）。

以下是一个 Gemm 操作的 MLIR 表示示例：

%result = linalg.matmul
  ins(%A, %B : tensor<1024x256xf32>, tensor<256x512xf32>)
  outs(%C : tensor<1024x512xf32>)
  -> tensor<1024x512xf32>

这里 linalg.matmul 保留了矩阵乘法语义，直到 lowering 到循环之前，编译器都能识别并应用 matmul 专用优化（分块、向量化、共享内存缓存）。

1.2 类型系统：结构化张量

MLIR 支持静态形状张量（如 tensor<4x8x128xf32>）、动态维度（tensor<?x?xf32>）、带地址空间的内存引用（memref<1024x1024xf32, 1>）和外部类型（如 !llvm.ptr）。类型系统与方言解耦，类型本身可以属于某个方言（如 !stablehlo.token），这种开放性是 LLVM IR 无法实现的。

1.3 属性和特性

属性承载编译期元数据，如索引映射（affine_map）和 tile 尺寸。特性（Traits）则提供结构化约束，如 SingleBlockTerminator、HasParent<"ModuleOp">、AttrSizedOperandSegments 等，让 Pass 基础设施能自动验证操作合法性。

二、核心方言体系：从高层到低层

MLIR 的精髓在于方言的分层组织，每个方言代表一个抽象层级：

2.1 算术层：arith 与 math

arith 方言是高层到循环层的桥梁，提供整数/浮点运算（如 arith.addi、arith.mulf）和比较操作（arith.cmpi）。math 方言则提供超越函数（如 math.sqrt、math.exp）和近似计算支持。

2.2 结构化控制流：scf

scf 方言使用结构化控制流（Structured Control Flow），保证每个循环/条件分支都有单个入口和出口，极大简化了后续分析和转换。典型构造包括 scf.for（计数循环）、scf.if（条件分支）、scf.while（通用循环）等。

2.3 张量运算层：linalg

linalg 方言是 MLIR 在 AI 领域最重要的贡献，它用索引映射（Indexing Maps） 统一描述张量运算。一个 linalg.generic 操作包含：

• indexing_maps：定义每个循环维度在不同张量中的使用方式
• iterator_types：标记并行（parallel）或归约（reduction）维度
• 循环体内的计算逻辑（通过基本块描述）

这种表示让编译器自动知道：可以对并行维度做分块并行，对归约维度做向量化优化。

2.4 向量抽象层：vector

vector 方言是硬件向量指令（AVX-512、NEON、SVE）的统一抽象，提供 transfer_read/write（内存到向量寄存器映射）和 vector.contract（向量矩阵乘法）等操作，是分层 lowering 的关键节点。

三、Lowering 流水线：层级渐进式转换

MLIR 的设计目标是让每一层优化都在最合适的信息密度下执行。以下是一个典型的 AI 工作负载 lowering 路径：

1. Frontend 层（TensorFlow/XLA/PyTorch）：将计算图转换为稳定方言
2. StableHLO/MHLO 方言：保持高层算子语义不变
3. Linalg 具名算子（如 linalg.batch_matmul）：将高层映射到近硬件层
4. Linalg Generic：统一索引映射表示，进入算子优化核心层
5. Affine 方言：多面体模型优化（循环交换、分块、展开）
6. SCF + Memref：结构化控制流和内存访问
7. Vector 方言：自动向量化
8. LLVM IR：标量优化和硬件无关优化
9. 机器代码生成

融合优化实例

稀疏注意力层涉及多个非线性算子（matmul、scale、softmax、matmul、layernorm），传统编译器需要多级 lowering 才能完整优化。MLIR 在 linalg 方言层能一次性看到整个算子图，执行逐元素算子融合和归约融合，将整个 attention layer 合并为单个 linalg.generic 操作，大幅减少内存带宽占用。

四、Pass 基础设施：变换引擎

MLIR 的 Pass 系统是编译优化的执行框架，支持多种粒度的 Pass：

• FunctionPass：按函数粒度处理（如循环融合）
• ModulePass：按模块粒度处理（如链接优化）
• OperationPass：按特定操作类型处理（如 LinalgTilingPass）
• 转换 Pass：方言间转换（如 ConvertLinalgToLLVM）

MLIR 使用声明式重写系统（DRR），通过 TableGen 模式匹配声明 Pass 逻辑，自动生成匹配和替换代码。这种声明式风格让优化 Pass 的编写更加可靠和可维护。

自定义 Pass 开发：批处理矩阵乘法拆分

当 NPU 对小 batch 矩阵乘法的处理效率较低时，可以编写一个自定义 Pass 自动拆分为多个独立的 matmul 再拼接结果。核心逻辑为：遍历函数内的所有 linalg.batch_matmul 操作，判断 batch 维度和矩阵大小，若 batch 过小而矩阵过大，则沿 batch 维度切片并逐个计算后拼接。

这种 Pass 通常只需 200-300 行 C++ 代码就能完成，编译后通过 mlir-opt 的 --pass-name 参数即可集成到现有编译流水线中。

五、生产部署：MLIR 在真实 AI 系统中的位置

5.1 TensorFlow/XLA

XLA 使用 StableHLO 方言作为输入 IR，通过 MLIR 流水线降级到 HLO Target IR，最终经 LLVM 生成 GPU/TPU/CPU 代码。完整路径为：TF Graph → StableHLO → MHLO→Linalg → Affine/SCF → Vector → LLVM IR → GPU PTX → CUBIN。

5.2 PyTorch 2.x Inductor

PyTorch 的 torch.compile 直接通过 Inductor 生成 halide/triton 方言的 MLIR IR，绕过传统计算图。对于 scaled_dot_product_attention 等高层算子，内部会直接生成包含 flash_attention 等操作的 MLIR IR，充分利用 MLIR 的优化 Pass。

5.3 Triton-to-MLIR 编译路径

OpenAI Triton 编译器实际上是一个 MLIR 前端，把用户定义的 Triton DSL 翻译到 triton 方言，再经 triton_gpu 到 PTX 的 lowering。用户编写的核函数自动获得 MLIR 提供的自动内存合并优化和共享内存管理，而 PTX codegen 则直接复用 LLVM 的 NVPTX 后端。

六、面向自定义加速器的方言设计

设计一个 AI 加速器的 MLIR 方言需要考虑四个维度：操作集定义（用 TableGen TD 文件描述操作名、参数类型、结果类型、 verifier 验证逻辑）、自定义 lowering（将高层算子映射到专用操作）、完整工具链构建（前端→方言→优化→后端→Runtime）。

通过这种方式，新硬件可以在数周内完成从概念到可运行编译器后端的开发，而传统方法通常需要数月的时间。

七、实战案例：NPU 卷积切片优化 Pass

为自研 NPU 设计卷积切片优化 Pass 的典型流程如下：

1. 收集目标：在模块内遍历所有 linalg.Conv2DNhwcHwcfOp 操作
2. Tile 分割：根据 SRAM 容量（如 256x256 Byte）设置 tile 尺寸
3. 循环生成：通过 scf::tileReductionLoopWithScf 生成嵌套循环
4. 流水线调度：添加 NPU 特定的双缓冲或三级流水调度原语
5. DMA 传输：在循环外层插入 DMA 搬运原语，实现计算与访存重叠

经过这些步骤，原本在 NPU 上运行缓慢的卷积层可以获得数倍的性能提升。整个过程可以注册为 mlir-opt 的一个标准 Pass，与其他 Pass 自由组合使用。

八、MLIR 与 GPU 编程模型的对比分析

MLIR 的独特定位是作为"编译器的基础设施"——它本身不直接面向开发者编程，而是为 Triton、XLA、IREE 等前端提供统一的中间层，让不同的硬件后端可以共享相同的高层优化框架。

总结：MLIR 在 AI 基础设施中的角色演进

回顾编译器史，每一次中间表示的重大革新都催生了技术浪潮：GCC RTL（1987）开启了开源编译器革命，LLVM IR（2003）推动了编译器工程化时代，而 MLIR（2020）正在构建 AI 和异构计算的编译基础设施。

MLIR 的核心贡献不是发明了新的优化算法，而是创建了一个可扩展的编译基础设施，让领域专家可以在自己的抽象层级上表达优化知识，而不必从头构建整个编译器。

对于系统工程师而言，掌握 MLIR 意味着：能够为自己的硬件设计专用的编译路径、在 AI 推理服务中实现算子级别的极致优化、参与构建下一代 AI 编译器栈。这不只是编译技术问题——在 AI 越来越依赖硬件效率的时代，谁能更好地连接高层算子底层硬件，谁就能掌握 MLIR 时代的核心竞争力。

关键要点： - MLIR 通过可组合方言解决语义下降悬崖问题 - linalg 方言通过索引映射统一描述张量运算 - 分层 lowering 让每层优化在合适的信息密度下执行 - 自定义 Pass 平均只需 200-300 行 C++ 就能实现专业级优化 - MLIR 正在成为 AI 编译器的事实标准，覆盖从 TPU 到自定义 NPU 的全硬件生态