AICF_fw / core 역할 문서

목표

• 1-step train 을 IR 로 컴파일 - lowering - warmup - CUDA graph capture - replay 까지 닫음
• 동일한 IR 을 IRExecutor 로도 실행 가능
• 캡처 안정성 규칙 : capture 중 새 버퍼 할당 금지, leaf grad 포인터 고정

 

파일별 책임

1) ir.py - IR 데이터 모델 

책임

• IRGraph / IRNode / IRValue 정의
• value id / node id 발급
• name uniquify

불변식

• IRValue.id, IRValue.name 그래프 내 유일

 

2) functional.py - Ops 구현 

책임

• forward / backward 에서 쓰는 primitive op 들 구현
  • linear, relu, mse_grad
  • optimizer helpers
• 실행 경로 : backend op_call_out 호출 + BufferPool 로 out-buffer 고정
• tracing 경로 : IR 에 op emit + SSA 바인딩 유지

불변식

• in_capture() == True 동안 새 텐서 할당 금지
• BufferPool 은 capture 전에 warmup 으로 전부 materialize 되어야 함
• tracing 시 output SSA 는 명확해야 함

 

3) autograd.py - Tensor / Node + backward 엔진 + capture guard 

책임

• Tensor, Parameter, Node 정의
• backward() 구현
  • capture-safe leaf grad 정책 ( 포인터 고정 overwrite / add )
  • non-leaf grad 는 local gmap 만 사용
• tracing 중에는 실제 autodiff 대신 IR 에 Backward 노드 emit

하지 않는 것

• 개별 op 수학
• capture begin / end
• lowering

불변식

• capture 중 leaf.grad 새로 만들기 금지
• accumulate = True 는 capture 중 금지

 

4) compile.py - 오케스트레이션

compile - validate - lower - warmup - capture - trace - artifact

책임

• compile_ir(step_fn)
  • tracing context 에서 step_fn 실행 - IR 생성
• lower_to_backend_ops(ir)
  • IR 을 backend op list 로 변환 ( 정해진 순서 / 패치 포함 )
• compile_and_capture(step_fn, ...)
  • validate / lower
  • warmup_capture_safe 호출
  • backend capture_begin / end + trace 수집
  • CompileArtifact 생성

하지 않는 것

• Tensor / autograd 구현
• runtime executor 구현
• env exact binding 생성

 

5) runtime.py - 실행기 / 캡처 안전 유틸

책임

• warmup_capture_safe
  • capture 전에 모든 버퍼 / leaf grad / 커널 lazy init materialize
• IRExecutor
  • lowerd 리스트를 env 로 실행
  • adam_step 강제 in-place semantics

 

6) moduld.py - 사용자 facing API + env exact binding

책임

• Module 베이스
  • parameter / module registration
  • compile / replay / get_artifact