[triton] AMD gfx1250에서 Async Copy와 TDM 경로의 Padded Layout 휴리스틱 통합

PR 링크: triton-lang/triton#9747 상태: Merged | 변경: +123 / -41

들어가며

AMD gfx1250 GPU에서 shared memory의 bank conflict를 줄이기 위한 padded layout 선택은 async copy 경로와 TDM(Tensor Data Movement) 경로에서 각각 다른 함수를 사용했습니다. 이 PR은 두 경로의 휴리스틱을 composePaddedLayout이라는 단일 함수로 통합합니다.

핵심 코드 분석

Before - 별도 함수 호출

// async copy 경로
tempAttr = composePaddedLayout(targetInfo, dotOpEnc, srcTy, sharedOrder,
                               canUseAsyncCopy);

// TDM 경로
tempAttr = getPaddedEncodingForDotOp(
    loadedValue.getContext(), dotOpEnc.getOpIdx(), srcTy.getShape(),
    sharedOrder, cgaLayout, bitWidth, dotOpEnc.getKWidth(), targetInfo);

After - 통합 함수

// 두 경로 모두 같은 함수 사용
tempAttr = composePaddedLayout(targetInfo, dotOpEnc.getOpIdx(),
                               dotOpEnc.getKWidth(), srcTy, sharedOrder,
                               dotOpEnc, canUseAsyncCopy);

// TDM 경로도 동일
tempAttr = composePaddedLayout(targetInfo, dotOpEnc.getOpIdx(),
                               dotOpEnc.getKWidth(), srcTy, sharedOrder);

왜 이게 좋은가

1. 코드 통합: 두 경로에서 padding 값이 달라질 위험을 제거했습니다.
2. 인터페이스 간소화: getPaddedEncodingForDotOp의 많은 파라미터를 srcTy에서 자동 추출하도록 개선했습니다.
3. MLIR 테스트 추가: f16과 f8E4M3FN에 대한 기대값을 명시한 FileCheck 테스트가 추가되었습니다.

정리

shared memory padding 휴리스틱의 코드 중복을 제거한 정리 작업입니다. async copy와 TDM이 같은 padding 로직을 공유하여 일관된 bank conflict 회피가 보장됩니다.

참고 자료

• PR #9747

---

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