[Triton] Concurrency Sanitizer에 TMA Store 검증 추가

PR 링크: triton-lang/triton#8672 상태: Merged | 변경: +156 / -57

들어가며

Triton의 **Concurrency Sanitizer(CONSAN)**는 GPU 커널에서 발생할 수 있는 데이터 레이스를 감지하는 도구다. 기존에는 cp.async와 WGMMA만 추적했지만, TMA(Tensor Memory Accelerator) Store도 비동기 연산이므로 동일한 검증이 필요했다. 이 PR은 commit/wait 기반의 동시성 추적을 일반화하여 TMA Store도 지원한다.

핵심 코드 분석

Commit 추적 일반화

Before:

RegionToValueMap asyncCpCommits;
RegionToValueMap wgmmaCommits;

각 비동기 연산마다 별도의 맵을 관리했다.

After:

namespace CommitKind {
enum Kind { None = -1, AsyncCp = 0, Wgmma, TmaStore, NumCommitKinds };
}
RegionToValueMap commits[CommitKind::NumCommitKinds];

enum 배열로 일반화하여 새로운 비동기 연산 타입을 쉽게 추가할 수 있다.

TMA Store Wait 처리

if (auto tmaStoreWaitOp = dyn_cast<ttng::TMAStoreWaitOp>(op)) {
    funcBuilder.createClearOutstandingCommitsTransferReadsCall(
        b, baseThread, getThreadPeersMask(thread),
        tmaStoreWaitOp.getPendings(), nullptr,
        auxData.commits[CommitKind::TmaStore][op],
        auxData.readVisibility[(int)MemType::SHARED_MEM][op], op);
}

왜 이게 좋은가

1. 일반화된 설계: 새로운 비동기 연산 추가 시 enum에 항목 하나만 추가하면 됨
2. TMA Store 검증: shared memory에서 global memory로의 비동기 쓰기도 레이스 감지 가능
3. baseThread 사용: commit-count 추적에서 helper thread가 아닌 base thread만 사용하도록 수정

정리

GPU 프로그래밍에서 비동기 연산의 동시성 버그는 찾기 가장 어려운 종류의 버그다. Sanitizer가 더 많은 연산 타입을 추적할수록 개발자가 안전한 코드를 작성하기 쉬워진다. 일반화된 commit tracking 패턴은 향후 새로운 비동기 연산이 추가될 때도 확장이 용이하다.

참고 자료

• PR #8672

---

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