본문으로 건너뛰기

[Triton] Gluon에 coalesced layout 추가 — 메모리 접근 효율 최적화

PR 링크: triton-lang/triton#8604 상태: Merged | 변경: +740 / -327

들어가며

GPU에서 글로벌 메모리 접근 성능은 coalescing(병합)에 크게 좌우된다. 같은 warp의 스레드들이 연속된 메모리 주소에 접근하면 하나의 트랜잭션으로 병합되어 대역폭을 최대로 활용할 수 있다.

Gluon DSL에서는 사용자가 텐서의 레이아웃을 명시적으로 지정하는데, 이 PR은 coalesced layout을 추가하여 글로벌 메모리 접근이 자동으로 coalescing되도록 보장한다.

핵심 코드 분석

Coalesced Layout 정의

# Gluon Python API
layout = ttgl.CoalescedLayout(
    element_bitwidth=16,
    rank=2,
    order=[1, 0],  # 마지막 차원이 가장 빠르게 변하는 순서
)

C++ 구현: resolveCoalescedEncoding

기존 resolveAutoEncoding 패스와 유사한 구조로, coalesced layout을 concrete encoding으로 해석하는 패스가 추가되었다:

// Before: auto encoding만 지원
void resolveAutoEncoding(ModuleOp moduleOp) {
  // blocked layout 등으로 자동 해석
}

// After: coalesced encoding도 해석
void resolveCoalescedEncoding(ModuleOp moduleOp) {
  moduleOp.walk([](Operation *op) {
    if (auto loadOp = dyn_cast<LoadOp>(op)) {
      auto resultType = loadOp.getResult().getType();
      if (hasCoalescedEncoding(resultType)) {
        // 메모리 접근 패턴 분석 후
        // coalescing이 보장되는 blocked layout으로 변환
        auto coalescedLayout = computeCoalescedBlockedLayout(
            resultType, op);
        loadOp.getResult().setType(
            resultType.cloneWithEncoding(coalescedLayout));
      }
    }
  });
}

핵심 로직: coalescing 보장 레이아웃 계산

BlockedEncodingAttr computeCoalescedBlockedLayout(
    RankedTensorType tensorType, Operation *op) {
  auto shape = tensorType.getShape();
  auto elementBitWidth = getElementBitWidth(tensorType);

  // 128바이트 캐시라인에 맞춰 스레드당 연속 요소 수 결정
  unsigned elemsPerThread = 128 / (elementBitWidth / 8);

  // 가장 안쪽 차원(contiguous dim)을 따라 스레드 배치
  SmallVector<unsigned> sizePerThread(rank, 1);
  sizePerThread[order[0]] = elemsPerThread;

  return BlockedEncodingAttr::get(
      ctx, sizePerThread, threadsPerWarp, warpsPerCTA, order);
}

왜 이게 좋은가

  1. 자동 최적화: 사용자가 CoalescedLayout만 지정하면 컴파일러가 아키텍처에 맞는 최적의 blocked layout을 자동 계산한다.
  2. 추상화 레벨 적절: "메모리 접근이 coalesced되어야 한다"는 의도를 표현하면서, 구체적인 스레드 배치는 컴파일러에 위임한다.
  3. 기존 인프라 재사용: blocked layout으로 변환되므로, 이후 파이프라인의 모든 최적화가 그대로 적용된다.

정리

이 PR은 Gluon DSL에 coalesced layout을 추가하여, 글로벌 메모리 접근의 coalescing을 추상적으로 표현하고 컴파일러가 최적의 concrete layout으로 해석하도록 한다. 740줄 추가의 대규모 변경으로, layout 해석 패스와 Python API를 포함한다.

참고 자료

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

댓글

관련 포스트

PR Analysis 의 다른글