[Paper Review] Compiling ONNX Neural Network Models Using MLIR

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

Parallel Computing and Optimization Techniques3 references39 citations

TL;DR

This paper presents onnx-mlir, a compiler that translates ONNX models into native code using MLIR, introducing ONNX and krnl dialects and several optimization passes. It demonstrates preliminary results on MNIST and ResNet50 across multiple architectures.

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.

Motivation & Objective

• Motivate model portability by enabling ONNX models to run efficiently in target environments.
• Describe the architecture and design of the onnx-mlir compiler.
• Introduce two new MLIR dialects (onnx and krnl) to bridge ONNX semantics to lower-level code.
• Present optimization passes and their impact on neural network compilation.

Proposed method

• Use MLIR to build a multi-dialect compiler pipeline that lowers ONNX graphs to LLVM IR and native code.
• Introduce two dialects: an ONNX dialect that encodes ONNX semantics and a krnl dialect for loop-based lowering.
• Apply optimization passes including graph rewriting, shape inference, and constant propagation.
• Provide an end-to-end flow from ONNX model import to a dynamic entry point function for inference.
• Demonstrate the pipeline by compiling and running MNIST and ResNet50 models on POWER systems.

Experimental results

Research questions

• RQ1 Can ONNX models be compiled into efficient native code using MLIR with dedicated dialects?
• RQ2 What are the required dialects and passes to lower ONNX to LLVM/Native code?
• RQ3 How do initial optimizations affect compilation and inference performance on representative models?

Key findings

• The onnx-mlir compiler can translate ONNX models to a native code library with an entry function _dyn_entry_point_main_graph.
• ONNX and krnl dialects enable structured lowering and loop optimizations, with subsequent conversion to affine, std, and llvm dialects.
• Graph rewriting and constant propagation are effective early optimizations in the MLIR-based pipeline.
• MNIST can compile quickly (0.237 s compilation; 0.001 s inference) on POWER9, while ResNet50 compiles in 7.661 s and infers in 7.540 s on POWER9, illustrating feasibility for large models.
• The current implementation focuses on correctness and base performance, with plans to add polyhedral, SIMD, and loop-fusion optimizations to improve speed.
• The system faces cross-architecture challenges (e.g., big-endian systems) that required patches to ONNX and MLIR; support now extends to x86, Power, Z, and Windows.