본문 바로가기

AI Compiler framework

Pattern matching vs Compiler - 어떤 점이 달라져야 하는가 ( 일련의 컴포넌트를 따라 변환 )

패턴이 의미 + 제약 + 합법성을 가진 객체가 되어야 한다.

pattern = IR rewrite rule with legality & cost

 

이전의 인식, pattern matching

연산 그래프의 형태를 보고, 특정 커널로 변환

이것은 형태 인식이지 컴파일이 아니다.

부족한 내용

  • 의미 보존 : 수치적 동일성 / 안정성 보장
  • 제약 검증 : shape, layout, dtype, mask
  • 대체 실패 시 fallback : 항상 가능한지
  • 비용 모델 : 언제 쓰면 손해인지

 

Patten -> Compiler 로 바뀌는 핵심 전환점

패턴은 반드시 Rule 이 되어야 한다.

컴파일러에서 패턴은 이렇게 생긴다

Pattern:
  IR Subgraph A

Preconditions:
  - shape multiple of tile
  - causal_mask == true
  - softmax is row-wise
  - no dropout (or supported)

Rewrite:
  IR Op B (AttentionFused)

Guarantees:
  - semantic equivalence (within fp tolerance)
  - determinism level: relaxed_fp

위 순간부터 pattern = 컴파일 규칙

 

실제 compiler 로 만들기 위해 필요한 것들 ( 핵심 목록 )

IR 의 정규화 normal form

패턴 매칭이 되려면 IR 이 같은 의미면 같은 형태여야 한다.

  • Softmax
  • exp( QK - max ) / sum( exp(...) )

위 둘이 서로 다른 형태면 매칭 불가능

필수 개념

  • canonical IR
  • algebraic normalization
  • commutativity / associativity 정리

없으면 패턴 수 폭증

 

Pattern 의 의미적 전제 조건 명시

이 패턴 + 이 조건이면 가능

필수 전제들

  • reduction axis
  • mask 위치
  • broadcasting 규칙
  • numerical tolerance

이 커널이 틀렸는지 설명하는 컴파일러 위해

 

Rewrite 이후 IR 가 더 추상적이어야 함

두 가지 rewirte 비교

QK → Softmax → V
↓
FlashAttentionKernel(tile=16, warp=2)

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

QK → Softmax → V
↓
AttentionFused(streaming_softmax=true, causal=true)

패턴 매칭의 결과는 CUDA 가 아닌 IR 여야 한다.

 

Legality checker (이게 compiler 를 만든다)

패턴 매칭 이후 반드시 거쳐야 하는 단계

Can this rewrite be applied?

검증 항목

  • shape legality
  • dtype support
  • layout compatibility
  • determinism requirement 충족 여부
  • memory footprint 가능 여부

 

Cost / preference 개념 (초기엔 룰루 충분)

Compilier 는 항상 rewrite 를 적용하지 않는다.

  • sequence length 작음 - naive 가 더 빠름
  • training mode + strict determinism -> fallback
if seq_len >= 128:
  prefer fused
else:
  fallback

 

pattern matching 을 컴파일러화 화는 최소 구조

[IR Graph]
   ↓
[Normalize]
   ↓
[Pattern Matcher]
   ↓
[Legality Checker]
   ↓
[IR Rewrite]
   ↓
[Kernel Matcher]