[triton] Async TMA Lowering에서 Cluster Barrier 로직 수정

PR 링크: triton-lang/triton#9510 상태: Merged | 변경: +39 / -4

들어가며

NVIDIA GPU의 TMA(Tensor Memory Accelerator) 비동기 복사에서, cluster 내 여러 CTA가 barrier를 공유할 때 올바른 barrier 타입(shared::cta vs shared::cluster)을 선택해야 합니다. 이 PR은 barrier mask가 0이 아닐 때만 cluster barrier를 사용하도록 조건을 수정합니다.

핵심 코드 분석

Before

// barrier mask와 multicast mask의 교집합으로 판단
bool clusterBarrier = barrierMask & ~maskCGABroadcast;

After

// barrier mask가 0이 아니면 cluster barrier 사용
bool clusterBarrier = barrierMask != 0;

추가된 테스트 케이스

// Barrier mask zero -> shared::cta 사용
// CHECK: cp.async.bulk.tensor.2d.shared::cta.global.mbarrier
tt.func @tma_copy_barrier_mask_zero(...) {
  ttng.async_tma_copy_global_to_local %tma[%x, %x] %alloc, %barrier, %pred
}

// Barrier mask non-zero -> shared::cluster 사용
// CHECK: cp.async.bulk.tensor.2d.shared::cluster.global.mbarrier
tt.func @tma_copy_barrier_mask_nonzero(...) {
  ttng.async_tma_copy_global_to_local %tma[%x, %x] %alloc, %barrier, %pred
}

왜 이게 좋은가

1. 정확한 의미론: barrier mask의 의미를 정확히 반영하여 불필요한 cluster barrier 사용을 방지합니다.
2. 테스트 커버리지: barrier mask가 0인 경우와 아닌 경우를 모두 명시적으로 테스트합니다.
3. Multi-CTA 안정성: cross-CTA mbarrier init 동기화도 함께 개선하여 다중 CTA 시나리오의 안정성을 높였습니다.

정리

TMA cluster barrier 선택 로직의 미묘한 버그를 수정한 PR입니다. barrier mask의 의미를 단순화하여 코드 가독성도 향상되었습니다.

참고 자료

• PR #9510

---

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