[Triton] AMD GFX9에서 AsyncCopy shared layout order 수정

PR 링크: triton-lang/triton#9373 상태: Merged | 변경: +59 / -5

들어가며

AMD GFX9 아키텍처에서 async copy는 Global Memory에서 LDS(Local Data Share)로 데이터를 비동기적으로 복사한다. 이때 shared memory의 layout order가 잘못 설정되면 coalesced direct-to-LDS 쓰기가 불가능해져 성능이 크게 저하된다. 이 PR은 layout order 결정 시 사용하는 contiguity 계산 방식과 vector size clamping을 수정한다.

핵심 코드 분석

Before: 잘못된 contiguity 계산

auto contig = llEnc.getElemsPerThread(srcTy.getShape());
SetVector<unsigned> orderSet;
if (contig[regOrder[0]] > 1)
  orderSet.insert(regOrder[0]);
orderSet.insert(threadOrder.begin(), threadOrder.end());
order = orderSet.takeVector();

getElemsPerThread는 스레드당 전체 요소 수를 반환하므로, 실제 메모리 연속성(contiguity)과 다를 수 있다.

After: 올바른 contiguity + 하드웨어 벡터 크기 clamp

auto regContig = llEnc.getContigPerThread()[regOrder[0]];
unsigned elemBitWidth = srcTy.getElementType().getIntOrFloatBitWidth();
unsigned finalRegContig =
    fitToValidDirectToLdsVecSize(regContig, elemBitWidth, targetInfo);
// 하드웨어가 지원하는 벡터 크기인 경우에만 order 변경
if (finalRegContig > 0) {
  if (finalRegContig > 1)
    orderSet.insert(regOrder[0]);
  orderSet.insert(threadOrder.begin(), threadOrder.end());
  order = orderSet.takeVector();
}

변경 포인트:

1. getElemsPerThread -> getContigPerThread: 실제 연속 접근 가능한 요소 수를 사용
2. fitToValidDirectToLdsVecSize로 하드웨어가 지원하는 벡터 크기로 clamp
3. 유효한 벡터 크기가 없으면 order 변경을 하지 않음

왜 이게 좋은가

1. 정확한 contiguity 정보 사용: getContigPerThread는 메모리에서 실제로 연속인 요소 수를 반환하므로, 올바른 vectorization 판단이 가능하다.
2. 하드웨어 호환성 보장: fitToValidDirectToLdsVecSize를 통해 GFX9가 지원하는 direct-to-LDS 벡터 크기만 허용한다.
3. 안전한 fallback: 유효한 벡터 크기가 없으면 order를 변경하지 않아, 최소한 기존 동작을 보장한다.

정리

이 PR은 AMD GFX9에서 async copy의 shared layout order 결정 로직을 수정했다. getContigPerThread로 정확한 연속성 정보를 사용하고, 하드웨어가 지원하는 벡터 크기로 clamp하여 coalesced direct-to-LDS 쓰기를 보장한다. 관련 lit 테스트도 추가되었다.

참고 자료

• triton-lang/triton#9373

---

이 글은 AI를 활용하여 PR의 핵심 변경사항을 분석하고 정리한 것입니다. 실제 코드의 맥락은 원본 PR을 참고해 주세요.