AI/Deep Learning

[DL] TVM 이란?

운호(Noah) 2022. 1. 12. 19:52

들어가기 앞서,

  • 딥러닝 모델을 특정 디바이스 상에서 효율적으로 동작시키기 위해서는,
  • 동작시킬 딥러닝 모델을, 타겟 디바이스에서 최적의 속도와 정확도를 낼 수 있는 머신 코드로 변환해야합니다.
  • 이러한 코드 변환 작업을 자동으로 지원해주는 도구를, ‘딥러닝 컴파일러’라고 합니다.

딥러닝 컴파일러 개념도

  • 딥러닝 컴파일러는, 다양한 딥러닝 플랫폼에서 학습된 딥러닝 모델을 입력으로 받아,
  • 특정 하드웨어에서 동작 가능한 머신 코드(또는, 백엔드 코드)를 자동으로 생성합니다.
  • 최근 제안된 XLA, TVM, Glow 와 같은 딥러닝 컴파일러들은,
  • TensorFlow, Pytorch, MxNet, ONNX 등의 프레임워크로 작성된 모델을 입력으로 하여,
  • CPU 및 GPU 용 백엔드 코드를 생성하고 있습니다.
    • CPU 용 백엔드 코드 : llvm
    • GPU 용 백엔드 코드 : CUDA, OpenCL, Metal 등

딥러닝 컴파일러 구조

  • 계산 그래프 생성기

    • 계산 그래프 생성기는, 다양한 딥러닝 플랫폼에서 생성된 딥러닝 모델을 로딩하여, 그래프 구조로 재구성한 뒤,
    • 모델에서 명시한 연산자들을, 프리미티브 연산자로 변환하여, 계산 그래프를 생성합니다.
    • 이후, 불필요한 노드 제거, 연산자 결합(Operator Fusion), 양자화(Quantization), 상수 폴딩(constant-folding)과 같은, 그래프 최적화를 기법들을 적용하여,
    • 최종적으로 최적화 된 계산 그래프를 생성합니다.

  • IR 생성기

    • 계산 그래프 생성기로부터 생성된 그래프는, 하드웨어에 독립적인 정보만 표현 할 수 있습니다.
      • ex) 연산자 간 입출력 관계, 연산자의 입출력 데이터 값, etc
    • 그러나 컴파일러는, 딥러닝 모델의 실행 성능을 높이기 위해,
    • 모델이 동작할 하드웨어의 아키텍처를 고려한, ‘하드웨어 의존적 최적화’를 필수적으로 수행해야합니다.
    • 따라서, 계산 그래프를 ‘하드웨어 의존적 최적화’에 적용 가능한 형태인,
    • 중간 표현 (IR: Intermediate Representation)으로 변환하는 과정이 필요합니다.
      • 참고로, IR 은 입출력 변수(또는 Tensor)들에 대한 선언과,
      • 변수들을 사용하는 연산식에 대한 정보를 포함하고 있습니다.
    • 컴파일러는, 생성된 IR 을 바탕으로 메모리 할당, 병렬화, 연산식 실행 순서 결정 등의 ‘하드웨어 의존적 최적화’ 작업을 수행한 뒤, 코드를 생성합니다.
    • 요약하면, IR 생성기는 컴파일러가 하드웨어 의존적 최적화를 용이하게 수행하도록, 계산 그래프를 IR 로 변환하는 역할을 합니다.

  • 백엔드 코드 생성기

    • 백엔드 코드 생성기는, IR 을 기반으로, 딥러닝 모델 워크로드가 배치될 타겟 하드웨어(CPU, GPU, TPU 등)의 아키텍처에 최적화 된 백엔드 코드를 생성합니다.
    • 즉, 타겟 하드웨어 별로, 병렬로 처리 가능한 데이터의 크기, 캐시 크기, 효율적으로 수행 가능한 딥러닝 오퍼레이터의 타입, 실행 순서 등을 참조하여,
    • 스케쥴 정보를 생성한 뒤, 해당 하드웨어에 최적화 된 백엔드 코드를 생성합니다.

TVM 실행 흐름

  • DL framework로부터 생성한 모델을 import 한 뒤, relay 를 통해, 해당 모델을 TVM 의 IRModule 로 변환합니다.
    • IRModule에는 모델을 나타낼 수 있는 다양한 함수들이 포함되어 있습니다.
  • IRModule 을 점점 Lowering 하는 Transformation 과정을 거칩니다.
    • IRModule 이 원하는 하드웨어(메모리 구조, 데이터 타입 등)에 알맞게 최적화될 수 있도록 Tensorize 시킵니다.
    • 이를 위해 Quantization, Graph-Packing 등, 다양한 작업들이 포함되며,
    • 이때, AutoTVM 이 사용됩니다.
  • 타깃 하드웨어에서 잘 작동하는 코드를 생성하기 위해, Translation 과정을 거칩니다.
    • 예를 들어, 타깃 하드웨어가 GPU 라면 [CUDA/OpenCL/Metal 등], CPU라면 [llvm] 코드를 생성합니다.
    • 이는, TVM 의 runtime.Module 에 encapsulate 되어있고,
    • 이를 통해, 타깃 하드웨어에 쉽게 export / load / execute 시킬 수 있습니다.
  • runtime.Module 를 로드하여, 타깃 하드웨어에서 최적화 된 모델을 동작시킵니다.

TFLite, TensorRT, TVM 의 차이

  • TFLite
    • TensorFlow(Keras)의 모델만 최적화가 가능합니다.
    • 여러 하드웨어에서 추론이 가능합니다.
    • 타겟 하드웨어에 따른 추가적인 최적화는 없습니다.
  • TensorRT
    • 다양한 프레임워크의 모델 최적화가 가능합니다.
    • NVIDIA GPU 만 추론 하드웨어로 사용 가능합니다.
    • 추론에 사용하려는 NVIDIA GPU 종류에 따라 추가적으로 최적화할 수 있습니다.
  • TVM
    • 다양한 프레임워크의 모델 최적화가 가능합니다.
    • 여러 하드웨어서 추론이 가능합니다.
    • 추론에 사용하려는 하드웨어 종류에 따라 추가적으로 최적화할 수 있습니다.

추가

참고