MLIR：摩尔定律终结的编译器基础结构论文解读-CSDN博客

本文链接：https://blog.csdn.net/just_sort/article/details/122786078

【GiantPandaCV导语】本文是对MLIR的论文解读以及实践，这里的实践指的是把MLIR的要点在OneFlow Dialect中进行了对应，并解释了每个要点的实现方法以及这些要点的相关性，算是对MLIR学习过程的一个阶段总结。本文分为2大部分，第一部分为1-6节，主要是阅读MLIR论文，第7节是根据OneFlow Dialect解释论文中提到的MLIR基础架构中的要点如Type，Attribute，Operation，Trait，Interfaces，Region，Block等等。本文只是想起到抛砖引玉的效果让更多小伙伴了解MLIR这个编译架构，如果对你有帮助欢迎关注一下我这个从零开始学深度学习编译器的github仓库：https://github.com/BBuf/tvm_mlir_learn。

0x0. 前言

之前以MLIR的Toy Tutorials教程为起点了解了一点MLIR，然后又对MLIR的ODS，DRR要点以及Interfaces等知识进行了整理。在继续学习分享MLIR的相关知识前，我想对MLIR做一个总结。而要了解MLIR的全貌，阅读MLIR论文是一个不错的方式。这篇文章在论文阅读的基础上我还做了一个思维导图把MLIR实现Dialect的组件画出来了，再以OneFlow的Dialect为例子详解了这些组件是如何实现的以及它们的关系。相信看完本文会对不熟悉MLIR的小伙伴有一些帮助和启发。希望起一个入门效果。

本文阅读方法步骤大概是（数字代表先后顺序）：

标题
摘要
引言
结论
相关工作
MLIR设计相关
评论
参考文章

MLIR论文链接：https://arxiv.org/pdf/2002.11054.pdf

0x1. 标题

MLIR论文标题

论文标题翻译为，MLIR: 摩尔定律终结的编译器基础结构 。从题目可以知道MLIR是一个编译器架构，终结摩尔定律这个不太好理解，我们需要往后看看。我们还可以发现MLIR是Chris大神（LLVM，CLang、Swift项目的发起人）领衔的，这让MLIR项目的质量有很大的保障，相信这也是目前MLIR编译架构非常流行的原因之一？

0x2. 摘要

这篇文章提出了MLIR，这是一种构建可重用、可扩展编译器基础结构的新方法。MLIR旨在解决软件碎片化，改进异构硬件的编译过程，大大降低了构建领域特定编译器的成本，并有助于和现有的其它编译器互相连接。MLIR还有助于在不同抽象级别、不同跨应用程序域、不同硬件目标和执行环境下改善code generators、translators和optimizers的设计和实现。贡献包括：(1) 讨论MLIR作为文本研究成果可能的扩展和进化，并指出这个新方法在设计、语义、优化规范、系统和工程等方面带来的挑战和机遇。(2) 评估MLIR作为可减少构建编译器成本的通用架构-通过描述各种用例，显示本文研究成果在未来编程语言、编译器、执行环境和计算机体系结构方面的研究和教学机会。然后还介绍了MLIR设计基本原理、结构和语义。

这一节主要是讲了一下MLIR的卖点，即MLIR是一个新的编译器架构，它着力于解决软件碎片化并降低了构建特定领域编译器的成本。

其实今天来看软件碎片化问题MLIR是没有完全解决的，它只是把软件碎片化问题转移为各个Dialect间的碎片化，然后这些Dialect又属于同一种语言可以混用以此缓解了软件碎片化带来的影响。这里为什么是缓解而不是完全解决呢？首先我理解软件的碎片化应该就是针对N种前端框架（如TensorFlow，PyTorch…）和M种后端（GPU，CPU…）的适配问题，如果没有一个中间的IR表示那么这个适配的工作量是 $N * M$ ，然后微软提出的ONNX尝试作为一个中间的IR使得这个 $N * M$ 的问题变成M，即所有的前端框架都可以转换到ONNX，只需要适配ONNX的后端就可以了。但理想和现实往往不一样，ONNX为了适配各种前端框架捏了一系列更加通用的算子（opset）来匹配各个前端框架的算子语意，但这样做的后果就是前端框架和ONNX互转的时候往往引入了一些新的胶水Op使得IR变得更加复杂。说回MLIR，各个前端框架把自己的IR对接为MLIR的Dialect上之后要走相当多的DialectConversion才可以到可以做代码生成的LLVM IR。虽然各个Dialect可以混用这样就不会出现ONNX里面那种多出胶水Op的情况，但Dialect可以混用不代表DialectConversion的畅通无阻。假设Dialect A下有一个Op X我们要将其转换为Dialect B下的Op，并且Dialect B下面没有对应Op X语意的Op或Dialect B下对应Op X语意的Op和X的语意有一些差距那么我们必须对Dialect B进行扩展以满足需求，这和ONNX不断增加Opset似乎没什么两样，并且MLIR的Dialect链路可能会很长，所以这种情况下感觉会比ONNX更麻烦。但乐观的想，MLIR开源到现在就2-3年，相信随着各个Dialect的丰富，这种碎片化风险真的会逐渐变小。

而降低构建特定领域编译器成本应该指的是在MLIR的生态更加完善之后，理论上我们只需要在对应硬件上实现一个边界Dialect，然后在这个Dialect中定义硬件的Operation，之后就可以选取生态中已有的Dialect来构建一个完整的编译流程即可。

0x3. 引言

编译器设计是一个成熟的领域，包括许多广为人知的算法，可用于代码生成、静态分析、程序转换等。编译器设计领域已发展出许多成熟技术平台，这些平台现在已经在整个编译器社区大规模应用，包括LLVM编译器基础结构[25]、Java虚拟机（JVM）[26]等系统。这些流行系统的一个共同特征是它们的“one size fits all”方法，即与系统接口的是单一抽象级别，例如LLVM中间表示（IR）大致是“C with vectors”，JVM提供了一个“具有垃圾收集器的面向对象类型系统（object-oriented type system with a garbage collector）”抽象。这种“one size fits all”的方法非常有价值，因为从源语言（C/C ++和Java）到这些抽象领域的映射非常直接。

同时，一些问题在更高或者更低的抽象层级建模会更好，比如在LLVM IR上对C ++代码进行源代码级分析十分困难。注意到，许多语言（例如Swift，Rust，Julia，Fortran）都开发了自己的IR，以解决这些语言领域特定的问题，例如语言/库相关的优化、flow-sensitive 类型检查（例如线性类型）和优化lowering过程的实现。类似地，机器学习系统通常将“ML graphs”用作领域特定的抽象。

尽管领域特定IR的开发是一项已经被充分研究的技术，但其工程和实现成本仍然很高。对于这些系统的实现者而言，有时候基础结构的质量不一定是优先考虑的因素。这可能导致编译器系统的实现质量降低，包括一些用户常见的问题，例如编译时间慢、错误的实现、诊断质量欠佳、优化代码的调试体验差等等。

MLIR项目的目的就是要应对这些编程语言设计和实现方面的挑战—通过非常方便的定义和引入新的抽象级别，并提供“in the box”基础架构来解决常见的编译器工程问题。 MLIR的做法是：（1）标准化基于静态单赋值（SSA）的IR数据结构（2）提供用于定义IR dialect的声明系统，（3）提供广泛的通用基础结构（包括文档、解析和打印逻辑、位置跟踪、多线程编译支持、pass管理等）。

论文探讨了MLIR系统的各个设计要点，将作者们的经验应用于不同的问题，并讨论了这项工作可能对编程语言设计和教学产生的影响。

论文的贡献可以总结为如下几点：

描述了一种对工业界和学术界有重要应用价值的新型编译器基础结构。
提出了一种构建可扩展和模块化编译器系统的新方法。
选择了一些MLIR在不同领域的应用，说明了系统的通用性。
分享了在MLIR基础架构上开发编译系统的经验。

这一节还提到了MLIR的产生动机：

我们首先意识到现代机器学习框架由许多不同的编译器、图技术和运行时系统组成（请参见Figure 1），但是这些部分没有共享公共的基础结构或设计观点，而且有些部分没有遵循最佳编译器设计实践，导致的后果是用户可以明显感觉到不便，包括不完善的错误消息、边界情况下的错误、不可预测的性能，以及难以支持新硬件。

我们很快意识到，整个编译器行业都存在一个类似的问题，那就是，诸如LLVM之类的现有编译系统在跨多语言实现的统一和集成方面非常成功，但是现代高级语言通常最终会构建自己的高级IR，并重复发明许多相同的更高层抽象技术（请参见Figure2）。同时，在LLVM社区经常出现一些争论，比如，如何最好地表示并行结构，如何共享常见的前端Lowering基础架构实现（例如，用于C调用约定或诸如OpenMP之类的跨语言功能），但都没有得出令人满意的解决方案。

面对这些挑战，我们认为我们无法承担实现N个改进编译器的工作量，因此我们需要构建一个更通用的解决方案。我们可以花精力开发一套高质量的基础架构，这会让多个领域受益，会让我们能够逐步升级现有系统，让我们能够更轻松地解决眼下紧迫的问题，例如专用加速器的异构编译。现在，我们在构建和部署基于MLIR的系统方面积累了大量经验，可以回顾一下MLIR基础架构的原理和设计，并讨论为什么要朝这个方向发展。

这一节列举了一下相关工作以及MLIR的产生动机，以此加强说明MLIR的创新点和贡献。

0x4. 结论

本文介绍了MLIR，可用作构造编译器的灵活且可扩展的基础结构。本文描述了MLIR的具体设计，展示了其在一系列重要领域中的适用性，并描述了许多原创研究和工程意义。

展望未来，我们希望看到编译器社区（例如Clang C和C ++编译器）和不同领域的专家如何能从更高级的、语言特定IR中受益。我们也想知道，MLIR是否能为教授编译器和IR设计技术提供新的方法，并希望看到这种基础设施加速新领域的研究。

这里介绍了一系列未来的工作方向，感兴趣的可以自行看一下。由于我对这部分不太了解，这里就不继续看了。

0x5. 相关工作

MLIR是一个涵盖多个不同领域的项目。虽然其基础设施提供了一个新的系统，但组成基础设施的各个组件在相关文献中都已有类似模块。

MLIR是类似于LLVM[25]的编译器基础结构，但LLVM在标量优化和同构编译做得很好，而MLIR的目标是将各种数据结构和算法建模为第一优先级的值和Operations，包括张量代数和算法、图表示以及异构编译。MLIR 允许混合匹配优化将编译pass分解为组件并重新定义lowering。这主要归功于模式重写基础设施，将完整的变换捕获为小型局部模式的组合，并控制在单个操作的粒度上应用哪些模式进行重写。自动扩展、形式化和验证重写逻辑将是重要的下一步 [9, 27]。在后端，MLIR 的 DDR 类似于 LLVM 的指令选择基础设施，支持以多结果模式和规范作为约束的可扩展操作[49]。

许多编程语言和模型都解决了硬件异构问题。同构编程模型OpenMP基于StarSs和OpenACC[34,31]等较早的建议，增加了对卸载（offloading）任务和加速器并行区域[32]的支持。 C++ AMP、HCC和SyCL利用传统的Clang/LLVM流程和C++为硬件加速提供高级抽象[46]。但是，所有这些例子都依赖于宿主语言（通常为C++）中的已有优化来减轻抽象造成的损失，从而将高级构造快速lower到对运行时执行环境的调用。扩展 LLVM IR 的并行中间表示解决了部分问题，但传统上专注于同构设置 [23, 42] 。迄今为止，最有雄心的工作可能是Liquid Metal[3]，其中提供了协同设计的领域特定语言（DSL），以及将被管理对象的语义转换为静态的、向量的或可重配置硬件的编译流程。然而，在其Lime编译器中，大部分工作量都放在将round对象装配到square硬件中（Kou和Palsberg [24]）。 MLIR通过可扩展的Operation和Type集合，为包含异构特性的高级语言提供直接嵌入手段，同时提供了一个通用基础结构，可逐步lowering这些结构，并最大程度地在不同目标之间重用通用组件。

解决语言异构性已成为元编程系统，尤其是多阶段编程的长期目标。Lightweight Modular Staging（LMS）[39]是最新的技术框架和运行时代码生成器，提供了可生成高效代码并将DSL嵌入Scala的核心组件库。 Delite[45]声称可以大幅提高DSL开发者的效率，同时支持并行和异构执行。我们认为这种方法是对MLIR的补充，为嵌入DSL提供了更高层次的抽象，并通过通用元编程构造实现了优化。

在语言语法上更进一步，ANTLR [33] 是一类解析器生成器，旨在使开发新的编译器前端变得容易。 MLIR 目前没有通用解析器生成，没有 AST 构造或建模功能。将 MLIR 与 ANTLR 等系统相结合，可以生成从用户输入到代码生成的可重用编译器库。

XLA[57]、Glow[40]和TVM[11]通过在机器学习中的应用，解决类似的异构编译目标。但是这些技术都是很具体的代码生成实例，从图形抽象开始，针对的是加速器的多维矢量抽象。这些技术都可以将MLIR用作基础架构，在使用各自现有的代码生成策略的同时，充分利用MLIR的通用功能。同样，来自Halide[36]和TVM的循环嵌套元编程技术，较早的循环嵌套元编程文献[19,41,5,14]，和全自动流程，如PolyMage[28]、Tensor Com-Phenhension[52]、Stripe[58]、Diesel[16]、Tiramisu[4]及其底层多面体编译技术[17,54,8,55]，可以在基于MLIR的编译框架中以不同的代码生成路径共存。序列化和互操作性格式有不同的方法解决ML前端的多样性问题，例如，ONNX[48]的方法是通过提供不同框架都可以映射的通用op集合。ONNX会成为MLIR的一种dialect选择，其他op可以被降级为该dialect。

0x6. MLIR设计相关

0x6.1 设计原则

内置少，一切可定制（Little builtin, everything customizable） MLIR系统基于最少量的基本概念，大部分IR都完全可定制。在设计时，应当用少量抽象（类型、操作和属性，这是IR中最常见的）表示其它所有内容，从而可以使抽象更少、更一致，也让这些抽象易于理解、扩展和使用。广义上讲，可定制性确保编译系统可以适应不断变化的需求，并且更有可能适用于未来的问题。从这个意义上讲，我们应该将IR构建为支持其中间语言的语法和语义、具有可重用组件和编程抽象的基础结构。定制化成功的标准是可以表达多种抽象，包括机器学习图、ASTs、数学抽象（例如多面体）、控制流图（CFGs）和指令级IR（例如LLVM IR），而且从这些抽象到编译系统无需使用任何硬编码的概念。当然，由于兼容性不佳，可定制性会带来内部碎片化的风险。 虽然不可能有一种纯粹的技术解决方案来解决生态系统碎片化问题，但系统应鼓励设计可重用抽象，并假定这些抽象会在设计的预料范围之外被使用。

SSA and regions 静态单赋值形式[15]是编译器IR中广泛使用的表示形式。它提供了许多优点，包括使数据流分析简单和稀疏，因其与continuation-passing风格的关系而被编译器社区广泛理解，并在主要框架中应用。尽管许多现有的IR使用扁平的，线性CFG，但代表更高级别的抽象却推动将嵌套区域（nested regions）作为IR中的第一概念。这超越了传统的region形式，提升了抽象级别（例如，loop trees），加快了编译过程、指令提取或SIMD并行性[22，21，37]。为了支持异构编译，系统必须支持结构化控制流、并发构造、源语言中的闭包等等。一个具体的挑战就是在嵌套区域之上构造基于CFG的分析和转换。

为了这样做，会牺牲LLVM的归一化（normalization），有时甚至牺牲其规范化（canonicalization）属性。能够将各种数据和控制结构降级为更小的归一化（normalized）表示集合，这对于控制编译器的复杂性至为重要。具有pre-header、header、latch、body的规范循环（canonical loop）结构是前端语言中各种循环构造的线性化控制流表示的典型情况。MLIR的目的是为用户提供一种选择，即，根据编译流程中pass的编译算法，可以将嵌套循环捕获为嵌套region或线性化控制流。通过提供这种选择，我们可以脱离LLVM的normalization-only方向，同时保留了在必要时处理更高级别抽象的能力。反过来，采用MLIR的这些方法也产生了如何控制抽象规范化（normalization）的问题，这是下一段的主题。

渐进式降级（Progressive lowering） 编译系统应支持渐进式lower，即，以较小的步幅，依次经过多个抽象级别，从较高级别的表示降低到最低级别。需要多层抽象是因为通用编译器基础结构必须支持多种平台和编程模型。以前的编译器已经在其pipeline中引入了多个固定的抽象级别，例如Open64 WHIRL表示[30]具有五个级别，Clang/LLVM编译器从AST降级到LLVM IR、SelectionDAG、MachineInstr和MCInst。上述降级实现方式较为僵化，因而需要更灵活的设计来支持抽象级别的可扩展性。这对转换的相位排序有深刻的影响。随着编译器专家们实现越来越多的变换pass，这些pass之间开始出现复杂交互。实际情况表明，将优化pass结合起来运行可以使编译器发现更多的程序有用信息。能说明组合pass好处的例子有混合常量传播、值编号（value numbering）和死代码消除的尝试[13]。一般而言，编译器pass可大致分为四个角色：（1）优化变换（2）使能变换（3）lowering（4）cleanup。编译系统应该允许在单个操作的粒度上混合和匹配这些角色，而不是在整个编译单元上顺序执行这些pass。

保持高层级语意（Maintain higher-level semantics） 系统需要保留分析或优化性能所需的高级语义和计算结构。一旦降低语义再试图提高语义会很难成功，并且将这种信息强行塞进一个低层次IR的环境中通常都有破坏性（例如，在使用调试信息来记录结构的情况下，所有pass都需要进行验证/重新访问）。相反，系统应保持计算结构并逐步lowering到硬件抽象。这时，可以有意识的丢弃结构信息，并且这种丢弃只在不再需要此结构来匹配基础执行模型的情况下才会发生。例如，系统应在整个相关转换过程中保留结构化的控制流，例如循环结构。删除此结构，即转到基于CFG的控制流，实质上意味着将不再在此级别上执行任何变换。在编译器开发中对并行计算结构进行建模的最新技术突出了该任务通常可能是多么困难[23, 42]。

为了允许编译系统的一部分IR保留较高层级的抽象，而另一部分被降低IR层级，在同一IR中混合不同级别的抽象和不同概念必然成为系统的关键属性。比如，自定义加速器的编译器可以在IR中复用系统定义的一些高级结构和抽象，IR同时也可表达加速器特有的基本标量/矢量指令。

IR验证（IR validation） 生态系统的开放性要求有宽泛的验证机制。验证和测试不仅对于检测编译器错误很有用，而且在可扩展的系统中，对验证方法和工具健壮性的需求也在不断提高。验证机制应使得定义简洁和实用，并可以作为正确结果的唯一来源。一个长期目标是复现成功的变换验证 [35、29、50、51] 和现代编译器测试方法 [12] 。在可扩展的编译器生态系统中，验证和测试都还是有待解决的两个问题。

声明式重写模式（Declarative rewrite patterns） 定义表示修饰符应该和定义新抽象一样简单。通用变换应实现为声明式表达的重写规则，并以机器可分析的格式推理出重写的属性，例如复杂性和完成度。重写系统的健全性和效率很高，因此被广泛研究，并已被应用于从类型系统（type systems）到指令选择的众多编译问题。我们（MLIR）的目标是实现前所未有的可扩展性和渐进lowering功能，可以通过许多途径将程序变换建模为重写系统。它还提出了有关如何表示重写规则和策略，以及如何构建能够通过多个抽象级别引导重写策略的机器描述的有趣问题。系统需要在解决这些问题的同时，保持可扩展性并执行合理、单调和可复制的行为。

源位置跟踪和可追溯性（Source location tracking and traceability） 操作的来源（包括其原始位置和应用的变换࿰