본문 바로가기

AI Compiler framework

lower.py 의 기능 추가를 통한 AI compiler 로서의 역할 수행하도록 만들기~

현재 lower.py 는 컴파일이 아닌 translator + 즉시 커널 박제의 역할,

AI compiler 로서의 실질 역할 (IR 재구성) 을 하기 위해선 다음과 같이 바꿔야 함

  • IR - 더 좋은 IR 로 바꾸는 단계
  • 그 결과에 맞춘 커널 선택 단계 kernel mapping, choice
  • 실행을 위한 바인딩 / 캡처

 

IR 재구성이 필요한 이유

현재는

  • linear -> gemm + bias_add 로만 분해
  • linearBwd - gemm + gemm + reduce_sum
  • 각 primitive 에 대해 dtype 만 보고 kernel_id 박제

최적화 여지가 없음

진짜 컴파일러의 역할

fusion

  • gemm + bias_add + relu - fused epilogue gemm
  • reduce_sum + sgd_step 같이 update 쪽도 fusion 

 

lower.py 에서 op- 커널 문자열로 직결하는 것이 아닌, 중간 단계 필요