본문으로 건너뛰기

[triton] [AMD Triton] LLVM InstCombine의 함정을 피하는 법: TDM 텐서 클램핑 최적화

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

들어가며

GPU 프로그래밍, 특히 Triton과 같은 커널 DSL(Domain Specific Language)을 개발할 때 가장 까다로운 부분 중 하나는 컴파일러 백엔드(LLVM)의 최적화 동작을 예측하는 것입니다. 컴파일러는 코드를 더 효율적으로 만들기 위해 다양한 패스(Pass)를 거치지만, 때로는 특정 하드웨어 아키텍처의 특성을 고려하지 못한 '지나친 최적화'가 오히려 성능 저하를 불러오기도 합니다.

이번 글에서는 AMD GPU 백엔드에서 TDM(Tensor Data Mover) 디스크립터를 생성할 때 발생했던 비효율적인 코드 생성(Suboptimal Codegen) 문제를 해결한 PR을 분석합니다. 핵심은 LLVM의 InstCombine 패스가 특정 비교 연산 패턴을 만났을 때, 이를 벡터 연산 장치(VALU) 전용 명령어로 변환하면서 발생하는 불필요한 레지스터 복사 비용을 제거하는 것입니다.

문제의 배경: SGPR vs VGPR

AMD GPU 아키텍처(CDNA/RDNA)에는 두 가지 주요 레지스터 타입이 있습니다.

  1. SGPR (Scalar General Purpose Register): 모든 워크아이템이 공유하는 스칼라 연산용 레지스터.
  2. VGPR (Vector General Purpose Register): 각 워크아이템마다 개별적으로 가지는 벡터 연산용 레지스터.

TDM(Tensor Data Mover)은 메모리 간 데이터를 효율적으로 옮기는 하드웨어 유닛이며, 이를 제어하는 Descriptor 정보는 보통 모든 스레드에 공통적인 값이기 때문에 SGPR에 위치하는 것이 훨씬 효율적입니다. 만약 이 값이 VGPR로 넘어가게 되면, 다시 SGPR로 가져오기 위해 v_readfirstlane_b32와 같은 추가적인 명령어가 필요하며 이는 곧 오버헤드가 됩니다.

코드 분석: 무엇이 바뀌었나?

1. TDMUtility.cpp: 텐서 쉐이프 클램핑 로직 변경

기존 코드는 텐서의 크기(shape)에서 오프셋(offset)을 뺀 값이 유효한 범위 내에 있는지 확인하기 위해 icmp_ule(Unsigned Less or Equal)를 사용했습니다.

[Before]

// Update tensor shapes based on offset
for (size_t i = 0; i < numDims; ++i) {
  auto diff = b.sub(tensorShape[i], offset[i]);
  Value inBounds = b.icmp_ule(diff, tensorShape[i]);
  tensorShape[i] = b.select(inBounds, diff, b.i32_val(0));
}

위의 select(icmp_ule(sub(a, b), a), sub(a, b), 0) 패턴은 LLVM의 InstCombine 패스에서 특정 최적화 타겟이 됩니다. LLVM은 이 패턴을 보고 "아, 이건 부호 없는 뺄셈의 포화 연산(Saturation)이구나"라고 판단하여 더 단순한 형태로 합치려 시도합니다. 문제는 이렇게 합쳐진 형태가 AMDGPU 백엔드에서 VALU(Vector ALU) 명령어로만 매핑된다는 점입니다.

결과적으로 스칼라 연산으로 충분한 작업이 벡터 연산으로 수행되고, 이를 다시 디스크립터로 쓰기 위해 SGPR로 복사하는 v_readfirstlane 명령어가 생성되었습니다.

[After]

// Update tensor shapes based on offset
Value zero = b.i32_val(0);
for (size_t i = 0; i < numDims; ++i) {
  // Clamp in this specific way to avoid suboptimal codegen by LLVM.
  Value diff = b.sub(tensorShape[i], offset[i]);
  Value clamped = b.smax(diff, zero);
  Value isNeg = b.icmp_slt(offset[i], zero);
  tensorShape[i] = b.select(isNeg, zero, clamped);
}

개선된 코드는 smax(Signed Max)와 icmp_slt(Signed Less Than)를 조합하여 동일한 클램핑 로직을 수행합니다. 이 방식은 LLVM InstCombine이 VALU 전용 패턴으로 오인하지 않게 유도하며, 결과적으로 전체 연산이 SALU(Scalar ALU) 내에서 완결되도록 합니다.

2. Gather/Scatter 디스크립터 적용

동일한 문제가 Gather/Scatter 연산을 위한 디스크립터 생성 로직에서도 발견되어 수정되었습니다.

[Before]

// Adjust column tensor shape for OOB handling
tensorShape[1] = b.smax(b.i32_val(0), b.sub(tensorShape[1], globalColOffset));

[After]

{
  Value zero = b.i32_val(0);
  Value diff = b.sub(tensorShape[1], globalColOffset);
  Value clamped = b.smax(diff, zero);
  Value isNeg = b.icmp_slt(globalColOffset, zero);
  tensorShape[1] = b.select(isNeg, zero, clamped);
}

단순히 smax(0, sub)를 쓰는 대신, 오프셋이 음수인 경우를 명시적으로 select로 처리함으로써 컴파일러가 최적화 과정에서 의도치 않은 경로로 빠지지 않도록 방어적인 코드를 작성했습니다.

왜 이게 좋은 최적화인가?

1. 불필요한 명령어 제거 (Instruction Reduction)

기존 방식에서는 VALU 연산 + v_readfirstlane이 발생했습니다. v_readfirstlane은 벡터 레지스터의 첫 번째 레인 값을 스칼라 레지스터로 복사하는 명령어로, 실행 유닛 간의 데이터 이동을 수반하므로 일반적인 산술 연산보다 비용이 큽니다. 이 PR을 통해 이러한 데이터 이동 자체가 사라졌습니다.

2. 레지스터 압박(Register Pressure) 완화

값이 VGPR에 머물게 되면, 커널이 사용하는 전체 VGPR 개수가 늘어날 수 있습니다. GPU에서 VGPR 사용량은 점유율(Occupancy)과 직결되는데, 스칼라 연산으로 처리할 수 있는 부분을 SGPR로 돌림으로써 VGPR 자원을 아끼고 더 높은 병렬성을 확보할 수 있는 여지를 만듭니다.

3. 컴파일러 유도(Compiler Hinting)의 정석

고수준 언어에서 작성한 코드가 저수준 어셈블리로 어떻게 변하는지 이해하고, 컴파일러의 중간 단계(IR) 최적화가 하드웨어 특성과 충돌할 때 이를 우회하는 코드를 작성하는 것은 시니어 엔지니어의 핵심 역량입니다. 이 PR은 단순히 '결과가 같은 코드'를 쓴 것이 아니라, '컴파일러가 올바른 선택을 하도록 가이드하는 코드'를 작성했다는 점에서 훌륭한 사례입니다.

결론

이번 변경사항은 코드의 외형상으로는 조금 더 복잡해 보일 수 있지만, 실제 생성되는 어셈블리 관점에서는 훨씬 효율적입니다. LLVM과 같은 강력한 컴파일러를 사용할 때도, 우리는 항상 "컴파일러가 내 의도대로 하드웨어 자원을 활용하고 있는가?"를 의심하고 검증해야 합니다. 특히 GPU와 같이 아키텍처 특성이 강한 환경에서는 이러한 미세한 조정이 모여 전체 성능의 차이를 만듭니다.

요약

  • 문제: LLVM InstCombine이 특정 클램핑 패턴을 VALU 명령어로 변환하여 SGPR 복사 오버헤드(v_readfirstlane) 발생.
  • 해결: smax와 select 조합의 명시적 로직으로 변경하여 SALU 연산 유도.
  • 효과: 불필요한 명령어 제거 및 레지스터 효율성 향상.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글