aicf_v2 현재 구조 정리

목표

• Layer 가 IR 을 직접 emit 한다.
• IR 을 그대로 LowredOp 로 변환해 op_call 로 실행한다
• rewrite/pass는 나중. 지금은 “IR → 실행” 파이프를 먼저 고정.

 

디렉토리 구조와 역할

src / aicf_v2 

핵심 모델 / IR 계층

• tensor_spec.py
  • TensorSpec 정의
  • 런타임에서 텐서 allocate / 검증 기준이 됨
• graph.py
  • Value : vid, name, spec, prod / users 같은 메타
  • Op : kind, inputs, outputs, attrs, constraints, name
  • dump / debug 출력 포맷
• builder.py
  • 값 / 오퍼레이션을 누적 관리
  • input / value / emit 같은 low-level API
  • input_vids, output_vids, values, ops 들고 있음
• model.py
  • 사용자 - facing DSL
    • m.input
    • m.add
    • m.output
  • 내부적으로 Builder 를 조작함

 

레이어 계층

• layers
  • base.py
  • linear.py
  • relu.py
  • add.py

 

backend / runtime

 

• backends/cuda/
  • registry.py
    • IR kind(str) → KernelSpec(kind_id, attr_schema) 매핑
    • kind_id는 C++ OpKind enum과 반드시 일치해야 함
    • attr_schema는 op마다 다름
      • 예: BiasAdd는 fourcc("BADD")
      • Gemm은 0 (테스트 코드로 확정)
  • attrs.py
    • op.attrs(dict) → attr_blob(bytes) packer
    • ABI(바이트 레이아웃)는 C++ attr decode와 반드시 일치해야 함
      • BiasAdd: <q axis
      • Gemm: <ii transA transB
      • Relu/Add: 빈 blob
  • bridge.py
    • from aicf_v2 import _C 로 확장 모듈 로드
    • op_call(...) thin wrapper
    • current_stream_u64() 제공
• runtime/cuda_exec.py
  • CudaExecutor.compile(model):
    • Builder.ops를 순회하며 LoweredOp 생성 (registry+attrs 사용)
    • (옵션) alias/inplace 플랜 적용 가능
  • CompiledCudaProgram.run(feed):
    • input/param 텐서 주입
    • 나머지 value는 spec으로 torch.empty allocate
    • lowered 순서대로 _C.op_call(...) 실행
    • output vid들 반환

 

 

새 op를 추가할 때: 무엇을 건드려야 하나

아래 체크리스트 순서대로 하면 “추가 → 단독 테스트 → 모델 조합” 흐름으로 안전하게 확장됨.

---

0) C++ 쪽에 op가 존재하는지 확인

필수:

• OpKind enum에 항목 존재
• register_all()에서 해당 kind가 등록됨
• _C.op_call(OpKind::<X>, ...) 단독 호출 테스트가 통과해야 함

단독 호출에서 NotImplemented면 v2를 건드려도 안 됨. C++ 등록/구현부터.

---

1) IR kind 이름을 정한다 (Python 문자열)

예:

• "layernorm_fwd"
• "reduce_sum"
• "copy"

이 문자열이 이후 모든 매핑의 키가 됨.

---

2) Layer에서 emit한다 (layers/*.py)

새 레이어가 필요하면:

• 입력/출력 Value(vid) 생성: shape/spec 결정
• ctx.emit_op(kind, inputs, outputs, attrs, constraints, hints...) 형태로 IR op를 추가

예시 패턴(Linear처럼 여러 op emit도 가능):

• fwd: op1 -> op2 -> op3 체인
• training: saved 정책을 op.saved/메타로 기록(실제 copy op를 넣을지 정책 결정)

---

3) backends/cuda/registry.py에 매핑 추가

IR kind → C++ enum kind_id + schema

• kind_id: C++ OpKind::<X> 숫자와 동일해야 함
• attr_schema: C++이 요구하는 schema id
  • 어떤 op는 fourcc("....")
  • 어떤 op는 0
  • 이건 C++ 단독 테스트 코드가 “정답”임

예:

 

"reduce_sum": KernelSpec(kind_id=4, attr_schema=fourcc("RSUM"))

---

4) backends/cuda/attrs.py에 packer 추가

IR attrs(dict) → bytes

C++ 테스트 코드(또는 C++ struct 정의)를 기준으로 정확히 동일한 레이아웃을 pack 해야 함.

예:

 

if kind == "reduce_sum": axis = int(attrs["axis"]) return struct.pack("<q", axis)

---

5) 런타임이 allocate 가능한지 확인 (spec/shape)

op가 요구하는 출력 shape를 Layer emit에서 정확히 잡아야 함.

실수 나는 케이스:

• broadcast 규칙
• reduce 후 shape 변화
• keepdim 옵션
• batch dimension(B) 추론

---

6) “단독 테스트 스크립트”를 먼저 만든다 (필수)

scratch/01x_run_<op>_only.py 형태로:

• v2로 IR을 최소 구성
• _C.op_call 경로 실행
• torch reference랑 비교

BiasAdd/Relu/Add/Gemm처럼 단독 0.0 또는 허용오차를 확보한 다음에,
그 다음에 모델(010) 같은 조합 테스트로 올려야 디버깅이 안 터짐.

---

ops를 “증가”시키는 경우별 가이드

케이스 A) 새로운 Layer를 추가해서 ops가 늘어남

건드릴 것:

1. layers/newlayer.py 생성 (emit 구현)
2. layers/__init__.py export
3. model.py에서 m.add(NewLayer(...), ...) 사용 가능하게

이 단계에서는 backend 지원이 없어도 dump()로 IR은 늘어남.

---

케이스 B) IR에 새 op kind가 추가되어 ops가 늘어남 (실행까지)

건드릴 것:

1. Layer emit에서 kind="new_op" 추가
2. registry.py에 "new_op" -> (kind_id, schema) 추가
3. attrs.py에 "new_op" pack 추가
4. 단독 실행 테스트 추가

---

케이스 C) 같은 op지만 attrs/constraints가 늘어남

예: Gemm에 algo 선택, epilogue, workspace 등

건드릴 것:

1. Layer emit에서 attrs dict 확장
2. attrs.py pack 확장(ABI 변경)
3. C++ decode struct와 동기화
4. 단독 테스트 업데이트

---

케이스 D) inplace/alias 같은 “플랜”이 추가되어 ops는 같지만 메모리 동작이 바뀜

건드릴 것:

• CudaExecutor.compile()에서 alias 플랜 생성
• CompiledCudaProgram.run()에서 슬롯 alias 적용
• (추가) ptr debug / slot dump 기능

예:

• bias_add: out_vid → in_vid alias

---

디버깅 루틴 (빠른 순서)

1. dump()로 IR이 기대대로 나왔는지 확인
2. 단독 op 테스트에서:
   • NotImplemented → C++ 등록/구현 문제
   • InvalidArgument / BadSchema / BadAttrSize → registry/attrs ABI 문제
   • shape mismatch → Layer emit(spec) 문제
3. 모델 조합 테스트로 올리기

---

현재 남은 큰 확장 포인트 (다음 단계 후보)

• compile()/run() 분리 확정 + 캐시
• alias/inplace 플랜 확장
• CUDA Graph capture/replay
• rewrite/fuse(pass): gemm + bias_add -> gemm_epilogue 등