[triton] AMD GFX9 AsyncCopy를 위한 Padded Layout 선택 확장

PR 링크: triton-lang/triton#9544 상태: Merged | 변경: +71 / -30

들어가며

AMD CDNA4 GPU에서 async copy를 사용할 때 shared memory의 bank conflict를 줄이기 위한 padded layout 선택이 기존에는 16비트(f16) 데이터 타입에만 적용되었습니다. 이 PR은 8비트(f8) 타입과 kWidth 16까지 확장하고, ds_read 명령어 종류에 따른 bank conflict 분석을 정교화합니다.

핵심 코드 분석

Before - 16비트만 지원

if (elemByteWidth != 2) {
  return {};  // f8 등은 padded layout 미적용
}
if (!llvm::is_contained({4, 8}, kWidth)) {
  return {};
}

After - 8비트 및 넓은 kWidth 지원

if (!llvm::is_contained({1, 2}, elemByteWidth)) {
  return {};  // 8비트, 16비트 모두 지원
}
if (!llvm::is_contained({4, 8, 16}, kWidth)) {
  return {};
}

// ds_read 명령어 종류에 따른 bank 수 결정
bool useDsReadB128 = isKContig && kWidthBytes == 16;
bool useDsReadB64Tr = !isKContig && kWidthBytes >= 8;
unsigned numberOfBanks = (useDsReadB128 || useDsReadB64Tr) ? 64 : 32;

왜 이게 좋은가

1. f8 지원: FP8 학습/추론에서 async copy + padded layout 조합의 성능이 개선됩니다.
2. 정교한 분석: ds_read_b128, ds_read_b64_tr, 32-bank 모드를 구분하여 최적의 padding을 계산합니다.
3. x-way conflict 허용: 2-way conflict까지 허용하는 휴리스틱으로 불필요한 padding 실패를 방지합니다.

정리

CDNA4의 async copy에서 padded layout 적용 범위를 8비트 타입으로 확장하고, bank conflict 분석을 ds_read 명령어 종류별로 정교화한 개선입니다.

참고 자료

• PR #9544

---

이 글은 AI(Claude)의 도움을 받아 작성되었으며, 원본 PR의 코드 변경 사항을 기반으로 분석한 내용입니다.