[Triton] gfx950에서 PaddedLayout + AsyncCopy 파이프라이닝 지원

PR 링크: triton-lang/triton#8365 상태: Merged | 변경: +300 / -40

들어가며

AMD CDNA GPU에서 matmul 성능을 높이려면 shared memory의 bank conflict를 줄여야 한다. Padded layout은 행 사이에 padding을 넣어 bank conflict를 방지하는 기법이다. 기존에는 padded layout을 사용하면 AsyncCopy(direct-to-LDS load)를 사용할 수 없었는데, 이 PR은 둘을 함께 사용할 수 있도록 파이프라인 lowering을 확장한다.

핵심 코드 분석

Shared encoding 결정 로직 확장

Before:

tempAttr = ttg::SwizzledSharedEncodingAttr::get(
    loadedValue.getContext(), dotOpEnc, srcTy.getShape(),
    sharedOrder, ctaLayout, bitWidth, /*needTrans=*/false);

항상 SwizzledSharedEncoding만 사용했다.

After:

bool canUseAsyncCopy = false;
if (useAsyncCopy && isa<tt::LoadOp>(loadOp)) {
    canUseAsyncCopy = canBeConvertedToAsyncLoad(
        2, cast<tt::LoadOp>(loadOp), {}, axisInfoAnalysis, targetInfo);
}
tempAttr = composePaddedLayout(targetInfo, dotOpEnc, srcTy,
                               sharedOrder, canUseAsyncCopy);
if (!tempAttr) {
    tempAttr = ttg::SwizzledSharedEncodingAttr::get(...);
}

AsyncCopy가 가능한지 먼저 확인하고, 가능하면 padded layout을 시도한다. padded layout이 적합하지 않으면 swizzled로 fallback한다.

AsyncCopy 호환성 검사

bool canBeConvertedToAsyncLoad(unsigned numBuffers, tt::LoadOp loadOp,
                               ttg::SharedEncodingTrait sharedEnc, ...) {
    if (paddedEnc) {
        sharedLayout = paddedEnc.getLinearComponent();
    } else {
        sharedLayout = triton::gpu::toLinearLayout(srcShape, sharedEnc);
    }
    auto regToSharedLayout = regLayout.invertAndCompose(sharedLayout);
    // vec size가 타겟이 지원하는 범위인지 확인
}

왜 이게 좋은가

1. 성능 향상: padding + async copy 조합으로 bank conflict 없이 빠른 로드 가능
2. 자동 최적화: 컴파일러가 최적의 shared memory 레이아웃을 자동 선택
3. 안전한 fallback: padded layout이 불가능하면 swizzled로 자동 전환

정리

Shared memory 최적화는 GPU matmul 성능의 핵심이다. padding과 async copy를 결합할 수 있게 됨으로써, AMD CDNA 아키텍처에서 Triton의 matmul 성능이 한 단계 더 개선될 수 있다.

참고 자료

• PR #8365

---

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