[triton] AMD: padded shared layout을 더 작은 block size에도 적용하여 bank conflict 제거

PR 링크: triton-lang/triton#9074 상태: Merged | 변경: +124 / -97

들어가며

AMD GPU의 LDS(Local Data Share)에서 bank conflict는 성능에 큰 영향을 미칩니다. 기존 Triton의 padded shared layout은 16KB 이상의 블록에서만 동작했는데, 이번 PR은 wrap-around 패딩 기법을 활용하여 32x64와 같은 작은 블록에서도 bank conflict 없는 레이아웃을 생성할 수 있도록 개선합니다.

핵심 코드 분석

Before: 16KB 최소 크기 제한

// Utility.cpp - 기존 코드
unsigned bitWidth = getIntOrFloatOrPtrBitWidth(srcTy.getElementType());
unsigned elemByteWidth = std::max(bitWidth / 8u, 1u);
auto loadBytes = shape[0] * shape[1] * elemByteWidth;
if (loadBytes < 16384) {  // 16KB 미만이면 padded layout 불가
  return {};
}
// 최소 16행 필요
if (std::min(shape[0], shape[1]) < 16) {
  return {};
}

After: wrap-around을 활용한 동적 계산

// Utility.cpp - 개선된 코드
unsigned padding = 0;
if (isKContig) {
  padding = mfmaNonKDim == 16 ? (kWidth * 2) : kWidth;
} else {
  padding = mfmaNonKDim == 16 ? 16 : 32;
}
constexpr unsigned vecSize = 8;
unsigned contigLanes = contigDim / vecSize;
unsigned wrap = std::min(contigDim, 128u) / padding;
unsigned requiredDim = warpSize / contigLanes * wrap;
if (nonContigDim < requiredDim) {
  return {};
}
// 16행 wrap이 가능하면 최적 wrap 사용
bool useBestWrap = false;
unsigned bestWrap = 16;
if (nonContigDim >= warpSize / contigLanes * bestWrap && bestWrap > wrap) {
  useBestWrap = true;
  wrap = bestWrap;
}

핵심 아이디어는 행을 wrap-around 방식으로 재배치하는 것입니다. 예를 들어 32x64 블록에서:

r0,  r4,  r8,  r12, r16, r20, r24, r28
pad, r1,  r5,  r9,  r13, r17, r21, r25
r29, pad, r2,  r6,  r10, r14, r18, r22
r26, r30, pad, r3,  r7,  r11, r15, r19

왜 이게 좋은가

1. 작은 블록 지원: 32x64 같은 작은 타일에서도 bank conflict 없는 접근이 가능해져 더 다양한 커널 구성을 효율적으로 지원합니다.
2. 동적 padding 계산: warpSize와 실제 텐서 차원에 기반한 동적 계산으로 하드코딩된 제한을 제거했습니다.
3. 하위 호환성: 기존 16KB 이상 블록의 동작은 그대로 유지됩니다.

정리

GPU 공유 메모리의 bank conflict는 실질적인 성능 병목입니다. 이 PR은 고정된 크기 제한을 제거하고, 수학적으로 bank 패턴의 wrap-around 주기를 계산하여 더 넓은 범위의 블록 크기에서 최적의 메모리 접근 패턴을 보장합니다.

참고 자료

• triton-lang/triton#9074

---

이 글은 AI(Claude)의 도움을 받아 작성되었으며, PR의 실제 diff를 기반으로 분석한 내용입니다.