深度学习编译MLIR初步

深度模型的推理引擎

目前深度模型的推理引擎按照实现方式大体分为两类：解释型推理引擎和编译型推理引擎。

解释型推理引擎

一般包含模型解析器，模型解释器，模型优化器。

模型解析器负责读取和解析模型文件，转换为适用于解释器处理的内存格式；
模型优化器负责将原始模型变换为等价的、但具有更快的推理速度的模型；
模型解释器分析内存格式的模型并接受模型的输入数据，然后根据模型的结构依次执行相应的模型内部的算子，最后产生模型的输出。

编译型推理引擎

编译型推理引擎一般包含模型解析器和模型编译器。

模型解析器的作用与解释型推理引擎相同
模型编译器负责将模型编译为计算设备（CPU、GPU 等）可直接处理的机器码，并且可能在编译的过程中应用各种优化方法来提高生成的机器码的效率。

编译型推理引擎的优势

由于机器码的模型可以直接被计算设备处理，无需额外的解释器的参与，其消除了解释器调度的开销。相对于解释型推理引擎，由于生成机器码的过程更加靠底层，编译器有更多的优化机会以达到更高的执行效率。由于现在业界对于推理引擎的执行速度有了更高的需求，编译型推理引擎也逐渐成为高速推理引擎的发展方向。编译型推理引擎有 Apache TVM、oneDNN、PlaidML、TensorFlow XLA、TensorFlow Runtime 等。

中间表示

为了便于优化，一般来说推理引擎会把模型转换为中间表示，然后对中间表示进行优化和变换，最终生成目标模型（对于解释型推理引擎）或目标机器码（对于编译型推理引擎）。此外，除了深度学习领域，在很早以前编程语言领域就引入了中间表示来做优化和变换。而新的编程语言层出不穷，因此就出现了各种各样的中间表示：不同的推理引擎或者编译器都会有自己的中间表示和优化方案，而每种中间表示和优化方案可能都需要从头实现，最终可能会导致软件的碎片化和重复的开发工作。

MLIR

听到这个名字的时候，下意识觉得这应该是个关于机器学习的中间表达结构，毕竟这年头听到ML都会认为是Machine Learning，但是MLIR还真不是Machine Learning Intermediate Representation，而是Multi-Level Intermediate Representation。

1. 什么是MLIR

MLIR（Multi-Level Intermediate Representation，多级中间表示）是一种用来构建可重用和可扩展编译基础设施的新方法。MLIR旨在解决软件碎片化，改进异构硬件的编译，显著减少构建特定领域编译器的成本，并帮助连接现有的编译器。

具体来说，MLIR是通过定义一个通用的中间表示（IR），将在TensorFlow和类似的ML框架中执行高性能机器学习模型所需的基础设施进行统一，包括高性能计算技术应用或强化学习这类搜索算法的集成。MLIR旨在降低开发新硬件的成本，并提高现有TensorFlow用户的可用性。

在这里插入图片描述

MLIR通过一种混合的中间表示，实现在一个统一的架构基础上支持以下多个不同的需求：

表征数据流图的能力，包括动态形状（dynamic shape），用户可扩展的操作集，TensorFlow变量，等等。
支持典型的图优化和图变换，例如TensorFlow Grappler。
可以以一种易于优化的形式来表征ML操作的Kernel。
具备掌控跨Kernel的高性能计算风格（high-performance-computing-style）循环优化（例如fusion，loop interchange，tiling等polyhedral compilation）的能力，以及变换数据内存排布的能力。
代码生成“下降lowering”转换，显式的插入DMA，显示的缓存cache管理，内存分块tiling，为适应1D和2D寄存器架构的自动向量化
表达特定后端操作的能力，比如特定加速器的高层操作。
在深度学习图上已有的量化和其他图变化。

MLIR 并非万能，不支持底层机器代码生成的算法（比如寄存器分配和指令调度），这些更适合于低层优化器（如llvm），它也不支持用户手写Kernel（不同于CUDA）。MLIR是提供一种框架，使自定义的 DSL（Domain-Specific Language，领域专用语言）表达可以融入到一个统一的生态。

2. MLIR的作用

虽然 MLIR 中的 ML 不是Machine Learning，但是Machine Learning确实是是MLIR的一个重要应用领域。我们接下来主要看一下在机器学习领域，MLIR可以做哪些事情。在了解MLIR是怎么工作之前，我先得弄明白这个IR在解决什么问题。

说到机器学习，我们就用TensorFlow这个框架来举例。我们知道TensorFlow是使用数据流图作为数据结构来进行各种数值计算，要让这些计算运行在硬件上，我们需要一个TensorFlow的编译生态系统：

在这里插入图片描述

如图中所示，TensorFlow图能够以多种不同的方式运行，包括：

将其发送至调用手写运算内核的 TensorFlow 执行器
将图转化为 XLA 高级优化器 (XLA HLO) 表示，反之，这种表示亦可调用适合 CPU 或 GPU 的 LLVM 编辑器，或者继续使用适合 TPU 的 XLA。（或者将二者结合！）
将图转化为 TensorRT、nGraph 或另一种适合特定硬件指令集的编译器格式
将图转化为 TensorFlow Lite 格式，然后在 TensorFlow Lite 运行时内部执行此图，或者通过 Android 神经网络 API (NNAPI) 或相关技术将其进一步转化，以在 GPU 或 DSP 上运行

此外，我们有时甚至会采用更复杂的途径，包括在每层中执行多轮优化。例如，Grappler 框架现在便能优化 TensorFlow 中的张量布局和运算。

通常来说，整个编译流程先将TensorFlow的图转化为XLA HLO，即一种类似高级语言的图的中间表达形式，可以基于此进行一些 High-Level 的优化。接着将XLA HLO翻译为LLVM IR，使用LLVM编译到各种硬件的汇编语言，从而运行在硬件上进行数值计算。

下图的蓝色阴影部分是 基于图的IR，绿色阴影部分是基于 SSA（static single-assignment，静态单一赋值）的IR，然而这样的编译方式的缺点在于构建这样的编译系统的开销比较大，每一层的设计实现会有重复部分，同一个层次的IR彼此之间虽然相似，但是存在天生的“生殖隔离”，升级优化缺乏迁移性，即改变优化一个模块，并不能惠及到同层次的其他模块。因此，目前存在的问题就在于各种IR之间转换的效率和可迁移性不高。

在这里插入图片描述

对于上述问题，MLIR希望为各种DSL提供一种中间表达形式，将他们集成为一套生态系统，使用一种一致性强的方式编译到特定硬件平台的汇编语言上。利用这样的形式，MLIR就可以利用它模块化、可扩展的特点来解决IR之间相互配合的问题。

在这里插入图片描述

到此为止，我们大致知道了MLIR的诞生是为了解决什么问题。目前它对我来说还是一个黑盒子，下面的工作就是要去看看MLIR内部究竟是一个什么样的结构，看看它是怎么把各层IR整合到一起，又是如何实现扩展功能的。

3. MLIR中的 “方言”：dialect

为什么要有方言

之前我们说到当前的编译结构的问题在于各种IR之间转换的效率和可迁移性不高。MLIR试图使用一种一致性强的方式，为各种DSL提供一种中间表达形式，将他们集成为一套生态系统，编译到特定硬件平台的汇编语言上。这样的目标是通过什么手段实现的呢？

从源程序到目标程序，要经过一系列的抽象以及分析，通过Lowering Pass来实现从一个IR到另一个IR的转换，这样的过程中会存在有些操作重复实现的情况，也就导致了转换效率低的问题。

这就好比，IR们组成一个流水线要合起伙来干一个大买卖，但是互相配合不默契，谁也明白不了对方究竟干了啥，为了保险起见，每个IR拿到上一个IR的产品之后只能多干点活保证不出错，这样一来效率自然就低了。MLIR面对这种IR群雄割据的现状，打算一统天下！打天下容易，守天下难呀，不让我们用各种IR了，你倒是给我们一条活路呀，怎么才能让源语言变成汇编语言然后跑在机器上呀？于是，统一IR的第一步就是要统一“语言”，各个IR原来配合不默契，谁也理解不了谁，就是因为“语言”不通，没法用统一的“语言”指挥流水线干活。MLIR看准时机拿出了自己的法宝：Dialects！让各个IR学习Dialects这个“语言”，这样一来，不光能指挥流水线高效干活了，还能随意扩展更改分工，从此IR们就可以完美地分工协作。

为区分不同的硬件与软件受众，MLIR 提供 “方言”，其中包括：

TensorFlow IR，代表 TensorFlow 图中可能存在的一切
XLA HLO IR，旨在利用 XLA 的编译功能（输出到 TPU 等）
实验性仿射方言，侧重于多面表示与优化
LLVM IR，与 LLVM 自我表示之间存在 1:1 映射，可使 MLIR 通过 LLVM 发出 GPU 与 CPU 代码
TensorFlow Lite，将会转换以在移动平台上运行代码

每种方言均由一组存在不变性的已定义操作组成，如：“这是一个二进制运算符，输入与输出拥有相同类型。”

将方言添加至 MLIR

MLIR 没有众所周知的固定或内置的操作列表（无 “内联函数”）。方言可完全定义自定义类型，即 MLIR 如何对 LLVM IR 类型系统（拥有一流汇总）、域抽象（对量化类型等经机器学习 (ML) 优化的加速器有着重要意义），乃至未来的 Swift 或 Clang 类型系统（围绕 Swift 或 Clang 声明节点而构建）进行建模。

如果我们想要连接新的低级编译器，则需要创建新方言，以及 TensorFlow 图方言与我们的方言之间的降阶。如此一来，硬件及编译器制造商便可畅通无阻。我们甚至可以在同一个模型中定位不同级别的方言；高级优化器将保留 IR 中不熟悉的部分，并等待较低级别的优化器来处理此类部分。

对于编译器研究者和框架制造者，则可以借助 MLIR 在每个级别进行转换，甚至是在 IR 中定义自己的操作和抽象，从而针对试图解决的问题领域构建最佳模型。由此看来，MLIR 比 LLVM 更像是纯编译器基础设施。

虽然 MLIR 充当 ML 的编译器，但我们也看到，MLIR 同样支持在编译器内部使用机器学习技术！这一点尤为重要，因为在进行扩展时，开发数字库的工程师无法跟上 ML 模型或硬件的多样化速度。MLIR 的扩展性有助于探索代码降阶策略，并在抽象之间执行逐步降阶。

dialect是如何工作的

dialects是将所有的IR放在了同一个命名空间中，分别对每个IR定义对应的产生式以及绑定相应的操作，从而生成一个MLIR的模型。整个的编译过程，从源语言生成AST（Abstract Syntax Tree，抽象语法树），借助Dialects遍历AST，产生MLIR的表达式，此处可为多层IR通过Lowering Pass依次进行分析，最后经过MLIR分析器，生成目标语言。

在这里插入图片描述