본문으로 건너뛰기

[Triton] tcgen05 op의 cga_layout barrier 검증 강화

들어가며

NVIDIA Blackwell GPU의 tcgen05 명령어(MMA, commit, CLC 등)는 completion barrier를 사용하여 비동기 연산의 완료를 추적한다. 이 barrier의 CGA(Cooperative Grid Array) layout이 잘못되면 런타임에 정의되지 않은 동작이 발생한다. 이 PR은 completion barrier의 CGA layout을 컴파일 타임에 검증하는 verifier를 추가한다.

핵심 코드 분석

검증 로직

static LogicalResult verifyCompletionBarrierLayout(Operation *op, Value barrier) {
    auto barrierTy = cast<MemDescType>(barrier.getType());
    auto expectedCGALayout =
        CGAEncodingAttr::get1DLayout(op->getContext(), gpu::lookupNumCTAs(op));
    auto actualCGALayout = getCGALayout(barrierTy.getEncoding());
    if (actualCGALayout != expectedCGALayout)
        return op->emitOpError("completion barrier cga_layout must be ")
               << formatCGALayout(expectedCGALayout) << ", got "
               << formatCGALayout(actualCGALayout);
    return success();
}

이 검증이 TCGen5MMAOp, TCGen5CommitOp, TCGen5MMAScaledOp, CLCTryCancelOp에 모두 적용된다.

테스트 코드에서의 변경

// Before: CGA layout이 [[0]]인 barrier가 2-CTA 환경에서 사용됨 (잘못됨)
#barrier = #ttg.swizzled_shared<{..., CGALayout = [[0]]}>

// After: 올바른 2-CTA barrier layout
#barrier = #ttg.swizzled_shared<{..., CGALayout = [[1]]}>

[[0]]은 모든 CTA가 같은 barrier를 공유한다는 의미이고, [[1]]은 각 CTA가 별도 barrier를 갖는다는 의미다. 2-CTA 모드에서는 [[1]]이 올바르다.

왜 이게 좋은가

  • 조기 오류 감지: 런타임 정의되지 않은 동작 대신 컴파일 타임 오류를 발생시킨다.
  • 일관된 검증: 모든 tcgen05 관련 op에 동일한 검증이 적용된다.
  • CLC 검증 추가: CLC(Cluster-Level Communication) op에 대해서도 barrier와 result buffer의 CGA layout을 검증한다.

정리

+169/-20 변경으로, 다가오는 multi-CTA 기능에 대한 안전망을 강화한 PR이다. GPU 컴파일러에서 "유효하지 않은 설정은 컴파일에서 차단"하는 원칙의 좋은 사례다.

참고 자료

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

댓글