본문 바로가기

AI Compiler framework

Core v2 실행/캡처 문서 (Direct aicf_cuda._C, op_call only)

 

  • Python 레벨에서 backend adapter(aicf_backend.AICFBackend)를 제거한다.
  • 모든 실행은 단일 바인딩 aicf_cuda._C의 op_call() 하나로만 수행한다.
  • CUDA Graph 캡처/리플레이도 동일하게 _C.graph_* API로만 제어한다.
  • plan 기반으로 포인터 안정성(pointer-stable) 을 보장해 capture/replay에서 동일 메모리로 실행되게 한다.

 

구성 요소

1) IRGraph (trace 결과)

  • trace_ir(build_fn)로 연산 그래프(IRGraph) 생성.
  • IRGraph.values: vid → (name/shape/dtype/device) 메타
  • IRGraph.nodes: op + (inputs vids → outputs vids) + attrs

 

2) LoweredOps (backend 단위 ops)

  • lower_to_backend_ops(ir)로 IR을 backend 실행 단위로 내린다.
  • lowered op는 문자열 op + vids + attrs 형태의 리스트.

예: Linear → gemm + bias_add
예: Save → copy
예: LinearBwd → gemm(dx) + gemm(dW) + reduce_sum(db)

 

3) BindingPlan (메모리 계획 + 바인딩 규칙)

  • build_binding_plan(ir)로 plan 생성.
  • plan은 vid를 역할별로 분류:
    • inputs: 사용자 입력(x,t)
    • params: 사용자 파라미터(W,b 등)
    • statics: 실행 중 생성/재사용되는 텐서(activation, grad, optimizer state 등)

중요: core_v2는 “실행 중 임시 텐서”도 capture-safe 하려면 모두 static으로 선할당해야 한다.
(그래야 replay에서도 동일 포인터를 유지함)

 

실행 모델 (PlannedExecutor)

핵심 원칙: “env[vid] = torch.Tensor”

  • 실행 중 모든 텐서는 env: Dict[int, torch.Tensor]에 저장된다.
  • lowered op는 env에서 입력/출력 텐서를 꺼내서 _C.op_call(...)에 넘긴다.

 

build_env: 포인터 안정성의 핵심

단계

static 선할당

  • allocate_static_env(ir, plan)이 plan.statics의 모든 vid에 대해 텐서를 만든다.
  • 여기서 만들어진 텐서는 capture/replay 전체 동안 동일 포인터로 재사용된다.
  • reuse_static=True면 매 run마다 다시 만들지 않고 유지한다.

inputs/params 바인딩

  • 사용자가 inputs={"x": x, "t": t}, params={...}로 넘긴 텐서를
  • plan.specs의 (name/shape/dtype/device)와 매칭되는 vid에 바인딩한다.

sanity 체크

  • plan.inputs + plan.params + plan.statics가 모두 env에 존재하는지 확인.

왜 이게 중요한가

CUDA Graph 캡처는 “캡처 당시의 메모리 주소로 커널이 기록” 된다.
따라서 replay 때 텐서 포인터가 바뀌면 그래프가 깨지거나 잘못된 메모리를 참조한다.

 

실행 경로: eager run

개념

  • lowered ops를 순회하면서 _C.op_call을 호출한다.

실제 호출 형태

  • kind = OpKind 매핑(opname → enum int)
  • ins_t = [env[vid] for vid in input_vids]
  • outs_t = [env[vid] for vid in output_vids]
  • _C.op_call(kind, ins_t, outs_t, attrs)

in-place / alias 처리

  • Stage6 lowering에서 흔히 output이 input과 같은 vid인 경우가 있다.
    • bias_add (v003, v002) -> (v003)
  • 이런 경우 outs_t는 동일 텐서 객체를 가리키며, 커널은 in-place로 업데이트한다.

 

“aicf_backend 제거”의 의미 (가장 큰 변화)

이전 구조

  • Python backend adapter가
    • op name → kind 변환
    • 출력 텐서 할당
    • capture guard(in_capture) 처리
    • warmup lr=0, step no-op 같은 정책
      을 담당했다.

지금 구조 (Direct binding)

  • Python은 더 이상 “연산 의미”를 모르고,
  • 모든 실행은 plan이 만든 텐서를 그대로 들고 _C.op_call만 호출한다.
  • 결과적으로 파이썬은:
    • IR/Lower/Plan 생성
    • env 바인딩
    • op_call dispatch
    • graph_begin/end/launch/reset orchestration
      만 책임진다.

장점

  • 실행 경로가 단순해지고 디버깅 지점이 줄어듦
  • “backend 정책”이 사라져서 capture-safe를 구조적으로 보장(= plan 기반)
  • op 추가 시 Python backend 수정 범위가 줄어듦 (거의 OpKind 매핑만)

주의점

  • warmup 정책(예: lr=0, step no-op)이 backend에 있었다면,
    이제는
    • (1) lowering 단계에서 attrs를 조정하거나
    • (2) executor가 특정 op를 스킵/치환하거나
    • (3) C++ op_call 내부에서 env var 기반으로 처리
      중 하나로 옮겨야 한다.