[Triton] gfx1250에서 TDM Store 지원 추가

PR 링크: triton-lang/triton#8392 상태: Merged | 변경: +611 / -224

들어가며

AMD gfx1250 아키텍처의 **TDM(Tensor Data Mover)**은 텐서 디스크립터를 사용하여 global memory와 shared memory 간 비동기 데이터 전송을 수행한다. 기존에는 global-to-shared load만 지원했는데, 이 PR은 shared-to-global store도 추가한다.

핵심 코드 분석

Gluon API 추가 (tdm.py)

After:

@builtin
def async_store(dest: tensor_descriptor, offsets: List[...], src: shared_memory_descriptor,
                _semantic=None) -> None:
    offset_handles = _semantic._convert_to_ir_values(offsets, require_i64=False)
    _semantic.builder.create_async_tdm_copy_local_to_global(
        dest.handle, offset_handles, src.handle)

SwizzledSharedLayout 지원

layout = _unwrap_if_constexpr(layout)
assert isinstance(layout, (PaddedSharedLayout, SwizzledSharedLayout)), \
    "Expected layout to be a PaddedSharedLayout or SwizzledSharedLayout"

TDM store는 PaddedSharedLayout뿐 아니라 SwizzledSharedLayout에서도 동작하도록 확장되었다.

MLIR 연산 등록 (gluon_ir.cc)

.def("create_async_tdm_copy_local_to_global",
     [](GluonOpBuilder &self, Value descPtr, std::vector<Value> &indices,
        Value src) {
       self.create<ttag::AsyncTDMCopyLocalToGlobalOp>(
           descPtr, indices, src);
     })

왜 이게 좋은가

1. 양방향 비동기 전송: load와 store 모두 TDM으로 처리하여 CPU 개입 최소화
2. 레이아웃 유연성: Padded와 Swizzled 두 가지 shared memory 레이아웃 모두 지원
3. Gluon 통합: 고수준 Python API로 저수준 하드웨어 기능에 접근 가능

정리

TDM은 AMD의 차세대 GPU에서 텐서 연산의 데이터 이동을 최적화하는 핵심 하드웨어 유닛이다. Store 연산의 추가로 TDM을 활용한 완전한 데이터 파이프라인을 구축할 수 있게 되었다.

참고 자료

• PR #8392

---

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