[triton] Proton 커널 내 프로파일러 Global Memory 지원

PR 링크: triton-lang/triton#8641 상태: Merged | 변경: +217 / -51

들어가며

Triton Proton의 intra-kernel profiler는 커널 실행 중 타이밍 데이터를 수집합니다. 기존에는 shared memory에만 버퍼를 할당했는데, shared memory가 부족한 커널에서는 프로파일링이 불가능했습니다. 이 PR은 global memory를 대안 버퍼로 사용할 수 있도록 확장합니다.

핵심 코드 분석

1. SegmentAllocOp의 입력 타입 확장

Before:

let arguments = (ins TTG_MemDescType:$buffer);

After:

let arguments = (ins AnyTypeOf<[TTG_MemDescType, TT_Ptr]>:$buffer);

Segment alloc이 shared memory memdesc뿐 아니라 global memory pointer도 받을 수 있게 되었습니다.

2. Global memory 버퍼에서의 CircularStore

if (addrSpace == 1) {
  // Global memory: PTX st.global.cg 명령어 사용
  auto stInst = builder.create<>("st")->o("global").o("cg").v(2).b(32);
  auto *ptrOpr = builder.newAddrOperand(dataPack.ptr, "l");
  // vectorized 2x32-bit store for num_warps > 1
  // non-vectorized for num_warps == 1 (alignment 문제 방지)
}

3. Finalize에서 global memory 조기 반환

if (objBaseTy.getAddressSpace() == 1) {
  // Global memory 버퍼는 이미 global에 있으므로
  // shared -> global 복사가 불필요
  writeBackPostFinalTime(b, rewriter, op, isFirstThread, scratchPtr);
  rewriter.eraseOp(op);
  return success();
}

왜 이게 좋은가

• 범용성: Shared memory 부족 시에도 프로파일링이 가능합니다.
• 성능 최적화: st.global.cg(cache-global) 명령어와 vectorized store로 write 오버헤드를 최소화합니다.
• 단일 워프 호환: num_warps=1일 때 alignment 문제를 비벡터화 store로 우회합니다.

정리

Proton intra-kernel profiler의 버퍼를 global memory에 배치할 수 있도록 확장한 PR입니다. 타입 시스템 변경, PTX store 구현, finalize 최적화가 핵심입니다.

참고 자료

• triton-lang/triton#8641

---

이 글은 AI의 도움을 받아 작성되었으며, 원본 PR의 코드 변경 사항을 기반으로 분석한 내용입니다.