[논문 리뷰] Compiling ONNX Neural Network Models Using MLIR

Jin Tian, Gheorghe-Teodor Bercea외 9명|arXiv (Cornell University)|2020. 08. 19.

Parallel Computing and Optimization Techniques참고 문헌 3인용 수 39

한 줄 요약

이 논문은 ONNX 모델을 MLIR을 사용해 네이티브 코드로 변환하는 컴파일러인 onnx-mlir을 소개하며, ONNX와 krnl 다이얼렉트 및 여러 최적화 패스를 도입합니다. 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 컴파일러의 아키텍처와 설계를 설명한다.
• ONNX 의미를 하위 레벨 코드로 연결하기 위한 두 개의 새로운 MLIR 다이얼렉트(on nx와 krnl)를 도입한다.
• 신경망 컴파일에 대한 최적화 패스와 그 영향력을 제시한다.

제안 방법

• 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 기반 파이프라인에서 효과적인 초기 최적화이다.
• MNIST는 POWER9에서 빠르게 컴파일될 수 있으며(0.237 s 컴파일; 0.001 s 추론), ResNet50은 POWER9에서 7.661 s에 컴파일되고 7.540 s에 추론되어 대형 모델의 실행 가능성을 보여준다.
• 현재 구현은 정확성 및 기본 성능에 초점을 두고 있으며, 속도 향상을 위한 폴리헤드럴, SIMD, 루프 퓨전 최적화를 추가할 계획이다.
• 시스템은 빅 엔디안 시스템과 같은 교차 아키텍처 도전에 직면해 있으며, 이를 해결하기 위해 ONNX와 MLIR에 패치를 적용했다. 현재는 x86, Power, Z, Windows를 지원한다.