AI 컴파일러의 정의와 Frame_work 내 사용 방법 고민

AI 컴파일러란?

AI 모델을 효율적으로 실행하기 위해 중간 표현(Intermediate Representation, IR)으로 변환하고, 하드웨어별 최적화된 실행 코드로 컴파일하는 시스템을 의미

전통적인 컴파일러가 “C/C++ 같은 소스 코드를 기계어로 변환”한다면, AI 컴파일러는 “딥러닝 모델 그래프를 GPU, TPU, CPU 같은 다양한 하드웨어에서 최적 성능을 내도록 변환”하는 역할

 

1. 일반적인 정의

• 입력: 딥러닝 프레임워크(Pytorch, TensorFlow 등)에서 정의된 연산 그래프(예: matmul, conv, relu …).
• 중간 표현(IR): 하드웨어 독립적인 중간 언어로 모델을 추상화 (ex: XLA HLO, MLIR, TVM IR).
• 출력: 특정 하드웨어용 최적화된 코드 (CUDA kernel, LLVM IR → 기계어, etc).

즉, 모델 정의 → IR 변환 → 최적화 → 하드웨어 코드 생성을 수행합니다.

 

2. 주요 역할

1. 그래프 최적화 (Graph Optimization)
   • 불필요한 연산 제거 (fusion, constant folding 등)
   • 커널 합치기 (fusion)로 메모리 접근 최소화
   • 레이아웃 변환(NHWC ↔ NCHW)
2. 하드웨어 추상화
   • 같은 모델을 GPU, TPU, CPU, NPU 등 다양한 기기에서 실행 가능
   • 벤더별 라이브러리(CuBLAS, MKLDNN, ROCm 등) 자동 호출
3. 성능 향상
   • 연산 병렬화, 메모리 사용 최적화
   • 연산자 단위보다 큰 “연산 그래프 단위”로 최적화 → 기존 프레임워크 실행보다 빠름

 

예시 (Dense + ReLU)

• 모델 그래프: Z = ReLU(X @ W + b)
• 기본 실행:
  • Kernel 1: MatMul(X, W)
  • Kernel 2: Add(b)
  • Kernel 3: ReLU(Z)
• AI 컴파일러 최적화 후:
  • Kernel Fusion → 하나의 CUDA 커널에서 MatMul + Add + ReLU까지 수행
  • 메모리 왕복 최소화 → 속도 증가

 

AI 컴파일러가 모델 그래프를 해석 → 최적 커널 선택까지 가려면, 결국 “선택지”가 다양해야 하니까 여러 종류의 커널 구현이 필수

 

1. 왜 다양한 커널이 필요한가?

• 연산 특화: 같은 행렬 곱이라도 입력 크기/배치 크기/레이아웃에 따라 다른 알고리즘이 더 효율적임.
  • 작은 행렬 → 단순 GEMM보다 타일링 없는 direct kernel이 더 빠름
  • 큰 행렬 → cuBLAS GEMM 같은 고도로 최적화된 라이브러리가 유리
• 하드웨어 차이: GPU, CPU, TPU, NPU마다 최적의 연산 방식이 다름.
• 연산 조합(fusion): Conv + BiasAdd + ReLU를 따로 하면 느린데, fused kernel을 구현하면 훨씬 빠름.

즉, AI 컴파일러가 “이 상황에선 어떤 커널을 써야 제일 빠른가?”를 고르려면 다양한 후보 커널이 필요

 

2. 필요한 커널 종류 예시

1. 기본 연산 커널 (Primitive Ops)
   • MatMul / GEMM / GEMV
   • Convolution (direct, Winograd, FFT 기반 등)
   • Element-wise (Add, Mul, ReLU, Sigmoid, Softmax …)
2. Fused 커널 (연산 결합 버전)
   • Conv + Bias + ReLU
   • MatMul + Add + Activation
   • LayerNorm + Dropout
3. 메모리 최적화 커널
   • inplace Add / inplace Activation
   • layout transform (NCHW ↔ NHWC)
   • reduction ops (sum, mean, max) with shared memory 최적화
4. 하드웨어 특화 커널
   • GPU: CUDA warp-level GEMM, Tensor Core 활용 커널
   • CPU: AVX/AVX512, OpenMP 벡터화된 GEMM
   • TPU/NPU: systolic array에 맞는 특수 커널

 

현재 나는 기본 연산 단위만 실행, 여러 커널 후보 등록 ㅣㅍㄹ요