跳转至

【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 的小伙伴有一些帮助和启发。希望起一个入门效果。

本文阅读方法步骤大概是(数字代表先后顺序):

  1. 标题
  2. 摘要
  3. 引言
  4. 结论
  5. 相关工作
  6. MLIR 设计相关
  7. 评论
  8. 参考文章

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 表示那么这个适配的工作量是NM,然后微软提出的 ONNX 尝试作为一个中间的 IR 使得这个NM 的问题变成 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),但是这些部分没有共享公共的基础结构或设计观点,而且有些部分没有遵循最佳编译器设计实践,导致的后果是用户可以明显感觉到不便,包括不完善的错误消息、边界情况下的错误、不可预测的性能,以及难以支持新硬件。

Figure1

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

Figure2

面对这些挑战,我们认为我们无法承担实现 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) 操作的来源(包括其原始位置和应用的变换)应易于在系统中追溯。这是为了解决在复杂编译系统中常见的缺乏透明性问题,而在复杂编译系统中,很难了解最终表示是如何从原始表示中构造出来的完整过程。在编译安全性至关重要的敏感应用程序时,这是一个突出的问题,在这类程序中,跟踪 lowering 和优化步骤是软件认证程序的重要组成部分 [43]。当使用安全代码(例如加密协议,或对隐私敏感的数据进行操作的算法)进行操作时,编译器常会碰到看似冗余或繁琐的计算,这些计算会嵌入未被源程序的功能语义完全捕获的安全性或私有属性,而安全代码可以防止旁路暴露或加强代码以防止网络攻击或故障攻击。优化可能会改变或使此类保护完全失效 [56];这种缺乏透明性在安全编译中称为 WYSINWYX[6]。准确地将高层次信息传播到较低层的一个间接目标就是帮助实现安全且可回溯的编译过程。

这一小节实际上说明 MLIR 具有的宏观特性,它是一个具有多层 IR 结构的编译架构,实际上就是多层 Dialect,各个 Dialect 分别对不同的层级概念进行建模。比如 LLVM Dialect 负责系统级别的转换,Linalg,Tensor,Vector 等 Dialect 负责协同生成代码,而 Affine,Math 等 Dialect 用来描述底层计算。

0x6.2 IR 设计细节

本节根据上一节中阐述的原理,介绍 MLIR 中 IR 的设计。

Operations(操作)

MLIR 中的语义单位是一个 “操作”,称为 Op。在 MLIR 系统中,从指令到函数再到模块,一切都建模为 Op。 MLIR 没有固定的 Op 集合,因此允许并鼓励用户自定义扩展 Op。编译器 pass 会保守地对待未知 Op,并且 MLIR 支持通过特征(traits)、特殊的 Operation hooks 和 Interfaces 等方式为 pass 描述 Op 语义。

Op(见 Figure3)具有唯一的操作码(opcode)。从字面上看,操作码是一个字符串,用于标识它所在的 dialect 和操作。Op 可以有零个或多个值作为操作数和结果,并以静态单赋值的形式(SSA)维护操作数和结果。所有值都有一个类型,类似于 LLVM IR。除了操作码、操作数和结果外,Op 还可能具有属性、区域、块参数和位置信息(Attributes, Regions, Block Arguments, and Location Information)。Figure4 说明了值和 Op,标识符是命名值(包),如果包中有多个值,后指定包中值的数量(注:如 Figure3 中的%results:2,表示返回值有 2 个),而 “#” 表示特定值。在一般的文本表示形式中,操作名称是用引号括起来的字符串,后跟括号括起来的操作数。

Figure3

Figure4

Attributes(属性)

MLIR 属性是结构化的编译期静态信息,例如整数常量值、字符串数据或常量浮点值列表。属性有类型,每个 Op 实例都有一个从字符串名称到属性值的开放键值对字典映射。在通用语法描述中,属性在 Op 操作数和其类型之间,键值对列表中的不同键值对用逗号分隔,并用大括号将整个键值对列表括起来。(如 Figure3 中的{attribute="value" : !d.type}以及 Figure4 的{lower_bound = () -> (0), step = 1 : index, upper_bound = #map3})。其中,lower_boundstepupper_bound是属性名称。() -> (0)标识用于内联仿射形式,在这个例子中是产生常数 0 的仿射函数。#map3标识用于属性别名,该属性别名允许将属性值与标签预先关联,并可以在任何需要属性值的地方使用标签。与操作码一样,MLIR 没有固定的属性集。属性的含义由 Op 语义或与属性相关的 dialect 中得出。属性也是可扩展的,允许直接引用外部数据结构,这对于和现有系统集成很有帮助。例如,某个属性可以引用 ML 系统中(在编译期已知的)数据存储的内容。

Location information (位置信息)

MLIR 提供了位置信息的紧凑表示形式,并鼓励在整个系统中处理和传播位置信息。位置信息可用于保留产生 Op 的源程序堆栈踪迹,用以生成调试信息。位置信息使编译器产生诊断信息的方式变得标准化,并可用于各种测试工具。位置信息也是可扩展的,允许编译器引用现有的位置跟踪系统、高级 AST 节点、LLVM 风格的文件 - 行 - 列(file-line-column )地址、DWARF 调试信息或其它高质量编译实现所需的信息。

上面三个要点我们可以基于 Toy 语言的 transpose Op 来加深理解:

%t_tensor = "toy.transpose"(%tensor) {inplace = true} : (tensor<2x3xf64>) -> tensor<3x2xf64> loc("example/file/path":12:1)

结构拆分解释:

  • %t_tensor:这个 Operation 定义的结果的名字,前面的%是避免冲突,见 https://mlir.llvm.org/docs/LangRef/#identifiers-and-keywords 。一个 Operation 可以定义 0 或者多个结果(在 Toy 语言中,只有单结果的 Operation),它们是 SSA 值。该名称在解析期间使用,但不是持久的(例如,它不会在 SSA 值的内存表示中进行跟踪)。
  • "toy.transpose" :Operation 的名字。它应该是一个唯一的字符串,Dialect 的命名空间前缀为 “.”。 这可以理解为 Toy Dialect 中的 transpose Operation。
  • (%tensor):零个或多个输入操作数(或参数)的列表,它们是由其它操作定义的 SSA 值或 block 参数的引用。
  • { inplace = true }:零个或多个属性的字典,这些属性是始终为常量的特殊操作数。 在这里,我们定义了一个名为 “inplace” 的布尔属性,它的常量值为 true。
  • (tensor<2x3xf64>) -> tensor<3x2xf64>:函数形式表示的操作类型,前者是输入,后者是输出。<2x3xf64>号中间的内容描述了张量的尺寸2x3和张量中存储的数据类型f64,中间使用x连接。
  • loc("example/file/path":12:1):此操作的源代码中的位置。

Regions and Blocks(区域和块)

Op 的实例可能附有一系列附加区域。区域为 MLIR 中的嵌套结构提供了实现机制:一个区域包含一系列块,一个块包含一系列操作(操作中可能又包含区域,如 Figure3 所示)。与属性一样,区域的语义由其附加的操作定义,但是区域内部的块(如果有多个)可形成控制流图(CFG)。例如,Figure4 中的affine.for操作是一个循环,其中位于({})定界符之间的单独块是一个区域。 Op 指定了跨区域的控制流。在这个例子中,重复执行主体直到达到循环上限。每个区域的主体是一系列块,每个块以终止符(terminator)操作结尾,终止符操作可能具有后继块,控制流可以转移到后继块。每个终止符(例如 “switch”,“conditional branch” 或“unwind”)定义自己的语义。终止符可以选择将控制流转移到同一区域中的另一个块,或将其返回到包含该区域的 Op。后续块的图定义了 CFG,从而允许区域内有基于标准静态单赋值(SSA)的控制流。MLIR 不使用ϕ 节点,而是使用静态单赋值(SSA)的函数形式。其中,终止符将值传给后继块定义的块参数(block arguments)。每个块都有一个(可能为空)类型化的块参数列表,这些参数是常规值并符合静态单赋值。终止符 Op 的语义定义了在控制权转移后该块的参数会采用的值。对于该区域的第一个(入口)块,值由包含 Op 的语义定义。例如,affine.for使用入口块参数%arg4作为循环归纳变量。

这里表达的意思就是一个 Operation 可能有多个 Region,然后 Region 又是由一系列 Block 组成,然后 Block 又包含一系列 Op。这样就形成了一个嵌套的关系,可以表达作用域和控制流关系。

Value dominance and visibility

Op 只能使用作用域内的值,即根据 SSA 支配、嵌套和包含 Operation 的语义限制可见的值。如果值遵循标准的 SSA 支配关系,则在 CFG 中可以看到这些值,在这些值中,可以确保控件在使用前先经过定义。

基于区域的可见性是根据区域的简单嵌套来定义的:如果 Op 的操作数在当前区域之外,则必须在使用的区域上方用外部词法对其进行定义,这允许affine.for操作中的 Op 使用外部作用域中定义的值。

MLIR 还允许将操作定义为与上方隔离,表明该操作是作用域 barrier(scope barrier)。例如, “std.func” Op 定义了一个函数,该函数内的操作不能引用该函数外定义的值。除了提供有用的语义检查之外,由于没有 use-def 链可以跨过隔离障碍(isolation barriers),包含与上方隔离(isolated-from-above)的 Op 的 Module 也可以由 ML 编译器并行处理。这对于利用多核计算机进行的编译很重要。

Symbols and symbol tables

Op 还可以附加一个符号表。这个符号表是将名称(以字符串表示)与 IR 对象(称为符号)相关联的标准方法。 IR 没有规定符号的用途,而是交由 Op 定义。对于无需遵守静态单赋值规则的命名实体,符号很有用。符号不能在同一表中重复定义,但是可以在定义之前使用符号。例如,全局变量、函数或命名模块可以表示为符号。没有这种机制,就不可能定义递归函数(在定义中引用自己)。如果附带符号表的 Op 的关联区域包含相似的 Op,那么符号表可以嵌套。 MLIR 提供了一种机制来引用 Op 中的符号,包括嵌套符号。

Dialects

MLIR 使用 Dialect 管理可扩展性。Dialect 在一个唯一的命名空间下提供 Ops、属性和类型的逻辑分组。Dialect 本身并未引入任何新的语义,而是用作逻辑分组机制,并且可用于提供 Dialect 通用 Op 支持(例如,dialect 中所有 op 的常量折叠行为)。Dialect 命名空间在操作码中是以 “.” 分隔的前缀,例如,Figure 4 使用的affinestd dialect。

概念上可以将 Ops、类型和属性抽象为 Dialect,这类似于设计一组模块化库。例如,某种 Dialect 可以包含用于对硬件向量进行操作的 Op 和类型(例如,shuffle、insert/extract 元素、掩码等),而另一种 Dialect 可以包含用于对代数向量进行操作的 Op 和类型(例如,绝对值、点积等 )。两种 dialect 是否使用相同的向量类型以及该类型属于哪一个,可以由 MLIR 用户在设计时决定。

我们也可以将所有 Op、类型和属性放在一个 dialect 中,但容易想到,这必然很快就会因为大量概念和名称冲突等问题,导致 Dialect 变得难以管理。尽管每个 Op、类型和属性都只属于一个 dialect,但是 MLIR 明确支持多种 dialect 的混合以便能实现渐进式 lowering。来自不同 dialect 的 Op 可以在 IR 的任何级别共存,可以使用在不同 dialect 中定义的类型,等等。Dialect 的混合可以加强重用性、可扩展性和灵活性。

类型系统

MLIR 中的每个值都有类型,该类型在产生该值的 Op 或将值定义为参数的 Block 中指定。类型为 IR 提供了编译期语义。 MLIR 中的类型系统是用户可扩展的,并且可以引用已有外部类型系统(例如 llvm::Type 或 clang::Type)。 MLIR 强制执行严格的类型等价检查,并且不提供类型转换规则。Op 使用类似尾函数的语法列出其输入和结果类型。Figure4 中,affine.load从内存引用和索引类型映射到加载的值的类型。从类型理论的角度来看,MLIR 仅支持非依赖类型,包括 trivial 类型、参数类型、函数类型、求和和乘积类型。

标准类型

此外,MLIR 提供了一组标准化的常用类型,包括任意精度整数、标准浮点类型和简单的通用容器,如元组(tuple)、多维矢量和张量。这些类型仅是方便 Dialect 开发者,但是不要求一定使用。

Functions and modules(函数和模块)

与常规 IR 相似,MLIR 通常被构造为函数和模块,这些不是 MLIR 的新概念。函数和模块在 builtin dialect 中作为 Op 实现。模块是一个具有单独区域的 Op,这个区域包含了一个单独的块。模块被一个不转移控制流的 dummy Op 终止。

模块定义了一个可以被引用的符号。像任何块一样,其主体包含一系列 Op,这些 Op 可以是函数、全局变量、编译器元数据或其它顶级构造。函数是具有单个区域的 Op,其参数对应于函数参数。

函数定义了一个可以按名称引用的符号。使用函数调用 Op 可以将控制流转移到函数中。 一旦进入内部,控制流程遵循区域中各个块的 CFG。 “return” 终止符没有后继,而是终止区域执行,从而将控制流转移回函数的调用方。 “return” 终止符 Op 的任何操作数都是函数的返回值。

上面介绍了 MLIR 的 IR 设计细节,可以结合 MLIR 官方文档的语法规则来更好的熟悉:https://mlir.llvm.org/docs/LangRef/

0x6.3 IR 基础设施

除了 IR 本身之外,MLIR 还提供了用于定义 IR 元素(如 dialect、Ops、 pattern rewrite、验证和可重用 passes)的基础结构。 当定义新的抽象并将 MLIR 用作优化工具包时,MLIR 的基础结构对于提供可扩展性和易用性至关重要。

0x6.3.1 Operation description(操作描述)

MLIR 使用 TableGen[47] 规范定义操作描述(Operation Descriptions, ODS),以声明的方式定义 Op 的结构及其验证程序组件。 TableGen 是一种在 LLVM 中广泛使用的数据建模工具,目的是帮助定义和维护领域特定( domain-specific)信息的记录。 ODS 可以看作是嵌入 TableGen 语言并用来定义 MLIR Op 的 DSL。因此 ODS 语法由 TableGen 规定,但 MLIR 特定的语义由 ODS 规定。 ODS 定义最终会转换为 C++ 代码,这些代码可以与编译系统的其余部分互操作。

MLIR 使用 TableGen Op 类在 ODS 中对 Op 进行建模。Figure 5 显示了 Op 用 ODS 定义的示例。每个 Op 定义都有一个名称,该名称是唯一标识符。Op 的特征(trait)列表描述了 Op 属性。Op 的 argument(参数)列表指定 Op 的操作数和属性。Op 定义中还有一个 result(结果)列表。Op 的参数和结果具有名称和类型约束(例如 float 或 int32 的固定形状张量)。 Op 定义还可以指定人类可读的 Op 描述。当 Op 需要定义比 ODS 提供的更精细的控制时,可以通过 builder、printer、parser、verifier 语句注入额外 C++ 代码。Op trait 可以是通用的,例如 “has no side-effects”,也可以是特定于 Dialect 或 ODS 的,例如“has custom exporter”。 ODS 中的 traits 可以由定义 trait 行为的 C++ 类支持。MLIR 没有固定的 trait 集合,但是有些 trait 或者 optimizer(对应论文的 6.1 节)对 ODS 来说是已知的(例如,“shape result and operand type” 表示对于给定输入类型完全捕获输出类型的约束)。

类型约束会检查参数 / 结果类型的属性,并且由用户 / dialect 扩展。 MLIR 基础结构还提供了许多预定义的类型约束,例如 “any type”、““tensor with element satisfying the given constraint”、““vector of given rank” 等。ODS 对自动推断操作数结果的返回类型的支持很有限,这些操作数使用了由特征带来的约束。更多信息请参见下一节(对应论文的 4.2 节)。

Op的ODS定义

0x6.3.2 Declarative rewrites(声明式重写)

许多 MLIR 变换涉及 Op 操作,尽管某些变换需要对 IR 进行复杂的修改,但许多其它转换可以表示为对静态单赋值 use-def 关系定义 DAG 的简单重写。 MLIR 提供了一个图重写框架,并辅以声明性重写规则(Declarative Rewrite Rule, DRR)系统,使得模式(pattern)表达变得简单。

与 ODS 相似,DRR 是嵌入到 TableGen 语言中的 DSL。 DRR 表示源和目标 DAG pattern 以及约束(包括动态约束 [49])并从 pattern 优先优先级中受益。pattern 可以捕获和重用 Op 的参数。从概念上讲,DRR 表示在特定约束下 DAG 的等效性。Figure 6 给出了 DRR 模式的示例,该模式将 Fiugure 5 中定义的 Op 转换为由compareselect组成的通用低级别实现。

DRR图改写规则

DRR 被转换为 C++ 代码,可以使用通用图重写框架将其与直接在 C++ 中定义的更复杂的模式混合。通过这项功能,MLIR 可以使常见用例保持简洁,且不会限制框架的通用性。

0x6.3.3 Pass Manager

MLIR pass 管理器以各种粒度组织并处理 IR pass 序列,保证 pass 的高效执行。现有编译系统中的 pass 管理通常是按照固定的粒度(例如,模块、函数或循环 pass 管理器)定义的。但在 MLIR 中,模块和函数并无特殊,它们只是具有区域的 Ops,并且有多种变体。因此,MLIR pass 管理器也不专门针对固定的 Op 集合,而是针对任意嵌套级别的任意 Op。

并行编译 MLIR 的一个重要需求是利用多核计算机来加快编译速度。pass 管理器支持并发遍历和修改 IR,这可以通过 Op 的 “与上方隔离(isolated-from-above)” 属性提供的不变量来实现,因为静态单赋值 use-def 链无法跨越这些 op 的区域边界,因此具有这种行为的 Op(例如“ std.func” Op)定义了可以并行处理的区域树。

这个需求也是 MLIR 不具有 whole-module use-def 链的原因(这与 LLVM 相反)。全局对象通过符号表条目进行引用,而常量则由具有关联属性的 Op 实现。

0x6.4.4 可相互变换的 IR 文本表示形式

MLIR 中的 IR 和 Op 具有文本表示形式,可以完全反映内存中的 IR 表示,这对于调试、理解变换期间的 IR 以及编写测试用例至关重要。Figure4 所示的原始 IR 表示冗长且难以理解,因此 MLIR 允许用户为 Op 定义定制的打印和解析格式,这使得示例可以如 Figure 8 所示进行打印和解析,这更容易使用。两种形式可以完全相互转换,并且可以使用文本形式作为输入和输出,分别测试每个编译器 pass。由于没有隐藏状态,因此运行单个 pass 的结果与在完整 pass pipeline 中运行相同 pass 的结果相同。这种方法对用户友好,因为可以手动创建 IR 格式,并可方便跟踪 IR 转换。

自定义解析格式的Affine Dialect IR

0x6.4.5 文档

Dialect、Op 和 Interfaces 都有从其对应 ODS 描述生成的文档。除了 summary 和更易读懂的 description 之外,生成的文档还包括参数和结果类型约束。由于验证代码和文档使用相同的来源,因此文档可以与运行时行为保持同步。

0x6.4.6 验证器

验证器用于增强 IR 的结构正确性和 Op 的不变性 让 pass 确定已验证的 IR 不变式是经过检查的,并且还可以用作调试工具。验证过程以 MLIR 总体结构属性检查开始,比如,检查类型必须完全匹配,值仅定义一次且遵守支配规则和可见性,符号名称在符号表中是唯一的,所有块均以终结符 Op 结尾,等等。之后,应用各个 Op 和属性的验证器。每个 Op 可以定义一组检查结构和语义有效性规则。例如,二元 Op 会检查是否有两个操作数,一些 Op 只接受特定类型的值,而一些 Op 需要附加特定的属性或区域。同样,Dialect 属性只能在特定的 Op 上被允许使用,或者通过这些属性对其所附加的 Op 做进一步的限制。例如,Dialect 属性可以要求 Op 仅使用 Dialect 中定义的类型,即使 Op 本身更通用。验证失败被视为 invariant violation 并中止编译。

0x6.5 评估:MLIR 的应用

MLIR 系统的目的是统一和驱动各种不同类型的编译器项目,因此我们的主要评估指标是展示 MLIR 已被哪些项目采用。本节提供了用户社区活动的简介,并详细描述了一些用例,突出说明 MLIR 的通用性和可扩展性,并展示 MLIR 如何能很好地实现定制设计原则。

目前,MLIR 还是一个不断发展的开源项目,其用户社区遍布学术界和工业界。来自 4 个不同国家的 4 个国家实验室和 16 个大学的人士参加了在高性能计算(HPC)中使用 MLIR 的学术研讨会。 MLIR 还得到了 14 家跨国公司的认可。在 LLVM Developer Meeting 上,超过 100 个业界开发人员参加了有关 MLIR 的圆桌会议。有超过 26 种 dialect 正在开发中,并且来自不同公司的 7 个项目正在用 MLIR 替换自定义编译器基础结构。这表明了对 MLIR 的真实需求,并认可了 MLIR 的可用性。

0x6.5.1 TensorFlow graphs

尽管大多数编译器开发人员也都熟悉其它表示形式,但是 MLIR 的关键用例之一是支持机器学习框架的开发。机器学习框架的内部表示通常基于具有动态执行语义的数据流图 [53]。

TensorFlow[1] 是这种框架的一个例子。TensorFlow 的表示是高级数据流计算,其中的节点是可以放置在各种设备(包括特定的硬件加速器)上的各种计算过程。

TensorFlow 使用 MLIR 对该内部表示进行建模,并针对 Figure1 所示的用例进行转换,将简单的代数优化转换为能在(硬件加速器的)数据中心集群上并行执行的、新形式的图,并将 IR lowering 为能使用 XLA[57] 这类工具生成高效本地代码、适合移动端部署的表示。 MLIR 中的 TensorFlow Graph 表示如图 7 所示:

TensorFlow Graph对应的MLIR表示

0x6.5.2 Polyhedral code generation 多面体代码生成

MLIR 的最初动机之一是探索加速器的多面体代码生成。affine dialect 是简化的多面体表示形式,设计目的是实现渐进式 IR lowering。尽管对设计要点的全面探讨不在本文的讨论范围之内,本文还是说明了 affine dialect 的几个方面,以展示 MLIR 的建模能力,并将 affine dialect 与过去的一些表示形式进行了对比 [17、19、54、55、52] 。

共同点

MLIR affine dialect 可以对所有内存访问的结构化多维类型做操作 。在默认情况下,这些结构化类型是注入的(injective),保证不同的索引不会因构造而混叠,这是多面体依赖分析的常见前提。

Affine modeling 可分为两个部分。属性用于在编译时对仿射图和整数集建模,而 Op 则用于对代码应用仿射约束。即,affine.for Op 是一个 “for” 循环,其边界表示为值的仿射图,并且这些值要求在函数中保持不变。因此,循环具有静态控制流。与此类似,affine.if是受仿射整数集限制的条件语句。循环和条件语句的主体是区域,这些区域使用affine.loadaffine.store将索引限制为循环迭代器的仿射形式。这样可以进行精确的仿射依赖分析,同时避免了从低级表示中推断仿射形式。

区别

MLIR 与现有多面体代码生成框架之间的差异很多,可以将其分为以下四类: (1)丰富的类型:MLIR 结构化的内存引用类型包含了一个将缓冲区索引空间连接到实际地址空间的布局图。这两种空间的分隔可以改善循环和数据转换的组合,因为对数据布局的修改不会影响到代码,也不会污染依赖关系分析。文献 [38] 已经探讨过这种转换混合,但并不常见。 (2)抽象的混合:MLIR 中的仿射循环体可以通过类型化(typed )静态单赋值的 Op 来表示。因此,所有传统的编译器分析和转换过程仍然适用,并且可以与多面体转换交错使用。相反,多面体编译器经常将这些细节完全抽象掉,这使得多面体编译器难以操作某些对象,例如向量类型。 (3)较小的表示差异:多面体模型的主要特征之一是能够表示类型系统中循环迭代的顺序。但是,多面体转换会将 IR 提升为与原始 IR 完全不同的表示形式[20,10]。此外,从变换后的多面体到循环的转换在计算上很困难[7]。基于 MLIR 的表示在低级表示中保持了高级循环结构,因而不再需要提升 IR。 (4)如第 0x6.3.3 Pass Manager 节所述,编译速度是 MLIR 的关键目标,但现有大多数多面体方法并不关注编译速度。这些多面体方法严重依赖指数复杂度的算法:依赖整数线性编程自动推导出循环顺序和依赖多面体扫描算法将 IR 转换回循环。 MLIR 采用的方法不依赖多面体扫描,因为循环保留在 IR 中。

论文还举了一些例子用来说明 MLIR 在领域特定编译器的应用以及基于 MLIR 开发的 Fortran IR,这里就不再阅读了,感兴趣的可以看下原文了解。

0x6.6 MLIR 设计的成果

MLIR 设计有助于对新语言和编译抽象进行建模,同时有助于重用现有的、通用的相关编译方法。MLIR 对很多问题的有效解决方法是 “添加新操作、新类型”,如果可能,将其收集到“某个新 dialect” 中。对于编译器工程而言,这是重大的设计转变,产生了新的机遇,挑战和见解。本节将探讨其中部分观点。

0x6.6.1 可重用的编译器 Pass

在一个 IR 中可以表示多个抽象级别的能力自然产生了编写跨多个抽象级别工作的 pass 的想法。 关于 MLIR 的一个常见问题是,既然 MLIR 具有可扩展的操作和类型系统,那么如何编写编译器 pass?虽然编译器 pass 可能总是以保守、正确的方式处理未知结构,但 MLIR 的目标是生成高性能代码,主要有四种方法:

基本操作特征 一些 “bread and butter” 编译器 pass(如 “死代码消除” 和“通用子表达式消除”)只依赖我们定义为 Op traits 简单属性(例如 “has no side effect” 或“is commutative”)。ODS 中 Op 的定义允许 Op 的开发者指定这些特征,并且 pass 可以使用此信息来保持操作在许多不同抽象域都适用。

MLIR 的可扩展性体现为包含一些结构属性,其中包括下述信息:是否知道某个操作是控制流终止符是否知道某个操作包含的区域是与上方隔离的(isolated-from-above)等等。这些信息可用于函数、闭包、模块和其他代码结构的建模和处理。

Privileged operation hooks(Op 的特殊钩子)虽然某些特征可以用单比特建模,但是其它很多特征则需要 C++ 代码实现,例如常量折叠逻辑。 MLIR 对适用于大量 pass 的某些 hook 提供了最好的支持。这些 hook 可以基于每个操作实现,也可以在 dialect 对象中实现。后一种方法对支持诸如 TensorFlow ops 的常量折叠之类 pass 很方便,在这种情况下,很容易实现对现有逻辑的委托。

尽管常量折叠是非常重要的功能,但更有意思的 hook 是getCanonicalizationPatterns,这个 hook 允许指定应用于操作的折叠模式。这使得重要的代数简化形式(例如 x − x→0,min(x,y,y)→min(x,y)等)具有可扩展性,并可帮助将普通 “规范化(Canonicalization)”pass 应用到所有 dialect 。这些都使得单一的可扩展系统可以包含像“InstCombine”、“DAGCombine”、“PeepholeOptimizer”、“SILCombine” 这类 pass,以及 LLVM 生态系统(和其它编译器)中的其它特殊用途 pass。

Optimization interfaces (优化接口) MLIR 的主要目标是可扩展性,不仅在 Op 和类型方面,而且在转换方面也要有可扩展性。虽然规范化(canonicalization)和常量折叠是关键操作,但仍需要以某些方式对许多标准转换进行参数化设置,才能描述转换的特定属性,才能实现代码模型等。

问题的解决方案是称为 “优化接口” 的子系统。考虑一下 MLIR 内联 pass,我们希望 inliner 可以处理 TensorFlow 图、Flang 函数、函数语言的闭包等,但是 inliner 不知道调用方是什么,甚至不知道被调用方是什么。inliner 需要了解的核心特性是:

  • 将给定操作内联到给定区域是否有效;
  • 如何处理内联后终止于块中间的终止符操作。

为了了解这些属性,Inliner pass 定义了 Figure 10 中的接口。各个操作和 dialect 可以向 MLIR 注册该接口在操作和 dialect 中的实现,并从通用的 Innerer pass 中获益。如果某个操作或 dialect 没有提供接口,则相应的优化 pass 将会保守地对待该操作。这种设计让 dialect 的开发者能快速启动开发并运行 dialect。随着时间的推移,通过将更多的精力投入到接口的开发,可以从系统中获得更多收益。

内联Pass接口

优化接口还为核心编译器提供了模块化优势,因为 dialect 特定的逻辑是在 dialect 自身内部实现,而不是在核心转换中实现。

Dialect 特定 pass 最后,定义特定 dialect 可以定义专用 pass,MLIR 系统中的这些 pass 和在其它编译器系统中的 pass 一样都很有用。比如说,如果想让代码生成器根据特定的机器约束对机器指令进行自定义调度,就可以通过专用 pass 达到目的。这可当作开发新转换 pass 的起点,不需要考虑 pass 的通用性。

0x6.6.2 Dialect 的混合

MLIR 中一个最根本(也是最难理解)的部分是允许并鼓励将来自不同 dialect 的操作混合在一个程序中。尽管在某些情况下(例如,将主机和加速器的计算保存在同一模块中),这样做很容易理解,但最有趣的情况是,在 MLIR 中可以将 dialect 直接混合(因为这样可以实现整个类的重用),这在其它系统中是见不到的。

考虑第 0x6.5.2 节中描述的 affine dialect。affine 控制流和 affine 映射的定义与 affine 区域中包含的操作的语义无关。在我们的案例中,我们将 affine dialect 与 “standard” dialect 结合起来,以目标无关的形式(如同 LLVM IR)表示简单算术,也可以针对内部加速器,将 affine dialect 与多个目标相关机器指令 dialect 结合。也有人将 affine dialect 与其它问题领域的抽象相结合。

重用通用多面体变换(使用 Op Interface 获取特定转换中操作的语义)的能力是分解编译器基础结构的一种有力方法。另一个例子是,可以在各种源语言 IR 中使用和重用 OpenMP dialect。

0x6.6.3 互操作性

本文的工作涉及与大量现有系统的互操作,例如,protobuff 格式的机器学习 graphs、包括 LLVM IR 在内的编译器 IR、各种专有指令集等。任何一种表示形式不可避免都有各种缺陷,虽然这些缺陷在某个现有系统的适用场景下是合理的,但是 MLIR 的表达能力使 MLIR 成为一种更好的表示形式。因为 importer 和 exporters 的测试难度很大(测试用例通常是二进制格式),因此我们希望确保其复杂性最低。

问题的解决方案是尽可能定义与外部系统直接相对应的 dialect,从而能以一种简单且可预测的方式来回转换该格式。一旦将 IR 导入 MLIR 格式中,就可以使用 MLIR 基础结构中所有转换,将导入的 IR 升级或降级为某种更适合的 IR 格式,并允许对这些转换 pass 进行类似于所有其它 MLIR pass 的测试。

这类 dialect 的例子很多,包括:a)LLVM dialect,可将 LLVM IR 映射为 MLIR; b)TensorFlow 的图表示形式,提出这种表示是为了简化 TensorFlow 中 “切换和合并(switch and merge)” 节点相关的分析和转换;c )函数式控制流运算符。“functional while”和 “functional if” 在机器学习图中很常见,在这种情况下,将其代码主体作为区域而不是外联(out-of-line)函数更方便。

这种方法对我们来说效果很好,并且 MLIR 工具对于编写外来二进制文件格式的测试用例也很有用。

0x6.6.4 非标准化设计带来了新挑战

虽然 MLIR 允许开发者定义几乎任意的抽象,但 MLIR 也几乎没有提供相关指导,比如,在实践中哪种方法效果更好或更差?现在,一些工程师和研究人员已经有这方面的经验,并且已经意识到,编译器 IR 设计和抽象设计的 “艺术” 在编译器和语言领域并未得到很好的理解。许多人在已建立系统的约束下工作,但是相对而言,很少人有机会自己定义抽象。

这是一个挑战,但也是未来研究的机遇。MLIR 社区正在通过这些抽象设计积累专业知识,随着时间的推移,这将是一个硕果累累的研究领域。

0x6.6.5 期望

在构建并将 MLIR 应用于许多不同的系统之后,可以发现 MLIR 的设计与其它编译器基础结构有很大不同。我们相信仍有很多应用领域有待发现,完全理解 MLIR 所有设计要点并建立最佳实践,需要更多的研究时间。例如,out-of-tree dialect 的兴起、前端使用 MLIR 的源语言数量的增加、在抽象语法树上的可能应用,以及对结构化数据(如 JSON,协议缓冲区等)的应用,这些都还处于很早期,可能会从中发现许多有趣的新挑战和机遇。

0x7. 评论(以 OneFlow Dialect 为例)

以上就是 MLIR 论文的大致内容。把 MLIR 论文中的提到的组件画成一张思维导图大概为:

Dialect的组成元素

下面我以 OneFlow Dialect 为例来讲解这张图。

正如论文提到的,在 MLIR 中 Operation 是 MLIR 的一个基本语意单位。定义了一个新的 Dialect 后我们首先就要考虑 Operation 的定义,要定义 Operation 必须先定义 Attribute 和 Type。OneFlow Dialect 的定义在oneflow/ir/include/OneFlow/OneFlowDialect.td这个文件中,它基于 ODS 规则设置了descriptioncppNamespce等关键信息,然后依靠mlir-tblgen可执行文件(MLIR 提供)自动生成了 OneFlow Dialect 的 C++ 代码。

def OneFlow_Dialect : Dialect {
    let name = "oneflow";
    let summary = "OneFlow MLIR dialect.";
    let description = [{
        This dialect is the IR of OneFlow.
    }];
    let cppNamespace = "::mlir::oneflow";
    let dependentDialects = [
        "StandardOpsDialect"
    ];
}

然后 Type 的定义在oneflow/ir/include/OneFlow/OneFlowBase.tdoneflow/ir/include/OneFlow/OneFlowEnums.td这两个文件中,分别对 OneFlow 的 Tensor 类型以及后续 Operation 指定 Attribute 需要的 Type 进行定义,需要说明的是 OneFlow 的 Operation 定义中除了下面定义的 Type 类型还大量使用了 MLIR 中提供的基础 Type:

def OneFlow_Tensor : TensorOf<[AnyType]>;
def SI32ArrayAttr : TypedArrayAttrBase<SI32Attr, "signed 32-bit integer array attribute"> {}

def SI64ArrayAttr : TypedArrayAttrBase<SI64Attr, "signed 64-bit integer array attribute"> {}

def ShapeAttr : TypedArrayAttrBase<SI64Attr, ""> {}
...

Attribute 的定义在每个 Operation 定义中,使用let attrs=来指定。下面以 LeakyReLU 为例看一下 OneFlow Dialect 的 Operation 定义(在oneflow/ir/include/OneFlow/OneFlowUserOps.td):

def OneFlow_LeakyReluOp : OneFlow_BaseOp<"leaky_relu", [NoSideEffect, DeclareOpInterfaceMethods<UserOpCompatibleInterface>]> {
  let input = (ins
    OneFlow_Tensor:$x
  );
  let output = (outs
    OneFlow_Tensor:$y
  );
  let attrs = (ins
    DefaultValuedAttr<F32Attr, "0.">:$alpha
  );
  let has_logical_tensor_desc_infer_fn = 1;
  let has_physical_tensor_desc_infer_fn = 1;
  let has_get_sbp_fn = 1;
  let has_data_type_infer_fn = 1;
}

可以看到OneFlow_LeakyReluOp继承了OneFlow_BaseOp,并声明了输入输出和 Attribute,最下面的 4 个标记是 OneFlow 在 llvm 的table-gen上做的一点扩展方便自动生成 Op 一些信息推导的接口这里可以不关心。

上面讲到了 Attribute,Type 和 Interface,接下来我们讲一讲 OneFlow Dialect 中 Operation 的 Trait 和 Constrait。在 MLIR 中 Trait(特征)和 Constrait(约束)的基类为OpTrait类,特征和约束通常用来指定 Operation 的特殊属性和约束,比如 Operation 是否具有副作用,Op 的输出是否与输入是否具有相同的形状等。

在 OneFlow 的 Operation 定义中不仅使用了 MLIR 提供的特征如 Leaky-ReLU 中的 NoSideEffect,还自定义了特征如IsOpConfCompatible。在oneflow/ir/include/OneFlow/OneFlowBase.tddef OneFlow_IsOpConfCompatible : NativeOpTrait<"IsOpConfCompatible">;的这句话就是使用 MLIR 提供的 ODS 方法NativeOpTrait声明了一个自定义的特征用来检查 OneFlow Dialect 定义的 Op 是否有某些共用属性例如 OpName,DeviceDagAttr 等等。这里只是在 ODS 中声明了自定义的属性,它真正定义在oneflow/ir/include/OneFlow/OneFlowOpTraits.h。这里简单摘出来看一下:

template<typename ConcreteType>
class IsOpConfCompatible : public TraitBase<ConcreteType, IsOpConfCompatible> {
 public:
  static StringRef getOpNameAttr() { return "op_name"; }
  static StringRef getDeviceTagAttr() { return "device_tag"; }
  static StringRef getDeviceNameAttr() { return "device_name"; }
  static StringRef getScopeSymbolIDAttr() { return "scope_symbol_id"; }
  static StringRef getHierarchyAttr() { return "hierarchy"; }
  static LogicalResult verifyTrait(Operation* op) { return impl::VerifyIsOpConfCompatible(op); }
};

LogicalResult VerifyIsOpConfCompatible(Operation* op) {
  for (auto attr : {
           IsOpConfCompatible<void>::getOpNameAttr(),
           IsOpConfCompatible<void>::getDeviceTagAttr(),
       }) {
    if (!op->hasAttrOfType<StringAttr>(attr)) {
      return op->emitError("expected operation to have attribute: " + attr);
    }
  }
  if (!op->hasAttrOfType<ArrayAttr>(IsOpConfCompatible<void>::getDeviceNameAttr())) {
    return op->emitError("expected operation to have attribute: "
                         + IsOpConfCompatible<void>::getDeviceNameAttr());
  }
  return success();
}

除了 Trait 之外,OneFlow 还使用了 MLIR 提供的一些特征如SameOperandsAndResultType。在oneflow/ir/include/OneFlow/OneFlowBase.tdOneFlow_UnaryBaseOp定义这里:

class OneFlow_UnaryBaseOp<string mnemonic, list<Trait> traits = []> :
        OneFlow_BaseOp<mnemonic, !listconcat(traits, [SameOperandsAndResultType, NoSideEffect])> {
  let summary = "";
  let input = (ins AnyType:$x);
  let output = (outs AnyType:$y);
  let has_logical_tensor_desc_infer_fn = 1;
  let has_physical_tensor_desc_infer_fn = 1;
  let has_get_sbp_fn = 1;
  let has_data_type_infer_fn = 1;
}

这个特征表达的意思就是继承了 UnaryBaseOp 的 Operation 的操作数和结果的类型都是相同的。当然特征也可以和约束一样自定义,也是用NativeOpTraittd文件中指明,然后实现也在oneflow/ir/include/OneFlow/OneFlowOpTraits.h

经过上面的讲解,相信大家对 MLIR 里面的 Type,Attribute,Operation,Trait,Constrait 都有所认识了。接下来我们要说一说 Interfaces,Interfaces 可以翻译成接口,MLIR 的 Interfaces 提供了和 IR 交互的通用方式。Interfaces 的设计目标是可以不用侵入到具体某个 Dialect 下的特定 Operation 和 Dialect 的特定知识就达到可以转换和分析 MLIR 表达式。这样就可以将转换,分析和新增一个 Dialect 和对应的 Operation 进行解耦,大大增强 MLIR 的可扩展性。为了说明 Interfaces 的重要性,我再这个专题实际上结合官方文档专门写了一篇文章来介绍,可以参考:【从零开始学深度学习编译器】十八,MLIR 中的 Interfaces

在 OneFlow 中,各个自定的 Interfaces 在oneflow/ir/include/OneFlow/OneFlowInterfaces.td这里。我们以UserOpCompatibleInterface为例来看一下 Interface 的具体实现:

def UserOpCompatibleInterface : OpInterface<"UserOpCompatible"> {
  let description = [{
    Interface to getting the hard-coded bn
  }];

  let methods = [
    StaticInterfaceMethod<"",
        "const std::vector<std::string>*", "inputKeys", (ins), [{
        static std::vector<std::string> val(mlir::oneflow::support::GetInputKeys(ConcreteOp::getOperationName().split('.').second.str()));
        return &val;
    }]>,
    StaticInterfaceMethod<"",
        "const std::vector<std::string>*", "outputKeys", (ins), [{
        static std::vector<std::string> val(mlir::oneflow::support::GetOutputKeys(ConcreteOp::getOperationName().split('.').second.str()));
        return &val;
    }]>,
    InterfaceMethod<"",
        "std::pair<unsigned, unsigned>", "getODSOperandIndexAndLength", (ins "unsigned":$index), [{
        return $_op.getODSOperandIndexAndLength(index);
    }]>,
    InterfaceMethod<"",
        "std::pair<unsigned, unsigned>", "getODSResultIndexAndLength", (ins "unsigned":$index), [{
        return $_op.getODSResultIndexAndLength(index);
    }]>
  ];
}

可以看到UserOpCompatibleInterface使用了 Interface ODS 规范中的 StaticInterfaceMethod 和 InterfaceMethod 为这个 Interface 指定了获取 Operation 输入操作数名字,输出操作数名字,操作数以及长度,结果以及长度等方法。然后在 OneFlow 的oneflow/ir/include/OneFlow/OneFlowUserOps.td中使用DeclareOpInterfaceMethods<UserOpCompatibleInterface>来将其指定为 Operation 层次的 Interface,在生成的 Operation 代码中就会带上这个 Interface 声明。

那么这样做有什么好处吗?第一点就是由于 OneFlow 的 UserOp 都带上了 UserOpCompatibleInterface,只要我们为 OneFlow 的 UserOp 实现一个通用的GetInputKeys函数,那么所有 UserOp 派生出来的 Operation 都拥有了这个函数的功能,因为它们都带上了 UserOpCompatibleInterface 这个接口。

关于 Interface 更加通用和经典的一个例子是基于 Interface 来开发一些通用 pass,比如内联和形状推导 pass,这个属于 Dialect 层次的 Interface 应用。具体见【从零开始学深度学习编译器】十三,如何在 MLIR 里面写 Pass?

思维导图中还剩下 Block 和 Region 没有讲了,实际上 MLIR 论文中对 Region 和 Block 的解释我觉得已经到位了。一个 Op 会附加一系列 Region,Region 为 MLIR 的嵌套结构提供了实现机制:一个 Operation 有一系列 Region,然后 Region 又是由一系列 Block 组成,然后 Block 又包含一系列 Op。这样就形成了一个嵌套的关系,可以表达作用域和控制流关系。在 OneFlow 的 Dialect 中对 Region 和 Block 的应用目前主要是在函数相关的语意中,例如在oneflow/ir/lib/OneFlow/Passes.cpp里面实现了一个OutlineMulCast的 Pass 可以将 IR 中指令的 op 模式外联到一个 FuncOp 类型的 Operation 中进行执行,就使用到了 Block 来确定这个 FuncOp 要插入到 IR 中的位置。再举一个例子,要访问 FuncOp 的参数时也需要用到 Block,如oneflow/ir/lib/OneFlow/OneFlowOps.cpp里为 Job Op 实现了一个 verify 函数,来验证函数的参数列表和入口 Block 的参数列表是否对齐:

static LogicalResult verify(Job op) {
  // If this function is external there is nothing to do.
  if (op.isExternal()) return success();

  // Verify that the argument list of the function and the arg list of the entry
  // block line up.  The trait already verified that the number of arguments is
  // the same between the signature and the block.
  auto fnInputTypes = op.getType().getInputs();
  Block& entryBlock = op.front();
  for (unsigned i = 0, e = entryBlock.getNumArguments(); i != e; ++i)
    if (fnInputTypes[i] != entryBlock.getArgument(i).getType())
      return op.emitOpError("type of entry block argument #")
             << i << '(' << entryBlock.getArgument(i).getType()
             << ") must match the type of the corresponding argument in "
             << "function signature(" << fnInputTypes[i] << ')';

  return success();
}

我们在 Job 这个 Op 的定义也可以发现,它通过let regions = (region AnyRegion:$body);绑定一个名为 body 的 Region,所以上面 verify 函数中访问 FuncOp 的 Block 时隐式访问了 Block 对应的 Region,即将op.front()改成op.body().front()是等价的效果。

Region 和 Block 相当于 MLIR 中的作用域,可以通过它们来实现 MLI 中复杂结构的控制流关系,并且将普通 Op 和 FuncOp,ModuleOp 区分开,实现在 MLIR 中 Operation 统一性原则。

至于 OneFlow Dialect 中涉及到的 Pass 机制我在以 OneFlow 为例探索 MLIR 的实际开发流程 这里已经介绍过了,这里暂时不重复写了。

总的来说 MLIR 是一个可重用、可扩展性都比较好的编译基础设施,至少从工程开发的角度来看是值得跟进的。这部分我介绍了 OneFlow Dialect 中的各种组件以及它们的关系,但还没有介绍 OneFlow Dialect 和其它各种 Dialect 的关系,以及目前 Dialect 的 lowering 流程。实际上 MLIR 还有相当多的技术细节需要深入研究和学习,所以这里主要起到一个小结和启发的作用。

0x8. 相关参考文章

本文参考了 https://arxiv.org/pdf/2002.11054.pdf MLIR 原论文和这篇 https://zhuanlan.zhihu.com/p/336543238 翻译。


本文总阅读量1282