[triton] AMD ds_read_tr 명령어 제한 완화로 더 유연한 레이아웃 지원

PR 링크: triton-lang/triton#8442 상태: Merged | 변경: +349 / -263

들어가며

AMD GPU의 ds_read_tr 명령어는 shared memory(LDS)에서 데이터를 읽으면서 동시에 레인 간 전치(transpose)를 수행합니다. 기존에는 DotOperand 레이아웃에서만 사용 가능했지만, 실제로는 임의의 linear layout에서도 충분히 활용할 수 있습니다. 이 PR은 불필요한 제한을 제거하고 API를 일반화합니다.

핵심 코드 분석

Before - 고정된 파라미터

LinearLayout chooseDsReadB64TrLayout(Attribute enc, ArrayRef<int64_t> shape,
                                     int32_t elemBitWidth);
// 64비트, 16레인 그룹으로 하드코딩
unsigned needContigReg = 64 / bitwidth;
unsigned needContigLane = 16;

After - 파라미터화된 유연한 API

std::optional<LinearLayout>
chooseDsReadTrLayout(Attribute enc, ArrayRef<int64_t> shape,
                     int32_t elemBitWidth, unsigned instBitWidth,
                     unsigned numLanesInShuffleGroup);

// 유효하지 않은 레이아웃에 대해 std::nullopt 반환
if (ldsTransLayout.getInDimSizeLog2(kRegister) < numRegBases ||
    ldsTransLayout.getInDimSizeLog2(kLane) < numLaneBases) {
  return std::nullopt;
}

왜 이게 좋은가

1. 일반화: instBitWidth와 numLanesInShuffleGroup을 파라미터로 받아 향후 새로운 ds_read 변형을 쉽게 지원합니다.
2. 안전한 실패: std::optional을 반환하여 유효하지 않은 레이아웃에서 크래시 대신 우아하게 fallback합니다.
3. 임의 layout 지원: DotOperand뿐 아니라 linear layout에서도 ds_read_tr을 활용할 수 있어 최적화 기회가 확대됩니다.

정리

ds_read_tr 활용 범위를 넓히면서도 API를 더 안전하고 유연하게 만든 개선입니다. 특히 std::optional 반환을 통한 안전한 실패 처리는 컴파일러 코드에서 본받을 만한 패턴입니다.

참고 자료

• PR #8442

---

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