Deep Learning Compiler 1장: TVM 시작하기와 개요

이동욱

2026/08/16

Categories: 개발 Tags: TVM 컴파일러 딥러닝 인공지능

시작하며

최근 인공지능과 딥러닝 모델의 크기가 비대해지고 연산량이 급증함에 따라, 단순히 모델을 설계하는 것을 넘어 다양한 하드웨어 백엔드(CPU, GPU, NPU 등)에서 어떻게 효율적으로 모델을 실행하고 최적화할 것인가에 대한 고민이 깊어졌다. 이 과정에서 모델을 최적의 기계어로 자동 컴파일하여 각 하드웨어의 최대 성능을 끌어내는 딥러닝 컴파일러(Deep Learning Compiler) 분야에 큰 매력과 관심을 느끼게 되었다.

특히 오픈소스이자 대표적인 딥러닝 컴파일러인 Apache TVM은 이 분야의 핵심 기술로 자리 잡고 있다. 이에 대해 깊이 있게 이해하고 학습하기 위해 Dive into Deep Learning Compiler 교재를 바탕으로 공부를 시작했다. 이 글은 그 첫 단추인 ‘TVM 시작하기(Getting Started)’ 장을 공부하며 이해한 바를 정리하고 기록한 내용이다.

해당 책은 현재 개발중에 있고 전체 목차는 다음 링크 를 참고하면 된다.

해당 책은 독자들이 NumPy 사용 경험 정도만 갖춘 최소한의 배경 지식을 가지고 있다고 가정합니다. 이를 고려하여, 내용은 기초부터 차근차근 설명하고, 필요할 때마다 관련 배경 지식을 소개할 것입니다. 하지만 경험이 풍부한 독자분들께도 이 내용이 유용할 것입니다.

1장 : 시작하기

이제 여정을 시작해보자. 배워야 할 내용은 많지만 모든 여정은 어딘가에서 시작된다. 이 부분에서는 다음 내용을 다룰 것이다.

  • (거의) 모든 챕터를 직접 실행할 수 있도록 필요한 라이브러리를 설치하는 방법
  • 두 벡터를 더하는 연산자 작성하기
  • 추론을 실행하기 위해서 신경망 모델을 컴파일하고, 컴파일된 라이브러리를 저장하는 방법

1.1 설치

이 책의 각 섹션은 Jupyter 노트북으로 구성되어 있다. 이를 실행하는 가장 쉬운 방법은 HTML 페이지 우측 상단에 있는 COLAB 버튼을 클릭하는 것이다. 그러면 해당 노트북이 열린 상태로 Google Colab 으로 바로 이동한다.

첫 번째 코드 셀을 실행하면 호스트 런타임에 연결되며 다음과 같은 경고 메시지가 표시된다. RUN ANYWAY 를 클릭하여 계속 진행할 수 있다.

이 섹션의 나머지 부분에서는 파이썬 환경 설정 방법, Jupyter 의 대화형 노트북, 관련 라이브러리, 그리고 이 책을 따라가며 여러분의 컴퓨터에서 직접 실행할 수 있는 코드를 다룰 것이다.

1.1.1 소스 코드

모든 노트북이 포함된 소스코드 패키지는 아래에서 확인할 수 있다. 이 패키지를 다운로드 받아서 폴더에 압축을 풀어주세요. 예를 들어 Linux/macOS 에서 wget 과 unzip 이 모두 설치되어 있다면 다음 명령어를 통해서 수행할 수 있다.

wget http://tvm.d2l.ai/d2l-tvm.zip
unzip d2l-tvm.zip -d d2l-tvm

1.1.2 실행 환경 설치

Python 3.5 이상 버전이 설치되어 있고 pip도 설치되어 있다면, 실행 환경을 설치하는 가장 쉬운 방법은 다음과 같습니다.

  • Jupyter 및 저장된 코드 블록과 같은 모든 종속성을 위한 d2ltvm
  • 우리가 사용할 딥러닝 컴파일러를 위한 tvm
  • 일부 장에서 기준 모델로 사용할 mxnet

먼저 d2ltvm 패키지를 설치하기 위해서 아래 명령어를 수행합니다.

pip install git+https://github.com/d2l-ai/d2l-tvm

그런 다음에 소스코드를 사용하여 tvm 을 컴파일 합니다. tvm 은 사용자의 시스템에 설치된 라이브러리에 크게 의존 하기 때문에 pip 패키지가 제공되지 않는다. tvm 을 설치하려면 tvm.ai 에 있는 지침을 따르세

이 책에서 사용하는 config.cmake 파일의 구성에는 최소한

set(USE_LLVM ON)

를 포함해야 합니다. 또한 만약 Nvidia GPU 를 사용한다면 다음과 같이 USE_CUDA 를 ON 으로 설정해야 합니다.

set(USE_CUDA ON)

또한 성능을 향상시켜주는 cpython 을 활성화 하는 것을 잊지 마세요. 단지 TVM 소스 폴더에서 make cpython 명령어를 실행하기만 하면 됩니다.

현재 저희 소스코드는 tvm-0.7-dev 에서 실행됩니다. 마지막으로 GPU 를 사용할 수 있다면 MXNet 의 CUDA 버전을 설치하세요

pip install mxnet-cu101

사용중인 CUDA 버전에 맞춰서 101을 변경할 수 있습니다. 모든 패키지가 설치되면 다음 명령어를 실행하여 Jupyter 노트북을 열 수 있습니다.

jupyter notebook

이 단계에서 브라우저에서 http://localhost:8888 에 접속하여 Jupyter 노트북을 사용할 수 있습니다.

1.1.3 코드

이 책 전체에 걸쳐서, 재사용 가능한 코드 블록은 코드 앞에 # dwltvm 패키지에 저장 이라는 주석을 추가하여 d2ltvm 패키지에 저장합니다.

다음 코드는 d2ltvm 패키지에 저장된 코드의 예시입니다.

# Save to the d2ltvm package.
import tvm
from tvm import te
import time
import timeit
import numpy as np
from matplotlib import pyplot as plt
from IPython import display
try:
  import mxnet as mx
except:
pass

1.2 벡터 덧셈

이제 모든 라이브러리를 설치했으니, 첫 번째 프로그램을 작성해 볼 것이다. n 차원 벡터 a 와 b 를 더하는 프로그램이다. NumPy 에서는 이 작업이 매우 간단하며 c = a + b 와 같이 간단하게 작성할 수 있다.

import numpy as np
np.random.seed(0)

n = 100
a = np.random.normal(size=n).astype(np.float32)
b = np.random.normal(size=n).astype(np.float32)

c = a + b

여기서는 길이가 100 인 두 개의 임의 벡터를 생성하고, 요소별로 합산한다. Numpy 는 기본적으로 64비트 부동소수점 또는 64비트 정수를 사용하는데, 이는 딥러닝에서 일반적으로 사용되는 32비트 부동 소수점과 다르므로, 데이터 형을 명시적으로 변환한다.

Numpy 의 내장 ‘+’ 연산자를 사용하여 요소별 덧셈을 구현할 수도 있지만, 이번에는 스칼라 연산자만 사용하여 이를 구현해볼 것이다. 이렇게 하면 TVM 을 이용한 구현 방식을 이해하는데 도움이 될 것이다. 다음 함수는 for-루프 를 사용하여 벡터의 모든 요소를 순회한 다음에, 매번 스칼라 ‘+’ 연산자를 사용하여 두 요소를 더한다.

def vector_add(a, b, c):
    for i in range(n):
        c[i] = a[i] + b[i]

d = np.empty(shape=n, dtype=np.float32)
vector_add(a, b, d)
np.testing.assert_array_equal(c, d)

다음 장들에서 결과를 저장하기 위해 두 개의 임의의 ndarray 와 또 다른 빈 ndarray 를 자주 생성하게 될 것이므로, 나중에 재사용할 수 있도록 이 루틴을 저장해 둔다.

# SAVE TO d2ltvm package.
def get_abc(shape, constructor=None):
    """
    Return random, a, b and empty c with sampe shape.
    """
    np.random_seed(0)
    a = np.random.normal(size=shape).astype(np.float32)
    b = np.random.normal(size=shape).astype(np.float32)
    c = np.empty_like(a)
    if constructor:
        a, b, c = [constructor(x) for x in (a, b, c)]
    return a, b, c

Numpy, TVM 및 기타 라이브러리 간의 비교를 용이하게 하기 위해서 항상 동일한 결과가 나오도록 난수 시드를 고정해 두었음을 유의하시기 바랍니다. 또한, 데이터를 다른 형식으로 변환하기 위한 선택적 생성자를 지원합니다.

1.2.2 TVM 계산 정의

이제 TVM 에서 vector_add 를 구현해볼 것이다. TVM 구현은 위와 같은 두 가지 점에서 다르다.

  1. 함수 전체를 작성할 필요는 없으며, 출력의 각 요소 즉 c[i] 가 어떻게 계산되는지만 저장하면 된다.
    1. TVM 은 기호적이므로, 기호 변수의 차원을 지정하여 생성하고, 프로그램이 어떻게 계산될지 정의한다.

다음 프로그램에서는 먼저 tvm.te.placeholder 를 통해 두 입력에 대한 플레이스 홀더 A 와 B 의 모양 (n,) 을 지정하여 선언합니다. A 와 B 는 모두 Tensor 객체이며, 나중에 여기에 데이터를 입력할 수 있다. 나중에 프로그램을 읽기 쉽게 출력할 수 있도록 이들에게 이름을 지정한다.

A, B, C 는 모두 텐서 (Tensor) 객체이며, 이는 Numpy 의 ndarray 를 기호적으로 표현한 것으로 볼 수 있다. 데이터 유형이나 모양과 같은 변수의 속성에 접근할 수 있다. 하지만 이러한 속성들은 현재로서는 구체적인 값을 가지고 있지 않다.

텐서 객체를 생성하는 연산은 A.op 를 통해서 접근할 수 있다.

A, C 에 대한 연산의 유형은 다르지만, 둘 다 텐서 객체를 생성하는 연산을 나타내는 기본 클래스 Operation 을 공유하고 있음을 알 수 있다.

1.2.3 스케줄링 작성하기

계산을 실행하려면 프로그램의 실행 방식, 예를 들어 데이터에 접근하는 순서나 멀티 스레드 병렬화를 수행하는 방법 등을 지정해야 한다. 이러한 실행 계획을 ‘스케줄’ 이라고 합니다. C 가 출력 텐서이므로, 해당 연산자에 대한 기본 스케줄을 생성하고 의사 코드를 출력해보자.

s = te.create_schedule(C.op)

참고: 해당 구식 Tensor Expression API의 한계로 인해서 제거 되었다. 최신 방식은 TVM Script 로 구현한다.

스케줄은 여러 단계로 구성된다. 각 단계는 스케줄이 어떻게 편성되는지 설명하는 작업에 해당한다. 특정 단계에는 s[C] 또는 s[C.op] 를 통해서 접근할 수 있다.

나중에 하드웨어 리소스를 더 효율적으로 활용하여 성능을 향상 시키기 위해서 실행 계획을 변경하는 방법을 살펴볼 것이다. 여기선느 C 언어와 유사한 의사 코드를 출력하여 기본 실행 계획을 확인해볼 것이다.

tvm.lower(s, [A, B, C], simple_mode=True)

produce c {
  for (i, 0, 100) {
  c[i] = (a[i] + b[i])
  }
}

아래의 메서드는 스케줄과 입력 및 출력 텐서를 인수로 받는다. simple_mode=True 를 사용하여 간단한 의사 코드를 출력한다. 해당 프로그램은 출력 모양에 따라서 적절한 for 루프를 생성했다는 점에 유의하세요. 전반적으로, 이는 이전에 봤던 함수인 vector_add 와 유사한 방식으로 동작한다.

이제 TVM 이 계산과 스케줄을 분리한다는 것을 알 수 있었다. 계산은 결과가 어떻게 산출되는지를 정의하며, 이는 프로그램을 어떤 하드웨어 플랫폼에서 실행하든 변하지 않는다. 반면 효율적인 스케줄은 종종 하드웨어에 의존적이지만, 스케줄을 변경하더라도 정확성에는 영향을 미치지 않는다. 계산과 스케줄을 분리하는 이 개념은 TVM 이 Halide 에서 계승한 것이다.

1.2.4 컴파일 및 실행

계산과 스케줄이 모두 정의되면 tvm.build 를 사용하여 이를 실행 가능한 모듈로 컴파일 할 수 있습니다. 이 함수는 tvm.lower 와 동일한 인수를 받는다. 사실, 이 함수는 먼저 tvm.lower 를 호출하여 프로그램을 생성한 다음에, 이를 기계어로 컴파일 한다.

mod = tvm.build(s, [A, B, C])
type(mod)

tvm.runtime.module.Module

실행 가능한 모듈 객체를 반환한다. 이제 A, B, C 에 데이터를 입력하여 실행할 수 있습니다. 텐서 데이터는 tvm.ndarray.NDArray 객체여야 합니다. 가장 쉬운 방법은 먼저 NumPy ndarray 객체를 생성한 다음에 tvm.nd.array 를 사용하여 이를 TVM ndarray 로 변환하는 것이다. asnumpy 메서드를 사용하면 이를 다시 NumPy 로 변환할 수 있다.

x = np.ones(2)
y = tvm.nd.array(x)
type(y), y.asnumpy()

(tvm.runtime.ndarray.NDArray, array([1., 1.]))

이제 데이터를 생성하고 이를 TVM ndarray 로 반환해 봅시다.

a, b, c = get_abc(100, tvm.nd.array)

이제 계산을 수행하고 결과를 확인한다.

mod(a, b, c)
np.testing.assert_array_equal(a.asnumpy() + b.asnumpy(), c.asnumpy())

1.2.5 인자 제약 조건

A 와 B 를 선언할 때 두 입력 모두 길이가 100 인 벡터로 지정했다는 점을 기억하십시오.

A.shape, B.shape, C.shape

> ([100], [100], [100])

TVM 은 입력이 해당 사양을 충족하는지 확인한다.

try:
  a, b, c = get_abc(200, tvm.nd.array)
  mod(a, b, c)
except tvm.TVMError as e:
  print(e)

TVM 의 기본 데이터 타입은 float32 입니다.

A.dtype, B.dtype, C.dtype
> ('float32', 'float32' ,'float32')

데이터 유형이 다른 값을 입력하면 오류가 발생합니다.

try:
  a, b, c = get_abc(100, tvm.nd.array)
  a = tvm.nd.array(a.asnumpy().astype('float64'))
  mod(a, b, c)
except tvm.TVMError as e:
  print(e)

1.2.6 모듈 저장 및 불러오기

컴파일된 모듈은 디스크에 저장할 수 있다.

mod_fname = 'vector-add.tar'
mod.export_library(mod_fname)

그리고 나중에 다시 불러올 수 있다.

loaded_mod = tvm.runtime.load_module(mod_fname)

결과를 확인하면 같다.

a, b, c = get_abc(100, tvm.nd.array)
loaded_mod(a, b, c)
np.testing.assert_array_equal(a.asnumpy() + b.asnumpy(), c.asnumpy())

1.2.7 마무리

TVM 을 사용하여 오퍼레이터를 구현하는 과정은 아래와 같다.

  1. 입력 및 출력 형태와 각 출력 요소의 계산 방식을 지정하여 계산을 선언한다.
  2. (가급적) 시스템 자원을 최대한 활용할 수 있도록 스케줄을 생성한다.
  3. 하드웨어 대상 플랫폼으로 컴파일 합니다. 또한, 컴파일된 모듈을 디스크에 저장해 두었다가 나중에 다시 불러올 수 있다.

1.3 신경망 추론

1.2 절에서는 간단한 벡터 덧셈에 대해서 오퍼레이터를 구현하고 컴파일 하는 방법을 살펴보았다. 이제 우리는 더 나아가서 일련의 연산자로 구성된 사전 학습된 신경망 전체를 컴파일 하여 추론을 해볼 것이다.

import numpy as np
import mxnet as mx
from PIL import Image
import tvm
from tvm import relay

이 장에서는 이전 장에 비해 세 가지 모듈이 추가로 도입됩니다. PIL 을 사용하여 이미지를 읽고, mxnet 을 통해서 사전 훈련된 신경망을 확보하며, TVM 의 릴레이 모듈을 사용하여 신경망을 변환하고 최적화 할 것이다. 릴레이는 TVM 에서 신경망을 표현하기 위한 고수준 중간 언어(IR) 입니다.

1.3.1 사전 학습된 모델 확보

사전 훈련된 모델이란 데이터 세트를 통해서 매개변수가 훈련된 신경망을 의미한다. 여기서는 MxNet 의 모델 동물원에서 pretrained=True 옵션을 사용하여 사전 훈련된 모델을 확보한다. 이 모델에 대한 자세한 내용을 알고 싶다면 다른 부분을 참고하세요. 더 많은 모델은 MXNet 모델 동물원에서 찾을 수 있습니다.

model = mx.gluon.model_zoo.vision.resnet18_v2(pretrained=True)
len(model.features), model.output

로드된 모델은 약 100만장의 자연물 이미지를 포함하는 ImageNet 1K 데이터 셋을 기반으로 훈련되었다. 이 데이터 셋은 1,000 개의 클래스로 구성되어 있으며, 모델은 두 부분을으로 이루어져 있으며, 본체인 model.feature 에는 13개의 블록이 포함되어 있고, 출력층은 1,000 개의 출력을 갖는 밀집층이다. 다음 코드 블록은 ImageNet 데이터셋의 각 클래스에 대한 텍스트 레이블을 불러온다.

with open('../data/imagenet1k_labels.txt') as f:
labels = eval(f.read())
>> Home