[triton] AMD Pipelined Loop에서 TDM Load의 Buffer Race 수정

PR 링크: triton-lang/triton#9720 상태: Merged | 변경: +12 / -8

들어가며

AMD GPU의 pipelined loop에서 TDM(Tensor Data Movement) load를 사용할 때, async copy와 마찬가지로 레지스터 버퍼를 거치지 않고 직접 shared memory에 쓰기 때문에 추가 버퍼가 필요합니다. 기존에는 async copy에만 추가 버퍼를 할당하고 TDM load에는 하지 않아 데이터 경쟁이 발생했습니다.

핵심 코드 분석

Before

// async copy에만 추가 버퍼 할당
if (useAsyncCopy) {
  numBuffers += 1;
}

After

// async copy 또는 TDM load 사용 시 추가 버퍼 할당
if (useAsyncCopy || hasTDMLoad) {
  numBuffers += 1;
}

TDM load 감지 로직

bool hasTDMLoad = false;
for (const auto &[load, info] : loadOpToIndLevel) {
  auto useTDM = isa<tt::DescriptorLoadOp>(load);
  if (useTDM) {
    hasTDMLoad = true;
    // ...
  }
}

왜 이게 좋은가

1. 데이터 경쟁 수정: producer가 아직 사용 중인 버퍼에 consumer가 덮어쓰는 버그를 방지합니다.
2. 일관성: async copy와 TDM load가 동일한 버퍼링 정책을 따릅니다.
3. 명확한 조건 분기: descriptor load를 TDM 경로로, 일반 load를 async copy/stream copy 경로로 분리합니다.

정리

Pipelined loop의 버퍼 할당 로직에서 TDM load를 누락한 버그를 수정한 PR입니다. 레지스터 버퍼를 거치지 않는 비동기 메모리 연산은 모두 추가 버퍼가 필요하다는 원칙을 확인시켜 줍니다.

참고 자료

• PR #9720

---

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