MLIR: Scaling Compiler Infrastructure for Domain Specific Computation

作者/机构: Chris Lattner, Mehdi Amini, Albert Cohen, Andy Davis, Jacques Pienaar, River Riddle, Tatiana Shpeisman, Nicolas Vasilache, Oleksandr Zinenko (Google, USA/France), Uday Bondhugula (Indian Institute of Science, India)

A1 主要贡献

本文介绍了一种名为 MLIR (Multi-Level Intermediate Representation) 的新型编译器基础设施，旨在解决软件碎片化问题、异构硬件编译挑战，并显著降低构建领域特定编译器的成本。

核心问题与研究目标:
传统的编译器基础设施如 LLVM 和 JVM 采用“一刀切”的方法，提供单一的抽象级别。虽然这种方法在某些领域非常成功，但许多问题在更高或更低的抽象级别上能被更好地建模。因此，许多现代编程语言（如 Swift, Rust）和机器学习系统最终都开发了自己的高级中间表示（IR），这导致了高昂的工程成本、重复实现以及潜在的低质量编译器系统（编译时间长、实现有 bug、诊断信息差等）。特别是现代机器学习框架，它们由许多不同的编译器、图技术和运行时系统组成（如图 1），缺乏共享的基础设施，导致错误信息差、性能不可预测和难以支持新硬件等问题。整个编译器行业都面临类似的问题，即高级语言重复发明用于高级抽象的技术（如图 2）。

MLIR 项目的目标是通过提供一个高质量、可重用和可扩展的基础设施来直接应对这些挑战。它旨在降低定义和引入新抽象级别的成本，并提供“开箱即用”的基础设施来解决常见的编译器工程问题，从而促进代码生成器、翻译器和优化器的设计与实现，跨越不同的应用领域、硬件目标和执行环境。

创新点与设计原则:
MLIR 的设计基于三个核心原则：
1. 简约性 (Parsimony): 采用奥卡姆剃刀原则，内置的语义、概念和编程接口都非常少。通过抽象操作和类型的属性来驾驭内在和偶然的复杂性。这为系统的可扩展性和定制化打开了大门。
2. 可追溯性 (Traceability): 强调保留信息而非恢复信息。通过声明式规则和属性来支持转换，并提供通用的信息追踪方法。这对于安全关键和安全编译等领域至关重要。
3. 渐进性 (Progressivity): 核心理念是“过早降低是万恶之源”。MLIR 允许IR中混合不同层次的抽象，并支持将计算结构逐步降低到硬件抽象，只在不再需要结构信息时才进行有意识的结构丢失。

本文的主要贡献包括:
1. 定位问题: 将构建可扩展和模块化编译器系统的问题，用经过验证的设计和工程原则（简约性、可追-溯性、渐进性）来进行阐述。
2. 新颖的基础设施: 详细描述了一个遵循这些原则的新型编译器基础设施 MLIR，它具有重要的工业和研究应用价值。
3. 应用探索: 通过展示 MLIR 在多个不同领域的应用案例，说明该方法的通用性，并分享基于 MLIR 构建系统的经验。

图 1. 跨越不同框架的 TensorFlow 执行流程。
图 1. 跨越不同框架的 TensorFlow 执行流程。
图 2. 带有中级语言 IR 的编译流水线。
图 2. 带有中级语言 IR 的编译流水线。

A3 设计原则

本文探讨了指导 MLIR 设计的需求及其与总体原则的关系。

少量内置，一切皆可定制 [简约性]。该系统基于最少数量的基本概念，使得中间表示的绝大部分内容都可以完全定制。仅使用少数几个抽象——类型、操作和属性，这些是 IR 中最常见的元素——来表达其他一切，从而实现了更少、更一致且易于理解、扩展和采用的抽象。广泛而言，可定制性确保了系统能够适应不断变化的需求，并更有可能适用于未来的问题。从这个意义上说，我们应该将 IR 构建为一个富有可重用组件和编程抽象的丰富基础设施，以支持其自身的中间语言的语法和语义。衡量可定制性成功的一个标准是，能否表达多种多样的抽象，包括机器学习图、抽象语法树（AST）、多元主面体等数学抽象、控制流图（CFG）以及像 LLVM IR 这样的指令级 IR，而所有这些都无需将这些抽象的概念硬编码到系统中。当然，可定制性也带来了因抽象兼容性差而导致内部分裂的风险。虽然可能没有纯粹的技术解决方案，但该系统应鼓励设计可重用的抽象，并假设它们将在其初始范围之外被使用。

SSA 和区域 [简约性]。静态单赋值（SSA）形式【3，Efficiently computing static single assignment form and the control dependence graph, 1991, ACM Trans. Program. Lang. Syst.】是编译器 IR 中广泛使用的表示法。它提供了众多优势，包括使数据流分析变得简单和稀疏，因其与续体传递风格（continuation-passing style）的关系而被编译器社区广泛理解，并在主要框架中得到确立。因此，IR 强制执行 SSA 基于值的语义、其引用透明性和算法效率，所有这些都被视为现代编译器基础设施的基石。然而，尽管许多现有 IR 使用扁平、线性的 CFG，但表示更高级别的抽象推动了将嵌套区域作为 IR 中头等概念的引入。这超越了传统的区域形成，以提升更高级别的抽象（例如，循环树），从而加速编译过程或提取指令级或 SIMD 并行性【4，The program structure tree: Computing control regions in linear time, 1994, PLDI ’94】【5，Treegion scheduling for wide issue processors, 1998, Fourth International Symposium on High-Performance Computer Architecture】【6，On loops, dominators, and dominance frontiers, 2002, ACM Trans. Program. Lang. Syst.】。为了支持异构编译，系统必须支持结构化控制流、并发构造、源语言中的闭包以及许多其他用途的表达。一个具体的挑战是使基于 CFG 的分析和转换能够在嵌套区域上组合。

对 LLVM 规范化属性的取舍。这样做，我们同意牺牲 LLVM 的规范化，有时甚至是其规范化属性。能够将各种数据和控制结构降低到一个较小的规范化表示集合中，是控制编译器复杂性的关键。典型的例子是规范的循环结构，其带有前置头（pre-header）、头（header）、闩（latch）和体（body），这是前端语言中各种循环构造的线性化控制流表示。我们的目标是为用户提供一种选择：根据感兴趣的编译算法或编译流程中的传递，嵌套循环可以被捕获为嵌套区域，或者作为线性化的控制流。通过提供这样的选择，我们偏离了 LLVM 仅规范化的方向，同时保留了在需要时处理更高级别抽象的能力。反过来，利用这种选择也引发了如何控制抽象规范化的问题，这正是下一段的目的。

维持更高级别的语义 [渐进性]。系统需要保留分析或优化性能所需的信息和结构。一旦降低后试图恢复抽象语义的尝试是脆弱的，并且在低级别上硬塞这些信息通常具有侵入性（例如，在使用调试信息记录结构的情况下，所有传递都需要重新审视）。相反，系统应该维持计算的结构，并逐步将其降低到硬件抽象。结构的丢失是经过深思熟虑的，并且只在结构不再需要匹配底层执行模型时发生。例如，系统应该在相关转换的全过程中保留结构化控制流，如循环结构；移除这种结构，即降低到 CFG，本质上意味着将不再执行利用该结构的其他转换。在生产编译器中建模并行计算结构的现有技术水平凸显了这项任务在一般情况下的困难程度【7，LLVM parallel intermediate representation: Design and evaluation using OpenSHMEM communications, 2015, LLVM ’15】【8，Tapir: Embedding forkjoin parallelism into LLVM’s intermediate representation, 2017, SIGPLAN Not.】。

支持渐进式降低。一个必然的推论是，在同一个 IR 中混合不同级别的抽象和不同概念是关键，这使得一部分表示可以保持在较高级别的抽象，而另一部分则被降低。例如，这将使一个为定制加速器设计的编译器能够重用系统定义的某些高级结构和抽象，同时使用该加速器特有的原始标量/向量指令。另一个推论是，系统应支持渐进式降低，即从最高级别的表示逐步降低到最低级别，这个过程沿着多个抽象分小步进行。对多层次抽象的需求源于编译器基础设施必须支持的各种平台和编程模型。以前的编译器在其流水线中引入了多个固定的抽象级别——例如，Open64 的 WHIRL 表示【9，Open64 compiler and tools, 2001】有五个级别，Clang 编译器也是如此，它从 AST 降低到 LLVM IR，再到 SelectionDAG、MachineInstr 和 MCInst。为了支持可扩展性，需要更灵活的设计。这对转换的阶段排序有深远的影响。随着编译器专家实现越来越多的转换传递，这些传递之间开始出现复杂的交互。早期研究表明，结合优化传递可以让编译器发现更多关于程序的事实。结合传递好处的最早例证之一是混合常量传播、值编号和不可达代码消除【10，Combining analyses, combining optimizations, 1995, ACM Trans. Program. Lang. Syst.】。

声明与验证 [简约性与可追溯性]。定义表示修饰符应该像引入新抽象一样简单；一个编译器基础设施的优劣取决于它支持的转换。常见的转换应该可以实现为以声明方式表达的重写规则，并采用机器可分析的格式，以便对重写的属性（如复杂性和完成度）进行推理。重写系统因其健全性和效率而被广泛研究，并应用于众多编译问题，从类型系统到指令选择。由于我们的目标是前所未有的可扩展性和增量降低能力，这为将程序转换建模为重写系统开辟了众多途径。它还提出了一些有趣的问题，关于如何表示重写规则和策略，以及如何构建能够通过多层次抽象引导重写策略的机器描述。系统需要在解决这些问题的同时，保持可扩展性并强制执行单调和可复现的行为。

对验证机制的需求。生态系统的开放性也要求一个广泛的验证机制。虽然验证和测试对于检测编译器错误和捕获 IR 不变性很有用，但在一个可扩展的系统中，对健壮验证方法和工具的需求被放大了。该机制应旨在使定义变得容易，并尽可能具有声明性，提供一个单一的真实来源。一个长期目标是重现翻译验证【11，Translation validation, 1998, Tools and Algorithms for Construction and Analysis of Systems, 4th International Conference, TACAS ’98】【12，Translation validation for an optimizing compiler, 2000, SIGPLAN Not.】【13，Formal verification of translation validators: a case study on instruction scheduling optimizations, 2008, POPL 2008】【14，Verified validation of lazy code motion, 2009, PLDI 2009】和现代编译器测试方法【15，Taming compiler fuzzers, 2013, PLDI ’13】的成功。在可扩展编译器的背景下，这两者目前都是开放问题。

源位置追踪 [可追溯性]。操作的来源——包括其原始位置和已应用的转换——应该在系统中易于追溯。这旨在解决复杂编译系统中常见的不透明问题，即几乎不可能理解最终表示是如何从原始表示构建而来的。这在编译安全关键和敏感应用程序时尤其成问题，因为追踪降低和优化步骤是软件认证程序的重要组成部分【16，Embedded Program Annotations for WCET Analysis, 2018, WCET 2018】。当处理安全代码（如加密协议或操作隐私敏感数据的算法）时，编译器经常面临看似冗余或繁琐的计算，这些计算中嵌入了源程序功能语义未完全捕获的安全或隐私属性：这些代码可能旨在防止侧信道暴露或加固代码以抵御网络或故障攻击。优化可能会改变或完全破坏这些保护措施【17，Secure delivery of program properties through optimizing compilation, 2020, ACM SIGPLAN 2020 International Conference on Compiler Construction (CC 2020)】；这种不透明性在安全编译中被称为 WYSINWYX（所见非所执行）【18，Wysinwyx: What you see is not what you execute, 2010, ACM Trans. Program. Lang. Syst.】。准确地将高级信息传播到低级别的一个间接目标是帮助支持安全和可追溯的编译。

A2 方法细节

我们的主要贡献是提出一个遵循上一节定义的原则的 IR。这就是 MLIR 所做的，我们将在本节回顾其主要设计要点。

通用文本表示。MLIR 有一个通用的文本表示（如图 3 所示），它支持 MLIR 的可扩展性，并完全反映了内存中的表示，这对于可追溯性、手动 IR 验证和测试至关重要。可扩展性带来了冗长的负担，但这可以通过 MLIR 支持的自定义语法来弥补；例如，图 7 展示了图 3 的用户定义语法。

// 属性别名可以前向声明。
#map1 = (d0, d1) -> (d0 + d1)
#map3 = ()[s0] -> (s0)

// 操作可以附加区域。
"affine.for"(%arg0) ({
  // 区域由一个带有参数的块的 CFG 组成。
  ^bb0(%arg4: index): // 块是操作的列表。
    "affine.for"(%arg0) ({
      ^bb0(%arg5: index): // 操作使用和定义类型化的值，这些值遵循 SSA。
        %0 = &quot;affine.load&quot;(%arg1, %arg4) {map = (d0) -> (d0)} : (memref<?xf32>, index) -> f32
        %1 = &quot;affine.load&quot;(%arg2, %arg5) {map = (d0) -> (d0)} : (memref<?xf32>, index) -> f32
        %2 = &quot;std.mulf&quot;(%0, %1) : (f32, f32) -> f32
        %3 = &quot;affine.load&quot;(%arg3, %arg4, %arg5) {map = #map1} : (memref<?xf32>, index, index) -> f32
        %4 = &quot;std.addf&quot;(%3, %2) : (f32, f32) -> f32
        &quot;affine.store&quot;(%4, %arg3, %arg4, %arg5) {map = #map1} : (f32, memref<?xf32>, index, index) -> ()
        // 块以一个终止符操作结束。
        &quot;affine.terminator&quot;() : () -> ()
      // 操作有一个属性列表。
    }) {lower_bound = () -> (0), step = 1 : index, upper_bound = #map3} : (index) -> ()
    &quot;affine.terminator&quot;() : () -> ()
}) {lower_bound = () -> (0), step = 1 : index, upper_bound = #map3} : (index) -> ()

图 3. 使用 affine 和 std 方言表示多项式乘法的 MLIR 通用表示。同样的 IR 在图 7 中以自定义语法显示。

操作（Operations）。MLIR 中的语义单元是“操作”，简称 Op。从“指令”到“函数”再到“模块”，所有东西都被建模为 Op。MLIR 没有固定的 Op 集合，而是允许（并鼓励）用户自定义扩展，这遵循了简约性和“一切可定制”的原则。该基础设施提供了一种基于 TableGen 【19，TableGen LLVM 10 Documentation, 2019】 的声明式语法来定义 Op，如图 5 所示。

操作的结构。操作（见图 4）有一个唯一的操作码，这是一个标识操作及其方言的字符串。操作接受并产生零个或多个值，分别称为操作数和结果，这些值以 SSA 形式维护。值表示运行时的数据，并有一个类型来编码关于数据的编译时知识。除了操作码、操作数和结果外，操作还可能有关联的属性、区域、后继块和位置信息。图 3 展示了值和操作，% 标识符是（一组）命名值，若多于一个则用 : 指定数量，# 表示一个特定的值。在通用文本表示中，操作名称是带引号的字符串字面量，后跟括号中的操作数。
图 4. 操作（Op）是 MLIR 中的主要实体；操作包含一个区域列表，区域包含一个块列表，块包含一个操作列表，从而实现递归结构。

未知操作的处理。编译器传递（pass）会保守地处理未知的操作，MLIR 通过特性（traits）和接口（interfaces）为向传递描述操作的语义提供了丰富的支持，具体如 V-A 节所述。操作的实现包含验证器，用于强制执行操作的不变性，并参与整个 IR 的验证过程。
图 5. 操作定义语法（ODS）提供了一种在 MLIR 中定义新操作的简洁方式。这里定义了一个 LeakyRelu 操作，它接受一个张量和一个浮点值，并返回一个与输入张量类型相同的张量。

属性（Attributes）。MLIR 属性包含除操作码之外的操作的编译时信息。属性是类型化的（例如，整数、字符串），每个操作实例都有一个开放的键值字典，将字符串名称映射到属性值。在通用语法中，属性位于一个由大括号括起来的、逗号分隔的键值对列表中。图 3 使用属性来定义一个循环的边界，这些边界已知是常数仿射形式：{lower_bound = () -> (0), step = 1 : index, upper_bound = #map3}，其中 lower_bound 是一个属性名的例子。() -> (0) 符号用于内联仿射形式，这里产生一个生成常数值 0 的仿射函数。#map3 符号用于属性别名，它允许预先将属性值与一个标签关联起来。属性的意义要么来自操作的语义，要么来自它们关联的方言（见 III 节）。与操作码一样，属性也没有固定的集合。属性可以引用外部数据结构，这对于与现有系统集成很有用，例如，在 ML 系统中（编译时已知的）数据存储的内容。

位置信息（Location Information）。MLIR 为位置信息提供了一种紧凑的表示方式，并鼓励在整个系统中处理和传播这些信息，遵循可追溯性原则。它可以用来保存产生某个操作的源程序堆栈跟踪，以便生成调试信息。它标准化了从编译器发出诊断信息的方式，并被广泛的测试工具使用。位置信息也是可扩展的，允许编译器引用现有的位置跟踪系统、高级 AST 节点、LLVM 风格的文件-行-列地址、DWARF 调试信息等。

区域（Regions）和块（Blocks）。一个操作实例可以附加一个区域列表。区域提供了 MLIR 中的嵌套机制：它包含一个块列表，每个块包含一个操作列表（这些操作可能包含更深层的区域）。与属性一样，区域的语义由附加它的操作定义，但是区域内的块（如果多于一个）会形成一个控制流图（CFG）。例如，图 3 中的 affine.for 操作是一个循环，其单块循环体作为区域附加，位于 ({ 和 }) 分隔符之间。该操作指定了跨区域的控制流。在这个例子中，循环体被重复执行，直到达到上界。

区域内的控制流。每个区域的主体是一个块列表，每个块都以一个终止符（terminator）操作结束，该操作可以有后继块，控制流可以转移到这些后继块。每个终止符（例如，“switch”、“条件分支”或“unwind”）都定义了自己的语义。它可以选择将控制流转移到同一区域中的另一个块，或者将其返回给包含该区域的操作。后继者的图定义了一个 CFG，允许在一个区域内进行标准的基于 SSA 的控制流。

函数式 SSA。MLIR 不使用 φ 节点，而是采用函数式的 SSA 形式【20，SSA is functional programming, 1998, ACM SIGPLAN NOTICES】，其中终止符将值传递给由后继块定义的块参数。每个块都有一个（可能为空的）类型化块参数列表，这些参数是常规值并遵循 SSA 规则。终止符操作的语义定义了在控制转移后块的参数将取什么值。对于区域的第一个（入口）块，这些值由包含它的操作的语义定义。例如，affine.for 使用入口块参数 %arg4 作为循环归纳变量。最后，这种显式的图设计和操作的可扩展性让人联想到节点之海（sea-of-nodes）表示【21，A simple graph-based intermediate representation, 1995, IR ’95】：这种联系是故意的，并且是选择 MLIR 的 SSA 风格的主要影响因素。

值的支配与可见性。操作只能使用在作用域内的值，即根据 SSA 支配关系、嵌套结构以及包含操作施加的语义限制可见的值。在一个 CFG 内，如果值遵循标准的 SSA 支配关系，即控制流保证在到达使用点之前通过定义点，那么这些值是可见的。

基于区域的可见性。基于区域的可见性是根据区域的简单嵌套定义的：如果一个操作的操作数位于当前区域之外，那么它必须在词法上定义在当前使用点区域的上方和外部。这使得 affine.for 操作内部的操作能够使用在外部作用域中定义的值。

隔离性。MLIR 还允许将操作定义为与上层隔离（isolated from above），这表明该操作是一个作用域屏障——例如，std.func 操作定义了一个函数，函数内部的操作引用函数外部定义的值是无效的。除了提供有用的语义检查外，包含与上层隔离的操作的模块可以由 MLIR 编译器并行处理，因为没有用-定义链可以跨越隔离屏障。这对于利用多核机器进行编译非常重要。所有这些设计选择都突出了渐进性原则，同时在某个概念看起来不够通用和基础以至于不能内置时，倾向于简约性。

符号（Symbols）和符号表（Symbol Tables）。操作可以附带一个符号表。这个表是一种标准化的方式，用于将以字符串表示的名称与称为符号的 IR 对象关联起来。IR 不规定符号的用途，将其留给操作定义。符号对于那些不需要遵循 SSA 规则的命名实体最为有用：它们不能在同一个表中被重新定义，但可以在定义之前被使用。例如，全局变量、函数或命名模块可以表示为符号。没有这种机制，就无法定义例如在其定义中引用自身的递归函数。如果一个带有符号表的操作有关联的区域，而区域中包含类似的操作，那么符号表可以嵌套。MLIR 提供了一种从操作引用符号（包括嵌套符号）的机制。

方言（Dialects）。MLIR 使用方言来管理可扩展性，它在唯一的命名空间下提供了操作、属性和类型的逻辑分组。方言本身不引入任何新的语义，而是作为一种逻辑分组机制，提供通用的操作功能（例如，方言中所有操作的常量折叠行为）。它们在遵循简约性原则的同时，组织了语言和领域特定语义的生态系统。方言命名空间在操作码中作为以点分隔的前缀出现，例如，图 3 使用了 affine 和 std 方言。

方言的模块化设计。将操作、类型和属性分离到不同的方言中是概念性的，类似于设计一组模块化的库。例如，一个方言可以包含用于操作硬件向量的操作和类型（例如，shuffle、insert/extract element、mask），而另一个方言可以包含用于操作代数向量的操作和类型（例如，absolute value、dot product 等）。两个方言是否使用相同的向量类型以及该类型属于哪里，是留给 MLIR 用户做出的设计决策。

方言的混合使用。虽然可以将所有操作、类型和属性放在一个单一的方言中，但这会因为大量同时存在的概念和名称冲突等问题而很快变得难以管理。尽管每个操作、类型和属性都恰好属于一个方言，但 MLIR 明确支持混合使用方言以实现渐进式降低。来自不同方言的操作可以在 IR 的任何级别、任何时间共存，它们可以使用在不同方言中定义的类型等。方言的混合使用带来了更强的可重用性、可扩展性，并提供了灵活性，否则开发者将不得不采用各种不可组合的变通方法。

类型系统（Type System）。MLIR 中的每个值都有一个类型，它在产生该值的操作中或在定义该值为参数的块中指定。类型编码了关于值的编译时信息。MLIR 中的类型系统是用户可扩展的，并且可以例如引用现有的外部类型系统。MLIR 强制执行严格的类型相等性检查，不提供类型转换规则。操作使用类似函数的尾随语法列出其输入和结果类型。在图 3 中，std.load 从内存引用和索引类型映射到它加载的值的类型。

类型理论的视角。从类型理论的角度来看，MLIR 只支持非依赖类型，包括平凡类型、参数化类型、函数类型、和类型和积类型。虽然通过结合操作与符号和用户定义类型可以实现依赖类型系统，但这类类型对 IR 来说将是不透明的。

内置类型。为方便起见，MLIR 提供了一套标准化的常用类型，包括任意精度整数、标准浮点类型以及简单的常用容器——元组、多维向量和张量。这些类型仅仅是一种实用工具，并非必须使用，这体现了简约性。

函数与模块。与传统 IR 类似，MLIR 通常被构造成函数和模块。然而，在 MLIR 中这些并不是独立的概念：它们是作为内置方言中的操作来实现的，这再次体现了设计的简约性。

模块的定义。模块是一个带有单个区域的操作，该区域包含单个块，并由一个不转移控制流的虚拟操作终止。像任何块一样，其主体包含一个操作列表，这些操作可能是函数、全局变量、编译器元数据或其他顶级构造。模块可以定义一个符号以便被引用。

函数的定义。类似地，函数是一个带有单个区域的操作，该区域可能包含零个（在声明的情况下）或多个块。内置函数与“std”方言的“call”和“return”操作兼容，这些操作分别将控制权转移到函数和从函数返回。其他方言可以自由定义自己的类函数操作。

A4 实验环境与结果

实验环境

本文的评估主要通过展示 MLIR 被多个不同项目采用和使用来证明其通用性和有效性，强调了其软件工程性质的贡献。
* 软件配置:
* 构建工具: 构建 MLIR 需要使用 LLVM C++ 工具链（包括 clang 和 lld）、Ninja 构建系统。
* 源代码管理: 通过 git 从 github.com/llvm/llvm-project.git 下载。
* 运行环境: 推荐在 Linux 操作系统上运行。
* 依赖项目: MLIR 作为 llvm-project 的一部分，需要启用 mlir 项目。
* 硬件配置: 论文没有指定统一的硬件配置，因为 MLIR 的应用场景非常广泛，涵盖了从数据中心集群、移动设备到专用硬件加速器（如 GPU、TPU）等各种异构硬件。
* 数据集与模型: 论文中提到的应用（如 TensorFlow、Lattice Regression）涉及机器学习模型，但没有具体说明用于评估的数据集。重点在于展示 MLIR 作为基础设施处理这些模型的能力。

实验结果

MLIR 作为一个旨在通用化并驱动各种编译器项目的系统，其主要评估指标是展示其在多样化项目中的采用和使用情况。

社区采纳情况：
MLIR 已经成为一个不断发展的开源项目，其社区横跨学术界和工业界。例如，首届关于在高性能计算中使用 MLIR 的学术研讨会吸引了来自 4 个国家 16 所大学和 4 个国家实验室的参与者。在 2019 年 LLVM 开发者会议上，MLIR 也得到了 14 家跨国公司的支持。社区的采纳和参与度是可用性和需求的代表性指标。目前，有超过 26 个方言正在公开或私下开发中，7 个跨公司的项目正在用 MLIR 替换其自定义基础设施。

A. TensorFlow 图：
MLIR 的一个关键用例是支持机器学习框架的开发。TensorFlow 的内部表示是一种高级数据流计算图，其节点是可放置在不同设备（包括特定硬件加速器）上的计算。
* 实验内容: 在 TensorFlow 中使用 MLIR 来建模其内部表示，并执行各种转换，如图 1 所示的用例。这包括从简单的代数优化到针对数据中心集群和异步硬件加速的并行与分布式执行的图重定向，以及通过 XLA 等领域特定代码生成器生成高效原生代码。
* 实验结果: MLIR 成功地模拟了 TensorFlow 图的异步并发模型（图 6），其中数据流图通过隐式 future 解除同步，而带副作用的操作通过显式控制信号串行化。MLIR 为 TensorFlow 模型和低级 LLVM IR 提供了相同的基础设施和转换能力，可表达如图死代码消除、常量折叠、公共子图消除、布局优化等图级转换，以及混合精度优化、操作融合等领域特定优化。

%0 = tf.graph (%arg0 : tensor<f32>, %arg1 : tensor<f32>, %arg2 : !tf.resource) {
  // 这些操作的执行是异步的，%control 返回值
  // 可用于施加额外的运行时顺序，例如下面显式地将
  // 对变量 %arg2 的赋值排在读取之后。
  %1, %control = tf.ReadVariableOp(%arg2) : (!tf.resource) -> (tensor<f32>, !tf.control)
  %2, %control_1 = tf.Add(%arg0, %1) : (tensor<f32>, tensor<f32>) -> (tensor<f32>, !tf.control)
  %control_2 = tf.AssignVariableOp(%arg2, %arg0, %control) : (!tf.resource, tensor<f32>) -> !tf.control
  %3, %control_3 = tf.Add(%2, %arg1) : (tensor<f32>, tensor<f32>) -> (tensor<f32>, !tf.control)
  tf.fetch %3, %control_2 : tensor<f32>, !tf.control
}

图 6. TensorFlow 图在 MLIR 中的表示。

B. 多面体代码生成：
affine 方言是为支持渐进式降低而设计的简化多面体表示。
* 实验内容: 使用 MLIR 的 affine 方言来探索用于加速器的多面体代码生成。图 7 展示了其自定义语法。
* 实验结果: affine 方言相比于现有方法【25-29】有几个关键区别和优势：
1. 丰富的类型: 结构化内存引用类型包含布局图，使循环和数据转换更好地组合。
2. 混合抽象: affine 循环体中可以使用基于 SSA 值的操作，使得传统编译器分析和转换可以与多面体转换交错进行。
3. 更小的表示鸿沟: MLIR 维持了高级循环结构，无需从低级表示中进行脆弱的“提升”操作。
4. 编译速度: MLIR 的方法不依赖于计算复杂度高的多面体扫描算法，因为循环结构在 IR 中得以保留，这对于即时编译（JIT）和预编译（AOT）场景都至关重要。

// 仿射循环是带有区域的操作。
affine.for %arg0 = 0 to %N {
  // 只允许循环不变量、循环迭代器以及它们的仿射函数。
  affine.for %arg1 = 0 to %N {
    // 仿射 for 循环的循环体遵循 SSA。
    %0 = affine.load %A[%arg0] : memref<? x f32> // 结构化内存引用 (memref) 类型可以有
                                                // 仿射布局图。
    %1 = affine.load %B[%arg1] : memref<? x f32, (d0)[s0] -> (d0 + s0)>
    %2 = mulf %0, %1 : f32
    // 仿射加载/存储可以使用仿射表达式作为下标。
    %3 = affine.load %C[%arg0 + %arg1] : memref<? x f32>
    %4 = addf %3, %2 : f32
    affine.store %4, %C[%arg0 + %arg1] : memref<? x f32>
  }
}

图 7. 使用 affine 方言自定义语法的多项式乘法。

C. Fortran IR (FIR)：
LLVM 的 Fortran 前端 "flang" 正在使用 MLIR。
* 实验内容: flang 使用 MLIR 支持 Fortran 特定的高级优化，如高级循环优化、数组拷贝消除、调用特化和去虚化。
* 实验结果: FIR 能够将 Fortran 虚拟分派表建模为头等概念（如图 8），这使得实现稳健的去虚化传递成为可能。使用 MLIR 让 flang 开发者能专注于领域 IR 设计而非重复实现基础设施。此外，MLIR 的设计使得重用其他方言成为可能，例如可以在 Fortran 和 C 前端之间共享一个语言无关的 OpenMP 方言，或通过共享 GPU 方言来支持 OpenACC。
图 8. FIR 对动态虚函数分派表提供头等支持。

D. 领域特定编译器：
* 优化 MLIR 模式重写:
* 实验内容: 将 MLIR 的模式重写系统本身表示为一个 MLIR 方言，从而利用 MLIR 基础设施来优化重写器。
* 实验结果: 成功地构建和优化了高效的有限状态机（FSM）匹配器和重写器，包含了其他系统（如 LLVM SelectionDAG）中看到的 FSM 优化。
* Lattice 回归编译器:
* 实验内容: 为一种名为 Lattice 回归的机器学习技术构建一个新的编译器，取代了原先基于 C++ 模板的实现。
* 实验结果: 仅投入 3 个人月的努力，就开发出了新的编译器。在一个生产模型上，性能提升高达 8 倍，同时还提高了编译过程的透明度。

A7 补充细节

MLIR 的设计在促进新语言和编译器抽象建模的同时，也重用了现有的通用抽象。许多问题的解决方案实际上是“添加新的操作、新的类型”，并可能将它们集合成“一个新的方言”。这对编译器工程而言是一个重大的设计转变，带来了新的机遇、挑战和见解。本节将探讨其中的一部分。

A. 可重用的编译器传递（Passes）

通过抽象操作属性实现通用传递。在一个 IR 中表示多层次抽象的能力，激励了那些跨越这些层次进行操作的传递。MLIR 通过反转通常的方法来处理可扩展性：因为操作（Op）的数量比传递（Pass）多，所以让操作了解传递比反过来更容易。这也提高了模块性，因为特定于方言的逻辑是在方言内部实现的，而不是在核心转换中。由于传递很少需要了解一个操作的所有方面，MLIR 依赖以下机制来实现通用的传递。

操作特性（Operation Traits）。许多常见的“基础”编译器传递，如死代码消除（Dead Code Elimination）或公共子表达式消除（Common Subexpression Elimination），依赖于诸如“是终止符”或“是可交换的”等简单属性。我们将这类属性定义为操作特性（Op Traits）。一个操作无条件地展现一个特性，例如，“标准分支”操作永远是一个终止符。对于许多传递来说，知道一个操作具有一组特性就足以对其进行操作，例如交换操作数或移除没有副作用且没有用户的操作。

特性作为验证钩子。特性可以作为验证钩子，从而在多个具有该特性的操作之间共享逻辑。例如，“与上层隔离（isolated from above）”特性会验证操作中的区域没有使用在包含该操作的区域中定义的值。它允许对函数、模块和其他自包含结构进行通用处理。

接口（Interfaces）。当无条件的、静态的行为表达能力不足时，可以通过接口来参数化处理，这是一个借鉴自面向对象编程的概念。接口定义了对 IR 对象行为的一种视图，它抽象掉了不必要的细节。与特性不同，接口是由 IR 对象实现的，使用任意的 C++ 代码，可以为不同的对象产生不同的结果。例如，“call”操作实现了一个“类调用（call-like）”接口，但该操作的不同实例调用不同的函数。

基于接口的传递实现。MLIR 传递可以依据接口来实现，与任何选择被该传递处理的操作建立一个契约。继续以类调用为例，考虑 MLIR 的内联传递，它可以作用于 TensorFlow 图、Flang 函数、函数式语言中的闭包等。这样的传递需要知道：（1）将一个操作内联到给定区域是否有效，以及（2）如何处理内联后出现在块中间的终止符操作。

接口作为通信契约。为了向一个操作查询这些属性，该传递定义了一个专门的接口，以便操作可以向 MLIR 注册它们的实现，从而从内联中受益。内联传递将保守地处理，即忽略任何未实现相应接口的操作。常量折叠也是通过相同的机制实现的：每个操作实现“fold”接口，提供一个函数，如果操作是可折叠的，该函数可能会产生一个持有该值的属性。更通用的规范化可以类似地实现：一个接口填充一个适用于模式重写的规范化模式列表。这种设计将通用逻辑与操作特定逻辑分离开来，并将后者置于操作自身之中，从而减少了 LLVM 中类似“InstCombine”、“PeepholeOptimizer”等众所周知的维护和复杂性负担。

接口的方言级实现。接口可以由方言而不是特定的操作来实现，这使得可以共享行为或委托给外部逻辑，例如在对 TensorFlow 操作进行常量折叠时。类型和属性也支持接口，例如，一个加法操作可能支持任何自声明为“类整数（integer-like）”并具有可查询的有符号性语义的类型。

B. 方言特定的传递

针对特定语义的传递。定义特定于某些方言的传递是有效且有用的，这些传递可以由它们所设计的方言中操作的完整语义来驱动。这些传递在 MLIR 系统中与在其他编译器系统中一样有用。例如，代码生成器可能希望根据特定的机器约束条件对机器指令进行自定义调度，或者使用其他不适合更广泛框架的技巧。对于新的转换来说，这是一个简单而有用的起点，此时并不需要通用化。

C. 混合使用方言

跨方言组合的威力。MLIR 最深刻（但也最难理解）的一个方面是，它允许并鼓励将来自不同方言的操作混合到同一个程序中。虽然某些情况相对容易理解（例如，在同一个模块中保存宿主机和加速器的计算），但最有趣的情况发生在方言直接混合时——因为这启用了一整类我们在其他系统中未曾见过的重用方式。

affine 方言的组合示例。考虑 IV-B 节中描述的 affine 方言。仿射控制流和仿射映射的定义与仿射区域中包含的操作的语义是独立的。在我们的例子中，我们将 affine 方言与代表简单算术的“标准（standard）”方言（其形式与 LLVM IR 类似，与目标无关）以及用于内部加速器的多个目标特定的机器指令方言结合起来。其他人则将其与来自其他问题领域的抽象结合起来。

基础设施的分解与重用。能够重用通用的多面体转换（通过操作接口获取特定转换中操作的语义）是一种强大（且令我们兴奋）的分解编译器基础设施的方式。另一个例子是，一个 OpenMP 方言可以在多种源语言 IR 中被使用和重用。

D. 并行编译

利用隔离性实现并行化。MLIR 的一个重要方面是能够使用多核机器来提高编译速度。特别是，“与上层隔离（isolated from above）”特性（V-A 节）允许诸如函数之类的操作选择加入 MLIR 传递管理器支持的并发 IR 遍历机制。实际上，这个特性保证了 SSA 的用-定义链不能跨越区域边界，因此可以被隔离处理。MLIR 也没有整个模块范围的用-定义链，而是通过符号表（III 节）引用全局对象，并将常量定义为带有属性的操作（III 节）。

E. 互操作性

与现有系统的集成挑战。我们的工作涉及与大量现有系统的互操作，例如，编码为协议缓冲区的机器学习图、包括 LLVM IR 在内的编译器 IR、专有指令集等。通常，这些表示法存在一些次优或不幸的设计决策，这些决策在现有系统的背景下是合理的，但 MLIR 的能力使得可以实现更具表达力的表示。因为导入器和导出器 notoriously 难以测试（测试用例通常是二进制的），我们希望确保它们的复杂性被最小化。

通过方言实现可靠的互操作。解决方案是定义一个尽可能直接对应于外部系统的方言——允许以一种简单且可预测的方式与该格式进行往返转换。一旦 IR 被导入到 MLIR 中，它就可以使用所有的 MLIR 基础设施被提升和降低到一个更方便的 IR，这使得这些转换可以像所有其他 MLIR 传递一样被测试。这类方言有很多例子，包括将 LLVM IR 映射到 MLIR 的 LLVM 方言。这种方法对我们来说效果很好，并且 MLIR 的工具在为这些外部文件格式编写测试时也很有用。

F. 无主见设计带来的新挑战

抽象设计的艺术。虽然 MLIR 允许人们定义几乎任意的抽象，但它对于应该做什么提供了很少的指导：在实践中什么做得更好或更差？我们现在有了一些工程师和研究人员将这些技术和方法应用于新问题领域的经验，并意识到编译器 IR 设计和抽象设计的“艺术”在编译器和语言领域并没有被很好地理解——许多人在已建立系统的约束下工作，但相对较少的人有机会自己定义抽象。

挑战与机遇。这是一个挑战，但也是未来研究的另一组机遇。更广泛的 MLIR 社区正在积累关于这些抽象设计权衡的专业知识，我们预计这会随着时间的推移成为一个富有成果的研究领域。

G. 展望未来

未知的探索领域。MLIR 的设计与其他编译器基础设施有很大的不同，以至于我们仍在学习中——即使在构建并将其应用于许多不同系统之后。我们相信还有很多东西有待发现，需要几年的研究才能更好地理解设计要点并建立最佳实践。例如，树外方言的兴起、越来越多使用 MLIR 的源语言前端、可能应用于抽象语法树（AST）以及应用于结构化数据（如 JSON、协议缓冲区等）的应用，这些都还处于非常早期的阶段，很可能会揭示有趣的新挑战和机遇。更好地支持即时编译和精确垃圾回收也将是有趣的方向，这可以利用 IR 的模块化和可编程性。

A5 结论

本文介绍了 MLIR，它是为应对设计灵活可扩展的编译器基础设施这一双重科学和工程挑战而提出的具体解决方案。其应用范围广泛，从后端代码生成和异构系统编排，到机器学习的图级建模，再到编程语言和领域特定框架的高级语言语义。我们展示了它在一系列领域中的适用性，并讨论了其研究意义。

受 LLVM 成功的激励并展望未来，我们渴望看到编程语言和高性能计算领域的成熟社区以及领域专家如何从引入更高级别、特定于语言的 IR 中受益。我们也相信，MLIR 能够催化新的研究领域，以及教授编译器和 IR 设计艺术的新方法。

A6 附录

A. 摘要

工件摘要。本文的工件包括 MLIR 系统、如何下载和构建它的说明，以及指向 TensorFlow 中与 MLIR 相关源代码的链接。

B. 工件清单（元信息）

工件的元信息。
* 程序: MLIR
* 编译: LLVM C++ 工具链
* 运行时环境: 推荐 Linux
* 公开可用?: 是
* 存档: DOI 10.5281/zenodo.4283090

C. 描述

1) 如何交付: 要下载 MLIR，请运行 git clone https://github.com/llvm/llvm-project.git。下载和构建 MLIR 的说明也可在 https://mlir.llvm.org/getting_started 上找到。更多信息可在 http://mlir.llvm.org 获取。

2) 软件依赖: 下载 MLIR 需要 git。构建 MLIR 需要 Ninja (https://ninja-build.org/) 和一个可用的 C++ 工具链，包括 clang 和 lld。

D. 安装

构建和测试 MLIR。要在 Linux 上构建和测试 MLIR，请执行以下命令：

mkdir llvm-project/build
cd llvm-project/build
cmake -G Ninja ../llvm \
    -DLLVM_ENABLE_PROJECTS=mlir \
    -DLLVM_BUILD_EXAMPLES=ON \
    -DLLVM_TARGETS_TO_BUILD=&quot;X86;NVPTX;AMDGPU&quot; \
    -DCMAKE_BUILD_TYPE=Release \
    -DLLVM_ENABLE_ASSERTIONS=ON \
    -DCMAKE_C_COMPILER=clang \
    -DCMAKE_CXX_COMPILER=clang++ \
    -DLLVM_ENABLE_LLD=ON
cmake --build . --target check-mlir

E. 应用

在 TensorFlow 中的应用。MLIR 在 TensorFlow 中的使用可以在位于 https://github.com/tensorflow/tensorflow/tree/master/tensorflow/compiler/mlir/ 的代码中观察到。位于 tensorflow/tests 子目录中的测试包含说明 TensorFlow 图表示和转换的 MLIR 代码片段。从源代码构建 TensorFlow 的说明可在 https://www.tensorflow.org/install/source 获得。