2020. 2. 14. 03:47ㆍResearch
TVM의 개념을 한 마디로 정리하자면, "주어진 하드웨어에 따라서 그 백엔드에 알맞게 high-level(graph level) 최적화와 low-level(operator level) 최적화를 동시에 제공하는 E2E 컴파일러"라고 할 수 있다. FPGA, ASIC 가속기들부터 여러 임베디드 edge device들까지 하드웨어 환경은 제각기 다르기 때문에, 각 환경의 특성에 알맞게 최적화를 하려면 서로 다른 방식으로 최적화 튜닝을 해주어야 한다. Tensorflow, PyTorch를 비롯한 여러 딥러닝 프레임워크들도 graph IR(Intermediate Representation)에 대해서만 최적화를 적용할 뿐, 하드웨어 특성에 따라 유동적으로 operator-level transformation을 해주지는 못 한다.
이와 같은 문제를 해결하기 위해서 TVM은 3가지 대안을 제시한다.
1. Tensor Expression Language
: 다양한 최적화 방식이 적용된 서로 다른 버전의 여러 프로그램을 생성하는 transformation primitives
2. Automated Program Optimization Framework
: 하드웨어 백엔드로부터 수집한 정보를 바탕으로 ML 기반의 cost model를 적용하여 자동으로 최적화된 tensor operator를 찾아낸다.
3. Graph Rewriter
: 기존 DL 프레임워크에서도 그러하듯이 high level(graph level)의 최적화를 제공한다.
TVM System Overview
글의 맨 앞에 제시한 그림이 TVM의 전체 시스템 workflow를 보여준다.
1. TVM은 우선 기존 딥러닝 프레임워크에서 프로그램된 모델을 입력으로 받은 뒤, computational graph representation으로 변환한다.
2. 최적화된 그래프를 생성하기 위해서 high level (graph level)의 최적화를 수행한다. (operator fusion 등)
3. 주어진 타겟 하드웨어(모델을 사용할 하드웨어)에 맞춰서 그럴 듯한 최적화 코드의 콜렉션을 뽑아낸다.
4. ML 기반의 cost model을 통해서 최적화된 operator를 찾아내고, 3에서 수집한 코드 콜렉션 중 가장 최적의 코드를 찾아낸다.
5. 생성된 코드를 packing하여 배포가능한 모듈로 만든다.
다음은 위에서 설명한 TVM workflow를 보여주는 사용자 코드의 예시이다.
import tvm as t
import tvm.runtime as rt
# keras로 작성된 이미 존재하는 모델을 로드하여 그래프 최적화를 수행한다.
graph, params = t.frontend.from_keras(keras_model)
# target HW를 인식
target = t.target.cuda()
# TVM 최적화를 통해서 다음을 얻을 수 있다.
## graph: 최적화된 계산 그래프
## lib: HW 백엔드에 최적화되어 생성된 operator
## params: module 파라미터
graph, lib, params = t.compiler.build(graph, target, params)
# 최적화 결과를 통해 배포용 모듈 생성
module = runtime.create(graph, lib, t.cuda(0))
module.set_input(**params)
module.run(data=data_array)
output = tvm.nd.empty(out_shape, ctx=t.cuda(0))
module.get_output(0, output)
Evaluation