본문 바로가기

AI Compiler framework

각 파일 별 역할 재정리, 실행 과정

📦python
 ┣ 📂aicf_fw
   ┣ 📂backend
   ┃ ┣ 📜aicf_backend.py
   ┃ ┣ 📜base.py
   ┃ ┣ 📜torch_backend.py
   ┃ ┗ 📜__init__.py
   ┣ 📂core
   ┃ ┣ 📜autograd.py
   ┃ ┣ 📜compile.py
   ┃ ┣ 📜functional.py
   ┃ ┣ 📜ir.py
   ┃ ┣ 📜module.py
   ┃ ┣ 📜tensor.py
   ┃ ┣ 📜trace.py
   ┃ ┗ 📜warmup.py
   ┣ 📂nn
   ┃ ┣ 📜linear.py
   ┃ ┣ 📜losses.py
   ┃ ┣ 📜relu.py
   ┃ ┗ 📜sequential.py
   ┗ 📂optim
     ┣ 📜adam.py
     ┣ 📜base.py
     ┣ 📜sgd.py
     ┗ 📜__init__.py

aicf_fw / backend

  • basy.py : 백엔드 공통 인터페이스 / 추상층
  • torch_backend.py : 기준 구현 백엔드, PyTorch eager 로 실행해서 정답 / 디버깅용 비교축
  • aicf_backend.py : AICF 실행 백엔드, 캡처 / 리플레이와 런타임 실행을 실제로 담당

aicf_fw / core

  • tensor.py : Tensor 래퍼 / 핸들 ( shape, dtype, device, requires_grad, grad 저장, storage 참조 등 )
  • module.py : Module 시스템 ( 파라미터 등록 / 버퍼 / forward, state 관리) Sequential / Linear 같은 nn 모듈이 이 규약 위에서 돌아감
  • functional.py : 연산 프리미티브( add, matmul, relu ) 를 모듈 / autograd / IR 이 공유하도록 묶은 레이어
  • autograd.py : backward 그래프 구성 / 미분 규칙, (노드 저장, grad 전파, 필요한 saved tensor 관리 )
  • ir.py : 추적 결과를 담는 IR 정의 ( op, value, shape, topo, 메타데이터 )
  • trace.py : Python 레벨 실행을 IR 로 기록하는 트레이서
  • compile.py : IR - 백엔드 실행체로 변환 ( 캡처 플랜 구성, 런타임 호출 시 어떤 순서로 실행할지 고정 )
  • warmup.py : 캡처 전 워밍업 / shape 고정 / 메모리 할당 안정화 ( CUDA Graph 캡처 조건 맞추기 위한 준비 단계 ) 

aicf_fw / nn

  • linear.py : Linear 모듈 정의
  • relu.py ...
  • losses.py : loss 함수들, train-step 담으려면 loss 도 trace / IR / compile 경로에 들어가야 함
  • sequential.py : 레이어 체이닝 컨테이너, 모듈 리스트를 순회하며 forward 를 구성

aicf_fw / optim

  • base.py : Optimizer 공통 규약 + 파라미터 그룹 관리
  • sgd.py : SGD 구현, 캡처 / 리플레이에서 상태 저장이 상대적으로 단순한 비교군
  • adam.py : Adam 구현

 

실행 과정 정리

Python 실행을 먼저 IR 로 기록하고 - 그 IR 을 기준으로 warmup 으로 shape/메모리/상태 고정한 후 - 같은 IR 시퀀스를 그대로 CUDA Graph 로 capture - 이후 replay 는 그 캡처된 그래프를 실하는 흐름

 

  • Trace - IR 생성 ( trcae, ir )
    • forward / loss / backward / optimizer 까지 연산 순서 + 텐서 메터 + 의존성을 IR 로 만듦
    • 여기서 실행 경로가 고정되어야 캡처가 가능해짐
  • Compile : IR - 실행 계획
    • IR 을 백엔드가 실행할 수 있는 형태로 정리
  • Warmup
    • 캡처 전에 IR 시퀀스를 실행해서
      • 모든 텐서 / 버퍼 할당이 끝나고
      • shape 가 변하지 않고
      • Adam state 같은 상태 텐서도 확정
      • 런타임 경로가 흔들리지 않게 만듦
    • 즉, 캡처 가능한 정적인 상태를 만든다
  • CUDA Graph Capture
    • warmup 으로 준비된 상태에서 IR 이 정의한 연산 호출을 그대로 캡처
    • 캡처 결과는 한 번에 실행되는 train-step 그래프가 됨
  • Replay
    • replay 때는 캡처 그래프 실행 전후로
      • 필요한 상태를 full-restore 규약에 맞게 복구 / 적용해서
      • 같은 입력이면 같은 결과가 나오게 유지

캡처는 IR 에 정의된 연산열을 캡처

warmup 은 그 연산열이 캡처 가능하도록 메모리 / shape / 상태를 미리 안정화시키는 단계