커널 선택 기반 방식 vs 수식적 최적화 기반 AI Compiler

1. 문제 의식

현재 커널 최적화 구현을 진행하면서 드러난 핵심적인 한계는 커널을 개별 엔티티로 취급하고, 연산 - 컨러로의 매핑 규칙을 누적하는 구조에 있다. 

현재 방식과 이를 대체할 수 있는 수식적, 최적화 문제 중심의 AI Compiler 설계를 비교 정리

 

2. 현재 방식 : 커널 선택 Selection 중심 구조

2.1 구조 요약

• 연산 타입에 따라 커널 후보군 정의
• shape, dtype, flags 에 따라 분기
• 조건 매칭 후 특정 CUDA kernel 호출

(op_type, shape, flags, device) - kernel_id

 

2.2 특징

• 커널은 이름을 가진 고정된 실행 단위
• Epilogue / Fusion 은 옵션 또는 별도 커널로 분기
• 성능 튜닝은 커널 내부에서만 발생

 

2.3 장점

• 초기 구현이 빠름
• 디버깅 용이
• 특정 workload 에 대해 빠른 최적화 가능

 

2.4 한계

• 커널 수 증가 - 선택 로직 폭증
• 하드웨어 변경 시 재작성
• 왜 특정 커널이 선택되었는지 설명 불가
• 커널 간 공통 구조를 추상화하기 어려움

결과적으로 룰 기반 dispatch 시스템에 가까워짐

 

 

3. 문제의 본질 재정의

핵심 전환

커널을 선택하는 문제가 아니라, 연산을 제약 조건 하에어 최적화하는 문제로 재정의

• 커널 = 결과물
• 컴파일러의 역할 = 최적화 문제 정의 + 해 탐색

 

4. 제안 방식 : 수식적 최적화 기반 구조

4.1 연산 표현의 변화

기존

C = GEMM(A, B) + Bias -> ReLU

제안

C[i,j] = Σ_k A[i,k] * B[k,j]
subject to:
shared_mem ≤ S_max
registers ≤ R_max
occupancy ≥ O_min

연산은 함수 호출이 아니라 수식 + 제약 조건으로 표현된다.

 

4.2 커널 공간의 재해석

기존

• tile = {64, 128}
• stages = {2, 3}
• warp = {2, 4}

이산적 선택

 

제안

• tile size - arithmetic intensity 함수
• stage 수 - memory latency / compute latency 비율
• warp 분배 - reduction depth 함수

performance = f(
FLOPs / bytes,
latency_compute / latency_memory,
reuse_distance,
parallelism
)

커널은 성능 공간 위의 한 점

 

4.3 Epilogue / Fusion 재정의

기존

• Bias, ReLU, LN 등 epilogue 옵션 분기

제안

• 연산 그래프 재작성 문제로 취급

minimize: global_memory_writes
subject to: register_pressure ≤ limit

Fusion 은 옵션이 아니라 cost 감소를 위한 graph rewrite 결과

 

5. Cost Model 중심 컴파일러 구조

5.1 Cost Function 예시

cost = α * memory_traffic
+ β * latency
+ γ * register_spill
- δ * parallelism

• alpha, beta, ramda, omega 는 hardware-specific
• 커널 이름과 무관

 

5.2 컴파일 파이프라인 비교

IR : 연산 단위 -> loop/ reduction / access 표현

Lowering : 커널 매핑 -> 파라미터화된 변환

선택 : rule 기반 -> cost 최소화

결과 : kernel_id -> 실행 계획 + 생성 코드

 

6. 구체적 적용 방향

6.1 IR 레벨

• 커널 이름 제거
• loop nest + reduction + memory access 명시

 

6.2 Lowering 단계

• tiling, unrolling, fusion 을 파라미터화
• 아직 커널 확정 X

 

6.3 Planning 단계

• cost model 평가
• heuristic / partial search 로 계획 선택

 

6.4 Codegen

• 선택된 계획을 CUDA kernel 로 물리화

 

 

7. 요약

• 현재 방식은 확정성 있는 커널 최적화에 구조적 불리
• 커널을 엔티티가 아닌 해(solution) 로 취급해야 함
• AI Compiler 의 본질은 커널 선택기가 아닌 최적화 문제 생성기
• 수식적 cost model 은 설명 가능성과 확장성을 동시에 제공