[Triton] swizzling=0 matrix descriptor 지원과 WGMMA lowering 일반화

PR 링크: triton-lang/triton#8378 상태: Merged | 변경: +81 / -72

들어가며

NVIDIA Hopper GPU의 WGMMA(Warp Group Matrix Multiply-Accumulate) 명령어는 shared memory의 matrix descriptor를 통해 operand를 참조한다. 기존에는 swizzling 값이 0인 경우를 제대로 지원하지 않았고, WGMMA의 operand 검증에서 SharedLinearEncodingAttr을 허용하지 않았다.

핵심 코드 분석

WGMMA 검증 확장 (Ops.cpp)

Before:

if (!isa<NVMMASharedEncodingAttr>(getB().getType().getEncoding()))
    return emitOpError("WGMMA B operand must have NVMMA shared layout");

After:

if (!isa<NVMMASharedEncodingAttr, SharedLinearEncodingAttr>(
        getB().getType().getEncoding()))
    return emitOpError("WGMMA B operand must have NVMMA shared layout");

SharedLinearEncodingAttr도 유효한 WGMMA operand 인코딩으로 인정한다.

MNdim 파라미터 변경 (MMAHelpers.h)

Before:

static DotOpMmaSmemLoader build(
    ..., int mmaVersion, bool isFp4 = false,
    std::optional<RankedTensorType> mmaTy = std::nullopt,
    std::optional<unsigned> MNdim = std::nullopt);

After:

static DotOpMmaSmemLoader build(
    ..., unsigned MNdim, int mmaVersion, bool isFp4 = false,
    std::optional<RankedTensorType> mmaTy = std::nullopt);

MNdim을 optional에서 필수 파라미터로 변경하여 타입 안전성을 높였다.

왜 이게 좋은가

1. swizzling=0 지원: 단순한(비 swizzled) shared memory 레이아웃에서도 WGMMA 사용 가능
2. 레이아웃 유연성: LinearLayout 기반 인코딩이 WGMMA와 호환
3. API 명확성: optional이던 MNdim을 필수로 만들어 실수 방지

정리

GPU 컴파일러에서 레이아웃 인코딩의 다양성을 지원하는 것은 성능 최적화의 기반이다. swizzling=0은 단순한 케이스지만, 이를 명시적으로 지원해야 LinearLayout 기반의 새로운 인코딩 전략을 활용할 수 있다.

참고 자료

• PR #8378

---

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