[triton] CGAEncodingAttr::getDefault를 get1CTALayout/get1DLayout로 분리하여 multi-CTA 지원

PR 링크: triton-lang/triton#9040 상태: Merged | 변경: +194 / -85

들어가며

Triton의 CGAEncodingAttr는 CGA(Cooperative Group Array) 레이아웃을 나타내는 속성입니다. 기존 getDefault 함수는 1CTA 모드에서만 올바르게 동작했지만, 이름에서 그 제한이 드러나지 않아 multi-CTA 환경에서 오용될 위험이 있었습니다. 이번 PR은 이를 명확히 분리하고, MemDesc 타입에 대한 레이아웃 검증도 추가합니다.

핵심 코드 분석

Before: 모호한 getDefault

// CGAEncodingAttr.td
static CGAEncodingAttr getDefault(MLIRContext *context, int rank);

// Dialect.cpp
CGAEncodingAttr CGAEncodingAttr::getDefault(MLIRContext *ctx, int rank) {
  auto kBlock = StringAttr::get(ctx, "block");
  LinearLayout::BasesT bases;
  bases[kBlock] = {};  // 빈 bases = 1CTA
  auto dims = standardOutDimNames(ctx, rank);
  return get(ctx, LinearLayout(std::move(bases), dims));
}

After: 명확한 이름의 두 함수

// CGAEncodingAttr.td
// 빈 bases로 1CTA 레이아웃 생성
static CGAEncodingAttr get1CTALayout(MLIRContext *context, int rank);
// identity layout으로 1D multi-CTA 레이아웃 생성
static CGAEncodingAttr get1DLayout(MLIRContext *context, int numCTAs);

// Dialect.cpp
CGAEncodingAttr CGAEncodingAttr::get1CTALayout(MLIRContext *ctx, int rank) {
  // 기존 getDefault와 동일
}

CGAEncodingAttr CGAEncodingAttr::get1DLayout(MLIRContext *ctx, int numCTAs) {
  auto kBlock = StringAttr::get(ctx, "block");
  auto dims = standardOutDimNames(ctx, /*rank=*/1);
  auto layout = LinearLayout::identity1D(numCTAs, kBlock, dims[0]);
  return get(ctx, std::move(layout));
}

추가로, verifyTensorLayouts가 MemDescType도 검증하도록 확장되었습니다:

// Traits.cpp
auto memDescTy = dyn_cast<MemDescType>(val.getType());
if (memDescTy) {
  return verifyLayoutInterface->verifyMemDescLayout(layout, memDescTy, op, makeErr);
}

왜 이게 좋은가

1. API 명확성: getDefault라는 모호한 이름 대신 get1CTALayout과 get1DLayout으로 의도가 명확해졌습니다.
2. Multi-CTA 지원: get1DLayout으로 다중 CTA 환경의 레이아웃을 올바르게 생성할 수 있습니다.
3. 검증 강화: MemDesc 타입에 대한 레이아웃 검증이 추가되어 잘못된 인코딩을 조기에 발견합니다.

정리

함수 이름은 코드의 의도를 전달하는 가장 중요한 수단입니다. getDefault처럼 모호한 이름을 구체적인 이름으로 바꾸는 것은 작은 변경이지만, multi-CTA 지원이라는 큰 기능의 기반이 됩니다.

참고 자료

• triton-lang/triton#9040

---

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