TVM: An Automated End-to-End Optimizing Compiler for Deep Learning

TVM은 딥러닝 컴파일러이다. 이번 글에서는 왜 TVM이 필요한지부터 TVM 내에서 일어나는 과정들에 대해서 소개하려고 한다.

 

딥러닝 모델을 학습하는 프레임워크는 너무 많고, 학습된 모델을 실행시킬 서버/디바이스도 너무 많다.

각각의 조합을 모두 만족하기 위해서 Tensorflow가 저 많은 서버/디바이스 별로 코드를 개발해야하고, Pytorch도 Pytorch 나름대로 코드를 개발하는 것은 비효율적이다.

그래서 딥러닝 컴파일러가 필요했고, 딥러닝 컴파일러 중에서 Apache에서 개발 중인 TVM이라는 컴파일러가 있다.

기본적으로 컴파일러란 사람이 이해할 수 있는 코드를 기계가 이해할 수 있는 언어(기계어)로 바꿔주는 것이다.

딥러닝 컴파일러가 중간에 들어가므로서 고려해야하는 가지 수가 크게 줄어들었다.

여기에서 TVM의 역할은 크게 두가지이다.

• 다양한 프레임워크로부터의 모델을 통합된 어떤 형태로 컴파일하기
• 통합된 모델을 타겟 하드웨어에 최적화된 형태로 실행시키기

TVM은 아래의 과정을 통해서 위 두 가지 역할을 수행한다.

 

Graph-Level로 최적화하는 것은 여러가지 딥러닝 프레임워크로부터 학습한 모델을 하나의 공통된 표현으로 바꾸기 위해서이다. 그 표현을 IR(Intermediate Representation)이라고 한다. 모델을 IR로 변형할 때 단순히 변형만 하지 않는다. 최적화를 한다.

| Optimizing Computational Graphs

1. Operator Fusion

Operator Fusion은 Layer Fusion과 같다고 생각하면 된다.

Layer Fusion을 수학적으로 이해하고 연산 결과가 100% 동일함을 예제 소스코드이다.

여기에서는 간단하게 Conv2D + BatchNorm을 단일 Conv2D 레이어로 Fusion하는 예제를 소개하려고 한다.

 

우선 이론적인 내용부터 알아보자. Conv2D + BatchNorm의 연산을 수식으로 전개하면 아래와 같다.

즉, Conv2D + BatchNorm 레이어는 단일 Conv2D 레이어로 변형 가능하다.

 

예제 코드를 봐보자. 아래의 SampleModel 레이어는 단일 Conv2D로 변형 가능하다.

class Conv_Bn_Activation(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, activation, bn=True, bias=False):
        super().__init__()
        pad = (kernel_size - 1) // 2

        self.conv = nn.ModuleList()
        if bias:
            self.conv.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, pad))
        else:
            self.conv.append(nn.Conv2d(in_channels, out_channels, kernel_size, stride, pad, bias=False))
        if bn:
            self.conv.append(nn.BatchNorm2d(out_channels))
            # self.conv.append(MyBatchNorm(out_channels))
        if activation == "mish":
            self.conv.append(Mish())
        elif activation == "relu":
            self.conv.append(nn.ReLU(inplace=True))
        elif activation == "leaky":
            self.conv.append(nn.LeakyReLU(0.1, inplace=True))
        elif activation == "linear":
            pass
        else:
            print("activate error !!! {} {} {}".format(sys._getframe().f_code.co_filename,
                                                       sys._getframe().f_code.co_name, sys._getframe().f_lineno))

    def forward(self, x):
        for l in self.conv:
            x = l(x)
        return x


class SampleModel(nn.Module):
    def __init__(self, in_channels, out_channels, kernel_size, stride, activation, bn=True, bias=False):
        super().__init__()
        conv_1 = Conv_Bn_Activation(in_channels, 1024, kernel_size, stride, activation)
        conv_2 = Conv_Bn_Activation(1024, 1024, kernel_size, stride, activation)
        conv_3 = Conv_Bn_Activation(1024, 1024, kernel_size, stride, activation)
        conv_4 = Conv_Bn_Activation(1024, 1024, kernel_size, stride, activation)
        conv_5 = Conv_Bn_Activation(1024, out_channels, kernel_size, stride, activation)
        self.convs = nn.ModuleList([conv_1, conv_2, conv_3, conv_4, conv_5])

    def forward(self, x):
        for conv in self.convs:
            x = conv(x)
        return x

 

위의 모델을 fuse하는 함수는 아래와 같다.

def fuse_layer_cbr(layer, inputs):
    c_layer = layer[0]
    b_layer = layer[1]

    assert type(c_layer) == nn.Conv2d, 'first layer is not nn.Conv2d'
    assert type(b_layer) == nn.BatchNorm2d, 'second layer is not nn.BatchNorm2d'

    kernel_size = c_layer.kernel_size[0]
    stride = c_layer.stride[0]
    pad = (kernel_size - 1) // 2
    conv = nn.Conv2d(
        c_layer.in_channels,
        c_layer.out_channels,
        kernel_size,
        stride,
        pad,
        bias=True)

    with torch.no_grad():
        for i in range(b_layer.num_features):
            # update new bias and weight
            precomputed = b_layer.weight[i] / (np.sqrt(b_layer.running_var[i] + b_layer.eps))
            new_bias = b_layer.bias[i] + precomputed * (0 - b_layer.running_mean[i])

            # update
            conv.weight[i,:,:,:] = c_layer.weight[i,:,:,:] * precomputed
            conv.bias[i] = new_bias

    return conv


def fuse_model(model, inputs):
    # print(f'original={model.conv.conv} length={len(model.conv.conv)}')

    if [type(m) for m in model.conv.conv] == [nn.Conv2d, nn.BatchNorm2d, nn.ReLU]:
        fused_conv = fuse_layer_cbr(model.conv.conv, inputs)

    return fused_conv

SampleModel을 하나 만들어서 그 함수에 (1, 3, 112, 112) 사이즈의 입력을 넣었을 때와 SampleModel을 fuse한 모델의 속도 연산속도 차를 비교해보면 아래와 같다. (100번 수행함)

x = torch.randn(1, 3, 112, 112)
run x 100 times
>>> SampleModel : duration=4.4151 params=30707200
>>> FusedModel  : duration=3.2266 params=30698496

 

이런 속도 차이가 나는 이유는 Conv2D 단일 레이어로 바꿈으로서 파라미터 수가 줄어들기 때문이 아니다. Conv2D와 BatchNorm을 실행시킬 때 GPU와 CPU의 메모리 복사 과정이 생략되기 때문이다.

2. Data Layout Transformation

모델의 연산을 그래프로 표현할 때 텐서들을 그래프에 여러가지 방법으로 표현할 수 있다.

예를 들어서 2D의 텐서가 있을 때 이 텐서의 데이터를 row-major로 표현할 지 column-major로 표현할지 정할 수 있다.

데이터의 레이아웃은 메모리에 접근하는 횟수를 결정할 수 있기 때문에 코드의 가속화에 큰 역할을 할 수 있다.

 

| Generating Tensor Operation

TVM의 핵심은 Pytorch/Tensorflow/ONNX 등등의 서로 다른 DL 프레임워크에서 학습된 모델을 Arm, Intel 등등의 서로 다른 하드웨어 아키텍쳐에 맞게 최적화 및 컴파일하는 것이다.

그렇게 하기 위해서 TVM에서는 모델 내에서 행해지는 연산들을 텐서로 만들어내는 과정을 거친다. 이때 메모리의 캐시를 조금 더 효율적으로 사용하기 위한 Cache-Friendly한 접근이 사용된다.

 

Cache-Friendly한 코드가 어떤 코드인지 간단하게 알아볼 수 있는 소스코드를 보자.

// cache-friendly approach
void matsum_v1(float** a, int n) {
    float sum = 0;
    for(int i = 0; i < n; i++)
        for(int j = 0; j < n; j++)
            sum += a[i][j];
    cout << sum << endl;
}

// cache-unfriendly approach
void matsum_v2(float** a, int n) {
    float sum = 0;
    for(int i = 0; i < n; i++)
        for(int j = 0; j < n; j++)
            sum += a[j][i];
    cout << sum << endl;
}

행렬의 모든 원소의 합을 구하는 함수이다. matsum_v1은 2차원에서 loop가 실행되기 때문에 2차원의 연속된 공간의 메모리가 캐시에 올라가게 된다. 따라서 캐시를 최대한 효율적으로 사용할 수 있다. 반면에 matsum_v2는 1차원의 공간을 loop가 iteration하고 있다. 1차원의 공간은 메모리에서 연속적으로 있지 않기 때문에 cache는 계속 miss하게 되어 결국 메모리로부터 계속 값을 읽어오므로서 속도가 느려진다. 위의 코드를 실행해보면 속도차이는 아래와 같다.

$ ./cache_friendly
3.60288e+16
matsum_v1: elapse time 201 ms
3.60288e+16
matsum_v2: elapse time 550 ms

matsum_v2가 두배 이상 느린 것을 확인할 수 있다. 캐시를 사용하는 방법만 바꿨을 뿐인데 속도차이는 어마어마하다.

 

Tensor Expression도 마찬가지이다. 같은 연산을 수행하는데 cache-friendly한 방식으로 코딩을 하면 속도 차이가 훨씬 빨라진다. 행렬 연산을 수행하는 tensor expression을 보면 그 차이를 명확하게 알 수 있다.

# run time: 1.832 sec
k = te.reduce_axis((0, K), "k")
A = te.placeholder((M, K), name="A")
B = te.placeholder((K, N), name="B")
C = te.compute((M, N), lambda x, y: te.sum(A[x, k] * B[k, y], axis=k), name="C")
evaluate_operation(s, [A, B, C], target=target, name="mmult", optimization="none", log=log)

# run time: 0.130 sec
xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)
(k,) = s[C].op.reduce_axis
ko, ki = s[C].split(k, factor=4)
s[C].reorder(xo, yo, ko, ki, xi, yi)
s[C].vectorize(yi)
evaluate_operation(s, [A, B, C], target=target, name="mmult", optimization="blocking & vectorization", log=log)

# run time: 0.030 sec
s = te.create_schedule(C.op)
xo, yo, xi, yi = s[C].tile(C.op.axis[0], C.op.axis[1], bn, bn)
(k,) = s[C].op.reduce_axis
ko, ki = s[C].split(k, factor=4)
s[C].reorder(xo, yo, ko, xi, ki, yi)
s[C].vectorize(yi)
evaluate_operation(s, [A, B, C], target=target, name="mmult", optimization="loop permutation", log=log)

똑같이 행렬 곱을 수행하는 Tensor Expression인데 수행하는 구현하는 방법에 따라서 1.83 → 0.13 → 0.03 초로 가속화시킬 수 있다.

 

Tensor Expression은 TVM에서 schedule primitive의 집합이다. 위 코드에서 tile, split, vectorize, reorder 등이 TVM의 schedule primitive이다. 하나의 모델을 schedule primitive로 표현할 때 여러가지 조합이 나올 수 있다. 이 중에서 TVM은 가장 효율적인 조합을 선택한다. 어떻게 할까? 

 

| Automating Optimization

주어진 schedule primitive들을 이용해서 각각의 레이어에 대해서 어떻게 최적화할 수 있을까?

가령, Conv2D를 schedule primitive로 구현할 수 있는 방법은 여러가지이다. 그 중에서 가장 빠른 방법을 찾는 과정을 알아보자.

 

1. Schedule Space Specification

TVM은 어떤 high-level tensor expression이 주어졌을 때, 그것과 연산 결과는 같지만 다른 expression들을 만들어내는 generic master template이 구현돼 있다.

generic master template을 통해서 생성된 tensor expression들이 하나의 큰 search space가 되는데, 이 search space에서 optimizer는 최적의 tensor expression을 찾아낸다.

쉽게 말해서, TE-1부터 TE-4까지의 tensor expression을 search space에 넣어두면 optimizer가 가장 빠른 연산을 찾아내는 것이다.

 

# TE-1
C = tvm.compute((n,), lambda i: A[i] + B[i])
s = tvm.create_schedule(C.op)
----------------------------------------
for (int i = 0; i < n; ++i) {
    C[i] = A[i] + B[i];
}


# TE-2
C = tvm.compute((n,), lambda i: A[i] + B[i])
s = tvm.create_schedule(C.op)
----------------------------------------
xo, xi = s[C].split(s[C].axis[0], factor=32)
for (int xo = 0; xo < ceil(n / 32); ++xo) {
    for (int xi = 0; xi < 32; ++xi) {
        int i = xo * 32 + xi;
        if (i < n) {
            C[i] = A[i] + B[i];
        }
    }
}


# TE-3
C = tvm.compute((n,), lambda i: A[i] + B[i])
s = tvm.create_schedule(C.op)
xo, xi = s[C].split(s[C].axis[0], factor=32)
s[C].recorder(xi, xo)
----------------------------------------
for (int xi = 0; xi < 32; ++xi) {
    for (int xo = 0; xo < ceil(n / 32); ++xo) {
        int i = xo * 32 + xi;
        if (i < n) {
            C[i] = A[i] + B[i];
        }
    }
}


# TE-4
C = tvm.compute((n,), lambda i: A[i] + B[i])
s = tvm.create_schedule(C.op)
xo, xi = s[C].split(s[C].axis[0], factor=32)
s[C].recorder(xi, xo)
s[C].bind(xo, tvm.thread_axis(“blockIdx.x”)
s[C].bind(xi, tvm.thread_axis(“threadIdx.x”)
----------------------------------------
int i = threadIdx.x * 32 + blockIdx.x;
if (i < n) {
    C[i] = A[i] + B[i];
}

 

2. ML-Based Cost Model

머신러닝 모델을 이용해서 최적의 tensor expression을 찾아낸다. 더 정확히 말하면 머신러닝 모델을 이용해서 최적의 후보 몇 개를 선택한다.

search space에 있는 tensor expression은 몇 억개 정도로 많다. 그 모든 것을 타겟 하드웨어에서 실행해서 제일 빠른 expression을 찾으면 너무 느릴 수 밖에 없다.

따라서 각각의 expression을 입력으로 해서 타겟 하드웨어에서의 실행시간을 예측하는 머신러닝 모델을 통해서 최적의 후보 몇 개를 선택한다.

TVM의 ML-Based Cost 모델은 XGBoost를 이용해서 학습돼 있다.

 

3. Schedule Exploration

ML-Based Cost 모델을 통해서 몇 개의 후보 expression을 구한 후에 각 expression의 실행시간을 측정한다.

그런데 expression을 실행할 때 주어진 설정대로 그대로 실행시키지 않는다. 처음에는 설정을 임의로 지정해서 실행한다. 그러면서 조금씩 주어진 설정대로 변경해가면서 실제로 실행시간이 단축되는지 확인해보는 과정을 거친다.

이 과정에서 실행시간이 실제로 단축이 되면 pass이고 그렇지 못하면 reject된다.

4. Distributed Device Pool and RPC

TVM은 RPC 프로토콜을 지원한다. 로컬에서 모델을 컴파일 한 후 컴파일 한 모델을 리모트(타겟 하드웨어)에 업로드해서 실행시키고, 그 결과를 다시 로컬로 가져올 수 있다.

우선 RPC가 무엇인지 알아보자. RPC는 Remote Procedure Call의 약자이다. 즉 원격 프로시져 콜이다. 내 컴퓨터에 있는 어떤 프로그램을 원격에 있는 다른 컴퓨터에서 실행시키고 결과값을 받는 것이다. 

# on local machine(alienware)
$ python rpc-example.py

# on remote target(RaspberryPi)
$ python -m tvm.exec.rpc_server --host 0.0.0.0 --port=9090
INFO:root:If you are running ROCM/Metal, fork will cause compiler internal error. Try to launch with arg ```--no-fork```
2022-12-06 18:00:36.393 INFO bind to 0.0.0.0:9090
2022-12-06 18:10:43.965 INFO connection from ('192.168.0.5', 53314)
2022-12-06 18:10:44.342 INFO load_module /tmp/tmp0e37s1ac/lib.tar
2022-12-06 18:10:44.440 INFO Finish serving ('192.168.0.5', 53314)

위의 예시는 노트북 컴퓨터를 로컬로 두고 라즈베리파이를 리모트 디바이스로 두고 서로를 RPC로 연결 시킨 후 노트북에서 모델을 컴파일 후 라즈베리에서 실행시킨 결과이다.

 

마지막으로 정리를 하고 마치려고 한다.

 

TVM은 딥러닝 컴파일러다. TVM은 우선 여러 딥러닝 프레임워크로부터 학습된 모델을 Intermediate Representation으로 변형을 한 후, 이 변형된 그래프 내에서 최적화를 진행한다. 최적화된 그래프는 다시 Tensor Expression으로 바뀌는데(schedule primitive의 집합) 변형 가능한 Tensor Expression들을 모두 하나의 search space에 모아두어서 머신러닝 모델로 제일 빠를 것 같은 Tensor Expression을 몇 개 찾고, 찾은 것들은 실제로 실행시켜보면서 최적의 속도를 보이는 Tensor Expression을 선택한다.

같은 결과를 내는 함수인데도 더 빠르게 실행시키기 위해서 코드 연산을 항상 Cache-Friendly하게 고치는 과정을 전체적인 TVM 컴파일 과정에서 보여주고 있다.

이렇게 컴파일된 모델은 RPC 등을 통해서 원격으로 실행도 시킬 수 있다.