본문 바로가기

AI Compiler framework

AICF v2 실행 파이프라인 - 함수 단위 호출 트레이스

0. 사용자 진입점

Model.add

파일
src/aicf_v2/model.py

호출 흐름
y = m.add(layer, *args)

실제 동작
return layer.emit(self.b, *args, ctx=self.ctx)
  • Builder
  • EmitContext 를 Layer.emit 에 전달하는 조율자 역할

 

1. Layer 단계 ( 고수준 의미 - IR 노드 생성 )

src/aicf_v2/layers/linear.py
src/aicf_v2/layers/relu.py
src/aicf_v2/layers/adam_step.py

역할

  • 수학적 의미를 IR 로 분해
  • shape / dtype / device 검증
  • Value 생성
  • 어떤 emitter 를 쓸지 결정

 

생성 데이터

value

vid = b.value(name, TensorSpec(...))

저장 위치
Builder.values[vid]

포함 정보

  • vid
  • name
  • TensorSpec(shape, dtype, device)
  • role (input / param / state / tmp / output)

아직 CUDA 나 커널 개념 없음

 

Emitter 호출

Layer 는 직접 b.emit 을 거의 안 부르고 대신 emitters / cuda 함수 호출

emitters.cuda.gemm(...)
emitters.cuda.relu(...)
emitters.cuda.adam_step(...)

 

 

2. Emitter 단계 ( IR - Backend-ready Op )

emitters/cuda/gemm.py
emitters/cuda/adam_step.py
emitters/cuda/bias_add.py

CUDA 호출에 필요한 불변 정보를 완성한다

  • kind_id : C++ OpKind enum 값
  • attr_schema : ABI 구분자
  • attr_blob : 커널에 전달할 packed bytes

 

예시 : gemm emitter

blob = struct.pack("<ii", ta, tb)

emit_resolved(
    b,
    kind="gemm",
    inputs=[A, B],
    outputs=[Y],
    kind_id=ctx.Gemm,
    attr_schema=0,
    attr_blob=blob,
)

 

 

3. Builder.emit ( IR 최종 수집 )

builder.py

def emit(...):
    self.ops.append(Op(...))

Builder 가 들고 있는 최종 IR

Builder.values

  • 모든 tensor slot 메타데이터

Builder.ops

각 Op 는 이미 다음 정보를 완전히 포함

  • inputs / outputs
  • kind_id
  • attra_schema
  • attr_blob
  • constraints / hints / saved

이 시점의 IR 은 backend-ready IR

 

4. Compile 단계 ( IR - ExecPlan )

def compile_cuda(m, registry):
    b0 = m.b
    b1 = optimize_ir(b0)        # 현재 identity
    lowered = lower_ir_cuda(b1)
    plan = make_exec_plan_cuda(b1, lowered)
    plan = optimize_plan(plan)  # 현재 identity
    return CompiledProgram(plan)

 

5. lower_ir_cuda ( serialize-only)

IR -> LoweredOp 의 역할

 

6. make_exec_plan_cuda ( 런타임 결정 )

compile / plan.py

  • 실행 순서 확정
  • alias / inplace 결정
  • CUDA Graph 에서 쓸 정책 데이터 준비
ExecPlan(
    lowered = [...],
    alias = { out_vid : in_vid }
)

동일 storage 를 쓰게 만드는 규칙

 

7. Runtime: eager 실행

runtime / cuda_exec.py

7.1 bind_and_alloc_slots

slots = bind_and_alloc_slots(b, feed)

생성

  • slots[vid] = torch.Tensor

 

7.2 alias 적용

slots[out_vid] = slots[in_vid]

 

7.3 커널 실행

op_call(
    kind_id,
    ins,
    outs,
    attr_schema,
    attr_blob,
)

 

8. CUDA Graph 경로

runtime / graph_capture.py

capture & replay

 

9. 전체 흐름 요약

Model.add
 → Layer.emit
   → emitter (불변 커널 스펙 완성)
     → Builder.emit (IR 수집)
       → compile_cuda
         → lower_ir_cuda (serialize)
         → make_exec_plan_cuda
           → CudaExecutor.run
             → bind/alias
               → op_call
               → (optional) cuda graph replay