본문 바로가기

AI Compiler framework

ge_v2 와는 다르게 forward / backward 를 같은 operation 내에 묶지 않는다?

같은 operation 패키지에 넣되, OpKind / registry entry 는 분리

 

OpKind 의 분리 이유

1) 입력 / 출력 시그니처가 다름

ReLU

  • forward : Y = relu(X)
  • backward : dX = relu_bwd(X, dY) or (Y, dY)

attrs 만으로 구분하는 순간 복잡해짐..

 

2) backward 는 저장값 / 캐시 정책이 개입된다.

saved tensor 가 필요한 backward 가 존재

  • ReLU : X, mask
  • LayerNorm : mean / bar
  • Softmax : sumexp / amax

같은 op 로 묶으면 forward 의 constract 가 오염됨

 

3) 커널 선택 / 우선순위가 달라진다

backward 와 forward 의 최적화 포인트가 달라짐

  • 메모리 패턴
  • reduction 필요 여부
  • fusion 가능성

 

4) 더 쉬운 IR Scheduling / memory planner 

IR Tensor 가 forward op / backward op 로 분리되어 있으면

  • 어떤 텐서가 saved 인지, grad 인지
  • 라이프타임이 어떻게 되는지가 명확해져서 planner 가 쉬워짐

 

✅ 파일 / 모듈 레벨에선 같이 둬도 괜찮, 하나의 launcher.cu 내에 fwd/bwd 둘 다 넣어도 됨

✅ registry 레벨에선 분리, ( OpKind::EltwiseRelu, OpKind::EltwiseReluBwd )