[论文阅读] 《Compiling ONNX Neural Network Models Using MLIR》

《Compiling ONNX Neural Network Models Using MLIR》

本文主要是对 ONNX-MLIR 这个编译框架的论文内容概括整理，速记论文中一些个人认为值得记录和分享的点，所以篇幅不长。

要解决的问题

当前模型加速技术有一些深度学习框架是通过使用一些高性能依赖库来实现加速的，在前向推理时让算子实现直接调用高性能库中已经实现的函数，这样的加速方式很直接，但同样会有一些问题：

• 能够加速的模型种类受限于高性能库能够支持的算子种类
• 让使用者需要额外的各种安装来确保高性能库的正确运行
• 库缺乏灵活性，同一个算子在不同情况下可能最优实现不同

名字由来

ONNX-MLIR 则是一款新的图编译器，能够将训练后的模型直接编译成目标硬件上的原生代码，并且可以在编译过程中使用各种成熟的优化技术。因为使用了 ONNX 做了输入模型的目标格式，并且使用了 Multi-level Intermediate Representation（MLIR）编译框架，所以项目直接取名，ONNX-MLIR 。

此处省略了论文中对 ONNX 框架的简单介绍及 MLIR 的简介。 MLIR 也有官方的论文，之后会发出相关的论文笔记，链接到这里。不过这里可以先给一张图，来表示 MLIR 的一些核心概念：

Operation 就可以理解为一个算子，layer，然后算子内部表示则是有属性（attributes），输入（SSA Operands)，输出，Regions 等。

一个或多个 Operation 组成了一个 Block，多个Block 组成了 Regions，而 Operation 里面有可以包含 Regions ，由此就可以实现嵌套，用于实现控制流等内容。

解决方法

这张图可以很好地归纳总结整个 ONNX-MLIR 的框架流程。为了实现图编译的功能，并且可以跟 ONNX 算子粒度想匹配，ONNX-MLIR 基于 MLIR 的 Dialect 的概念提出了 ONNX Dialect 和 Krnl Dialect 两个新的Dialect （这里的 Dialect 不了解的同学可以暂且理解为一层中间表示），其他部分则是直接使用 MLIR 实现好的 Dialect 来进行编译映射。

• 在模型原生算子到 ONNX Dialect 定义的算子的转化中主要可以进行类型推导（Type inference），算子的组合，分解，消除等优化操作。
• Krnl Dialect 则是主要用于对计算图中标量及循环操作的优化，从这层开始就主要集中在性能及内存方面的优化操作了。

主要是通过上图对 ONNX-MLIR 框架有了更多的了解，之前的了解仅局限在到 ONNX Dialect 这一层面，并不知道里面很多涉及到 Krnl Dialect 和 Affine Dialect 的内容都是干什么的，现在看论文一张图便很清晰了。