ONNX-Runtime에서 Custom Operator 구현 Study

서론

이 장은 ORT(ONNX-Runtime)에서 Custom Operator를 구현하는 방법에 대해 study한것을 설명한다.

ONNX Operation은 하나 이상의 커널을 갖는다. 각 커널은 연산이 어떻게 실행되는지를 정의하며 커널은 특정 하드웨어에 대해 최적화되어 있다. 따라서 ORT가 다양한 하드웨어 플랫폼에서 최적의 성능을 낼 수 있다. Core_ML을 바탕으로 연구했다.

목차

1. 프로세스
2. Create the Execution Provider
3. coreml_execution_factory.cc

프로세스

1. Custom Op 구현
2. Custom Op 등록
3. Custom Op 라이브러리 빌드
4. Custom Op 라이브러리 로드

Create the Execution Provider

1. Create a folder under onnxruntime/core/providers
2. Create a folder under include/onnxruntime/core/providers, it should has the same name as the first step.
3. Create a new class, which must inherit from IExecutionProvider. The source code should be put in ‘onnxruntime/core/providers/[your_provider_name]’
4. Create a new header file under include/onnxruntime/core/providers/[your_provider_name]. The file should provide one function for creating an OrtProviderFactoryInterface. You may use ‘include/onnxruntime/core/providers/cpu/cpu_provider_factory.h’ as a template. You don’t need to provide a function for creating MemoryInfo.
5. Put a symbols.txt under ‘onnxruntime/core/providers/[your_provider_name]’. The file should contain all the function names that would be exported from you provider. Usually, just a single function for creating provider factory is enough.
6. Add your provider in onnxruntime_providers.cmake. Build it as a static lib.
7. Add one line in cmake/onnxruntime.cmake, to the ‘target_link_libraries’ function call. Put your provider there.

coreml_execution_factory.cc

아래는 전체 코드다.

먼저 구조체인 CoreMLProviderFactory를 먼저 봐보자.

IExecutionProviderFactory를 상속받아 구조체를 생성하며 IExecutionProviderFactory의 구조체에 있는 CreateProvider라는 함수를 오버라이드하여 사용하려한다.

실제로 IExecutionProviderFactory를 봐보면 가상함수로CreateProvider가 존재한다.

그리고 그 오버라이드 한 CreateProvider함수를 재정의해준다. 해당 함수는 CoreMLExecutionProvider객체를 생성하고 초깃값으로 coreml_flags_값을 전달해주며 make_unique포인터로 wrapping하여 Return한다.

그럼 CoreMLExecutionPriovider객체를 생성하면 어떤 일이 일어나는지 봐보자.

CoreMLExecutionPriovider는 IExecutionProvider를 상속하여 class를 만든다.

class를 봐보면 GetCapability라는 함수가 있는데 이 함수는 IExecutionProvider에 속한 함수이며 오버라이드해서 사용하겠다는 뜻이다.

실제로 IExecutionProvider 클래스를 봐보면 GetCapability가 존재함을 알수 있다.

execution_provider.h 에서 좀 더 내려가보면 아래에 나와있다.

ORT는 다양한 Backend(CPU, TPU, GPU..)를 통해 ONNX 모델을 실행할 수 있도록 해주는 프레임워크인데, IExecutionProvider는 이러한 Backend에 대해 공통적인 기능을 정의하는 인터페이스 역할을 해준다. 따라서 IExecutionProvider를 상속받아 target device에 따른 특정 기능들을 구현할 수 있다.

IExecutionProvider 클래스를 봐보면 생성자 위임(Delegating Constructors)과 2개의 생성자를 통해 객체를 생성할 때 다양한 입력에 따라 대응한다.

GetCapability 함수를 봐보면 GraphView와 IKernelLookup 인스턴스를 인자로 받는데 Execution Provider가 지원되는 하드웨어 위에서 실행할 Graph를 target에 맞춰 할당하는 역할을 한다.

추측해보면 지원되는 하드웨어를 참조하는 것이 kernel_lookup인 것 같고, 모델의 layer를 할당받는 부분이 graph_view인것 같다.

virtual 함수로 구성되어 있는걸 봐선 상속해서 특정 하드웨어 플랫폼에 맞게 사용하라는 뜻이다.

IKernelLookup 객체는 아래와 같이 구현되어 있다.

GraphView객체는 아래와 같이 구현되어 있다. 주석을 보면 model의 graph를 읽어오는 역할을 하는 것 같다.

현재까지의 과정 상태

[참조]

https://onnxruntime.ai/docs/execution-providers/

다시 coreml_exeuction_provider.h로 돌아가보면, Compile 함수를 오버로드해서 사용한다.

Compile함수의 원형은 다음과 같다.