[triton] Triton AMD 백엔드 최적화: SGPR 활용과 루프 최적화를 통한 GEMM 성능 향상

PR 링크: triton-lang/triton#9849 상태: Merged | 변경: +4 / -1

들어가며

Triton은 고성능 GPU 커널을 작성하기 위한 강력한 도구이지만, 특정 하드웨어 아키텍처(특히 AMD GFX1250과 같은 최신 아키텍처)에서는 레지스터 활용 방식이 성능의 병목이 될 수 있습니다. 이번에 분석할 PR은 AMD GPU 환경에서 f16_gemm 커널의 TDM(Triton Data Mover) 로드 로직을 최적화하여, 불필요한 레지스터 간 데이터 이동을 줄이고 코드 크기를 최적화한 사례입니다. 이 변경사항은 크게 두 가지 핵심 최적화로 나뉩니다: VGPR 기반의 조건문을 SGPR 기반의 비트 연산으로 전환, 그리고 ttgl.assume을 통한 불필요한 분기 제거입니다.

코드 분석

1. VGPR에서 SGPR로의 조건문 전환

기존 코드에서는 루프 내에서 pred=i < epilogue_lb와 같은 비교 연산을 수행했습니다. 이는 GPU 내부에서 VGPR(Vector General Purpose Register)을 사용하여 평가될 가능성이 높으며, 결과적으로 v_readfirstlane과 같은 명령어를 유발하여 SGPR(Scalar General Purpose Register)로 값을 가져오는 추가적인 오버헤드가 발생했습니다.

Before:

producer = issue_loads(..., pred=i < epilogue_lb)

After:

pred = i - epilogue_lb
pred = (pred >> 31) & 1
producer = issue_loads(..., pred=pred)

위와 같이 변경함으로써 비교 연산을 산술적인 비트 시프트 연산으로 변환했습니다. (i - epilogue_lb) >> 31은 해당 값이 음수인지 양수인지(즉, i < epilogue_lb인지)를 판단하는 효율적인 방법입니다. 이를 통해 조건문이 SGPR 내에서 완전히 처리되어 v_readfirstlane 호출을 제거할 수 있게 되었습니다.

2. ttgl.assume을 통한 루프 분기 제거

컴파일러는 루프의 상한선(loop_ub)이 0보다 큰지 알 수 없을 때, 루프 진입 전후에 안전을 위한 분기(branch) 코드를 생성합니다. ttgl.assume(loop_ub > 0)을 추가함으로써 컴파일러에게 루프가 최소 한 번은 실행됨을 보장하여, 불필요한 진입부 분기 코드를 삭제하도록 유도했습니다.

Before:

loop_ub = ttgl.cdiv(K, BLOCK_K)
epilogue_lb = loop_ub - (NUM_BUFFERS - 1)
for i in range(0, loop_ub):

After:

loop_ub = ttgl.cdiv(K, BLOCK_K)
epilogue_lb = loop_ub - (NUM_BUFFERS - 1)
ttgl.assume(loop_ub > 0)
for i in range(0, loop_ub):

왜 이게 좋은가

이 최적화의 핵심은 **'하드웨어 레지스터 아키텍처의 이해'**에 있습니다. AMD GPU에서 SGPR은 제어 흐름과 조건 처리에 최적화되어 있으며, VGPR은 데이터 연산에 최적화되어 있습니다. VGPR에서 SGPR로 데이터를 옮기는 작업은 파이프라인 스톨(stall)을 유발할 수 있는 비용이 큰 작업입니다. 이를 비트 연산으로 대체함으로써 루프의 핫패스(hot path)에서 불필요한 동기화 명령을 제거했습니다.

또한, ttgl.assume을 통해 코드 크기를 약 17% 감소시켰다는 점은 매우 인상적입니다. 이는 명령어 캐시 효율을 높이고, 분기 예측기(branch predictor)의 부담을 줄여 전체적인 커널 실행 속도 향상으로 이어집니다. 일반적인 GPU 커널 최적화에서도 루프의 경계 조건을 명확히 컴파일러에게 전달하는 것이 얼마나 중요한지 보여주는 좋은 사례입니다.

참고 자료

• https://rocm.docs.amd.com/en/latest/reference/gpu-arch.html
• https://github.com/triton-lang/triton

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