AI Compiler 각 파이프라인 단계 역할 이해하기

전체 흐름 확인

[1] IR Graph 생성
      ↓
[2] IR 정규화 (Normalization)
      ↓
[3] Pattern Matching
      ↓
[4] Legality Checking
      ↓
[5] IR Rewrite
      ↓
[6] Kernel Matching / Selection

 

IR Graph 생성 - "의미를 잃지 않는 표현"

역할

•  모델 연산을 의미 단위로 표현
• 커널 / 하드웨어 정보 완전 배제

여기서 반드시 지켜야 할 규칙

• 연산은 수학적 의미 기준
• control flow / reduction axis 명시
• layout 은 속성이지 구현 아님

예시

Op: MatMul
Inputs: Q [B, H, M, D], K [B, H, N, D]
Attrs:
  transpose_b = true
  reduction_axis = D

 

여기선 다음 내용 없음

• tile size
• warp count
• shared memory

 

IR 정규화 (Normalization) - "같은 의미 = 같은 형태"

왜 필요한가

Pattern matching 은 형태를 본다.

형태가 다르면 의미가 같아도 매칭 실패.

하는 일

• algebraic rewrite
• broadcasting 명시화
• reduction 순서 표준화
• softmax canonical form 강제

예시

exp(x - max(x)) / sum(exp(x - max(x)))
↓
Softmax(x, axis=-1, stable=true)

해당 단계가 없으면

• 패턴 수가 폭발
• rule 유지 불가

 

Pattern Matching - "의미 있는 부분 그래프 인식"

역할

• 이 subgraph 는 무엇이다를 인식
• 아직 대체하는 과정은 아님

패턴의 정체

Pattern:
  MatMul(Q, K^T)
   → Mask(causal)
   → Softmax

중요

• shape 변수 허용
• attr match
• graph connectivity match

이 단계에서는

• 성능 판단 X
• 커널 선택 X

 

Legality Checking -이 rewrite 가 가능한가?

이것이 compiler 와 heuristic 의 분기점

Pattern 이 맞아도 못 바꾸는 경우가 반드시 존재

항목                                                         예시

Shape	seq_len % tile == 0
DType	fp16/bf16 only
Layout	contiguous
Numerics	determinism allowed?
Memory	smem ≤ limit

def is_legal(match, context):
    return (
        match.shape.K % 16 == 0 and
        context.device.supports_tensorcore and
        context.determinism != STRICT
    )

해당 코드가 없다면 compiler 가 틀린 코드를 낼 수 있음

 

IR Rewrite - 더 강한 의미 단위로 승격

핵심 원칙

Rewrite 결과는 항상 더 추상적이어야 한다.

두 가지 예시

Softmax + MatMul
→ FlashAttentionKernel


//////////////////////////////

Softmax + MatMul
→ AttentionFused

나쁜 예, 좋은 예 순

이유

• 아직 커널을 고르지 않음
• 다른 구현으로 갈 여지 유지

Rewrite 의 역할

• 연산 수 줄이기
• 의미 단위 압축
• 최적화 여지 생성

 

Kernel Matching / Selection - 구현으로의 연결

여기서 CUDA 등장

입력

• IR Op + requirements
• Kernel registry

동작

IR requirements ⊆ Kernel capabilities

예시

IR:
  AttentionFused:
    causal = true
    streaming_softmax = true
    determinism = relaxed

Kernel:
  FA_v2_tc:
    supports_causal = true
    streaming_softmax = true
    tensorcore = true

match!

결과

• 커널 선택
• template specialization
• CUDA graph capture 가능

 

요약

• IR : 의미 표현
• Normalize : 형태 통일
• Pattern : 의미 인식
• Legality : 컴파일러성 확보
• Rewrite : 추상화 상승
• Kernel match : 구현 선택