[triton] AMD GFX1250에서 TDM Software Pipelining 지원

PR 링크: triton-lang/triton#9302 상태: Merged | 변경: +388 / -33

들어가며

AMD GFX1250은 Tensor Descriptor Memory(TDM) 명령어를 지원하여, descriptor 기반 비동기 메모리 복사가 가능합니다. 이 PR은 TDM 복사를 Triton의 software pipelining 프레임워크에 통합하여, matmul 루프에서 데이터 로딩과 연산을 오버랩시킵니다.

핵심 코드 분석

1. TDMCopyChainOps 구조체

struct TDMCopyChainOps {
  triton::amdgpu::AsyncTDMCopyGlobalToLocalOp copyOp;
  ttg::AsyncCommitGroupOp commitOp;
  triton::amdgpu::AsyncTDMWait waitOp;
  ttg::LocalLoadOp maybeLocalLoadOp;
};

기존 StreamCopyChainOps/AsyncCopyChainOps와 동일한 구조로 TDM 체인을 정의합니다.

2. TDM 비동기 복사 생성

TDMCopyChainOps createTDMAsyncCopy(tt::DescriptorLoadOp loadOp, Value alloc,
                                   Value extractIdx) {
  Value pred = arith::ConstantIntOp::create(builder, loc, 1, 32);
  auto viewLoad = triton::createSingleBufferView(builder, alloc, extractIdx);
  auto copyOp = triton::amdgpu::AsyncTDMCopyGlobalToLocalOp::create(
      builder, loc, loadOp.getDesc(), loadOp.getIndices(), viewLoad, pred);
  auto commitOp = ttg::AsyncCommitGroupOp::create(builder, loc, copyOp->getResult(0));
  auto waitOp = triton::amdgpu::AsyncTDMWait::create(builder, loc, commitOp->getResult(0), 0);
  // ...
}

3. WMMA encoding과 padded shared layout 지원

} else if (auto dotEnc = getDotEncoding<ttg::AMDWmmaEncodingAttr>(
               userResult, &opIdx, &vecSize)) {
  tempAttr = dotEnc.composeSharedLayoutForOperand(
      cgaLayout, opIdx, srcTy.getShape(), order, vecSize, bitWidth,
      /*needTrans=*/false);
}

GFX1250의 WMMA 명령어에 맞는 padded shared encoding을 자동 추론합니다.

왜 이게 좋은가

• 성능 향상: 데이터 로딩과 연산이 오버랩되어 메모리 대기 시간이 은닉됩니다.
• 기존 프레임워크 통합: StreamOpVariant에 TDM을 추가하여 기존 pipelining 로직을 재활용합니다.
• 새 하드웨어 활용: GFX1250의 TDM 명령어를 직접 활용합니다.

정리

AMD GFX1250의 TDM 비동기 복사를 software pipelining에 통합한 PR입니다. TDMCopyChainOps 구조체, padded shared encoding 지원, WMMA layout 추론이 핵심 변경사항입니다.

참고 자료

• triton-lang/triton#9302

---

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