ONNX-MLIR 이란

※ 컴파일러에 대한 배경지식이 있어야 이해 가능한 글입니다 ※

1. 인공지능 컴파일러란?

인공지능 컴파일러는 일종의 그래프로 표현되어 있는 인공지능 모델을 하드웨어 명령어 집합으로 변환하는 일련의 과정을 말하며, 이 과정에서 그래프의 최적화 및 하드웨어별 자원 할당, 분배 과정이 동반된다고 한다. (이게 어려우므로 보통 인공지능 반도체 개발과 병행된다고 한다.)

그래프 형태로 인공지능 모델을 표현하고 바로 하드웨어 명령어 집합으로 변환하는 것이 아니라, 그 사이에 많은 과정이 존재한다. 그래프를 계산 그래프로 만들고, 중간 표현으로 바꾸고 백엔드 코드를 생성하고 하드웨어에 맞게 또 코드를 생성하고.. 이 과정은 또 High-level language와 하드웨어의 종류에 따라 다 다르게 적용되었다.

굉장히 복잡한 과정이어서 이런 과정들을 통합해보자! 하고 나온 프로젝트가 MLIR이다. MLIR은 다양한 딥러닝 프레임워크에 대한 컴파일러 Front-end 파트를 통합하는 프로젝트로, LLVM 프로젝트의 산하 프로젝트이다. LLVM은 Front-end, Middle-end, Back-end를 가지는 완전한 컴파일러인데 이 Front-end를 보완하고자 MLIR 프로젝트를 새로이 시작한듯 하다. (뇌피셜)

MLIR은 개발사 홈페이지의 소개에 따르면 딥러닝 프레임워크에서 고성능 머신러닝 모델을 실행하는 데 필요한 인프라를 통합하는 일반적인 중간 표현을 정의한다고 쓰여있다. 쉽게 말해서 다양한 프레임워크를 통합된 중간 표현으로 만드는 프로젝트라는 것!

MLIR 프로젝트가 Original로 존재하고.. 어떤 것을 입력으로 하는지에 따라 Torch-MLIR, ONNX-MLIR로 나뉘는거 같다. 이 글에서는 ONNX-MLIR에 대해 다뤄보겠다.

2. ONNX와 ONNX-MLIR

ONNX란?

Open Neural Network Exchange는 다양한 머신러닝 프레임워크로 작성된 모델들을 공통의 파일 형식으로 사용할 수 있게끔 하는 아주 유용한 도구이다. Tensorflow, PyTorch 등으로 작성된 모델을 .ONNX 라는 하나의 통합된 형식으로 바꿀 수 있다. 이렇게 변환된 ONNX 모델은 netron.app에서 그 내부를 그래프 형태로 살펴볼 수 있다.

이렇게 변환한 모델은 ONNX와 호환되는 ONNXRuntime이나 ONNXRuntime-extensions을 통해 실행할 수 있다.

ONNX-MLIR이란?

ONNX-MLIR은 개발자의 소개에 따르면 ONNX 그래프(= ONNX 모델)를 최소한의 런타임 지원으로 그래프를 구현하는 코드로 변환하는 컴파일러 기술을 제공한다고 한다. 그냥.. ONNX 전용 MLIR이라고 보면 되겠다. ONNX-MLIR을 잠깐 사용한 적이 있는데, ONNX-MLIR 하나로 LLVM IR까지 Lowering 가능하다.

ONNX-MLIR 깃허브의 예제를 함께 보면서 Lowering에 대한 감을 잡아보자.
ONNX-MLIR을 설치하는 방법은다음 깃허브를 참고 바란다.
리눅스에서 설치하는 방법
윈도우에서 설치하는 방법

ONNX-MLIR 사용 방법은 다음과 같다. 설치해서 확인해보면 ONNX-MLIR 말고도 최적화해서 Lowering하는 툴이 굉장히 많다. 필자가 추후에 기능을 다 정리해 글을 업로드할 예정이지만 급하다면 --help 옵션으로 하나하나 다 확인해보자.

OVERVIEW: ONNX-MLIR modular optimizer driver

USAGE: onnx-mlir [options] <input file>

OPTIONS:

Generic Options:

 --help        - Display available options (--help-hidden for more)
 --help-list   - Display list of available options (--help-list-hidden for more)
 --version     - Display the version of this program

ONNX-MLIR Options:
These are frontend options.

 Choose target to emit:
     --EmitONNXBasic - Ingest ONNX and emit the basic ONNX operations without inferred shapes.
     --EmitONNXIR    - Ingest ONNX and emit corresponding ONNX dialect.
     --EmitMLIR      - Lower the input to MLIR built-in transformation dialect.
     --EmitLLVMIR    - Lower the input to LLVM IR (LLVM MLIR dialect).
     --EmitObj       - Compile the input to an object file.
     --EmitLib       - Compile and link the input into a shared library (default).
     --EmitJNI       - Compile the input to a jar file.

 Optimization levels:
     --O0           - Optimization level 0 (default).
     --O1           - Optimization level 1.
     --O2           - Optimization level 2.
     --O3           - Optimization level 3.

간단한 예제를 통해 사용 방법을 익혀보자.
다음과 같은 방식으로 add.onnx라는 ONNX 모델을 ONNX IR로 낮출 수 있다.

./onnx-mlir --EmitONNXIR add.onnx

잘 낮추어 졌다면 다음과 같은 형식으로 출력될 것이다.

module {
  func.func @main_graph(%arg0: tensor<10x10x10xf32>, %arg1: tensor<10x10x10xf32>) -> tensor<10x10x10xf32> {
    %0 = "onnx.Add"(%arg0, %arg1) : (tensor<10x10x10xf32>, tensor<10x10x10xf32>) -> tensor<10x10x10xf32>
    return %0 : tensor<10x10x10xf32>
  }
}

ONNX-MLIR의 사용법을 간단히 해 보았는데..
현실은 이만큼 간단하지 않다
이것도 간단하긴 하지만 위 예제보다 더욱 실무적인 예제(?)를 경험해 보고 싶다면 MNIST ONNX 모델을 Lowering하는 예제를 수행해보자.

지금까지 인공지능 컴파일러부터 MLIR, ONNX-MLIR까지 간단하게 알아보았다. 본 글은 자세한 이해 보다는 MLIR에 대한 직관을 제공하는 글이다. 다음 글에서 ONNX-MLIR의 기능을 자세히 알아보겠다!

참고자료
(1) https://ettrends.etri.re.kr/ettrends/189/0905189004/032-042_%EA%B9%80%EC%A7%84%EA%B7%9C.pdf
(2) https://onnx.ai/
(3) https://github.com/onnx/onnx-mlir
(4) https://wikidocs.net/235922
(5) https://youtu.be/Y4SvqTtOIDk?si=jQnVLHBPocJIGVTw