본문으로 건너뛰기

[triton] AMD GPU에서 불필요한 워프 로드를 제거하여 성능을 최적화한 Triton PR 분석

PR 링크: triton-lang/triton#10332 상태: Merged | 변경: +160 / -50

들어가며

최신 GPU 아키텍처는 복잡한 연산을 효율적으로 처리하기 위해 다양한 최적화 기법을 활용합니다. 특히, AMD GPU의 경우 특정 조건에서 발생하는 불필요한 데이터 로드를 줄이는 것이 성능 향상의 핵심이 될 수 있습니다. 이번 PR (triton-lang/triton#1234)은 AMD GPU의 gfx9 아키텍처에서 BufferLoadToLocal 및 AsyncCopyGlobalToLocal 연산 시, 워프(warp) 단위로 동일한 데이터를 중복 로드하는 경우를 식별하고 이를 효과적으로 회피하는 방법을 제시합니다.

기존에는 이러한 중복 로드를 방지하기 위해 스레드 단위의 조건 분기(branch)를 사용했습니다. 하지만 이 방식은 LLVM 코드 생성 과정에서 VGPR(Vector General Purpose Register) 압력을 증가시키는 요인이 될 수 있습니다. 본 PR은 워프 내의 모든 스레드가 동일한 조건(warp-uniform predicate)을 만족할 경우, 분기 대신 공유 메모리(LDS - Local Data Share)에 접근할 때 유효하지 않은 주소(out-of-bounds address)를 사용하여 해당 워프 전체의 로드를 마스킹하는 새로운 최적화 기법을 도입합니다.

이 PR은 특히 작은 크기의 텐서를 로드할 때 발생하는 불필요한 로드를 제거함으로써, 일부 mxfp gemm 커널에서 VGPR 사용량을 최대 35%까지 감소시키는 효과를 가져옵니다. 이는 GPU 커널의 전반적인 효율성을 높이는 중요한 개선입니다.

코드 분석

이번 PR은 주로 test/Conversion/amd/ 디렉토리 내의 MLIR 테스트 파일과 third_party/amd/lib/TritonAMDGPUToLLVM/LoadStoreOpToLLVM.cpp 파일의 컴파일러 로직을 수정했습니다.

1. test/Conversion/amd/async_ops_to_llvm.mlir 및 test/Conversion/amd/buffer_load_to_local_to_llvm.mlir

이 파일들은 새로운 최적화 기법이 올바르게 적용되는지를 검증하기 위한 MLIR 테스트 케이스를 추가하거나 수정합니다. 핵심은 워프 단위로 동일한 조건이 만족될 때, 기존의 스레드 단위 분기 대신 공유 메모리 주소를 유효 범위를 벗어나게 설정하는(OOB address) 패턴이 생성되는지 확인하는 것입니다.

async_copy_warp_uniform_thread_pred 테스트 케이스 (추가됨):

이 테스트는 4개의 워프가 각각 64개의 요소를 로드하는 시나리오를 가정합니다. 이 경우, 일부 워프는 실제로는 데이터를 로드할 필요가 없지만 기존 방식으로는 분기 처리가 필요했습니다. PR 이후에는 다음과 같은 코드가 생성됩니다:

Before (기존 로직 예상):

// ... (분기 로직) ...
// rocdl.global.load.async.lds ...

After (PR 적용 후):

// CHECK: %[[OOB_I32:.*]] = llvm.mlir.constant(2147483647 : i32) : i32
// CHECK: %[[OOB_PTR:.*]] = llvm.inttoptr %[[OOB_I32]] : i32 to !llvm.ptr<3>
// CHECK: %[[PRED_ADDR:.*]] = llvm.select {{.*}}, {{.*}}, %[[OOB_PTR]] : i1, !llvm.ptr<3>
// CHECK: rocdl.global.load.async.lds {{.*}}, %[[PRED_ADDR]], {{.*}}

llvm.select를 사용하여 조건이 거짓일 경우, 공유 메모리 주소(shmemAddr) 대신 OOB_PTR (매우 큰 값으로 변환된 i32 주소)을 선택합니다. rocdl.global.load.async.lds는 이 유효하지 않은 주소로 인해 해당 워프의 로드를 무시하게 됩니다. 이는 분기문 자체를 제거하여 LLVM의 코드 생성 부담을 줄입니다.

buffer_load_warp_uniform_thread_pred 테스트 케이스에서도 유사한 패턴이 확인됩니다.

async_copy_warp_uniform_thread_pred_with_per_lane_mask 테스트 케이스 (추가됨):

이 테스트는 워프 단위 조건은 만족하지만, 스레드별 마스크(per-lane mask)가 존재하는 경우를 다룹니다. 이 경우에는 OOB 주소 최적화가 적용되지 않고, 기존처럼 스레드 단위의 분기(llvm.cond_br)가 유지됩니다. 이는 hasOther 조건이 존재할 때 분기가 필요한 이유를 명확히 보여줍니다.

// CHECK: %[[PRED:.*]] = llvm.and {{.*}} : i1
// CHECK: llvm.cond_br %[[PRED]], ^[[LOAD_BLOCK:bb[0-9]+]]
// CHECK-NEXT: ^[[LOAD_BLOCK]]:
// CHECK-NEXT: rocdl.global.load.async.lds

2. third_party/amd/lib/TritonAMDGPUToLLVM/LoadStoreOpToLLVM.cpp

이 파일은 Triton IR을 LLVM IR로 변환하는 로직을 포함하며, 이번 PR의 핵심 최적화가 구현된 곳입니다.

BufferLoadToLocalOpConversion::execute 함수 수정:

기존에는 emitRedundantThreadPredicateNonNull 함수를 호출하여 스레드 단위의 조건 분기를 생성했습니다. 이제는 이 함수 호출과 더불어, 해당 조건이 워프 단위로 균일한지(isThreadPredWarpUniform)를 확인하는 로직이 추가되었습니다.

// Before (기존 로직 일부)
Value threadPred = emitRedundantThreadPredicateNonNull(
 getFreeVariableMasks(ptrType), rewriter, loc, targetInfo);

// After (PR 적용 후)
auto freeVarMasks = getFreeVariableMasks(ptrType);
Value threadPred = emitRedundantThreadPredicateNonNull(
 freeVarMasks, rewriter, loc, targetInfo);
bool isThreadPredWarpUniform =
 isRedundantThreadPredWarpUniform(freeVarMasks, rewriter.getContext());

isRedundantThreadPredWarpUniform 함수는 lane 및 reg 마스크가 모두 0인지 확인하여 워프 균일성을 판단합니다. 이는 해당 조건이 스레드별로 다르지 않고 워프 전체에 동일하게 적용됨을 의미합니다.

이어서, emitBufferLoadLds 람다 함수 내부에서 isThreadPredWarpUniform 플래그와 hasOther (다른 값을 로드하는지 여부)를 사용하여 분기 로직을 결정합니다.

// Before (기존 로직 일부)
Value cond =
 hasOther ? b.and_(threadPred, maybeSwizzledMaskElem) : threadPred;

auto [loadBlock, afterLoadBlock] = emitBranch(rewriter, loc, cond);

auto bufferLoadToLds = bufferEmitter.emitLoadToLds(
 vecTy, vecBytesVal, rsrcDesc, offsetElem, shmemAddr,
 hasOther ? b.true_val() : maybeSwizzledMaskElem, op.getCache());

// After (PR 적용 후)
if (isThreadPredWarpUniform && !hasOther) {
 Value predicatedAddress =
 selectLdsAddressForPredicate(b, threadPred, shmemAddr);
 auto bufferLoadToLds = bufferEmitter.emitLoadToLds(
 vecTy, vecBytesVal, rsrcDesc, offsetElem, predicatedAddress,
 maybeSwizzledMaskElem, op.getCache());
 if (targetInfo.requiresAliasInfoForAsyncOps())
 AMD::addAsyncCopyAliasScope(bufferLoadToLds);
} else {
 // ... (기존 분기 로직 유지) ...
}

selectLdsAddressForPredicate 함수는 threadPred가 거짓일 때 shmemAddr 대신 OOB 주소를 반환합니다. 이 OOB 주소를 사용하여 bufferEmitter.emitLoadToLds를 호출하면, 해당 워프의 로드가 공유 메모리에 쓰여지지 않게 됩니다. 이는 불필요한 글로벌 메모리 접근 및 공유 메모리 쓰기를 방지합니다.

리뷰어 antiagainst가 제안한 캡처 리스트 간소화([&])는 코드 가독성을 높이기 위한 좋은 제안이며, 실제 코드에서는 이러한 스타일 개선이 이루어졌을 수 있습니다.

왜 이게 좋은가?

성능 향상

이 PR의 가장 큰 장점은 VGPR 사용량 감소입니다. 기존의 스레드 단위 조건 분기는 각 스레드마다 조건 검사 및 분기 로직을 위한 레지스터를 필요로 합니다. 특히 워프 내 모든 스레드가 동일한 조건을 만족하는 경우, 이러한 분기 로직은 불필요하게 많은 레지스터를 소모하게 됩니다.

새로운 기법은 워프 단위 조건이 균일할 때 분기 자체를 제거하고, 유효하지 않은 메모리 주소를 사용하여 로드를 무시하도록 합니다. 이는 LLVM 컴파일러가 더 적은 VGPR을 사용하여 코드를 생성하도록 돕습니다. PR 설명에 따르면, 일부 mxfp gemm 커널에서 VGPR 사용량이 최대 35%까지 감소하는 효과를 보였습니다. VGPR 사용량 감소는 더 많은 스레드를 동시에 실행할 수 있게 하거나, 더 복잡한 연산을 레지스터에 담을 수 있게 하여 전반적인 커널 성능을 향상시킵니다.

LLVM 코드 생성 최적화

GPU 컴파일러에서 조건 분기는 종종 복잡한 코드 생성을 유발합니다. 분기 예측 실패 시 성능 저하가 발생할 수 있으며, LLVM과 같은 컴파일러는 이러한 분기를 효율적으로 처리하기 위해 많은 노력을 기울입니다. 분기 자체를 제거하고 데이터 로드 메커니즘을 변경함으로써, LLVM의 코드 생성기가 더 단순하고 효율적인 코드를 생성할 수 있게 됩니다.

일반적인 교훈

  1. 하드웨어 특성 활용: AMD GPU의 gfx9 아키텍처가 특정 조건(워프 균일 조건)에서 공유 메모리 주소 유효성 검사를 통해 로드를 마스킹할 수 있다는 점을 적극적으로 활용했습니다.
  2. 분기 제거의 가치: 불필요한 조건 분기는 성능 병목의 주요 원인 중 하나입니다. 가능한 경우, 분기 대신 데이터 경로를 변경하거나 하드웨어 기능을 활용하여 분기를 제거하는 것이 유리합니다.
  3. 조건적 최적화: 모든 경우에 적용되는 최적화보다는, 특정 조건(예: 워프 균일성, hasOther 플래그의 부재)에서만 적용되는 최적화가 더 효과적일 수 있습니다. 이는 코드의 복잡성을 관리하면서도 최대의 성능 이득을 얻는 방법입니다.
  4. 테스트의 중요성: 새로운 최적화 기법이 올바르게 적용되는지, 그리고 기존 기능에 영향을 미치지 않는지를 검증하기 위해 MLIR 테스트 케이스를 추가하는 것이 매우 중요합니다.

References

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글