TVM 상세

개요

Apache TVM(Tensor Virtual Machine)은 다양한 하드웨어 가속기에서 고성능 머신러닝 실행을 가능하게 하는 오픈소스 딥 러닝 컴파일러 스택이다. 2017년에 소개된 이후 Apache 소프트웨어 재단 산하에서 활발히 개발되고 있으며, GPU(CUDA, OpenCL), CPU(LLVM), TPU, NPU 등 광범위한 하드웨어를 지원한다. TVM은 프레임워크 독립적인 IR(Intermediate Representation)을 통해 PyTorch, TensorFlow, ONNX 등 다양한 ML 프레임워크의 모델을 수용하고, 하드웨어에 특화된 고성능 커널을 자동으로 생성한다.

TVM의 핵심 설계 철학은 "엔드투엔드 최적화(end-to-end optimization)"이다. 모델 수준의 그래프 최적화부터 커널 수준의 루프 타일링, 벡터화, 스레드 바인딩까지 하나의 통합 파이프라인에서 처리한다. 최근 TVM은 고수준 그래프 추상화를 담당하는 Relax와 저수준 텐서 프로그램을 담당하는 TensorIR의 이중 IR 구조를 채택하여, 그래프-커널 수준 교차 최적화를 지원한다.

핵심 개념

TVM 스택 아키텍처

TVM 스택은 크게 네 가지 계층으로 구성된다:

IRModule: 핵심 데이터 구조

TVM의 모든 변환은 IRModule을 중심으로 이루어진다. IRModule은 두 종류의 함수를 포함한다:

• relax.Function: 고수준 함수형 프로그램 표현. 그래프 구조를 가지며, 제어 흐름과 복잡한 데이터 구조를 지원한다. 모델이나 서브그래프에 대응한다.
• tirx.PrimFunc: 저수준 텐서 프로그램 표현. 루프 중첩, 다차원 로드/스토어, 스레딩, 벡터/텐서 명령어를 포함한다. 융합된 커널 하나를 나타낸다.

컴파일 과정에서 모든 relax 연산은 tirx.PrimFunc이나 TVM PackedFunc으로 로어링(lowering)되며, relax 연산 호출은 저수준 함수 호출(예: R.call_tir)로 변환된다.

TVMScript

TVMScript는 Python 기반 DSL로, TVM IR를 기술할 수 있게 한다. @I.ir_module, @T.prim_func, @R.function 데코레이터를 사용하여 Python 문법으로 IRModule을 정의하며, 파이썬 인터프리터가 아닌 파서를 통해 TVM IR로 변환된다. 모든 IRModule은 TVMScript로 역직렬화(roundtrip)가 가능하다.

TensorIR (tirx + s_tir)

TensorIR은 텐서 프로그램을 위한 저수준 추상화이며, 두 모듈로 나뉜다:

• tirx (Tensor IR eXtensions): PrimFunc, Buffer, SBlock, 표현식, 문장 등 코어 IR 정의와 로어링 패스를 담당한다.
• s_tir (Schedulable TIR): 스케줄 프리미티브, MetaSchedule, DLight, 텐서 인트린시크를 담당한다.

TensorIR의 핵심 요소는 Block이다. Block은 반복문 겹침(loop nest) 내에서의 계산 단위를 나타내며, read/write region, 반복문 축(axis) 속성, 계산 차원 등 스케줄링에 필요한 구조적 정보를 명시적으로 포함한다.

Relax

Relax는 ML 모델의 그래프 수준 추상화를 담당하는 고수준 IR이다. 핵심 특징:

• Dataflow Pattern Language (DPL): 패턴 매칭과 재작성을 위한 언어. 연산 융합, 그래프 레이아웃 최적화 등에 활용된다.
• Graph-level Optimizations: Constant Folding, Dead Code Elimination, Operator Fusion
• Cross-level Optimization: Relax와 TensorIR 간의 교차 최적화. LegalizeOps 패스로 Relax 연산을 TensorIR PrimFunc으로 로어링하고, FuseOps + FuseTIR로 여러 연산을 하나의 커널로 융합한다.

비교/분석

TVM vs 다른 AI 컴파일러

TVM 내부 모듈 비교

정밀도 지원 비교

동작 원리

TVM 컴파일러 파이프라인

1. 모델 임포트 단계:
   - 프레임워크별 프론트엔드(Frontend)가 모델을 로드한다
   - ONNX: onnx.load() → relax.frontend.onnx.from_onnx()
   - PyTorch: torch.jit.trace() → relax.frontend.torch.from_fx()
   - 프레임워크 중간 표현을 Relax IR로 변환한다

2. Relax 그래프 최적화 단계:
   - Constant Folding: 컴파일 시점에 상수 연산 수행
   - Dead Code Elimination: 불필요한 연산 제거
   - ** FuseOps: 연산을 패턴 기반으로 융합 단위로 분리
   - LegalizeOps**: Relax 연산을 TensorIR PrimFunc으로 로어링

3. TensorIR 스케줄링 단계:
   - DLight (규칙 기반): 미리 정의된 고성능 스케줄 규칙 적용

   • Tiling: 루프를 타일로 분할하여 캐시 접근 최적화
   • Vectorization: SIMD 명령어 활용
   • Thread Binding: GPU 스레드 할당
   • MetaSchedule (검색 기반): 실제 성능 측정 기반 최적 조합 탐색
   • 공간 탐색(Space Explorer) + 성능 예측(Performance Model)

4. TensorIR 로어링 단계:
   - 멀티차원 접근을 1차원 포인터 접근으로 평탄화
   - 하드웨어 특화 인트린시크로 변환
   - 런타임 호출 규약 충족 함수 엔트리 데코레이션

5. 대상 번역(Target Translation) 단계:
   - LLVM IRBuilder를 사용하여 x86/ARM 기계어 생성
   - CUDA C, OpenCL 소스 코드 생성
   - BYOC(Bring Your Own Codegen)로 cuBLAS, CUTLASS, cuDNN 등 외부 라이브러리 디스패치

6. 런타임 실행 단계:
   - runtime.Module으로 컴파일 결과를 캡슐화
   - PackedFunc으로 타입 제거된 함수 인터페이스 제공
   - Python, C++, Rust, Go, Java, JavaScript 등 다양한 언어에서 호출 가능

런타임 아키텍처

TVM 런타임은 다음 핵심 구조로 구성된다:

Disco 분산 런타임

Disco는 모델이 단일 GPU에 들어가지 않을 때 사용되는 분산 실행 런타임이다. 핵심 개념:

• Session: 워커 그룹을 소유하고 SPMD 스타일 프로그래밍 인터페이스를 제공한다
• DRef: 워커에 상주하는 분산 참조. 각 워커의 로컬 샤드에 대응한다
• 세 가지 세션 백엔드:
• ThreadedSession: 단일 머신 멀티GPU 추론용 (가장 일반적)
• ProcessSession: 프로세스 격리가 필요한 경우
• SocketSession: 멀티노드 클러스터용 TCP 소켓 연결

MetaSchedule 자동 튜닝

MetaSchedule은 TensorIR 커널의 최적 스케줄을 검색 기반으로 찾는 프레임워크이다.

작동 방식:
1. 튜닝 가능한 작업(Tunable Task) 추출
2. 탐색 공간(Space) 정의: Tiling 크기, 벡터화 길이, 스레드 배치 등
3. 샘플링 및 성능 측정: 후보 스케줄을 실제 하드웨어에서 실행
4. 데이터베이스 저장: 최적 결과를 재사용 가능한 형태로 저장

DLight vs MetaSchedule:
- DLight: 규칙 기반. 빠르고 안정적이며, 범용 GPU 스케줄에 적합하다
- MetaSchedule: 검색 기반. 느리지만 더 나은 성능을 찾을 수 있으며, 특정 하드웨어/워크로드에 최적화할 때 유용하다

장단점

TVM 장단점

TVM 사용 시 고려사항

관련 기술

관련 문서

• AI 컴파일러 개요: AI 컴파일러 전체 지형
• CUDA 상세: TVM CUDA 백엔드와 관련된 CUDA 아키텍처
• ROCm 상세: AMD GPU 개발 플랫폼, TVM ROCm 백엔드
• TPU 소프트웨어 스택: JAX, XLA, TPU 런타임과 TVM 비교
• AI 가속기 개요: TVM이 지원하는 하드웨어 가속기 비교

참고 문헌

• Chen et al., "TVM: An Automated End-to-End Optimizing Compiler for Deep Learning," OSDI 2018
• Chen et al., "Learning to Optimize Tensor Programs," NeurIPS 2018
• Shen et al., "A Flexible Approach to Autotuning for Deep Learning Compiler," MLSys 2021
• Apache TVM Documentation: https://tvm.apache.org/docs/
• Apache TVM GitHub Repository: https://github.com/apache/tvm
• Relax Deep Dive: https://tvm.apache.org/docs/deep_dive/relax/index.html
• TensorIR Deep Dive: https://tvm.apache.org/docs/deep_dive/tensor_ir/index.html
• TIRx Overview: https://tvm.apache.org/docs/tirx/overview.html

핵심 정리

1. Apache TVM은 오픈소스 딥 러닝 컴파일러 스택으로, 다양한 하드웨어(GPU, CPU, NPU, TPU)에서 고성능 ML 실행을 가능하게 하며 Apache 재단 산하에서 활발히 개발되고 있다
2. 최신 TVM은 Relax(고수준 그래프 IR)와 TensorIR(저수준 텐서 IR)의 이중 IR 구조를 채택하여 엔드투엔드 교차 최적화를 지원하며, TVMScript로 직관적인 IR 기술이 가능하다
3. MetaSchedule(검색 기반)와 DLight(규칙 기반) 자동 튜닝을 통해 하드웨어 특화 최적화를 자동화하며, BYOC로 cuBLAS, CUTLASS 등 외부 라이브러리도 통합할 수 있다
4. Disco 분산 런타임은 NCCL/RCCL 기반 멀티GPU/멀티노드 실행을 지원하며, Relax VM은 레지스터 기반 인터프리터로 엔드투엔드 모델 실행을 담당한다
5. TVM의 주요 강점은 범용 하드웨어 지원과 자동 튜닝 기능이며, 단점은 학습 곡선과 컴파일 시간이다