跳转至

0x1. 前言

最近开始做一些 MLIR 的工作，所以开始学习 MLIR 的知识。这篇笔记是对 MLIR 的初步印象，并不深入，适合想初步了解 MLIR 是什么的同学阅读，后面会继续分享 MLIR 的一些项目。这里要大力感谢中科院的法斯特豪斯（知乎 ID）同学先前的一些分享，给了我入门 MLIR 的方向。

0x2. 什么是 IR？

IR 即中间表示（Intermediate Representation），可以看作是一种中介的数据格式，便于模型在框架间转换。在深度学习中可以表示计算图的数据结构就可以称作一种 IR，例如 ONNX，TorchScript，TVM Relay 等等。这里举几个例子介绍一下：

首先，ONNX 是微软和 FaceBook 提出的一种 IR，他持有了一套标准化算子格式。无论你使用哪种深度学习框架（Pytorch，TensorFlow，OneFlow）都可以将计算图转换成 ONNX 进行存储。然后各个部署框架只需要支持 ONNX 模型格式就可以简单的部署各个框架训练的模型了，解决了各个框架之间模型互转的复杂问题。

但 ONNX 设计没有考虑到一个问题，那就是各个框架的算子功能和实现并不是统一的。ONNX 要支持所有框架所有版本的算子实现是不现实的，目前 ONNX 的算子版本已经有 10 代以上，这让用户非常痛苦。IR 可以类比为计算机架构的指令集，但我们是肯定不能接受指令集频繁改动的。另外 ONNX 有一些控制流的算子如 If，但支持得也很有限。

其次，TorchScript 是 Pytorch 推出的一种 IR，它是用来解决动态图模式执行代码速度太慢的问题。因为动态图模式在每次执行计算时都需要重新构造计算图 (define by run)，使得性能和可移植性都比较差。为了解决这个问题，Pytorch 引入了即时编译（JIT）技术即 TorchScript 来解决这一问题。Pytorch 早在 1.0 版本就引入了 JIT 技术并开放了 C++ API，用户之后就可以使用 Python 编写的动态图代码训练模型然后利用 JIT 将模型（nn.Module）转换为语言无关的模型（TorchScript），使得 C++ API 可以方便的调用。并且 TorchScript 也很好的支持了控制流，即用户在 Python 层写的控制流可以在 TorchScript 模型中保存下来，是 Pytorch 主推的 IR。

最后，Relay IR 是一个函数式、可微的、静态的、针对机器学习的领域定制编程语言。Relay IR 解决了普通 DL 框架不支持 control flow（或者要借用 python 的 control flow，典型的比如 TorchScript）以及 dynamic shape 的特点，使用 lambda calculus 作为基准 IR。

Relay IR 可以看成一门编程语言，在灵活性上比 ONNX 更强。但 Relay IR 并不是一个独立的 IR，它和 TVM 相互耦合，这使得用户想使用 Relay IR 就需要基于 TVM 进行开发，这对一些用户来说是不可接受的。

这几个例子就是想要说明，深度学习中的 IR 只是一个深度学习框架，公司甚至是一个人定义的一种中介数据格式，它可以表示深度学习中的模型（由算子和数据构成）那么这种格式就是 IR。

0x3. 为什么要引入 MLIR？

目前深度学习领域的 IR 数量众多，很难有一个 IR 可以统一其它的 IR，这种百花齐放的局面就造成了一些困境。我认为中科院的法斯特豪斯同学 B 站视频举的例子非常好，建议大家去看一下。这里说下我的理解，以 TensorFlow Graph 为例，它可以直接被转换到 TensorRT 的 IR，nGraph IR，CoreML IR，TensorFlow Lite IR 来直接进行部署。或者 TensorFlow Graph 可以被转为 XLA HLO，然后用 XLA 编译器来对其进行 Graph 级别的优化得到优化后的 XLA HLO，这个 XLA HLO 被喂给 XLA 编译器的后端进行硬件绑定式优化和 Codegen。在这个过程中主要存在两个问题。

  • 第一，IR 的数量太多，开源要维护这么多套 IR，每种 IR 都有自己的图优化 Pass，这些 Pass 可能实现的功能是一样的，但无法在两种不同的 IR 中直接迁移。假设深度学习模型对应的 DAG 一共有 10 种图层优化 Pass，要是为每种 IR 都实现 10 种图层优化 Pass，那工作量是巨大的。
  • 第二，如果出现了一种新的 IR，开发者想把另外一种 IR 的图层优化 Pass 迁移过来，但由于这两种 IR 语法表示完全不同，除了借鉴优化 Pass 的思路之外，就丝毫不能从另外一种 IR 的 Pass 实现受益了，即互相迁移的难度比较大。此外，如果你想为一个 IR 添加一个 Pass，难度也是不小的。举个例子你可以尝试为 onnx 添加一个图优化 Pass，会发现这并不是一件简单的事，甚至需要我们去较为完整的学习 ONNX 源码。
  • 第三，在上面的例子中优化后的 XLA HLO 直接被喂给 XLA 编译器后端产生 LLVM IR 然后 Codegen，这个跨度是非常大的。这里怎么理解呢？我想到了一个例子。以优化 GEMM 来看，我们第一天学会三重 for 循环写一个 naive 的矩阵乘程序，然后第二天你就要求我用汇编做一个优化程度比较高的矩阵乘法程序？那我肯定是一脸懵逼的，只能 git clone 了，当然是学不会的。但如果你缓和一些，让我第二天去了解并行，第三天去了解分块，再给几天学习一下 SIMD，再给几个月学习下汇编，没准一年下来我就可以真正的用汇编优化一个矩阵乘法了。所以跨度太大最大的问题在于，我们这种新手玩家很难参与。我之前分享过 TVM 的 Codegen 流程，虽然看起来理清了 Codegen 的调用链，但让我现在自己去实现一个完整的 Codegen 流程，那我是很难做到的。【从零开始学深度学习编译器】九，TVM 的 CodeGen 流程

针对上面的问题，MLIR（Multi-Level Intermediate Representation）被提出。MLIR 是由 LLVM 团队开发和维护的一套编译器基础设施，它强调工具链的可重用性和可扩展性。下面我们具体分析一下：

针对第一个问题和第二个问题，造成这些深度学习领域 IR 的优化 Pass 不能统一的原因就是因为它们没有一个统一的表示，互转的难度高。因此 MLIR 提出了 Dialect，我们可以将其理解为各种 IR 需要学习的语言，一旦某种 IR 学会这种语言，就可以基于这种语言将其重写为 MLIR。Dialect 将所有 IR 都放在了同一个命名空间里面，分别对每个 IR 定义对应的产生式以及绑定对应的操作，从而生成 MLIR 模型。关于 Dialect 我们后面会细讲，这篇文章先提一下，它是 MLIR 的核心组件之一。

针对第三个问题，怎么解决 IR 跨度大的问题？MLIR 通过 Dialect 抽象出了多种不同级别的 MLIR，下面展示官方提供的一些 MLIR IR 抽象，我们可以看到 Dialect 是对某一类 IR 或者一些数据结构相关操作进行抽象，比如 llvm dialect 就是对 LLVM IR 的抽象，tensor dialect 就是对 Tensor 这种数据结构和操作进行抽象：

官网提供的MLIR Dialect

除了这些，各种深度学习框架都在接入 MLIR，比如 TensorFlow，Pytorch，OneFlow 以及 ONNX 等等，大家都能在 github 找到对应工程。

抽象了多个级别的 IR 好处是什么呢？这就要结合 MLIR 的编译流程来看，MLIR 的编译流程大致如下：

图源法斯特豪斯，侵删

对于一个源程序，首先经过语法树分析，然后通过 Dialect 将其下降为 MLIR 表达式，再经 MLIR 分析器得到目标程序。注意这个目标程序不一定是可运行的程序。比如假设第一次的目标程序是 C 语言程序，那么它可以作为下一次编译流程的源程序，通过 Dialect 下降为 LLVM MLIR。这个 LLVM MLIR 即可以被 MLIR 中的 JIT 执行，也可以通过 Dialect 继续下降，下降到三地址码 IR 对应的 MLIR，再被 MLIR 分析器解析获得可执行的机器码。

因此 MLIR 这个多级别的下降过程就类似于我们刚才介绍的可以渐进式学习，解决了 IR 到之间跨度太大的问题。比如我们不熟悉 LLVM IR 之后的层次，没有关系，我们交给 LLVM 编译器，我们去完成前面那部分的 Dialect 实现就可以了。

0x4. 总结

这篇文章简单聊了一下对 MLIR 的粗浅理解，欢迎大家批评指正。

本文总阅读量6199次