[Triton] AMD fine-grained cluster barrier 추가 및 Gluon 노출

PR 링크: triton-lang/triton#9206 상태: Merged | 변경: +227 / -1

들어가며

GPU에서 cluster는 여러 CTA(Cooperative Thread Array)를 하나의 그룹으로 묶어, CTA 간 데이터 공유와 동기화를 가능하게 하는 단위다. Cluster barrier는 클러스터 내 모든 CTA가 특정 지점에 도달했음을 보장하는 순수 실행 동기화 메커니즘이다. 이 PR은 AMD gfx1250 백엔드에 cluster barrier arrive/wait 연산을 추가하고, Gluon(Triton의 저수준 API) 인터페이스를 통해 커널 작성자에게 노출한다.

핵심 코드 분석

Gluon Python API (cluster.py)

@builtin
def arrive(_semantic=None):
    """클러스터 내 모든 CTA에 도달 신호를 보낸다."""
    _semantic.builder.create_amd_cluster_arrive()

@builtin
def wait(_semantic=None):
    """클러스터 내 모든 CTA의 도달을 대기한다."""
    _semantic.builder.create_amd_cluster_wait()

MLIR Op 정의 (TritonAMDGPUOps.td)

def ClusterBarrierArriveOp : TT_AMDGPU_Op<"cluster_barrier_arrive"> {
  let hasVerifier = 1;
  let assemblyFormat = "attr-dict";
}

def ClusterBarrierWaitOp : TT_AMDGPU_Op<"cluster_barrier_wait"> {
  let hasVerifier = 1;
  let assemblyFormat = "attr-dict";
}

LLVM 로우어링 (cluster_barrier_to_llvm.mlir)

// arrive → rocdl.s.barrier.signal id = -3
// wait   → rocdl.s.barrier.wait id = -3

왜 이게 좋은가

• 명시적 CTA 동기화: multicast load를 사용하는 커널에서 CTA 간 시간적 정렬을 보장하여, 데이터가 여러 CTA에 동시에 브로드캐스트될 확률을 높인다.
• CTA-scope barrier 분리: cluster barrier는 warp 하나만 신호하면 되므로, CTA 전체 동기화(async_wait 등)와 분리하여 불필요한 동기화 오버헤드를 방지한다.
• Verifier 포함: num-ctas > 1이 아닌 모듈에서 사용 시 컴파일 타임 에러를 발생시킨다.

정리

Multi-CTA 커널에서 CTA 간 동기화 granularity를 세밀하게 제어할 수 있는 primitive가 추가되었다. 특히 multicast load 패턴에서 CTA 간 시간 정렬이 성능에 직접적인 영향을 미치므로, 이 primitive는 고성능 커널 작성에 필수적이다.

참고 자료

• PR #9206 전체 diff
• AMD GFX1250 ISA Reference

---

이 글은 AI 도구의 도움을 받아 작성되었습니다.