Op 와 Kernel 의 개념적 분리와 Attribute 계층

1. 문제 인식

초기 설계 과정에서 op 와 kernel 을 동일한 단위로 인식하는 혼동이 있었음

특히 lowered ops 단계에서 각 연산이 이미 CUDA Kernel 과 1:1 로 대응되는 것처럼 보이기 때문에

op 단위에서 바로 kernel 이 선택된다는 인상을 주기 쉬웠음

=== LoweredOps(v2_kid_trace_by_id_test) ===
ops: 18

  #000 gemm       (0, 1) -> (3)  {transB=True}  kid=gemm_f16_tc_wmma_out_f16_v0
  #001 bias_add   (3, 2) -> (3)  kid=bias_add_f16_vec2_v0
  #002 relu       (3) -> (4)  kid=relu_f16_vec2_v0
  #003 copy_saved (4) -> (5)  kid=copy_f16_v0
  #004 gemm       (4, 6) -> (8)  {transB=True}  kid=gemm_f16_tc_wmma_out_f16_v0
  #005 bias_add   (8, 7) -> (8)  kid=bias_add_f16_vec2_v0
  #006 mse_grad   (8, 9) -> (10)  kid=mse_grad_f16_v0
  #007 gemm       (10, 6) -> (11)  {transA=False, transB=False}  kid=gemm_f16_tc_wmma_out_f16_v0
  #008 gemm       (10, 4) -> (12)  {transA=True, transB=False}  kid=gemm_f16_tc_wmma_out_f16_v0
  #009 reduce_sum (10) -> (13)  {axis=0, keepdim=False}  kid=reduce_sum_keep_lastdim_f16_v0
  #010 relu_bwd   (11, 5) -> (14)  kid=relu_bwd_f16_vec2_v0
  #011 gemm       (14, 1) -> (15)  {transA=False, transB=False}  kid=gemm_f16_tc_wmma_out_f16_v0
  #012 gemm       (14, 0) -> (16)  {transA=True, transB=False}  kid=gemm_f16_tc_wmma_out_f16_v0
  #013 reduce_sum (14) -> (17)  {axis=0, keepdim=False}  kid=reduce_sum_keep_lastdim_f16_v0
  #014 sgd_step   (1, 16) -> (1)  {lr=0.01}  kid=sgd_step_f16_half2_v0
  #015 sgd_step   (2, 17) -> (2)  {lr=0.01}  kid=sgd_step_f16_half2_v0
  #016 sgd_step   (6, 12) -> (6)  {lr=0.01}  kid=sgd_step_f16_half2_v0
  #017 sgd_step   (7, 13) -> (7)  {lr=0.01}  kid=sgd_step_f16_half2_v0

 

그러나 online kernel selection 및 op fusion 을 도입하는 순간 op 와 kernel 은 명확히 다른 역할을 갖는 계층임이 드러난다.

 

2. 핵심 개념 요약

• op : IR 수준의 의미적 연산자
• Lowered Op : 실행 친화적으로 분해된 primitive op
• Kernel : 실제 GPU 에서 실행되는 구현 단위
• Region : 여러 op 가 결합되어 하나의 kernel 로 실행되는 범위

Op 는 의미 단위, Kernel 은 구현 단위로 Kernel 은 하나 이상의 Op 를 포함할 수 있음

 

3. Op 는 Kernel 보다 더 작은 단위

3.1 Op 의 역할

• 어떤 수학적 / 의미적 연산을 수행하는가
• 입력 / 출력 텐서의 관계는 무엇인가
• 다른 op 와 결합할 때 의미가 어떻게 변하는가

독립적으로도 실행 가능하지만, 의미적으로는 서로 결합될 수 있다.

 

3.2 kernel 의 역할

• 어떤 연산 시그니처를 지원하는가
• dtype / layout / epilogue 조합이 가능한가
• 어떤 하드웨어 제약

Kernel 은 Op 들의 조합 결과를 만족할 수 있을 때만 선택된다.

 

4. 왜 Op Attribute 와 Kernel Meta 가 분리되어야 하는가

fusion / online selection 이 추가되었기 때문

4.1 실제 필요한 구조

• Op Attribute
  • 이 op 가 region 에 어떤 의미를 추가하는가
  • compose 의 입력
• Kernel Meta
  • 이 kernel 이 어떤 region attribute 를 지원하는가
  • cadidate filtering 의 대상

 

5. Attribute 계층 구조

5.1 Op Attribut 

• lowered op 에서 동적으로 생성
• 개별 op 의 의미를 표현

 

5.2 Region Attribute 

• 여러 OpAttr 를 compose 한 누적 상태
• canonial form 으로 수렴

 

5.3 Kernel Meta

• 사전에 정의된 정적 메타 데이터
• kernel 이 지원 가능한 RegionAttr 의 범위 명시

 

6. online kernel selection 에서의 올바른 흐름

1.  lowered op - OpAttr 생성
2. AttributeState 에 OpAttr compose
3. RegionAttr 갱신
4. KernelMeta 후보 교집합 갱신
5. 결합 불가 시 flush
6. Region 에 대해 kernel 하나 확정