[Triton] AMD gfx1250에 LDS 메모리 배리어 지원 추가

PR 링크: triton-lang/triton#8681 상태: Merged | 변경: +779 / -64

들어가며

AMD gfx1250 아키텍처는 LDS(Local Data Share)에 대한 새로운 memory barrier 메커니즘을 도입했다. 기존의 CTA 전체 동기화(s_barrier) 대신, LDS 접근에 특화된 세밀한 barrier를 사용할 수 있다.

이 PR은 gfx1250의 LDS barrier 관련 op들(allocate, init, arrive, wait, expect, invalidate)을 구현하고, LLVM lowering과 Gluon DSL 바인딩까지 전체 스택을 제공한다.

핵심 코드 분석

MLIR Op 정의: LDS Barrier 연산들

// TritonAMDGPUOps.td — 새로운 barrier op들
def TTAG_AllocBarrierOp : TTAG_Op<"alloc_barrier"> {
  let results = (outs TTG_MemDescType:$result);
}

def TTAG_InitBarrierOp : TTAG_Op<"init_barrier"> {
  let arguments = (ins TTG_MemDescType:$barrier);
}

def TTAG_ArriveBarrierOp : TTAG_Op<"arrive_barrier"> {
  let arguments = (ins TTG_MemDescType:$barrier, I32Attr:$count);
  let results = (outs TTAG_BarrierTokenType:$token);
}

def TTAG_WaitBarrierOp : TTAG_Op<"wait_barrier"> {
  let arguments = (ins TTAG_BarrierTokenType:$token);
}

Gluon Python API

# Gluon DSL에서 LDS barrier 사용
bar = ttgl.allocate_lds_barrier()
lds_barrier.init(bar)
token = lds_barrier.arrive(bar, count=1)
lds_barrier.wait(token)

LLVM Lowering: gfx1250 전용 intrinsic

// ArriveBarrierOp → LLVM IR
// gfx1250의 ds_arrive_barrier intrinsic을 생성
void lowerArriveBarrierOp(ArriveBarrierOp op) {
  auto barrierAddr = getSharedMemoryBase(op.getBarrier());
  auto count = op.getCount();
  // ds_arrive_barrier addr, count
  rewriter.create<ROCDL::DsArriveBarrierOp>(loc, barrierAddr, count);
}

Before vs After: barrier 인터페이스

// Before: CTA 전체 동기화만 가능
void barrier(Location loc, RewriterBase &rewriter) {
  rewriter.create<gpu::BarrierOp>(loc);  // s_barrier
}

// After: LDS에 특화된 세밀한 동기화 가능
auto token = rewriter.create<ttag::ArriveBarrierOp>(
    loc, barrier, /*count=*/1);
rewriter.create<ttag::WaitBarrierOp>(loc, token);

왜 이게 좋은가

1. 세밀한 동기화: CTA 전체가 아닌 LDS 접근에 특화된 barrier로 불필요한 동기화 오버헤드를 줄인다.
2. arrive/wait 분리: arrive와 wait를 분리하여, arrive 이후 wait 전까지 다른 연산을 수행하는 오버래핑이 가능하다.
3. 전체 스택 구현: MLIR op → LLVM lowering → Gluon API까지 한번에 제공하여 즉시 사용 가능하다.

정리

이 PR은 AMD gfx1250의 LDS memory barrier 기능을 Triton에 전체 스택으로 구현한다. allocate/init/arrive/wait/expect/invalidate의 6개 op, LLVM lowering, Gluon Python API를 포함하는 대규모 변경(+779/-64)으로, gfx1250에서의 파이프라인 병렬화 기반을 마련한다.

참고 자료

• AMD RDNA/CDNA Memory Model
• Triton AMD GPU Dialect

---

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