본문으로 건너뛰기

[Triton] TritonGPU Barrier 재설계 — 주소 공간별 메모리 가시성 보장

PR 링크: triton-lang/triton#9025 상태: Merged | 변경: +555 / -379

들어가며

GPU 커널에서 barrier는 스레드 간 동기화와 메모리 가시성을 보장하는 핵심 원시 연산(primitive)이다. 기존 Triton은 MLIR의 gpu.barrier를 사용했는데, 이는 단순한 CTA(Cooperative Thread Array) 수준 동기화만 제공하고 어떤 주소 공간의 메모리에 대한 가시성을 보장하는지 명시하지 않았다.

이 PR은 새로운 ttg.barrier op을 도입하여 shared memory, global memory 등 특정 주소 공간에 대한 메모리 가시성을 명시적으로 지정할 수 있게 한다.

핵심 코드 분석

Before: 단순 barrier 인터페이스

// TargetInfoBase.h
virtual void barrier(Location loc, RewriterBase &rewriter,
                     bool isWarpSync = false) const = 0;

bool isWarpSync라는 단일 플래그로만 barrier 종류를 구분했다.

After: 주소 공간 기반 barrier

// TargetInfoBase.h — 새로운 인터페이스
// Block/CTA level barrier with address space visibility
virtual void barrier(Location loc, RewriterBase &rewriter,
                     triton::gpu::AddrSpace targets) const = 0;
// Warp-level synchronization
virtual void warpSync(Location loc, RewriterBase &rewriter) const = 0;

새로운 AddrSpace enum이 도입되었다:

// TritonGPUEnums.td
def AddrSpace : I32EnumAttr<"AddrSpace", "address space", [
  I32EnumAttrCase<"Local", 0>,    // shared memory
  I32EnumAttrCase<"Global", 1>,   // global memory
]>;

MLIR op 정의에서 barrier가 어떤 주소 공간의 가시성을 보장하는지 명시한다:

// TritonGPUOps.td
def TTG_BarrierOp : TTG_Op<"barrier"> {
  let summary = "CTA-level barrier with memory visibility";
  let arguments = (ins
    AddrSpaceAttr:$addrspace
  );
}

유틸리티 헬퍼도 업데이트되었다:

// Before
mlir::gpu::BarrierOp barrier() {
  return mlir::gpu::BarrierOp::create(*builder, loc);
}

// After
mlir::triton::gpu::BarrierOp barrier(triton::gpu::AddrSpace addrspace) {
  return mlir::triton::gpu::BarrierOp::create(*builder, loc, addrspace);
}

왜 이게 좋은가

  1. 정확한 의미론: barrier가 보장하는 메모리 가시성 범위를 IR 수준에서 명시하므로, 백엔드가 최적의 명령어를 선택할 수 있다. shared memory만 필요한 곳에 global fence를 넣을 필요가 없다.
  2. warp sync 분리: warp 수준 동기화를 별도 인터페이스(warpSync)로 분리하여, CTA barrier와의 혼용을 방지한다.
  3. 백엔드 유연성: NVIDIA는 bar.sync와 fence.proxy.async, AMD는 s_barrier와 s_waitcnt 등 아키텍처별로 적절한 명령어를 매핑할 수 있다.

정리

이 PR은 Triton의 barrier 인프라를 gpu.barrier에서 ttg.barrier로 교체하며, 주소 공간별 메모리 가시성을 명시할 수 있도록 재설계했다. barrier → warpSync 분리와 AddrSpace enum 도입으로 더 세밀하고 정확한 동기화 제어가 가능해졌다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글