一文搞懂新一代AI编译器的基础设施——MLIR编译框架

MLIR的基本概念和特点

作为完全可扩展的编译器基础设施，MLIR具有高度模块化和开放的架构，其操作集、属性集和类型系统都是开放、可扩展的，并随着不同应用场景持续演进。这使得MLIR不仅适用于AI编译器的构建，还广泛适用于高性能计算、图形处理、量子计算等多个领域。

MLIR的设计思想支持将高阶IR逐步递降到特定目标形式，并涵盖从图形算法到低级代码生成的整个范围。除了处理循环嵌套和数据布局的高级转换这类典型的中级优化外，MLIR还可执行典型的针对后端IR的低级调度和映射决策，例如，将IR映射到专门的矢量指令、自动矢量化和软件流水线等。可以说，MLIR提供的多级抽象表示能力不但在表示能力方面提供了新的解决思路，而且以统一的设计框架提供多级抽象，以统一的标准IR和内在一致的概念表示各种形式的计算和源于各种框架的数据流图，为不同问题域的计算表示提供可共享和可重用的基础设施，并在不同层次上插入优化、转换和分析，经由各种高级优化和并行化方法递降得到高性能目标代码。

MLIR中的操作与方言

程序行为由操作描述。MLIR操作集通常用于描述程序的动态语义，其表示范围包括从硬件指令到函数，再到AI模型的构建块，涉及计算的方方面面。操作（简写作op）是MLIR的基本实体和最小代码单元，是抽象和计算的核心单元，开发者可将其视为硬件指令或其他高度抽象的操作的泛化表示。

操作使用值并定义新值。每个操作的使用是一组类型化的操作数，定值是一组类型化的结果。定值和使用的值代表了SSA形式的不可变数据单元。例如，在表达式“%3 = addi %0, %1”中，加法操作的使用为%0和%1，定值为%3。某个操作的输出结果与另一个操作的输入操作数之间的连接可描述为SSA形式的数据流。包含若干个操作的线性序列可组织成基本块（block），若干个基本块又可组织成区域（region），如循环体或函数体。

操作的属性（attribute）包含了操作在编译时已知的信息，每个操作可通过属性参数来描述操作的重要特征。与执行时才能获得操作数的值不同，属性的值在编译时就可确定且固定不变。MLIR中的值都有类型，值的类型中包含了值在编译时已知的信息。

MLIR的操作是完全可扩展的，并具有应用相关的语义，可用于表示LLVM中的所有核心IR结构，如指令、函数、模块等。为了支持可扩展性，MLIR仅将属性、类型、操作和区域等有限的几种概念作为其基本内建概念，其他概念则由基本内建概念衍生得到。例如，模块（module）和函数（function）被定义为具有特定语义、符号名称（symbol name）属性并以区域为主体的操作。不仅如此，从for循环这类高级结构到低级硬件指令都可以用操作定义。

模块化和重用性是MLIR的核心设计思想。MLIR可将不同的组件以独立库的形式分别实现，开发者可以根据需要链接这些库，这与LLVM的设计理念高度一致，同时更进一步强调了灵活性与工程可扩展性。虽然MLIR目前还部分依赖于LLVM的支持库，但在架构上与LLVM IR相互独立，并通过保持优化与转换逻辑与核心IR抽象的分离，实现了高度解耦的构建系统。构建系统可以独立选择编译链接特定的方言实现代码。

MLIR通过方言这一概念支持模块化和可扩展设计以及中间表示的多样化。方言是IR对象的逻辑组合，可在唯一的命名空间下，为不同应用领域的、逻辑上相关的各种抽象（操作、属性和类型）提供一种组合机制。MLIR方言可以笼统的被视为IR逻辑层次。这种通过一系列可组合的方言实现多阶段编译是MLIR不同于其他编译框架的重要特性。这一特性使开发者可以在针对不同异构硬件的编译栈上重用公共的方法和部件。

MLIR包含一系列特定领域的方言，这些方言描述了给定应用所支持的合法操作集，并根据各自工作的上下文与其他方言进行交互。不同的方言在代码生成对中间表示的处理过程中扮演不同角色。为了有效地与MLIR交互，开发者有时需要根据需求定义新的MLIR方言，为后续的分析和转换提供途径。因此，MLIR除了提供若干内置方言来表示通用功能，还为开发者提供开放的基础设施，允许开发者在不同粒度和抽象级别定义新的方言、自定义类型、操作和属性。

MLIR中的常用方言简介

MLIR方言提供了多个抽象级别，这有助于完成在单一抽象级别很难执行的转换和优化。例如，Torch、TOSA、MHLO等上层方言多与AI领域相关，这些方言通常负责导入高层AI框架的图并处理算子集合，捕获算子语义后将其转换为统一的表示形式，消除不同框架之间的不一致性，简化AI模型导入MLIR环境过程和代码生成任务。在统一表示的基础上，MLIR可以进行各种转换和优化，例如，融合算子、内存布局优化、并行化等。这些优化有助于生成高效的低层代码，并在不同的硬件平台上高效运行。

Affine、SCF（Structured Control Flow）、Linalg、Tensor、Vector等中间层方言负责对张量和缓冲区的操作及其他转换流程。中间层方言通常支持多样化的上层输入方言。例如，TOSA、MHLO等方言都可以作为Linalg方言的上层输入方言。

在中间层方言中，Affine方言使用多面体编译技术，主要用于表达和优化带约束的循环和内存访问模式。Affine方言提供了强大的循环变换能力，如循环拆分、循环融合和循环交换。

SCF方言表示IR中不受仿射分类规则约束的各种控制流结构，包括循环和条件。例如，scf.for循环操作接受整数值作为循环下、上限或步长，且不支持仿射映射。这些结构化的控制流有助于在高层次上进行控制流分析和优化，简化代码生成。

Linalg方言用于表达和优化与线性代数相关的张量操作。例如，矩阵乘法、卷积等操作可以通过Linalg方言描述，并利用分块（tiling）操作将大规模的计算问题分解成小规模的子问题，一方面可以提高数据局部性和缓存命中率，另一方面也可以更好地将计算映射到硬件资源，达到在目标硬件上高效执行的优化目标。

张量表示抽象的值类型数据序列，与此相关的Tensor方言专注于张量操作的抽象，并提供了处理多维数组的操作，可以在高层次上对张量进行操作，为后续的优化和转换提供基础。使用张量和Tensor方言操作可以提高算法开发者的开发效率，而MemRef方言表示较低级别的缓冲区访问，可用于构建与物理内存的桥梁。Vector方言提供了向量化的操作抽象，可以有效地进行数据并行优化，如向量化和SIMD化（单指令多数据）。综上所述，Affine、SCF、Math和Linalg等方言主要用于表示计算结构和控制流，Tenor、MemRef和Vector等方言则主要用于表示数据负载。

底层的NVVM、ROCDL等平台方言和LLVM方言则负责通过与其他编译器（特别是传统编译器和领域专用编译器）后端兼容的中间表示，实现MLIR与不同外部编译器后端之间的对接。针对特定硬件的底层方言，如NVGPU、AMDGPU等方言，实现了对特定硬件特性的抽象。这些方言可以视为MLIR中间表示在导出到外部编译器后端之前的过渡阶段。

MLIR完成优化和转换后，通常输出LLVM IR，然后借助LLVM框架，调用不同的处理器后端完成相应的代码生成等后端流程。LLVM方言提供了LLVM IR指令和类型到MLIR的直接映射。考虑到编译过程后续需要LLVM进一步处理，作为叶方言（leaf dialect）的LLVM方言可以简化MLIR和LLVM IR两种表示之间的翻译过程。

MLIR将域内的映射问题和其他问题限制在各自领域的方言内处理，意味着数据结构之间转换可以做到尽可能简单，得到正确转换结果的概率也更高。

MLIR的高层结构

如前所述，在MLIR中，基本块被定义为一系列不含控制流的顺序操作。由基本块构成的控制流图(CFG)在MLIR中被组织成区域，控制可以从区域中的一个块流向其后续块。MLIR支持递归结构，区域可以附加到一个操作上，该操作定义控制如何流入和流出这些区域，即操作可能包含块的关联区域，这些区域包含一系列块，每个块可能又包含一系列操作，这些操作可能也包含其他区域，从而允许IR在任意多个级别上实现嵌套。每个块以终止操作结束，控制流可能转移到该终止操作的后续块。图1总结了MLIR嵌套结构的组成。

图1.MLIR的高层嵌套结构

以下示例代码说明了MLIR的关键概念，包括操作、操作数和结果、属性（Attributes）、区域、基本块等的基本使用方法。

%value_definition = "dialect.operation"(%value_use) {attribute_name = #attr_kind<"value">} ({^block(%block_argument: !argument_type):"dialect.other_operation"() : () -> ()}) : (!operand_type) -> !dialect.result_type

在MLIR中，操作是MLIR的基本单位，由唯一的字符串（如tf.Conv2d、x86.repmovsb等）标识，可以接收多个操作数和返回多个结果，支持属性字典、后继（successor），以及零个或多个嵌套的区域。上述示例代码中的方言操作dialect.operation属于dialect方言，接受一个操作数%value_use，并返回一个类型为!dialect.result_type的结果。该操作包含一个命名属性attribute_name，使用#attr_kind<"value">的语法表示，这是MLIR中通用的属性机制。

操作表达范围可以非常广泛，既可以表示函数定义（如func.func操作）、函数调用（如func.call操作）、缓冲区分配（如memref.alloc操作）等高层抽象，也可以表示目标无关算术操作（如arith.addi操作）、特定硬件平台底层指令（如nvvm.mma.sync操作）等底层逻辑。MLIR中的pass是实现操作转换和优化的核心机制。借助pass及其管理框架，MLIR能够模块化、渐进式地对多层次的操作进行变换和优化，并支持开发者灵活地插入和定制新的转换逻辑。

上述dialect.operation操作通过大括号{ }引入了一个区域。区域是一个有序的基本块列表。区域内的语义由包含该区域的操作（即上述示例代码中的dialect.operation操作）定义。区域必须包含在操作中，且没有名称或地址，也没有类型或属性。

函数体可看作区域的一个例子。在函数体中，位于基本块末尾的基本块终止操作（block terminator）必须跳转到不同的块，或者从函数返回，且返回参数的类型必须与函数签名的结果类型匹配。同样，函数参数必须与区域参数的类型和数量匹配。

区域中不同基本块之间的引用或跳转必须局限在其所属的区域之内，即某基本块的终止操作无法跳转到不在同一个区域内的其他基本块。而某个值的可见性取决于其是否可以被后续操作引用，因此，区域的跳转限制意味着一个值定义在某区域中，则只能被该区域内部的操作访问。所以，值的作用域自然被限制在区域内。与此同时，MLIR默认允许区域内的操作引用外层区域中已存在的值，只要这些值本来就是封闭操作（enclosing operation）的合法操作数。若希望彻底隔离内层区域，阻止其引用外部值，可为封闭操作加上OpTrait::IsolatedFromAbove特征（trait）。

上述dialect.operation操作的区域中包含一个基本块^block。基本块是MLIR中控制流的基本单元，包含一组顺序执行的操作序列，只有一个入口和一个出口，没有显式的分支结构，且必须以终止操作（terminator operations）结束。终止操作实现基本块之间的控制流分支。区域中的第一个基本块被称为入口块（entry block）。入口块的块参数（block arguments）也是该区域的参数，而这些参数绑定的值由包含该区域的操作语义决定。对于区域中其他基本块的块参数，其绑定值则是由控制流终止操作决定，例如分支跳转（branch）操作，这些操作将该基本块作为其后继块之一。

在以下示例代码中的函数@simple中，其入口块^bb0的参数%a和%cond的值由func.func操作定义和绑定。而块^bb0作为入口块支配区域中的其他基本块，这意味着在其他基本块内可以引用%a和%cond，%a和%cond的作用域是整个区域，这也体现了入口块参数兼具区域参数的特点。

对于区域中其他基本块的块参数，如块^bb3的参数%c，其绑定值来自^bb3前驱块^bb1和^bb2的cf.br跳转操作（也是终止操作）。块^bb3将参数%c和%a一起传递给块^bb4，但%a不是^bb3的参数，而是直接引用自定义在^bb0的%a。

MLIR利用上述基本块结构和参数传递方法，替代了传统SSA通过复杂的显式PHI节点合并可能来自多个前驱基本块的变量值的机制，使得IR表达更简洁。

func.func @simple(i64, i1) -> i64 {^bb0(%a: i64, %cond: i1):cf.cond_br %cond, ^bb1, ^bb2^bb1:cf.br ^bb3(%a: i64)^bb2:%b = arith.addi %a, %a : i64cf.br ^bb3(%b: i64)^bb3(%c: i64):cf.br ^bb4(%c, %a : i64, i64)^bb4(%d : i64, %e : i64):%0 = arith.addi %d, %e : i64return %0 : i64 }

上述基本块^block中嵌套了另一个操作dialect.other_operation，该操作也可以拥有一个或多个区域，从而实现MLIR中间表示的递归层次化结构表示。其中，操作是语义与结构的基础单位，也是IR层级结构中的根节点。操作可以定义自身的语义，同时承载一个或多个区域。区域起到结构性作用域的角色，其本质是一个有序的基本块列表，用于封装控制流结构并限定值的可见性。基本块则是基本的控制流单元，包含按顺序排列的一组操作，这些操作又可进一步嵌套区域。

在MLIR的中间表示中，另一种重要的结构关系是值在操作之间传递与引用。MLIR中的值或者是从控制流前驱传入的块参数，或者是某个操作的输出结果。值的使用者是其他操作，通过操作的操作数引用。

图2展示了MLIR通过值的使用-定值（use-def）链在操作、基本块和操作数之间建立的数据流连接结构图。图2中的直角灰色框表示值的定义，可以来自块参数或某个操作的结果。直角白色框表示操作的操作数，是对已定义值的引用。虚线箭头表示操作数使用值的方向。

图2.MLIR中值的使用与定值

基本块的每个块参数是一个值。基本块中的每个操作或者引用块参数，或者引用前驱操作的结果作为操作数。每个操作可以返回多个值，并被其他操作使用，从而形成清晰的使用-定值链，便于分析、优化pass追踪值的来源和使用位置。每个值记录了该值的第一个使用（first use），即OpOperand，并可以从该使用出发，通过OpOperand的nextUse指针遍历后续所有OpOperand，或通过OpOperand的back指针向前回溯。显然，MLIR中值的使用是通过这种典型的双向链表结构进行组织。

以上内容节选自《MLIR编译框架开发指南：构建新一代AI编译器基础设施》作者：汪岩

▊《MLIR编译框架开发指南》

汪岩

融合了一线AI芯片工程师多年实战与沉淀，从MLIR核心设计思想出发，从底层机制到实际应用，全面解析MLIR的设计理念、实现细节和优化路径。循序渐进、由浅入深，书中配有二维码视频，使读者身临其境，迅速、深入地掌握各种经验和技巧。

撰  稿  人：计旭

责任编辑：郝建伟

审  核  人：曹新宇