[triton] WSSpecialize에서 초기화된 Barrier의 Invalidation 추가

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

들어가며

NVIDIA GPU의 mbarrier(memory barrier)는 warp 간 동기화를 위한 하드웨어 프리미티브입니다. WarpSpecialize 패스는 producer/consumer 파이프라인을 구축할 때 여러 mbarrier를 할당하고 초기화합니다. 그러나 사용이 끝난 barrier를 invalidate하지 않으면 하드웨어 리소스가 올바르게 해제되지 않아 후속 커널 실행에서 문제가 발생할 수 있습니다.

핵심 코드 분석

Before:

// WarpSpecialize 패스 완료 후
doTokenLowering(funcOp, numWarpGroups - 1);
// barrier invalidation 없음

After:

doTokenLowering(funcOp, numWarpGroups - 1);
invalidateWarpSpecializeBarriers(funcOp);  // 추가됨

invalidateWarpSpecializeBarriers의 구현:

void invalidateWarpSpecializeBarriers(triton::FuncOp funcOp) {
  for (ttg::WarpSpecializeOp wsOp : wsOps) {
    SetVector<Value> barrierAllocs;
    for (Value operand : partitionOp->getOperands()) {
      if (!isWarpSpecializeBarrierAlloc(operand))
        continue;
      barrierAllocs.insert(operand);
    }
    ImplicitLocOpBuilder builder(wsOp.getLoc(), wsOp);
    builder.setInsertionPointAfter(wsOp);
    for (Value barrierAlloc : barrierAllocs)
      invalidateBarrierAlloc(builder, barrierAlloc);
  }
}

Barrier를 식별하기 위해 kWarpSpecializeGeneratedBarrierAttrName 속성을 할당 시점에 표시합니다.

왜 이게 좋은가

mbarrier는 하드웨어 리소스이며, NVIDIA GPU에서는 mbarrier.inval을 호출하여 명시적으로 해제해야 합니다. 이를 생략하면 (1) 하드웨어가 stale barrier 상태를 유지하여 다음 커널 실행에 영향을 줄 수 있고, (2) barrier 리소스 누수가 발생할 수 있습니다. 특히 warp specialization으로 생성된 barrier만 선별적으로 invalidate하기 위해 속성 기반 태깅을 사용한 것은 깔끔한 설계입니다.

정리

WarpSpecialize 패스가 생성한 모든 mbarrier에 태그 속성을 추가하고, 패스 완료 후 이들을 자동으로 invalidate하는 로직을 추가하여 하드웨어 barrier 리소스 정합성을 보장합니다.

참고 자료

• triton-lang/triton#9592

이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.