본문 바로가기

AI Compiler framework

현재 IR 상태 요약 - IR 기반 실행은 아님, 변화 필요

  • IR 은 forward + optimizer 의 구조만 표현
  • 실제 연산은 IR 을 직접 해석해서 실행하지 않음
  • 실행은 Python step_fn - backend op dispatch - CUDA Graph capture / replay 로 이루어짐
  • backward 는 IR 을 따라서 실행되는 것이 아닌, forward 실행 중 autograd 가 런타임에서 호출한다

 

현재 IR 에는 무엇이 들어가 있나?

  • Linear
  • ReLU
  • MseGrad
  • StepInc
  • BiasCorr
  • AdamStep
  • Backward

즉 IR 은 훈련 step 의 논리적 구조를 담고 있음

  • IR 이 실행 엔진이 아님,
  • IR 을 순회하면서 실제 커널을 호출하는 건 아님
  • 단지 이 연산들이 이 순서대로 필요하다를 표현

 

그럼 실제 연산은 어떻게 실행되는지?

실행의 실제 주체

step_fn()

step_fn 내부에서

y = model(x)                  # forward
dY = F.mse_grad(y, t)         # loss grad 생성
autograd_backward(y, grad=dY) # backward 호출
optim.step_()                 # optimizer
  • model(x) 내부에서 Linear.forward, ReLU.forward 호출
  • 각 forward op 는 즉시 backend op_call 을 발생시킴
  • autograd_backward - backward 엔진을 즉시 실행
  • optim.step - step_inc, bias_corr, adam_step 을 즉시 실행

IR 을 따라 실행하는 것이 아닌 Python 코드가 실행을 주도, IR 은 단지 기록

 

IR 은 그럼 언제 / 어떻게 쓰이는지?

(1) tracing 중에 무슨 연산이 있었는지 기록

with tracing(ir):
    step_fn()
  • 이때 각 op_call 은 
    • 실제로 실행되고
    • 동시에 IRGraph 에 노드가 추가됨

IR 은 실행의 부산물

 

(2) lowering 계획 생성

lowered = lower_to_backend_ops(ir)
  • IR 을 보고
    • 이 step 에서는 gemm - relu - mse_grad - adam_step
    • 라는 기대 실행 시퀀스 생성
  • 컴파일 타임 분석
  • 검증용
  • 최적화 계획용

 

(3) runtime trace 와의 정합성 검증

trace_ops = backend.trace_get()
artifact.assert_runtime_matches_lowering(...)
  • 실제 실행 중 디스패치된 op_call 과 
  • IR 에서 기대한 lowering 결과가
  • 논리적으로 같은지 확인

 

backward 는 IR 을 따라서 실행되는 것이 아닌, 런타임에서 즉시 실행된다.

autograd_backward(y, grad=dY)

위 호출이

  • y 가 만들어질 때 쌓아둔 runtime tape
  • 각 Tensor 가 들고 있는 grad_fn
  • 의존성 역추적

을 기반으로 즉시 backward kernel 들을 실행한다 

  • IR 의 Backward 노드들을 해석해서 나온 게 아님
  • 실제 backward 구현이 런타임에서 분해되어 실행된 결과

 

하지만 컴파일러가 되려면 IR 이 실행의 근거가 되어야 함, 

현재 IR 이 복사인 구조 한계 

현재는

  • 실행 주체 : step_fn
  • IR : 실행 중에 기록되는 로그
  • 최적화 : IR 을 바꿔도 실제 실행은 여전히 step_fn 이 호출하는 op_call 이 기준

현재 상태에서 IR 최적화를 해도 

  • IR 최적화가 실행에 반영되지 않음
  • 검증 기준만 늘고 컴파일러 본질이 빠짐
  • 동적 요소가 계속 남음

최적화 후 바뀐 IR 을 따라서 실행해야 컴파일러임

 

IR 기반 실행으로 전환하려면 

기록이 아닌 프로그램이 되려면 4가지 필요

1. 값의 정보

  • 현재 IR 의 value 는 shape / dtype 정도의 의미
  • 실제 실행은 python 객체 / torch tensor 가 들고 있는 포인터가 기준

IR 실행하려면

  • 각 value 가 어느 버퍼에
  • in-place 인지
  • aliasing 여부

 

2. op 의 lowering 규칙이 실행 가능한 형태

단순 리스트가 아니라 value id 기반의 실행 그래프

 

3. 메모리 플래너

명시적 조건

  • 어떤 value
  • 어느 시점에 free / 재사용
  • workspace 어떻게 공유
  • grad buffer / optimizer state buffer 는 어디에

 

4. 실행 엔진

실제 IR 을 돌리는 엔진

두가지 방법

  • IR interpreter
    • IR 노드를 순서대로 읽고 backend op_call 호출
    • 구현이 쉽고 디버깅이 쉬움
  • IR - 캡처 가능한 실행으로 lowering
    • IR 을 실행해서 cuda graph 로 캡처
    • 이후 replay