MLIR 架构和方言

传统的编译器基础设施在优化大量机器学习 (machine learning)工作负载时常会遇到困难，因为它们强制立即从高级图表示转换为低级指令集。这种“语义丢失”发生在一个复杂运算（例如卷积）被直接降级为一组标量循环时。一旦这种转换发生，编译器就失去了对这是卷积这一原始意图的感知，优化就成为一项困难的标量指令模式匹配任务。

MLIR（多级中间表示）通过不作为一个单一的中间表示，而是作为一个构建中间表示的架构来解决这个问题。它使得编译器能够在同一编译单元中保留高级结构（如张量维度和填充），同时保持低级详情（如内存分配和向量 (vector)内联函数）。

运算的结构

在 MLIR 中，基本的执行单元是运算。与 LLVM IR 不同，后者指令集是固定并由语言规范定义的，MLIR 是开放的。用户可以定义新的运算，而无需修改核心编译器架构。

MLIR 中的一个运算是一个通用容器。它不与每个特定指令的 C++ 类一对一映射。相反，它采用一种轻量级存储机制，包含定义指令行为的必要组成部分。

每个运算包含以下主要组成部分：

1. 运算名称： 一个唯一的字符串标识符，通常带有命名空间（例如，tosa.matmul 或 scf.for）。
2. 操作数： 运算使用的一系列输入值（SSA 值）。
3. 结果： 运算生成的一系列输出值。
4. 属性： 在编译时确定的常量数据，例如整数文字、字符串名称或张量形状。
5. 区域： 嵌套的代码块，允许运算包含其他运算。
6. 后续者： 用于控制流运算（如分支）的、指向其他块的链接。

下图展示了一个 MLIR 模块的层级结构，呈现了运算如何存在于块中，而这些块又在其他运算或函数内部形成区域。

MLIR 模块的层级结构。运算可以定义区域，这些区域包含顺序运算块，从而形成一个可用于表示循环和函数的递归结构。

文本表示

MLIR 提供一种透明的文本格式，反映了其内存中的结构。能够读取这种 IR 对调试编译器过程来说是必要的。请看以下一个运算的通用语法：

$\%result:2 = \text{"dialect.op\_name"}(\%arg0, \%arg1) \{ attribute = 42 : i32 \} : (f32, f32) \to (f32, i1)$

在此结构中：

• %result:2 表明此运算生成两个结果。
• "dialect.op_name" 是唯一的标识符。
• (%arg0, %arg1) 是操作数（输入）。
• { ... } 包含一个属性字典（编译时常量）。
• 最后一部分 (f32, f32) -> (f32, i1) 定义了函数类型签名：两个 float32 输入，产生一个 float32 和一个 1 位整数。

方言体系

MLIR 的优势在于方言。方言是在一个独特命名空间下对运算、类型和属性进行的逻辑分组。如果说传统编译器像一个拥有固定工具集的车间，那么 MLIR 则是一个工厂车间，你可以在其中引入用于特定任务的专业机器。

方言允许不同抽象级别共存。在一个 MLIR 文件中，你可能会看到：

• tf (TensorFlow)：高级图节点，如 tf.Softmax。
• tosa (张量算子集架构)：标准化的张量代数，可用于硬件推理 (inference)。
• scf (结构化控制流)：循环（for、while）和条件语句（if），保持结构，与低级“goto”分支不同。
• arith (算术)：基本的标量数学，如整数加法或浮点乘法。
• llvm：与 LLVM IR 后端一对一映射的指令。

这种模块化使得渐进式降低成为可能。编译器不是将高级图直接翻译为机器代码，而是逐步降低它。一个 tf.MatMul 可能首先被降低为一个 linalg.matmul（一个结构化循环表示），然后降低为使用 scf.for 和 arith 运算的循环，最后降低为 llvm 方言进行代码生成。

特征与接口

由于 MLIR 运算是通用的，编译器需要一种方式来推断它们的行为，而无需确切了解每个具体运算的作用。这通过特征和接口来处理。

特征描述了一个运算的固有属性。例如，如果你定义一个自定义矩阵乘法运算，你可能会附带 NoSideEffect 特征。这告知编译器该运算是纯粹的，它仅依赖其输入并生成输出，而不修改全局内存或 I/O。因此，如果该运算的结果未被使用，死代码消除 (DCE) 过程可以自动移除此运算，即使 DCE 过程对矩阵乘法具体情况一无所知。

接口提供了一种通用查询或修改运算的方式。

• InferTypeOpInterface：允许运算根据输入形状计算其输出形状（形状推断）。
• LoopLikeInterface：允许优化过程将一个运算辨别为循环（无论是 scf.for 还是自定义的 my_accelerator.repeat），并应用循环不变代码移动。

类型系统与属性

MLIR 中的类型系统与运算集一样可扩展。虽然它包含 i32（32 位整数）或 f16（半精度浮点数）等标准类型，但方言可以定义复杂的特定方面类型。

在机器学习 (machine learning)中，RankedTensorType 普遍存在。它表示具有已知形状和元素类型的张量，表示为 tensor<4x4xf32>。MLIR 也支持动态维度，由 ? 表示。一个 tensor<?x128xf32> 类型意味着在编译时批次大小未知，但特征维度固定为 128。

属性与操作数的不同之处在于它们是静态的。它们在机器学习硬件编译器中大量使用以存储配置数据。例如，一个卷积运算需要填充、步长和扩张率。这些都作为属性存储，因为它们在编译过程中必须已知，才能有效地执行布局转换或内存分块。

验证器与不变量

在定义方言时，正确性通过验证器来保障。验证器是与一个运算关联的 C++ 钩子，用于检查不变量。

例如，一个 Reshape 运算可能要求输入张量中的元素总数等于输出张量中的元素总数。如果一个优化过程意外违反此规则（例如，通过错误地进行形状常量折叠），验证器将立即触发错误。这种快速失败机制在构建复杂的降低流水线时至关紧要，因为它在无效 IR 状态传播到后端之前捕获它们，而在后端调试会变得非常困难。

下图可视化了通用运算存储与定义特征和验证器的专用方言逻辑之间的相互作用。

通用运算存储与方言定义之间的相互作用。TableGen 定义和特征为原始的底层存储提供了语义。

通过将存储表示与语义定义解耦，MLIR 使得编译器能够高效地操作运算。过程可以遍历运算，检查特定特征（例如“这是否可交换？”），并执行转换，而无需将每个运算转换为特定的 C++ 类。正是这种架构实现了 MLIR 架构在多种硬件目标上的高可扩展性。