[Triton] AMD TDM L2 Prefetch 백엔드 지원 추가

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

들어가며

AMD GPU의 TDM(Tensor Data Movement) 엔진은 L2 캐시로 데이터를 미리 가져오는 prefetch 기능을 제공한다. 이 하드웨어 instrinsic은 256바이트 단위로 L2에 데이터를 적재하며, speculative/non-speculative 두 모드를 지원한다.

이 PR은 이 기능에 대한 TritonAMDGPU dialect op과 LLVM IR lowering을 구현하는 백엔드 레벨 작업이다. 후속 PR(#9148)에서 Gluon DSL API가 추가된다.

핵심 코드 분석

MLIR Op 정의: TDMPrefetchOp

// TritonAMDGPUOps.td
def TTAG_TDMPrefetchOp : TTAG_Op<"tdm_prefetch"> {
  let summary = "Prefetch a TDM tile into L2 cache";
  let description = [{
    The underlying intrinsic expects an address and prefetches 256 bytes
    into L2 (hardware enforces 256-byte alignment).
  }];
  let arguments = (ins
    TT_TensorDescType:$desc,
    Variadic<I32>:$coords,
    I1:$pred,
    BoolAttr:$speculative
  );
}

LLVM Lowering: speculative/non-speculative 분기

// LLVM lowering에서 speculative 모드에 따라 다른 비트 설정
// CHECK-DAG: %[[NON_SPECULATIVE_BITS:.*]] = llvm.mlir.constant(8 : i32)
// CHECK-DAG: %[[SPECULATIVE_BITS:.*]] = llvm.mlir.constant(9 : i32)

speculative 모드(비트 9)는 잘못된 주소에 대해서도 안전하게 무시되어 조건 분기 없이 사용할 수 있다. non-speculative 모드(비트 8)는 반드시 유효한 주소에만 사용해야 한다.

Tensor Descriptor 기반 주소 계산

Prefetch는 tensor descriptor에서 좌표를 기반으로 메모리 주소를 계산한다:

// LLVM IR 테스트에서 확인되는 주소 계산 패턴
// base_addr + coord[0] * stride[0] + coord[1] * stride[1]
// 결과 주소를 256바이트 경계로 정렬

유틸리티 수정: lookupNumCTAs 범용화

// Before
int triton::gpu::lookupNumCTAs(Operation *op) {
  auto mod = op->getParentOfType<ModuleOp>();
  // ...
}

// After — ModuleOp 자체가 전달될 수도 있도록 수정
int triton::gpu::lookupNumCTAs(Operation *op) {
  auto mod = dyn_cast<ModuleOp>(op);
  if (!mod)
    mod = op->getParentOfType<ModuleOp>();
  // ...
}

이 수정은 prefetch lowering에서 ModuleOp을 직접 전달하는 경우를 지원한다.

왜 이게 좋은가

1. 하드웨어 기능 활용: L2 prefetch를 컴파일러 수준에서 지원하여, 사용자가 메모리 접근 패턴에 맞는 prefetch를 삽입할 수 있는 기반을 마련한다.
2. 두 모드 지원: speculative(안전, 공격적) / non-speculative(정확, 보수적) 두 모드를 모두 지원하여 다양한 사용 시나리오에 대응한다.
3. 계층적 설계: 이 PR은 저수준(LLVM lowering)만 담당하고, 사용자 API는 후속 PR에서 추가하는 clean한 분리 구조다.

정리

이 PR은 AMD TDM L2 prefetch의 백엔드 구현으로, MLIR op 정의부터 LLVM lowering까지를 포함한다. 256바이트 단위 L2 프리페치를 speculative/non-speculative 모드로 지원하며, 후속 Gluon API PR의 기반이 된다.

참고 자료

• AMD CDNA Architecture — Memory Subsystem
• 후속 PR: Gluon API 노출 (#9148)

---

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