MLIR多层次中间表示与编译器优化基础设施：从Dialect体系到CIRCT硬件设计革命

引言：为什么需要多层次中间表示

传统编译器通常采用单一层次的中间表示（IR），最典型的是LLVM IR。然而随着AI框架、领域专用语言（DSL）和异构硬件的爆发式增长，单一IR面临根本性瓶颈：它要么过于底层以至于高层优化信息丢失，要么过于抽象以至于难以映射到多样化的硬件后端。

MLIR（Multi-Level Intermediate Representation）是Google于2019年开源的编译器基础设施，核心创新在于：不定义一种IR，而是定义一套可组合的IR构建块，让不同抽象层次的表示共存于同一个编译流水线中，各层次之间通过渐进式 lowering 相互衔接。

本文将深入解析MLIR的体系架构、核心概念（Dialect/Operation/Type/Attribute）、优化机制（Pass/Pattern Rewriter）、以及实际工程应用，展示它如何将编译器开发从"每次从头构建"转变为"搭积木式组装"。

一、MLIR 核心设计理念

1.1 Dialect：IR 的模块化单元

MLIR 中最核心的概念是 Dialect（方言）。每个 Dialect 定义了一组相关的 Operation、Type 和 Attribute，对应特定的抽象层次。例如：

• builtin Dialect：基础类型（IntegerType, IndexType, FunctionType, FloatType）和 Attribute（StringAttr, IntegerAttr, ArrayAttr, DenseElementsAttr）
• func Dialect：函数定义与调用（func.call, func.return, func.func）
• arith Dialect：算术运算（arith.addi, arith.muli, arith.cmpi 等基本整数/浮点操作）
• scf Dialect：结构化控制流（scf.for, scf.if, scf.while，作为高级循环表示）
• memref Dialect：内存引用（memref.alloca, memref.load, memref.store，抽象内存操作模型）
• affine Dialect：仿射约束（affine.for 带仿射边界，支持多面体编译优化）
• linalg Dialect：线性代数运算（linalg.matmul, linalg.conv_2d，张量级别的操作原语）
• tensor Dialect：不可变张量（tensor.empty, tensor.extract，SSA级别的张量值语义）
• vector Dialect：SIMD 向量类型（vector.contract, vector.transfer_read/write）
• gpu Dialect：GPU 编程模型（gpu.launch, gpu.barrier, thread/block/grid 层级）
• llvm Dialect：LLVM IR 的1:1映射，用于最终 lowering 为 LLVM bitcode

关键设计决策：Operation 是 SSA（Static Single Assignment）值，每个 Operation 产生零或多个结果值，消费零或多个操作数。这与 LLVM IR 的 SSA 形式一脉相承，但 MLIR 允许自定义任意复杂的 Operation——一个 Operation 可以有自己的 region（嵌套区域）、symbol table（符号表）、甚至自定义的解析/打印格式。

1.2 Operation：通用 IR 节点

MLIR 中一切计算表达为 Operation。与传统三地址码不同，MLIR Operation 是高度结构化的：

// MLIR Operation 的一般语法格式
%result = "dialect.op_name"(%operand1, %operand2) {
  // 可选的 Region（嵌套代码区域）
  ^bb0(%block_arg : type):
    "scf.yield"(%value) : (type) -> ()
} {attribute_name = value} : (input_types) -> (result_types)

具体示例：

// 一个带嵌套 region 的 affine.for 操作
%result = "affine.for"(%c0) = (0) to (128) {
  %i = "arith.index_cast"(%i64) : (index) -> i64
  %val = "memref.load"(%buffer, %i) : (memref<256xf32>, i64) -> f32
  %sq = "arith.mulf"(%val, %val) : (f32, f32) -> f32
  "memref.store"(%sq, %output, %i) : (f32, memref<256xf32>, i64) -> ()
} : (index) -> f32

这种统一表示使得 MLIR 能够同时容纳从高层张量运算到底层硬件指令的所有抽象层次——不同层次的 Operation 可以通过定义转换规则互相 lowering。

1.3 Type System：可扩展的静态类型

MLIR 的类型系统是可扩展的，分为以下几大类：

• 基础类型：IntegerType（支持任意位宽如 i1, i8, i32, i64）、FloatType（f16/f32/f64/bf16）、IndexType（目标平台指针宽度的无符号整数）
• 复合类型：MemRefType（内存引用，含布局映射和内存空间）、RankedTensorType（定秩张量）、VectorType（SIMD 向量）、TupleType（值元组）
• 函数类型：FunctionType（参数类型列表 + 返回类型列表）
• 不透明类型：通过 Dialect 扩展的任意自定义类型（如 llvm.ptr 表示指针）

Type 也可以携带参数——类型参数在编译期已知，值为运行时确定。例如 memref<4x8xf32> 中的 4 和 8 是编译期参数，这允许编译器基于类型信息进行静态内存分配和向量化决策。

1.4 Attribute：编译期元数据

Attribute 是编译期的静态元数据，与运行时的 Value 相对应。常见 Attribute 类型包括：

• StringAttr（字符串字面量）
• IntegerAttr（编译期整数值）
• FloatAttr（编译期浮点值）
• ArrayAttr（Attribute 数组）
• DictionaryAttr（键值对字典）
• DenseElementsAttr（密集张量常量值，如权重矩阵）
• SymbolsRefAttr（符号引用，用于链接）

Attribute 可以附加到 Operation 上作为编译期配置。例如 linalg.matmul 可以通过 Attribute 指定是否在转置后的操作数上执行，或者 gpu.launch 通过 Attribute 指定 thread/block 维度。

二、编译流水线：多阶段渐进式 Lowering

2.1 Lowering 流水线全览

MLIR 编译器的核心优势在于能够将一个高层计算表示通过多阶段 lowering 逐步转换为底层机器代码。以下是一个典型的 AI 模型编译流水线（以 TensorFlow/MLIR 为例）：

[图级表示 TF/XLA] 
    → [StableHLO Dialect]        // 高级算子语义（matmul, conv, softmax）
    → [Linalg Dialect + Tensor]  // 张量级别线性代数操作
    → [Linalg → SCF + MemRef]    // 循环化（tiling/fusion/unrolling）
    → [SCF → Affine]             // 仿射化（多面体优化）
    → [Affine → LLVM]            // 指针+循环的底层表示
    → [LLVM IR] → [目标代码]      // 传统 LLVM 后端优化+代码生成

每一步 lowering 不仅是格式转换，更是优化机会介入的缝合点。不同层次适合不同类别的优化：

• 高层（Linalg/Tensor）：算子融合（fusion）、内存规划、数据布局变换
• 中层（Affine/Loop）：循环交换（loop interchange）、分块（tiling）、多面体优化、向量化
• 底层（LLVM/Vector）：指令选择、寄存器分配、SIMD 指令映射、窥孔优化

2.2 Lowering Pass 与转换框架

MLIR 中 lowering 是通过 Pass 机制实现的。两种 Pass 类型：

• Operation Pass：针对特定 Operation 类型的局部转换（如 ConvertLinalgToLoops 仅处理 linalg.generic）
• Module Pass：模块级别的全局转换（如内联所有函数调用、全局死代码消除）

Pass 可以声明依赖关系和调度顺序，形成 Pass Pipeline。例如：

// 典型 pass pipeline 构造
pm.addPass(mlir::createConvertLinalgToLoopsPass());       // Linalg→SCF
pm.addPass(mlir::createConvertSCFToCFPass());            // SCF→CFG (分支/跳转)
pm.addPass(mlir::createConvertSCFToOpenMPPass());        // 循环并行化到 OpenMP
pm.addPass(mlir::createConvertVectorToLLVMPass());       // Vector类型 lowering
pm.addPass(mlir::createConvertFuncToLLVMPass());         // 函数调用约定转换
pm.addPass(mlir::createReconcileUnrealizedCastsPass());  // 清理过渡类型

三、模式重写系统：MLIR 优化的核心引擎

3.1 Pattern Rewriter 与 RewritePattern

MLIR 的优化操作通过 Pattern Rewriter 框架进行。每个优化规则是一个 RewritePattern，负责匹配特定的 Operation 子图并进行替换。重写分为两种模式：

• Greedy Pattern Rewrite：贪婪式全局扫描，反复应用模式直到收敛。可配置遍历顺序（从内到外/从外到内）、最大迭代次数、是否应用受益分析（benefit analysis）
• Conversion Pattern：一个 Operation Type A → Type B 的完整类型集转换（Dialect 级别的 lowering），由驱动框架保证类型一致性

3.2 多面体编译：affine Dialect 的独特优化

affine Dialect 是 MLIR 中最具特色的优化层次。它要求所有循环边界和数组索引表达为仿射函数（Affine Function）：

// affine.for：循环边界必须是仿射表达式
affine.for %i = #map0(%N) to #map1(%M) {
  // 数组索引也必须是仿射表达式
  %val = memref.load %A[%i, %c0] : memref<?x?xf32>
}

仿射约束使得编译器可以应用多面体模型（Polyhedral Model）进行精确的数据依赖分析。基于此可以实现：

• 自动循环交换（改变遍历顺序以优化缓存局部性）
• 自动循环分块（tile），改善 L1/L2 缓存重用
• 自动并行化（识别可并行 loop band 并标记为 parallel）
• 自动向量化（基于仿射向量化分析）
• 数组填充（padding）消除 bank conflict

虽然 affine Dialect 的表达力受限于仿射约束（无法处理间接索引或非线性边界），但对于科学计算和深度学习中占90%以上的规则循环模式，它提供了精确且高效的优化路径。

3.3 算子融合：Linalg 带来的组合式优化

linalg Dialect 的设计灵感来自"算子融合"的需求。在深度学习推理中，element-wise 操作（如ReLU、bias add）和 reduction 操作（如matmul）可以融合到同一个 kernel 中执行，避免中间结果的内存分配和读写。

linalg.generic 操作通过索引映射（indexing map）声明输入输出张量的映射关系：

// linalg.matmul: C[i,j] = sum_k(A[i,k] * B[k,j])
%result = linalg.matmul
  ins(%A, %B : tensor<MxKxf32>, tensor<KxNxf32>)
  outs(%C : tensor<MxNxf32>)
  -> tensor<MxNxf32>

// linalg.generic: element-wise 操作（如 ReLU）
%relu = linalg.generic {
  indexing_maps = [affine_map<(d0,d1) -> (d0,d1)>,  // input
                   affine_map<(d0,d1) -> (d0,d1)>], // output
  iterator_types = ["parallel", "parallel"]
} ins(%input : tensor<HxWxf32>) outs(%output : tensor<HxWxf32>)
{%arith_maxf(%in, %zero) : (f32, f32) -> f32
  linalg.yield %max : f32
}

基于索引映射的结构化表示，MLIR 可以自动推导融合规则：一个 element-wise 的 linalg.generic 可以融合进消费者（如matmul）的循环体中，共享外层循环索引。

四、符号表与模块化编译

4.1 Symbol Table 与 Symbol Reference

MLIR 内置了 Symbol Table 机制用于跨模块引用。一个 Operation 可以携带 SymbolTable 属性（如 ModuleOp、FuncOp），允许其他 Operation 通过 SymbolRefAttr 引用其中定义的符号。这在以下场景至关重要：

• 跨函数引用（函数指针、虚函数表）
• 跨模块链接（如链接外部定义的库函数）
• 硬件模块层次化实例化（如 CIRCT 中的硬件组件实例化）

4.2 部分求值与编译缓存

MLIR 支持部分求值（Partial Evaluation）：对已知编译期常量的计算子图提前求值。配合 MLIR 的 Module 缓存和 Pass 缓存机制，可以实现增量编译——只有变更的模块才需要重新编译和优化。

这对于 AI 模型的 JIT（Just-In-Time）推理引擎尤为重要：通常模型权重在推理时不变，可以被完全编译期绑定和优化，仅保留运行时输入维度作为动态参数。

五、CIRCT：硬件设计领域的 MLIR 革命

如果说 MLIR 在 AI 编译器领域的成功已广为人知，那么 CIRCT（Circuit IR Compilers and Tools）项目则将 MLIR 的 Dialect 思想带到了硬件设计自动化（EDA）领域。

CIRCT 定义了以下核心 Dialect：

• hw Dialect：硬件模块实例化（hw.module, hw.instance, hw.output）
• comb Dialect：组合逻辑运算（comb.add, comb.and, comb.xor, comb.mux）
• seq Dialect：时序逻辑（seq.compreg 表示寄存器，seq.fsm 表示有限状态机）
• sv Dialect：SystemVerilog 的直接映射（用于与现有 EDA 工具流集成）
• fsm Dialect：有限状态机的高层抽象（状态/转移/动作），可 lowering 为 seq 寄存器传输级
• esi Dialect：弹性 Streaming Interface，片上网络（NoC）通信层

通过 CIRCT，硬件设计中的高级综合（HLS）、形式验证、仿真和物理实现可以共享同一套 IR，消除传统流程中 Verilog/VHDL → 综合 → 布局布线 各阶段的信息丢失问题。例如 Chisel/SpinalHDL 可以编译为 Combone IR 再通过 CIRCT lowering 为可综合的 SystemVerilog。

六、实战：从零构建一个 MLIR Toy 编译器

6.1 Toy 语言定义

MLIR 项目官方提供了 Toy 语言教程，展示如何为一种简单的数组计算语言构建完整编译器。Toy 语言支持：

# 变量赋值
var a = [[1, 2, 3], [4, 5, 6]];  # 2D tensor 字面量
var b = a * 2.0;                    # broadcast 元素乘

# 转置
var c = transpose(b);

# 函数定义
def multiply_transpose(a, b) {
  return transpose(a) * transpose(b);
}

# 外部声明（运行时库）
extern func print(a);

6.2 自定义 Toy Dialect

编译器前端通过 MLIR 的方言机制定义 Toy Dialect，包含以下 Operation：

• toy.struct_access：结构体字段引用
• toy.mul：乘法运算（支持 broadcast）
• toy.transpose：转置操作
• toy.constant：常量张量（携带 DenseElementsAttr）
• toy.return：函数返回
• toy.generic_call / toy.call：函数调用
• toy.print：内建打印操作（用于调试输出）

6.3 Lowering 流程分析

Toy 编译器的 lowering 流水线展示了多阶段 lowered 的典型模式：

// Step 1: Toy Dialect（高层，保留 transpose、print 等高层语义）
func.func @multiply_transpose(%arg0: tensor<*xf64>, %arg1: tensor<*xf64>) {
  %0 = "toy.transpose"(%arg0) : (tensor<*xf64>) -> tensor<*xf64>
  %1 = "toy.transpose"(%arg1) : (tensor<*xf64>) -> tensor<*xf64>
  %2 = "toy.mul"(%0, %1) : (tensor<*xf64>, tensor<*xf64>) -> tensor<*xf64>
  "toy.return"(%2) : (tensor<*xf64>) -> ()
}

// Step 2: Lower to Affine（消除动态形状，转置变为显式循环）
func.func @multiply_transpose(%arg0: memref<3x2xf64>, %arg1: memref<3x2xf64>) {
  %alloca = memref.alloca() : memref<2x3xf64>
  // ... tiling 和 loop interchange 优化 ...
  affine.for %i = 0 to 3 {
    affine.for %j = 0 to 2 {
      %v = memref.load %arg0[%j, %i] : memref<3x2xf64>  // transpose
      memref.store %v, %alloca[%i, %j] : memref<2x3xf64>
    }
  }
  // ... 乘积循环化 ...
}

// Step 3: Lower to LLVM Dialect
llvm.func @multiply_transpose(%arg0: !llvm.ptr, %arg1: !llvm.ptr) {
  // pointer arithmetic + branch-based loop
  llvm.br ^bb1(%c0_i32 : i32)
^bb1(%i : i32):
  %cond = llvm.icmp "slt"(%i, %c3_i32) : (i32, i32) -> i1
  llvm.cond_br %br, ^bb2, ^bb5
  // ...
}

值得注意的是，Shape Inference Pass 在 lowering 过程中自动推导运行时未知维度：通过符号分析传播张量形状，将 tensor<*xf64> 转变为具体秩类型 tensor<3x2xf64>，使得后续内存预分配成为可能。

七、MLIR vs 传统编译器框架对比

维度	LLVM IR	MLIR
IR 层次	单一层次，面向通用CPU	多层次 Dialect 体系，支持任意抽象层级
扩展性	修改 IR 核心需重新编译整个LLVM	新增 Dialect 为独立模块，零侵入式扩展
高层语义保留	高层信息（循环、张量结构）在 lowering 后丢失	各层次独立保留语义，可随时视察
优化复用	优化 pass 硬编码于 C++ 模块	声明式模式重写，可组合成 pass pipeline
编译速度	LLVM 驱动，高度优化后的 lower 开销确定	多层 lowering 有额外开销，但可缓存部分编译结果
学习曲线	相对成熟稳定，生态丰富	概念更多（Dialect/Region/OpTrait），初始门槛更高

MLIR 并不替代 LLVM IR，而是在其之上提供更高层次的编译器构造能力。最终 MLIR 编译器通常仍会将 LLVM Dialect 转换为标准 LLVM IR，利用 LLVM 成熟的指令选择和后端代码生成能力。两者的关系可类比为：MLIR 是"编译器构造框架中的编译器"，LLVM IR 是"编译器后端的汇编语言"。

八、生态与前沿方向

8.1 IREE：可执行的运行时环境

IREE（Intermediate Representation Execution Environment）是基于 MLIR 的端到端编译器与运行时，专为 AI 推理/训练设计。其流程：TensorFlow/PyTorch → StableHLO → Linalg/Vectors → LLVM SPIR-V/Vulkan/CUDA → 原生机器代码。IREE 最终生成可直接部署的精简运行时镜像（从几MB缩至几十KB），嵌入式设备和移动端推理尤为重要。

8.2 Polygeist：C/C++ 到 MLIR 的编译桥梁

Polygeist 将 C/C++（通过 Clang AST）提升（raising）到 MLIR 的 Affine/Linalg 层次，使得传统循环密集型科学计算代码也能享受多面体优化和 MLIR优化流水线的优势。这种"旧代码新编译"策略让 HPC 和数值计算社区无需重写代码即可获得 MLIR 生态的优化能力。

8.3 ONNX-MLIR：AI 模型标准化入口

ONNX-MLIR 直接将 ONNX（Open Neural Network Exchange）格式的 AI 模型编译为优化的原生代码，绕过框架运行时开销。它通过 ONNX Dialect 定义了完整的神经网络算子集，并 lowering 到 Affine/Linalg/LLVM，支持 CPU/GPU/TPU 多后端。

8.4 自动调优与搜索：MLIR + 强化学习

MLIR 的可组合优化 pipeline 天然适合自动调优（AutoTuning）：将每个优化决策（tile size、展开因子、并行宽度、向量化维度）参数化，然后通过搜索算法（遗传算法/贝叶斯优化/强化学习）寻找最优组合。Google 的 Auto-TVM 和 IREE 内置的 heuristics 搜索即基于此原理。

九、总结与展望

MLIR 代表了编译器基础设施从"单次构建"到"可组合框架"的范式转变。其 Dialect 体系、可扩展类型系统和声明式重写框架使得领域专家可以在不深入了解整个编译器的情况下，为特定领域构建优化的编译流水线。

随着 AI 芯片架构持续分化（NPU、TPU、DPU、CXL-attached accelerator）和异构计算成为主流，MLIR 的多层次 IR 理念只会更加重要。未来方向包括：

• 动态形状编译的运行时优化：支持未知维度的张量运算，在 JIT 编译时结合具体输入形状触发超级线性优化
• 概率编程 Dialect：表达贝叶斯推断中的随机采样和变分推理原语
• 量子计算 Dialect：如 QIR 已在探索量子门操作和错误校正编码的 IR 表示
• 跨模块全程序优化：在 MLIR Module 边界执行全程序内联、死代码消除和数据布局统一

MLIR 不仅仅是一个编译器框架，更是一种思维的转变：承认计算的多样性，让 IR 的抽象层次适配问题本身——而不是迫使所有计算挤进同一条流水线。