Registry 의미 정리 = 시스템 커널 존재 목록, Selector

Registry 를 코드 컴파일러로 비유

Q : 컴파일러는 CPU 명령어를 어떻게 아는지?

A : 이 CPU 에는 이런 명령어들이 있다는 표가 존재

ADD r, r
MUL r, r
FMA r, r

• 판단이나 최적화가 아닌
• 그냥 존재 목록

instruction registry

 

AI Compiler 에서 registry 가 하는 일

AI Compiler 에서는 이런 CUDA 커널들이 있다를 적어두는 것

• gemm_f16_tc_128x128
• gemm_f16_generic
• ...

이걸 한 군데에 모아두는 것이 registry

 

registry 가 없으면 생기는 문제

if (op == GEMM) {
  if (dtype == FP16 && sm >= 80 && M % 128 == 0) {
    launch_kernel_A();
  } else if (dtype == FP16) {
    launch_kernel_B();
  } else {
    launch_kernel_C();
  }
}

• 커널이 추가될 때마다 if-else
• 조건이 코드 전체
• 해당 커널의 조건을 명확히 알 수 없음

 

registry 의 진짜 역할

해당 구조

[커널 이름]
[이 커널이 쓸 수 있는 조건]
[이 커널의 성격 힌트]
[이 커널을 실행하는 함수 포인터]

이 커널은 이 상황에서 쓸 수 있다를 선언해두는 곳

 

selector 와의 구분

registry

• 이런 커널들이 존재
• 각 커널은 이런 조건에서 가능

Selector

• 현재 상황에서 커널 판단
• 점수 계산
• 선택

 

Registry 를 실제로 쓰는 전체 흐름

Registry 는 selectyor 가 질문을 던질 대상이고, selector 는 registry 를 조회해서 답을 하나 고른다

 

registry 가정

Kernel A:
  - FP16
  - M % 128 == 0
  - sm80+
  - fused
  - priority = 100

Kernel B:
  - FP16
  - any shape
  - any sm
  - generic
  - priority = 40

Kernel C:
  - FP32
  - fallback
  - priority = 0

그냥 정보 목록

 

1. selector 가 던지는 질문

컴파일 시점에서 selector 의 질문

DispatchContext:
  op = GEMM
  dtype = FP16
  shape = (M=256, N=128, K=512)
  layout = row-major
  sm = 86

현재 상황

 

2. step1 - registry 에서 후보를 전부 꺼낸다

candidates = registry.lookup(op = GEMM);

/////////

[A, B, C]

GEMM 커널들 가져옴

 

3. Step 3 - 하드 조건 필터링 ( registry 정보 사용 )

selector 가 registry 에 적힌 조건을 그대로 써서 필터링

if (!kernel.supports(ctx)) drop;

 

• Kernel A
  • FP16 o
  • M=256, 256%128==0 o
  • sm86 ≥ sm80 o
    → 살아남음
• Kernel B
  • FP16 o
  • shape 제한 없음 o
    → 살아남음
• Kernel C
  • FP32 x
    → 탈락

 

filtered = [A, B]

논리 판단, 최적화 없음

 

4. Step 3 - 우선순위 계산 (selector 의 일)

selector 가 registry 에 적힌 힌트를 사용한다

score = kernel.base_priority;
if (kernel.is_fused) score += 50;
if (ctx.sm >= 80 && kernel.uses_tc) score += 30;

점수 계산

• Kernel A
  • base 100
  • fused +50
  • sm86 + tc +30
    → 180
• Kernel B
  • base 40
  • generic (추가 없음)
    → 40

 

Step 4 - 선택

best = argmax(score);

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

Kernel A 선택

 

Step 5 - Plan 에 기록

selector 의 선택을 Execution plan 에 넣는다

Plan:
  step 3:
    kernel_id = A
    launch_config = (grid, block, smem)

이후 Kernel A 실행