TVM Runtime System

TVM runtime의 주요 요구사항은 다음과 같다.

• Deployment: 컴파일된 함수를 python/javscript/c++ 언어에서 호출
• Debug: python으로 함수를 정의하고 이를 컴파일된 함수에서 호출
• Link: device specific code (CUDA등)을 호출하도록 드라이버 코드를 작성하고 이를 컴파일된 호스트 함수에서 호출
• Prototype: python에서 IR pass를 정의하고 이를 C++ backend에서 호출
• Expose: python같은 front-end를 가지는 c++로 개발된 컴파일러 스택
• Experiment: 컴파일된 함수를 타켓으로 전송하여 타겟에서 직접 실행
  

PackedFunc

PackedFunc 객체는 caller와 callee가 다른 언어가 될 수 있는 function call을 나타냄.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15

#include <tvm/runtime/packed_func.h>

void MyAdd(TVMArgs args, TVMRetValue* rv) {
  // automatically convert arguments to desired type.
  int a = args[0];
  int b = args[1];
  // automatically assign value return to rv
  *rv = a + b;
}

void CallPacked() {
  PackedFunc myadd = PackedFunc(MyAdd);
  // get back 3
  int c = myadd(1, 2);
}

위의 코드에서 서 PackedFunc는 MyAdd이고 다음과 같은 특징을 지님

• 두 개의 Argument(arg, rv)를 지니며 type-erased function 이다.
• 즉 함수의 입력 유형과 반환 유형을 제한하지 않음
• 내부적으로 PackedFunc를 호출 시, input argument를 스택의 TVMArg에 패킹하고 결과를 TVMRetValue을 통해 반환

C++의 Template 문법 덕택에 PackedFunc을 python 같은 dynamic language에서 glue 코드 없이 일반 함수 호출하듯이 콜할 수 있다.

1
2
3

// register a global packed function in c++
TVM_REGISTER_GLOBAL("myadd")
.set_body(MyAdd);

1
2
3
4
5
6

# packedfunction 호출
import tvm

myadd = tvm.get_global_func("myadd")
# prints 3
print(myadd(1, 2))

상기와 같은 호출 구조가 가능한 이유는 TVMArgs와 TVMRetValue 구조 덕분인데 TVM은 두 구조의 type을 아래와 같이 제한함

• int, float and string
• PackedFunc itself
• Module for compiled modules
• DLTensor* for tensor object exchange
• TVM Object to represent any object in IR

위의 제약들은 serialization이 필요 없이 구현을 간편하게 하기 위함이다. 그러나 이러한 제약에도 불구하고 딥러닝 deploy에서 대부분 DLTensor 또는 int/float만 사용하므로 사용하기 충분

하나의 PackedFunc가 다른 PackedFunc를 인수로 사용할 수 있으므로, 우리는 함수를 python(PackedFunc)에서 C++로 전달할 수 있습니다.

1
2
3
4
5

TVM_REGISTER_GLOBAL("callhello")
.set_body([](TVMArgs args, TVMRetValue* rv) {
  PackedFunc f = args[0];
  f("hello world");
});

 1
 2
 3
 4
 5
 6
 7
 8
 9
10

import tvm

def callback(msg):
  print(msg)

# convert to PackedFunc
f = tvm.convert(callback)
callhello = tvm.get_global_func("callhello")
# prints hello world
callhello(f)

TVM은 PackedFunc을 모든 언어에 임베디 할 수 있는 minimum C API을 제공 (python, java, javascript)

• 이는 Lua랑 비슷(새로운 언어 대신 C++을 이용한다는 것을 제외하면)

PackedFunc은 compiler와 deployment stack 양쪽에서 모두 쓰임

• TVM의 모든 compiler pass function은 frontend에 PackedFunc으로 노출됨
  • (역자주) TVM에서 내부적으로 사용하는 컴파일 패스 함수들 (예: Relay 최적화, 타입 추론 등)을 사용자가 Python 등에서 쓸 수 있도록 PackedFunc 형태로 포장해서 노출, C++에서 만든 내부 패스를 PackedFunc로 감싸면 Python에서 마치 일반 함수처럼 사용할 수 있게 됨

1
2
3
4
5
6
7
8
9

  from tvm import relay

mod = ...  # relay IRModule
mod = relay.transform.InferType()(mod)

# 여기서 InferType() 은 Relay의 타입 추론 패스입니다.
#내부적으로는 C++에 정의된 InferTypePass 같은 함수를
# → TVM_REGISTER_GLOBAL("relay._transform.InferType") 를 통해
# → PackedFunc 형태로 Python에 노출한 것입니다.

• 컴파일된 module도 컴파일된 함수를 PackedFunc로 반환
  • (역자주) TVM으로 모델을 컴파일하면, 그 결과물은 Module 객체로 나오는데, 그 안에 있는 각 함수(예: main, tvmgen_default_fused_…)도 PackedFunc 형태로 제공

1
2
3
4
5
6

lib = tvm.build(...)  # 모델 컴파일
f = lib["main"]       # "main" 함수 가져오기
f(input_tensor)       # 실행!

#여기서 lib["main"]은 실제로 PackedFunc 객체입니다.
#즉, lib["main"]을 호출하면 내부적으로 PackedFunc.__call__()이 실행되어 모델 추론이 일어납니다.

런타임 크기를 최소화 하기 위해 배포 런타임에서 IR Object 지원을 제외함. 이를 통해 배포 런타임의 크기를 200K~600K 정도로 만들 수 있었음

• (역자주) TVM은 크게 컴파일러 영역 (IR 사용)과 런타임 영역 (IR 제거)으로 나누는데 런타임영역에서는 IR기능을 분리, 예를 들어 라즈베리파이에 모델을 실행시키기 위해 libtvm_runtime.so만 넣음, 이 안에는 PackedFunc, DLTensor, CUDA 호출기 등 필요한 것만 포함하며 Relay 같은 IR 처리 로직은 빠져 있음

PackedFunc를 호출하는데 드는 오버헤드는 스택에 몇 가지 값만 저장하면 되기 때문에 적다.요약하자면, PackedFunc는 컴파일러와 배포를 지원하기 위해 광범위하게 사용하는 TVM의 범용 glue임

Module

TVM은 여러 유형의 장치를 지원하기 때문에 다양한 유형의 드라이버를 지원해야 하며 다음과 같은 일을 해야함

• 드라이버 API를 사용하여 kernel을 로드
• argument를 packed format으로 셋업하고 kernel을 실행
• 사용된 함수가 threadsafe하도록 kernel을 패치
• 이러한 driver glue를 C++로 구현하여 사용자에게 제공
  그러나 각 디바이스의 모든 함수에 대해 상기 작업을 할 수 없으므로 PackedFunc을 사용

TVM은 컴파일된 객체를 Module로 정의하며 사용자는 Module로 부터 컴파일된 함수를 PackedFunc형식으로 얻을 수 있음. 생성된 컴파일된 코드는 런타임에 Module에서 동적으로 함수를 가져올 수 있다. 첫 호출 시에 function handle을 캐시하고 subsequent call에서 이를 재사용한다. 우리는 이것을 생성된 코드에서 PackedFunc(e.g., python)과 디바이이스 코드를 연결하는데 사용할 수 있다.

ModuleNode는 각 유형의 장치에서 구현할 수 있는 abstract class 입니다. 지금까지는 CUDA, Metal, OpenCL 및 동적 공유 라이브러리 로드를 위한 모듈을 지원했습니다. 이 추상화를 통해 새로운 장치를 쉽게 도입할 수 있으며, 각 유형의 장치에 대해 Host Code 생성을 다시 할 필요가 없습니다.

Remote Deployment

PackedFunc / Module system은 argument를 직렬화하고 원격에서 계산을 수행하는 RPCModule을 사용하여 리모트 디바이스로 직접 전송이 가능.

RPC server는 runtime에 번들로 제공될 수 있으며 iPhone, Android, Raspberry pi, 브라우저에서 시작할 수 있다.

이러한 구조(즉각적인 피드백이 가능한)는 많은 이점을 제공하는데 예를 들어 RPC를 사용하여 iPhone에서 실행하고, 결과를 다시 복사한 후 numpy를 통해 호스트에서 검증이 가능하기에 host의 테스트 케이스를 다시 작성할 필요가 없다.

TVM Object and Compiler Stack

앞서 말했듯 compiler stack API 는 PackedFunc을 이용하여 구성됨. 새로운 primitive를 추가 할 때마다 새로운 Language object나 IR node가 필요했으나 TVM 개발자들은 API를 변경하는 대신 다음 내용을 원함

• 어떤 language object나 IR들을 직렬화 할 수 있을 것
• front-end language(ex. python)에서 IR object를 탐색, 출력, 조작 할 수 있을 것 상기 내용을 해결하기 위해 Base Class인 Object를 도입, 컴파일 스택의 모든 language object는 Object의 서브클래스 임
• (역자주) primitive는 TVM에서 새로운 연산(operators), IR 노드, 최적화 단위, 메모리 표현, 실행 스케줄링 전략을 의미
• (역자주) language object TVM 내부에서 모델을 표현하거나 최적화하는 데 사용되는 모든 추상 구조물(TIR, Relax, Type)을 의미
• (역자주) 새로운 연산(primitive)을 실험하고 싶다면 Relay, TIR, Relax 중 해당 영역의 새로운 IR 노드 또는 그 노드를 감싸는 “language object” (예: ObjectRef 타입)를 만들어야 한다

각 Object는 Object의 Type을 식별하기 위한 고유 식별자인 String 형식의 type_Key를 가지고 있음

• int 대신 String을 선택한 이유는 centeral repo에 코드를 추가하지 않고 분산 방식으로 새로운 Object 클래스를 추가할 수 있도록. 그러나 Runtime에서는 디스패치 속도를 높이기 위해 type_Key에 대응하는 integer type index를 할당

하나의 Object는 여러 곳에서 참조되므로 Object의 reference를 나타내는 ObjectRef 클래스를 사용

• ObjectRef는 Object 컨테이너의 shared_ptr 라고 볼 수 있음
• Object의 subtype을 유지하면서 ObjectRef의 서브클래스를 정의할 수 있으며 이 서브 클래스는 VisitAttr를 정의(가상함수 오버라이딩) 해야함

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19

class AttrVisitor {
public:
  virtual void Visit(const char* key, double* value) = 0;
  virtual void Visit(const char* key, int64_t* value) = 0;
  virtual void Visit(const char* key, uint64_t* value) = 0;
  virtual void Visit(const char* key, int* value) = 0;
  virtual void Visit(const char* key, bool* value) = 0;
  virtual void Visit(const char* key, std::string* value) = 0;
  virtual void Visit(const char* key, void** value) = 0;
  virtual void Visit(const char* key, Type* value) = 0;
  virtual void Visit(const char* key, ObjectRef* value) = 0;
  // ...
};

class BaseAttrsNode : public Object {
public:
  virtual void VisitAttrs(AttrVisitor* v) {}
  // ...
};

각 Object의 서브클래스는 자신의 각 멤버를 방문하기 위해 VisitAttr을 오버라이딩. 하기는 TensorNode 예제

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20

class TensorNode : public Object {
public:
  /*! \brief The shape of the tensor */
  Array<Expr> shape;
  /*! \brief data type in the content of the tensor */
  Type dtype;
  /*! \brief the source operation, can be None */
  Operation op;
  /*! \brief the output index from source operation */
  int value_index{0};
  /*! \brief constructor */
  TensorNode() {}

  void VisitAttrs(AttrVisitor* v) final {
    v->Visit("shape", &shape);
    v->Visit("dtype", &dtype);
    v->Visit("op", &op);
    v->Visit("value_index", &value_index);
  }
};

상기 예제에서 Operation, Array<Expr>은 ObjectRef이다. VisitAttrs는 각 멤버 object를 방문하기 위한 reflection API를 제공함

• VisitAttrs을 node를 방문하거나 language object를 재귀적으로 직렬화 하는데 사용 가능
• front-end언어에서 object 멤버를 얻는데 사용 가능(하기는 TensorNode의 예제임)

1
2
3
4
5
6

import tvm
from tvm import te

x = te.placeholder((3,4), name="x")
# access the op field of TensorNode
print(x.op.name)

front-end runtime을 변경하지 않고 새로운 Object를 C++에 추가할 수 있고, 추가된 Object를 위한 extention을 쉽게 컴파일 스택에 만들 수 있음, 이것은 object member를 expose하는 가장 빠른 방법은 아니나 가장 심플한 방법 중 하나임(TVM이 테스트를 위한 프로토타이핑에 python을 사용하고 내부적으로 무거운 작업에 C++를 사용하므로)

Implementation Details

PackedFunc의 각 argument는 union value인 TVMValue와 type code를 가지고 있음

• (역자주) TVM은 python에서 c++ 코드 호출을 위한 tvm ffi를 제공하며 PackedFunc에서는 argument로 ffi를 사용
  • Foreign Function Interface(FFI) : 한 프로그래밍 언어(이하 A)에서 다른 프로그래밍 언어(이하 B)의 코드를 호출하기 위한 인터페이스 이 디자인은 dynamically typed language를 이에 대응하는 type으로 직접 변환 할 수 있게, statically typed language는 변환 과정에서 runtime type checking을 수행할 수 있게함. 아래는 관련 파일임
• C++ API : packed_func.h
• C API and how to provide callback. : c_runtime_api.cc

Extension type을 지원하기 위해 type 정보를 등록하는 registry system을 사용, Extension types을 참조

• (역자주) 동적 타입 언어(dynamically typed language) : 변수를 선언할 때 타입을 명시하지 않아도 되며, 런타임 시에 타입이 결정되는 언어(파이썬), 동적 타입 언어는 값 자체에 타입 정보가 포함되어 있기 때문에, TVM의 TVMValue + type_code 구조로 쉽게 매핑됨

역자해설

TVM에서 Object와 PackedFunc는 런타임 시스템의 두 핵심 구성 요소이며 TVM이 유연하고 확장 가능한 컴파일러/런타임 환경을 제공하는 기반이 됨

• Object는 데이터를 표현하며 Relay/TIR/Relax 등 모든 TVM 언어의 노드, 타입, 표현식, 속성 등은 모두 Object의 서브클래스
• PackedFunc는 함수를 표현하며 TVM 런타임 상에서 함수를 호출할 수 있는 유니버설 함수 래퍼, 모든 함수는 TVMArgs → TVMRetValue 형태로 표현되며, 타입 정보 없이도 호출 가능 TVM에서 **사용자 정의 타입(custom object type)**을 등록하면, 이 타입도 Object 시스템 및 PackedFunc 시스템과 자연스럽게 통합되어 Python/C++ 경계를 넘을 수 있다. 이것은 하기의 예시임

1. 1단계 : 사용자 정의 객체 정의 (Object 기반)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27

// my_expr.h

#include <tvm/runtime/object.h>
#include <tvm/runtime/registry.h>

namespace tvm {
namespace myns {

class MyExprNode : public Object {
 public:
  int value;

  void VisitAttrs(AttrVisitor* v) {
    v->Visit("value", &value);
  }

  static constexpr const char* _type_key = "myns.MyExpr";
  TVM_DECLARE_FINAL_OBJECT_INFO(MyExprNode, Object);
};

class MyExpr : public ObjectRef {
 public:
  TVM_DEFINE_OBJECT_REF_METHODS(MyExpr, ObjectRef, MyExprNode);
};

}  // namespace myns
}  // namespace tvm

2. 2단계: 타입 등록 (TVM의 RTTI 시스템에)

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11

// my_expr.cc

#include "my_expr.h"

namespace tvm {
namespace myns {

TVM_REGISTER_OBJECT_TYPE(MyExprNode);

}  // namespace myns
}  // namespace tvm

3. 3단계: PackedFunc으로 전달 가능한 함수 등록

1
2
3
4

TVM_REGISTER_GLOBAL("myns.print_my_expr")
.set_body_typed([](tvm::myns::MyExpr e) {
  std::cout << "MyExpr.value = " << e->value << std::endl;
});

4. 4단계: Python에서 사용

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

from tvm.runtime import Object

# TVM 내부적으로 Object를 상속한 Python 래퍼 사용 가능
MyExpr = tvm._ffi.get_global_func("myns.MyExpr")  # 생성자 등록했을 경우

# 또는 C++에서 만든 객체를 받는 경우
print_func = tvm.get_global_func("myns.print_my_expr")

# 임의로 생성한 객체가 있다고 가정하고 전달
class MyExpr(Object):
    def __init__(self, value):
        self.__init_handle_by_constructor__(
            "myns.MyExpr",  # C++에서 등록한 type_key
            value
        )

e = MyExpr(42)
print_func(e)