在编译器的发展历程中，我们见证了从单一中间表示到多层次抽象的演进。MLIR（Multi-Level Intermediate Representation）代表了这一演进的最新成果——它不仅是一个新的编译器基础设施，更是对”编译器应该如何构建”这一根本问题的重新思考。本章将深入探讨 MLIR 的设计理念、核心机制，以及它如何解决现代编译器面临的复杂性挑战。
LLVM IR 在过去二十年取得了巨大成功，成为了编译器领域的事实标准之一。它提供了一个稳定、well-defined 的中间表示，使得前端和后端能够独立发展。然而，随着计算领域的快速演进，LLVM IR 的一些设计决策开始显现出局限性：
抽象层次的固定性：LLVM IR 被设计为接近汇编语言的低级表示。这种设计对于传统的 C/C++ 编译非常合适，但对于高级领域特定语言（DSL）来说，从源语言到 LLVM IR 的语义鸿沟过大。许多重要的高级信息（如张量操作、并行模式）在降级到 LLVM IR 时会丢失。
扩展性的挑战：向 LLVM IR 添加新的概念需要修改核心基础设施。例如，添加 coroutine 支持花费了多年时间，涉及大量的设计讨论和实现工作。这种僵化性使得 LLVM 难以快速适应新的编程模型。
优化机会的丢失：当高级抽象被过早地降级到低级表示时，许多领域特定的优化机会就会丢失。例如，矩阵乘法在降级为循环和内存访问后，就很难再识别和应用特定的优化算法。
2010 年代见证了领域特定编译器的爆炸式增长，特别是在机器学习领域：
传统路径：Python/Framework → Graph IR → Custom Optimizer → LLVM IR → Machine Code
↓ ↓ ↓
TensorFlow Graph XLA HLO 自定义优化器
PyTorch JIT Graph TVM IR 各自为战
JAX IR Glow IR 重复造轮子
每个框架都构建了自己的编译器栈，导致了大量的重复工作。更糟糕的是，这些编译器之间很难共享优化passes、分析算法或代码生成技术。一个在 TensorFlow 中实现的优化很难移植到 PyTorch，即使它们解决的是相同的问题。
MLIR 的核心洞察是：与其为每个领域构建独立的编译器，不如构建一个可重用的基础设施，让领域特定的编译器能够共享通用组件。这个愿景包括：
2017 年，Chris Lattner 离开 Apple 后短暂加入 Tesla，随后于同年加入 Google Brain 团队。作为 LLVM 的创始人，他对编译器基础设施有着深刻的理解，同时也清楚地看到了 LLVM 在新兴领域面临的挑战。在 Google，他面对的是一个完全不同的问题域：如何为 TensorFlow 这样的机器学习框架构建高效的编译器？
Lattner 观察到，TensorFlow 的编译器栈（包括 XLA、Grappler、TF Lite 等）存在大量的重复工作。每个组件都在解决类似的问题：图优化、内存规划、设备分配等，但它们之间很难共享代码。这种状况让他想起了 LLVM 项目开始前的编译器领域——每个编译器都是一个孤岛。
TensorFlow 在 2015 年开源时，主要依赖于运行时解释执行。随着模型规模和性能要求的提升，编译优化变得越来越重要。TensorFlow 面临的编译挑战包括：
多层次抽象的需求：从 Python API 到硬件指令，中间需要经过多个抽象层次：
每个层次都有自己的优化机会，但缺乏统一的框架来表示和转换这些不同层次的抽象。
异构硬件的支持：TensorFlow 需要支持 CPU、GPU、TPU 等多种硬件。每种硬件都有自己的编程模型和优化策略。传统的编译器架构很难优雅地处理这种异构性。
动态性与静态优化的平衡：Python 的动态特性与编译器的静态优化之间存在根本矛盾。如何在保持灵活性的同时获得编译优化的性能提升，是一个关键挑战。
XLA（Accelerated Linear Algebra）是 TensorFlow 的第一代编译器，它证明了编译优化对机器学习的重要性。XLA 可以将 TensorFlow 图编译为高效的机器码，在某些场景下获得数倍的性能提升。然而，XLA 也暴露了一些根本问题：
紧耦合的设计：XLA 与 TensorFlow 紧密耦合，其他框架很难重用 XLA 的技术。这导致 PyTorch、MXNet 等框架都需要构建自己的编译器。
有限的扩展性：向 XLA 添加新的操作或优化需要深入理解整个系统。这种高门槛限制了社区贡献。
单一的抽象层次：XLA HLO 是一个固定的抽象层次，不够灵活。某些优化需要更高级的信息，而另一些优化需要更低级的控制。
基于这些经验，Lattner 和团队开始构思一个全新的架构。2018 年初，MLIR 项目正式启动，目标是创建一个真正可扩展、可重用的编译器基础设施。
2019 年 4 月，Google 决定将 MLIR 开源，这个决策具有重要的战略意义：
建立生态系统：开源使得 MLIR 能够成为整个行业的标准，而不仅仅是 Google 的内部工具。这有助于建立一个健康的生态系统。
加速创新：通过开源，MLIR 可以吸引全球的贡献者，加速技术创新。事实证明，开源后的 MLIR 发展速度远超预期。
标准化的推动：MLIR 的开源推动了编译器中间表示的标准化。越来越多的项目开始采用 MLIR，形成了事实上的行业标准。
值得注意的是，MLIR 最终被贡献给了 LLVM 基金会，成为 LLVM 项目的一个子项目。这个决定体现了 MLIR 与 LLVM 的互补关系——MLIR 处理高级抽象，LLVM 处理低级代码生成，两者共同构成了完整的编译器栈。
MLIR 的设计哲学是”一切皆可扩展”。与 LLVM IR 的固定指令集不同，MLIR 提供了一个元框架，允许用户定义自己的操作、类型和属性。理解这三个核心概念是掌握 MLIR 的关键。
在 MLIR 中，Operation 是最基本的计算单元。每个 Operation 都有统一的结构：
%result = dialect.operation(%arg0, %arg1) {attr1 = value1} : (type0, type1) -> result_type
这种统一表示带来了几个重要优势：
可扩展性：任何 Dialect 都可以定义自己的 Operation，而不需要修改 MLIR 核心。例如：
arith.add - 算术加法tensor.extract - 张量元素提取gpu.launch - GPU kernel 启动async.execute - 异步执行区域语义完整性：每个 Operation 都携带完整的类型信息和属性，使得分析和转换更加可靠。Operation 可以有多个结果，支持复杂的计算模式。
区域和块：Operation 可以包含区域（Region），区域包含块（Block），块包含 Operation。这种递归结构可以表示控制流、闭包等复杂概念：
scf.if %condition -> (i32) {
%true_val = arith.constant 1 : i32
scf.yield %true_val : i32
} else {
%false_val = arith.constant 0 : i32
scf.yield %false_val : i32
}
验证机制：每个 Operation 都可以定义自己的验证规则，确保 IR 的正确性。这种”设计时验证”比运行时检查更加高效。
MLIR 的类型系统是完全可扩展的，这是它与 LLVM IR 的一个重要区别。LLVM IR 有固定的类型集合（整数、浮点、指针等），而 MLIR 允许每个 Dialect 定义自己的类型。
内建类型：MLIR 提供了一些基础类型：
i32、f64、indextensor<4x4xf32>、memref<10x20xi8>、vector<4xf32>Dialect 特定类型：每个 Dialect 可以定义自己的类型系统：
!tf.string - TensorFlow 字符串类型!llvm.ptr<i8> - LLVM 指针类型!gpu.async.token - GPU 异步令牌类型参数化：类型可以是参数化的，支持泛型编程：
!my_dialect.optional<i32> // 可选整数
!my_dialect.list<tensor<*xf32>> // 动态形状张量的列表
形状和秩：对于张量和内存引用，MLIR 支持静态和动态形状：
tensor<4x?xf32> - 第一维静态（4），第二维动态tensor<*xf32> - 完全动态的秩和形状这种灵活的类型系统使得 MLIR 可以精确地表示各种领域的概念，从机器学习的张量到量子计算的量子位。
Attribute 表示编译时已知的常量值和元数据。它们是 MLIR 中信息传递的重要机制。
基础属性类型：
42 : i32、3.14 : f64[1, 2, 3] : i32{key1 = "value1", key2 = 42}Dense 元素属性：用于表示大型常量数据：
%weights = arith.constant dense<[[1.0, 2.0], [3.0, 4.0]]> : tensor<2x2xf32>
符号引用：属性可以引用其他符号，支持模块化：
func.call @my_function(%arg) : (i32) -> i32
// @my_function 是一个符号引用属性
自定义属性：Dialect 可以定义自己的属性类型，携带领域特定信息：
#my_dialect.layout<"NHWC"> // 数据布局属性
#my_dialect.device<"GPU:0"> // 设备分配属性
属性的一个关键特性是它们是不可变的，这简化了编译器的实现并提高了性能。属性的哈希和比较操作都是高效的，使得基于属性的模式匹配变得实用。
Operation、Type 和 Attribute 这三个概念紧密协作，构成了 MLIR 的表达能力：
// Operation 使用 Type 定义接口，使用 Attribute 携带元数据
%result = my_dialect.convolution(%input, %kernel) {
strides = [1, 1],
padding = "SAME",
data_format = "NHWC"
} : (tensor<1x28x28x1xf32>, tensor<5x5x1x32xf32>) -> tensor<1x28x28x32xf32>
这个例子展示了：
my_dialect.convolution 定义了计算strides、padding 等提供了操作的配置这种设计使得 MLIR 能够精确地表示各种计算，同时保持足够的灵活性来适应不同的领域需求。
Dialect 是 MLIR 最具创新性的概念之一。它提供了一种模块化的方式来扩展编译器，使得不同的抽象层次和领域可以在同一个框架内共存。每个 Dialect 定义了一组相关的操作、类型和属性，形成了一个自包含的”语言”。
Dialect 可以理解为 MLIR 中的”命名空间”或”模块”。但它远不止于此——Dialect 是一个完整的抽象层次，包含了该层次所需的所有概念。这种设计带来了几个关键优势：
隔离性：不同 Dialect 之间是隔离的，避免了命名冲突和语义混淆。例如，arith.add 和 llvm.add 虽然都是加法，但它们有不同的语义和约束。
可组合性：多个 Dialect 可以在同一个模块中共存，允许渐进式转换。一个函数可以同时包含高级的 tensor 操作和低级的 llvm 操作。
演化能力：新的 Dialect 可以随时添加，现有 Dialect 可以独立演化，不会影响其他部分。
Dialect 的注册机制确保了系统的模块化和动态性：
// Dialect 定义
class MyDialect : public Dialect {
public:
explicit MyDialect(MLIRContext *context)
: Dialect("my_dialect", context, TypeID::get<MyDialect>()) {
// 注册操作、类型、属性
addOperations<MyAddOp, MyMulOp>();
addTypes<MyTensorType>();
addAttributes<MyLayoutAttr>();
}
};
// 使用时注册
MLIRContext context;
context.loadDialect<MyDialect>();
这种延迟加载机制意味着只有需要的 Dialect 才会被加载，减少了内存占用和启动时间。
MLIR 提供了一系列核心 Dialect，覆盖了从高级抽象到机器码生成的各个层次：
高级数学和张量操作：
tensor - 不可变张量操作linalg - 线性代数操作（类似 NumPy）sparse_tensor - 稀疏张量支持complex - 复数运算内存和缓冲区管理：
memref - 可变内存引用bufferization - 张量到缓冲区转换控制流和结构：
scf - 结构化控制流（for、if、while）cf - 控制流（branch、cond_br）func - 函数定义和调用并行和异步：
async - 异步执行gpu - GPU 编程模型omp - OpenMP 并行构造低级和目标相关：
llvm - LLVM IR 操作x86vector - x86 向量指令arm_neon - ARM NEON 指令spirv - SPIR-V for Vulkan/OpenCL每个 Dialect 都有明确的抽象层次和目标用途。例如，tensor Dialect 处理数学意义上的张量，而 memref Dialect 处理内存布局和访问模式。
开发自定义 Dialect 是扩展 MLIR 的主要方式。以一个简化的量子计算 Dialect 为例：
// 使用 TableGen 定义 Dialect
def Quantum_Dialect : Dialect {
let name = "quantum";
let summary = "Quantum computing operations";
let cppNamespace = "::mlir::quantum";
}
// 定义量子比特类型
def QubitType : TypeDef<Quantum_Dialect, "Qubit"> {
let summary = "Quantum bit type";
}
// 定义量子门操作
def HadamardOp : Quantum_Op<"hadamard"> {
let summary = "Hadamard gate";
let arguments = (ins QubitType:$qubit);
let results = (outs QubitType:$result);
let assemblyFormat = "$qubit attr-dict `:` type($qubit)";
let verifier = [{
// 自定义验证逻辑
return success();
}];
}
这个例子展示了 Dialect 开发的关键要素：
QubitType 表示量子比特HadamardOp 实现 Hadamard 门Dialect 之间的转换是 MLIR 编译流程的核心。转换通常遵循以下模式：
// 定义转换 Pass
struct TensorToMemrefPass
: public PassWrapper<TensorToMemrefPass, OperationPass<func::FuncOp>> {
void runOnOperation() override {
// 设置转换目标
ConversionTarget target(getContext());
target.addLegalDialect<memref::MemRefDialect>();
target.addIllegalDialect<tensor::TensorDialect>();
// 应用转换模式
RewritePatternSet patterns(&getContext());
tensor::populateTensorToMemrefPatterns(patterns);
if (failed(applyPartialConversion(getOperation(), target, patterns)))
signalPassFailure();
}
};
这种转换机制支持：
设计 Dialect 时需要考虑多个权衡：
抽象层次：太高级难以优化，太低级失去表达能力。好的 Dialect 在其目标领域找到合适的抽象层次。
操作粒度：粗粒度操作（如 linalg.matmul）易于优化但不够灵活；细粒度操作（如 arith.add）灵活但优化困难。
类型丰富度：丰富的类型系统提供更多信息，但也增加了复杂性。
语义明确性：操作的语义必须明确定义，避免歧义。例如，tensor.pad 明确指定了填充值和填充大小。
渐进式降级是 MLIR 的核心设计理念之一。与传统编译器的”全有或全无”转换不同，MLIR 支持逐步地、选择性地将高级抽象转换为低级表示。这种方法不仅使编译过程更加可控，也保留了更多的优化机会。
考虑一个矩阵乘法的编译过程，传统方法可能直接从高级表示跳到循环：
传统方法：
C = A @ B → for i: for j: for k: C[i,j] += A[i,k] * B[k,j]
（丢失了矩阵乘法的语义信息）
MLIR 的渐进式降级保留了中间层次：
// 层次 1：高级线性代数
%C = linalg.matmul ins(%A, %B : tensor<128x256xf32>, tensor<256x512xf32>)
outs(%C : tensor<128x512xf32>) -> tensor<128x512xf32>
// 层次 2：带 tiling 的结构化循环
scf.parallel (%i, %j) = (0, 0) to (128, 512) step (32, 32) {
%tile = linalg.matmul ins(%A_tile, %B_tile) outs(%C_tile)
}
// 层次 3：向量化的循环
scf.for %i = 0 to 128 step 4 {
%vec_a = vector.load %A[%i, %k]
%vec_b = vector.load %B[%k, %j]
%vec_c = vector.fma %vec_a, %vec_b, %vec_c
}
// 层次 4：目标相关的指令
llvm.inline_asm "vfmadd231ps %zmm1, %zmm2, %zmm0"
每个层次都保留了适合该抽象级别的优化机会。
MLIR 的转换框架提供了强大而灵活的机制来实现降级：
Type Converter：定义类型之间的映射关系
class TensorToMemrefTypeConverter : public TypeConverter {
TensorToMemrefTypeConverter() {
addConversion([](TensorType type) -> Optional<Type> {
return MemRefType::get(type.getShape(), type.getElementType());
});
}
};
Operation Converter：定义操作之间的转换规则
struct TensorLoadToMemrefLoad : public OpConversionPattern<tensor::LoadOp> {
using OpConversionPattern::OpConversionPattern;
LogicalResult matchAndRewrite(
tensor::LoadOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
rewriter.replaceOpWithNewOp<memref::LoadOp>(
op, adaptor.getMemref(), adaptor.getIndices());
return success();
}
};
Conversion Target：指定哪些操作是合法的
ConversionTarget target(context);
target.addLegalDialect<memref::MemRefDialect>();
target.addIllegalOp<tensor::LoadOp>();
target.addDynamicallyLegalOp<func::FuncOp>([&](func::FuncOp op) {
return typeConverter.isSignatureLegal(op.getFunctionType());
});
MLIR 支持部分转换，允许 IR 处于”混合”状态。这带来了几个重要优势：
增量开发：新的降级规则可以逐步添加，不需要一次性完成所有转换。
选择性优化：某些操作可以保持在高级形式以便优化，而其他操作已经降级。
调试友好：可以在任何中间状态检查 IR，理解转换过程。
// 部分降级的 IR：tensor 和 memref 操作共存
func.func @partially_lowered(%arg0: tensor<10xf32>) -> f32 {
%0 = tensor.extract %arg0[%c5] : tensor<10xf32> // 还是 tensor
%1 = memref.alloc() : memref<10xf32> // 已转为 memref
memref.store %0, %1[%c0] : memref<10xf32>
%2 = memref.load %1[%c0] : memref<10xf32>
return %2 : f32
}
成功的降级策略通常遵循一些设计模式：
1. 自顶向下降级：从最高级的抽象开始，逐步降级
linalg → scf → cf → llvm
2. 横向转换：在同一抽象层次内的转换
tensor → memref (都是中级抽象)
3. 局部精化：只降级特定的操作或区域
// 只降级特定的矩阵乘法，保留其他操作
transform.structured.tile_to_scf_for %matmul
4. 目标驱动降级：根据目标硬件选择降级路径
如果目标是 GPU：linalg → gpu → nvvm
如果目标是 CPU：linalg → vector → llvm
一个关键挑战是如何在降级过程中保留有用的信息。MLIR 提供了多种机制：
属性传播：高级操作的属性可以传播到低级操作
// 高级操作带有并行属性
%0 = linalg.matmul {parallel_dims = [0, 1]} ...
// 降级后保留并行信息
scf.parallel (%i, %j) ... {
// 并行循环
}
元数据保留：使用属性保存原始操作信息
// 降级后的操作仍然知道它来自矩阵乘法
%0 = llvm.fmuladd %a, %b, %c {original_op = "linalg.matmul"}
分析信息缓存：降级前的分析结果可以附加到 IR
// 在降级前进行依赖分析
auto deps = analyzeDependencies(op);
// 将分析结果作为属性附加
op->setAttr("deps", deps);