[triton] Triton PROTON: FinalizeOp 최적화를 통한 프로파일링 오버헤드 개선

PR 링크: triton-lang/triton#8635 상태: Merged | 변경: +394 / -125

들어가며

Triton의 프로파일링 도구인 PROTON은 커널 실행 중 데이터를 수집하고 이를 메모리에 기록하는 과정을 거칩니다. 기존의 FinalizeOp 구현은 단일 스레드 중심의 데이터 쓰기 방식을 사용하여, 커널 실행이 길어질수록 프로파일링 오버헤드가 누적되어 성능 저하를 유발하는 문제가 있었습니다. 특히, 복잡한 커널에서는 이 오버헤드가 2.6배까지 증가하는 현상이 관찰되었습니다. 본 PR은 FinalizeOp를 리팩토링하여 각 warp가 자신의 데이터를 병렬로 기록하도록 최적화함으로써 이 문제를 해결합니다.

코드 분석

1. PatternProtonGPUOpToLLVM.cpp에서의 제어 흐름 개선

가장 핵심적인 변경은 FinalizeOp의 LLVM 변환 로직입니다. 기존에는 모든 스레드가 순차적으로 혹은 비효율적으로 메모리에 접근했다면, 이제는 emitBlockLeaderPrologue, emitWarpIndexWriteback, emitWarpCopySection과 같은 함수들을 통해 제어 흐름을 세분화했습니다.

Before (기존 방식):

Value threadId = getThreadId(rewriter, loc);
Value warpId = b.udiv(threadId, b.i32_val(triton::gpu::TritonGPUDialect::getThreadsPerWarp(mod)));
// ... 단일 스레드 중심의 순차적 로직 ...
Value warpIndexOffset = b.add(warpId, b.i32_val(circularHeaderWordSize));

After (개선된 방식):

// Warp 단위로 병렬 처리를 수행하도록 구조화
Block *continuation = emitBlockLeaderPrologue(op, isBlockFirstThread, scratchPtr, ...);
continuation = emitWarpIndexWriteback(op, continuation, isWarpFirstThread, warpId, ...);
if (segmentBaseTy.getAddressSpace() == 3) {
  continuation = emitWarpCopySection(op, continuation, laneId, ...);
}
emitBlockLeaderEpilogue(op, continuation, isBlockFirstThread, ...);

이 구조는 각 warp가 자신의 데이터 영역을 독립적으로 처리하게 하여, 전체 쓰기 작업의 병렬성을 극대화합니다.

2. protongpu_to_llvm.mlir 테스트 케이스 업데이트

MLIR 수준의 테스트 코드에서도 기존의 단일 브랜치 구조가 조건부 브랜치(llvm.cond_br)를 포함한 복잡한 제어 흐름으로 변경되었음을 확인할 수 있습니다.

-  // CONVERT-BUILTIN: llvm.br ^bb{{.*}}(%{{.*}} : i32)
+  // CONVERT-BUILTIN: llvm.cond_br %{{.*}}, ^bb{{.*}}, ^bb{{.*}}

이는 각 warp가 자신의 실행 상태에 따라 분기하여 효율적으로 메모리 쓰기를 수행함을 의미합니다.

왜 이게 좋은가

이번 최적화의 핵심은 **'Warp 단위의 병렬 쓰기'**입니다. 기존에는 단일 스레드가 전체 버퍼를 처리하거나 비효율적인 동기화가 발생했지만, 이제는 각 warp가 자신의 데이터만 책임지게 함으로써 다음과 같은 성능 향상을 얻었습니다:

• Simple Kernels: 오버헤드 2x → 1.1~1.0x로 감소
• Long Running Kernels: 오버헤드 2.6x → 1.8~1.3x로 감소

리뷰 과정에서 warp specialization과 관련된 논의가 있었으나, 저자들은 이 최적화가 현재의 warp 구조 내에서 오버헤드를 줄이는 가장 효과적인 방법임을 입증했습니다. 특히, buffer_size를 튜닝하지 않고도 즉각적인 성능 이득을 볼 수 있다는 점이 큰 장점입니다.

교훈

1. 데이터 지역성 활용: GPU 프로그래밍에서는 스레드 간의 경합을 줄이는 것이 중요합니다. 각 warp가 자신의 메모리 영역을 독립적으로 처리하게 하는 것만으로도 오버헤드를 획기적으로 줄일 수 있습니다.
2. 컴파일러 패스 설계: FinalizeOp와 같은 후처리 작업은 커널의 전체 실행 흐름에 영향을 주지 않도록 정교한 제어 흐름(Control Flow) 설계가 필요합니다.

참고: 이번 변경은 buffer_size에 의존하지 않는 범용적인 최적화로, 향후 더 복잡한 warp 특화 시나리오에서도 확장 가능한 기반을 마련했습니다.

참고 자료

• https://mlir.llvm.org/docs/Dialects/LLVM/
• https://github.com/triton-lang/triton

⚠️ 알림: 이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.