[triton] MultiCTA Membar에 Fence + Cluster Relaxed 자동 삽입

PR 링크: triton-lang/triton#9615 상태: Merged | 변경: +539 / -128

들어가며

NVIDIA Multi-CTA 커널에서 cross-CTA mbarrier를 사용할 때, barrier init 후 첫 사용 전에 fence_mbarrier_init_release_cluster와 cluster_arrive/wait(relaxed=true)를 삽입해야 합니다. 이 PR은 이 과정을 자동화하여 프로그래머가 수동으로 동기화를 삽입하지 않아도 되게 합니다.

핵심 코드 분석

cross-CTA barrier 판별

static bool requiresCrossCTAMBarrierInitSync(
    ttng::InitBarrierOp initBarrierOp, FunctionOpInterface funcOp,
    Allocation *allocation, int numCTAs) {
  // 1. barrier 자체가 cross-CTA인 경우
  if (isCrossCTAMBarrier(initBarrierOp, numCTAs))
    return true;
  // 2. per-CTA barrier라도 multi-CTA consumer(TMA multicast, 2-CTA MMA)가 사용하면 동기화 필요
  return funcOp->walk([&](Operation *op) {
    if (usesTrackedBarrierInCrossCTAConsumerOp(op, initBarrierBuffers, allocation))
      return WalkResult::interrupt();
    return WalkResult::advance();
  }).wasInterrupted();
}

Verifier 추가

LogicalResult InitBarrierOp::verify() {
  bool crossCTA = barrierTy.getShape()[0] != gpu::lookupNumCTAs(getOperation());
  if (crossCTA &&
      getOperation()->getParentOfType<gpu::WarpSpecializeOp>())
    return emitOpError("cannot be used inside `ttg.warp_specialize`");
  return success();
}

왜 이게 좋은가

1. 자동 동기화: cross-CTA barrier의 fence/arrive/wait를 컴파일러가 자동 삽입하여 프로그래머 실수를 방지합니다.
2. 정확한 분석: allocation analysis를 활용해 alias 관계까지 추적하여 정확한 barrier 의존성을 판별합니다.
3. 안전장치: WarpSpecialize 영역 내에서 cross-CTA barrier 초기화를 금지하는 verifier를 추가했습니다.

정리

Multi-CTA 환경에서 barrier 동기화의 정확성을 컴파일러가 보장하도록 만든 중요한 개선입니다. allocation analysis 기반의 세밀한 분석으로 불필요한 동기화는 피하면서 필요한 곳에만 삽입합니다.

참고 자료

• PR #9615

---

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