ONNX-MLIR Guideline: AddCustomAccelerators.md

moze·2025년 2월 15일

이 글은 ONNX-MLIR Github에서 제공하는 커스텀 가속기 추가 가이드라인에 대한 한국어 해석을 담고 있습니다.

A guideline on adding a new custom accelerator

일반적으로 ONNX-MLIR 상에서는 커스텀 가속기를 플러그인 형태로 처리하며, ONNX-MLIR을 빌드하거나 ONNX 모델을 컴파일할 때 온-오프할 수 있는 방식으로 동작한다. 이는 보통 CMake를 통해서 이루어진다. NNPA 가속기 예제를 보면서 프로세스를 살펴보자.

1. 코드 폴더

ONNX-MLIR에서는 모든 가속기 관련 코드를 src/Accelerators 디렉토리 아래 별도 폴더로 저장한다. 즉 새로운 가속기를 추가하기 위해서는 해당 디렉토리에 가속기 전용 하위 디렉토리를 생성해야 한다.

생성한 폴더 이름은 ONNX-MLIR에서 해당 가속기의 이름으로 사용되고, 다음 용도로도 활용된다.

(1) CMake 빌드

CMake가 해당 가속기 디렉토리 아래의 코드들을 빌드한다.
CMake 변수 ONNX_MLIR_ACCELERATORS를 설정해서 아래와 같이 빌드할 수 있다.

$ cd build
$ cmake .. -DONNX_MLIR_ACCELERATORS='accel1;accel2'

잘 모르겠다면 AccelNNPAHowToUseAndTest.md 문서를 확인해보자.

(2) ONNX-MLIR을 이용한 모델 컴파일

onnx-mlir 툴을 사용할 때, 해당 가속기를 위한 모델 컴파일 옵션을 사용 가능하다.
--maccel 옵션을 사용해서 수행할 수 있다.

$ onnx-mlir --maccel='가속기명' model.onnx
# --maccel 옵션은 CMake에서 빌드된 가속기만 사용 가능

(3) ONNX-MLIR-OPT를 이용한 Pass 활성화

onnx-mlir-opt 명령어를 실행할 때, 가속기 관련 최적화 패스 활성화가 가능하다.
onnx-mlir 툴을 사용할 때와 마찬가지로 --maccel 옵션을 추가해서 수행 가능하다.

$ onnx-mlir-opt --maccel='가속기명' --'패스명' model.mlir
# --maccel 옵션은 onnx-mlir과 마찬가지로 CMake에서 빌드된 가속기만 사용 가능

폴더 내부 구조의 경우 각 가속기에 따라 유동적으로 설정할 수 있으나 ONNX-MLIR의 최상위 폴더 구조를 따르는 것을 권장하고 있다.

2. 코드 통합

(1) 매크로

각 가속기는 특정 매크로를 정의하고 이를 onnx_mlir::accel::Accelerator에 포함해야 한다. (잘 이해되지 않는다면 NNPA 예제를 보자. NNPAAccelerator.hpp를 src/Accelerators/NNPA/NNPAAccelerator.hpp에서 정의하고, 이를 Accelerator.hpp에서 include하고 있다.) 필수적으로 정의해야 하는 매크로는 다음과 같다.

  1. INSTRUMENTSTAGEENUM<accel_name>
  2. INSTRUMENTSTAGECL_ENUM<accel_name>
  3. PROFILEIRCL_ENUM<accel_name>
  4. OPTREPORTENUM<accel_name>
  5. OPTREPORTCL_ENUM<accel_name>

만약 가속기 이름이 ACCEL1이면 아래 예시와 같이 정의한다.

#define INSTRUMENTSTAGE_ENUM_ACCEL1
#define INSTRUMENTSTAGE_CL_ENUM_ACCEL1
#define PROFILEIR_CL_ENUM_ACCEL1
#define OPTREPORT_ENUM_ACCEL1
#define OPTREPORT_CL_ENUM_ACCEL1

(2) 다이얼렉트 및 패스

MLIR에서 코드를 작성할 때는 보통 다이얼렉트와 패스를 설계한다. ONNX-MLIR도 마찬가지로, 가속기 지원을 위해서는 ONNX-MLIR에 해당 가속기에 맞는 다이얼렉트와 패스를 등록해야 한다.

ONNX-MLIR에서는 가속기 통합(지금 이 프로세스)을 위해 기본 클래스 onnx_mlir::accel:Accelerator를 제공하며 이를 상속해 커스텀 가속기 코드(다이얼렉트와 패스)를 작성할 수 있다. 사용자는 다이얼렉트와 패스 작성을 마친 뒤, 이를 등록하는 훅을 정의해야 한다.

다이얼렉트와 패스를 등록하는 훅은 다음과 같다.

//===--------------------------------------------------------------------===//
// onnx-mlir 관련 훅
//===--------------------------------------------------------------------===//

/// 해당 가속기에 필요한 변환을 추가하는 역할
virtual void addPasses(mlir::OwningOpRef<mlir::ModuleOp> &module,
    mlir::PassManager &pm,
    onnx_mlir::EmissionTargetType &emissionTarget) const = 0;

//===--------------------------------------------------------------------===//
// onnx-mlir-opt 관련 훅
//===--------------------------------------------------------------------===//

/// 가속기에 필요한 MLIR 다이얼렉트 등록
virtual void registerDialects(mlir::DialectRegistry &registry) const = 0;

/// 가속기용 패스 등록 (명령어 옵션에서 사용 가능)
virtual void registerPasses(int optLevel) const = 0;

//===--------------------------------------------------------------------===//
// onnx-mlir 및 onnx-mlir-opt 공통 훅
//===--------------------------------------------------------------------===//

/// 가속기 패스 설정
virtual void configurePasses() const = 0;

//===--------------------------------------------------------------------===//
// onnx-to-krnl 변환 관련 훅
//===--------------------------------------------------------------------===//

/// TensorType을 MemRefType으로 변환
virtual mlir::MemRefType convertTensorTypeToMemRefType(
    const mlir::TensorType tensorType) const = 0;

/// ONNXToKrnl 변환 대상 설정
virtual void conversionTargetONNXToKrnl(
    mlir::ConversionTarget &target) const = 0;

/// ONNXToKrnl 변환 패턴 설정
virtual void rewritePatternONNXToKrnl(mlir::RewritePatternSet &patterns,
    mlir::TypeConverter &typeConverter, mlir::MLIRContext *ctx) const = 0;

//===--------------------------------------------------------------------===//
// krnl-to-llvm 변환 관련 훅
//===--------------------------------------------------------------------===//

/// KrnlToLLVM 변환 대상 설정
virtual void conversionTargetKrnlToLLVM(
    mlir::ConversionTarget &target) const = 0;

/// KrnlToLLVM 변환 패턴 설정
virtual void rewritePatternKrnlToLLVM(mlir::RewritePatternSet &patterns,
    mlir::LLVMTypeConverter &typeConverter, mlir::MLIRContext *ctx) const = 0;

ONNX-MLIR에는 다양한 패스가 존재하지만 가속기 통합에 대해서는 두 개의 핵심 패스에 대한 훅만 제공한다. onnx-to-krnl, krnl-to-llvm만 제공하는데, 이 두 패스가 onnx-mlir의 첫 번째 및 마지막 변환 패스이기 때문이다. (ONNX-MLIR은 크게 ONNX Model > ONNX Dialect > Krnl Dialeact > LLVM 순으로 변환을 제공한다.)

onnx-to-krnl 패스

  • 어떤 ONNX 연산자를 CPU에서 실행할지, 가속기에서 실행할지 결정
  • ONNX 연산자를 Krnl 다이얼렉트 또는 가속기 다이얼렉트로 변환

krnl-to-llvm 패스

  • krnl 다이얼렉트 및 가속기 연산자를 LLVM 다이얼렉트로 변환
  • 변환 결과로 어셈블리 코드를 생성하고 가속기 전용 API를 호출

예시로, ONNX-MLIR에서는 NNPA 가속기를 Built-in으로 지원한다. 이 가속기는 ZHigh, ZLow 다이얼렉트를 가지고 있다.

  • ZHigh Dialect: onnx-to-krnl 패스에서 사용
  • ZLow Dialect: krnl-to-llvm 패스에서 사용

즉, ONNX 연산자를 ZHigh 다이얼렉트로 변환하고 KrnlToLLVM 변환 과정에서 ZLow 다이얼렉트로 변환하는 구조이다. (ONNX Model > ZHigh Dialect > ZLow Dialect > LLVM)

3. 테스트

가속기 관련 테스트는 test 폴더 내에 배치해야 한다.
ONNX-MLIR에서는 두 가지 유형의 테스트를 제공한다.

(1) LIT 테스트

  • mlir/accelerators 내에 새로운 폴더 생성 후 작업
  • MLIR 패스를 검증하는 Low-Level Intergration Test 프레임워크를 활용하여 가속기 관련 기능을 테스트
  • example
    LIT 테스트는 다음과 같은 위치에 저장되어야 한다.
test/mlir/accelerators/'가속기명'/

(2) 기타 테스트

  • test/accelerators 내에 새로운 폴더 생성 후 작업
  • LIT 이외 일반적인 유닛 테스트 및 기능 테스트 수행
  • example
    기타 테스트는 다음과 같은 위치에 저장되어야 한다.
test/accelerators/'가속기명'/

한국인을 위한 3줄 요약

6줄이지만..

(1) 코드 폴더 생성 및 빌드 설정

src/Accelerators 디렉토리에 가속기 전용 폴더를 만들고, CMake 변수와 --maccel 옵션을 통해 가속기 코드를 빌드 및 컴파일 옵션으로 등록한다.

(2) 매크로 및 후크 구현

필수 매크로를 정의하고 onnx_mlir::accel::Accelerator를 상속받아 다이얼렉트와 패스 등록을 위한 후크를 구현하여 커스텀 기능을 통합한다.

(3) 테스트 배치

MLIR 패스 검증을 위한 LIT 테스트는 test/mlir/accelerators, 기타 유닛/기능 테스트는 test/accelerators 디렉토리에 배치하여 가속기 기능을 검증한다.

출처: https://github.com/onnx/onnx-mlir/blob/main/docs/AddCustomAccelerators.md

profile
moze
이전 포스트

Chapter 3: High-level Language-Specific Analysis and Transformation

다음 포스트

ONNX-MLIR의 NNPA 디렉토리를 톺아보자 (Accelerators)

0개의 댓글