[论文解读] Compiling ONNX Neural Network Models Using MLIR

Jin Tian, Gheorghe-Teodor Bercea+9 more|arXiv (Cornell University)|Aug 19, 2020

Parallel Computing and Optimization Techniques参考文献 3被引用 39

一句话总结

本文介绍 onnx-mlir，一种将 ONNX 模型翻译为本地代码的编译器，使用 MLIR，引入 ONNX 和 krnl dialect，以及若干优化阶段。它在多个体系结构上对 MNIST 和 ResNet50 展示了初步结果。

ABSTRACT

Deep neural network models are becoming increasingly popular and have been used in various tasks such as computer vision, speech recognition, and natural language processing. Machine learning models are commonly trained in a resource-rich environment and then deployed in a distinct environment such as high availability machines or edge devices. To assist the portability of models, the open-source community has proposed the Open Neural Network Exchange (ONNX) standard. In this paper, we present a high-level, preliminary report on our onnx-mlir compiler, which generates code for the inference of deep neural network models described in the ONNX format. Onnx-mlir is an open-source compiler implemented using the Multi-Level Intermediate Representation (MLIR) infrastructure recently integrated in the LLVM project. Onnx-mlir relies on the MLIR concept of dialects to implement its functionality. We propose here two new dialects: (1) an ONNX specific dialect that encodes the ONNX standard semantics, and (2) a loop-based dialect to provide for a common lowering point for all ONNX dialect operations. Each intermediate representation facilitates its own characteristic set of graph-level and loop-based optimizations respectively. We illustrate our approach by following several models through the proposed representations and we include some early optimization work and performance results.

研究动机与目标

• 通过使 ONNX 模型在目标环境中高效运行来提升模型的可移植性。
• 描述 onnx-mlir 编译器的架构与设计。
• 引入两种新的 MLIR 方言（onnx 和 krnl），以把 ONNX 语义桥接到底层代码。
• 展示优化阶段及其对神经网络编译的影响。

提出的方法

• 使用 MLIR 构建一个多方言编译管线，将 ONNX 图降低为 LLVM IR 并生成本地代码。
• 引入两种方言：一个编码 ONNX 语义的 ONNX 方言，以及一个用于基于循环的降级的 krnl 方言。
• 应用包括图重写、形状推断和常量传播在内的优化阶段。
• 提供从 ONNX 模型导入到用于推理的动态入口点函数的端到端流程。
• 通过在 POWER 系统上编译并运行 MNIST 和 ResNet50 模型来演示该管线。

实验结果

研究问题

• RQ1 是否可以利用带有专用方言的 MLIR，将 ONNX 模型编译为高效的本地代码？
• RQ2 将 ONNX 降到 LLVM/本地代码需要哪些方言和优化阶段？
• RQ3 初始优化如何影响代表性模型的编译和推理性能？

主要发现

• onnx-mlir 编译器能够将 ONNX 模型翻译为具有入口函数 _dyn_entry_point_main_graph 的本地代码库。
• ONNX 与 krnl 方言实现结构化降级和循环优化，随后转换为 affine、std 和 llvm 方言。
• 图重写和常量传播是基于 MLIR 的流程中的有效早期优化。
• 在 POWER9 上，MNIST 的编译很快（0.237 s 编译；0.001 s 推理），而 ResNet50 的编译为 7.661 s，推理为 7.540 s，显示了对大模型的可行性。
• 当前实现关注正确性和基础性能，计划添加多面体分析、多 SIMD、循环融合等优化以提升速度。
• 系统面临跨体系结构的挑战（如大端系统），需要对 ONNX 和 MLIR 进行修补；现在支持扩展到 x86、Power、Z 和 Windows。