파이프 라인 수정... - IR → 패스 → 커널선택(디스패치) → 스케줄/메모리플랜 → 실행으로 이어지는 컴파일러형 학습 파이프라인

전체 파이프라인(러프 스케치)

Frontend

• 입력: Python 모듈(네이티브 레이어 API), ONNX/TF/PyTorch trace 등
• 출력: High-Level IR (HLIR) 그래프 (op=Conv2D/GEMM/SDPA/LayerNorm/… + attr)

IR Passes (Optimizer/Lowering)

• 정적분석: dtype/shape/stride 추론, 정규화(브로드캐스트/컨티구어스), 패턴 매칭
• 최적화: 연산 병합/퓨전(GEMM+Bias+Act → regemm_epilogue), 상수폴딩, 레이아웃 선택(NHWC↔NCHW), AMP 정책 설정(fp16/bf16+fp32 Accum), 아토믹 회피/리덕션 전략 등
• 출력: Mid/Low-Level IR (MLIR/LLIR) — 백엔드에 가까운 op 집합(eg. GemmBiasAct, SdpaFused, EwiseAdd, ReduceMax …)과 레이아웃/타일링 힌트

커널선택(디스패치) & 코스트모델

• 각 op 노드마다 registry.select(op, problem, hw_caps, prefs) -> KernelHandle
• 규칙 기반 + 마이크로벤치 데이터(또는 런타임 오토튜닝) 혼합
• 예: GemmBiasAct → regemm_epilogue vs cublasLt epilogue vs naive + ewise
• 실패 시 폴백 커널 으로 보장

스케줄러 & 메모리 플래너

• 노드 토폴로지 정렬 → 스트림/이벤트 배치(겹치는 노드 병렬화)
• 라이프타임 분석 + 버퍼 재사용(아레나 할당) + 체크포인팅
• 오토그라드용 역전파 그래프와 공간-시간 트레이드오프 조절(activation recompute)

Executable 생성

• 선택된 커널 핸들 + 파라미터 + 텐서 레이아웃 + 런타임 의존성(스트림/이벤트/워크스페이스)로 실행플랜 고정
• 학습 루프에서 반복 사용(그래프 캐시/셰이프 키)

런타임(학습)

• forward() → 손실 커널 → backward()(규칙 기반 미분/커스텀 bwd 커널) → 옵티마이저(ewise/update)
• 배치마다 실행플랜만 주입 → 최소 오버헤드

 

최소 구현 로드맵(4 단계)

1. IR Core & Registry 연결
   • HLIR/LLIR 구조체 + shape/dtype 패스
   • registry.select() 스켈레톤 + _ops_gemm만 연결
   • 간단 네트워크(MLP) compile/run 성공
2. Fusion & AMP & Planner
   • GemmBiasAct 퓨전 + arena allocator + 1~2개 스트림 스케줄링
   • AMP 정책과 정확도 가드(ewise/reduce fp32)
3. 역전파 그래프 & 옵티마이저
   • rule-based bwd(mlp 기준) + SGD/Adam ewise 커널
   • compile once, train many 경로 성능/안정화
4. 확장
   • conv2d, sdpa, layernorm, softmax, reduction 번갈아 추가
   • 코스트모델/오토튜너(벤치 캐시 + 워밍업) 고도화
   • 멀티-GPU(DDP/AllReduce) 스텁

 

핵심 포인트 요약

• 지금의 플러그인식 바인딩 + 커널 레지스트리는 “IR 해석 → 커널선택” 컴파일러에 매우 적합.
• regemm_epilogue는 구조상 이미 분리된 커널 집합이고, 플러그인 모듈화만 하면 퓨전 에피로그를 디폴트 경로로 쓰도록 쉽게 바꿀 수 있어.
• 학습까지 포함하려면 역전파 규칙/커널과 메모리/스케줄 정책을 IR 패스와 함께 고정하는 것이 포인트.
• 이후에는 오토튜닝 DB와 코스트모델이 성능의 관건.