[triton] 클러스터 환경을 위한 Membar 패스 확장

PR 링크: triton-lang/triton#9318 상태: Merged | 변경: +861 / -31

들어가며

Multi-CTA 클러스터에서는 distributed shared memory를 통해 CTA 간 데이터가 이동하므로, 단순한 CTA 내 barrier로는 부족합니다. 이 PR은 membar 분석의 핵심 자료구조인 AllocationSlice에 buffer ID를 추가하고, MembarSliceFilterFn을 도입하여 slice 레벨의 세분화된 barrier 판단을 지원합니다.

핵심 코드 분석

1. AllocationSlice에 buffer ID 추가

Before:

struct AllocationSlice {
  AllocationSlice(Value value, Interval<size_t> allocationInterval);
  AllocationSlice(Interval<size_t> interval)
      : allocationInterval(interval), accessTy(nullptr) {}
};

After:

struct AllocationSlice {
  AllocationSlice(Value value, Interval<size_t> allocationInterval,
                  Allocation::BufferId bufferId);
  AllocationSlice(Interval<size_t> interval)
      : allocationInterval(interval), accessTy(nullptr),
        bufferId(Allocation::InvalidBufferId) {}
  Allocation::BufferId getBufferId() const { return bufferId; }
};

2. MembarSliceFilterFn 도입

using MembarSliceFilterFn =
    std::function<bool(const AllocationSlice &, const AllocationSlice &,
                       bool /*lhsIsRead*/, bool /*rhsIsRead*/, Allocation *)>;

기존 MembarFilterFn(op 레벨)에 더해, slice 레벨 filter를 추가하여 더 정밀한 판단이 가능합니다.

3. isIntersected의 read/write 방향 전달

Before:

return isIntersected(syncWriteSlices, other.syncReadSlices, filter,
                     allocation);

After:

return isIntersected(syncWriteSlices, other.syncReadSlices,
                     /*lhsIsRead=*/false, /*rhsIsRead=*/true,
                     filter, sliceFilter, allocation);

RAW/WAR/WAW를 구분하여 filter에 전달함으로써, 방향별로 다른 barrier 정책을 적용할 수 있습니다.

왜 이게 좋은가

• 정밀한 분석: Buffer ID와 read/write 방향 정보로 불필요한 barrier를 더 많이 제거할 수 있습니다.
• 클러스터 지원: ClusterBarrierInsertion과 통합되어 CTA 간 동기화를 정확히 처리합니다.
• 확장 가능한 설계: Slice filter와 op filter의 2단계 구조로 다양한 백엔드 요구사항을 수용합니다.

정리

Membar 분석의 기반 자료구조를 클러스터 환경에 맞게 강화한 대규모 PR입니다. Buffer ID 추가, slice-level filter, read/write 방향 구분이 핵심 변경사항입니다.

참고 자료

• triton-lang/triton#9318

---

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