[triton] AMD: BufferLoadToLocal을 UpdateAsyncWaitCount에 포함하여 성능 회귀 수정

PR 링크: triton-lang/triton#8621 상태: Merged | 변경: +47 / -19

들어가며

Triton의 AMD 백엔드에서 UpdateAsyncWaitCount 패스는 비동기 로드 명령어 사이의 wait 카운트를 계산합니다. 이전 PR(#8575)에서 이 패스가 TTIR->TTGIR에서 TTGIR->LLVM 단계로 이동하면서, 새로 도입된 amdgpu.buffer_load_to_local 명령어가 카운트에 포함되지 않는 문제가 생겼습니다. 이로 인해 보수적인 wait이 발행되어 성능이 저하되었습니다.

핵심 코드 분석

Before: buffer_load_to_local 미포함

// UpdateAsyncWaitCount.cpp
int getNumberOfLoadInstructions(TypedValue<RankedTensorType> ptrs,
                                ttg::MemDescType dstTy, Value mask,
                                ModuleAxisInfoAnalysis &axisInfo) {
  // AsyncCopyGlobalToLocal만 처리
  int contig = srcToSharedLayout.getNumConsecutiveInOut();
  contig = std::min<int>(contig, LLVM::AMD::getVectorSize(ptrs, axisInfo));
  // ...
}

After: buffer_load_to_local 포함

// UpdateAsyncWaitCount.cpp
int getNumberOfLoadInstructions(RankedTensorType srcTy, ttg::MemDescType dstTy,
                                Value mask, int contig,
                                ModuleAxisInfoAnalysis &axisInfo) {
  // 일반화된 인터페이스
  contig = std::min(contig, srcToSharedLayout.getNumConsecutiveInOut());
  // ...
}

int getOpNumberOfAsyncLoadInstructions(Operation *op, ...) {
  if (auto copyOp = dyn_cast<ttg::AsyncCopyGlobalToLocalOp>(op)) {
    int contig = LLVM::AMD::getVectorSize(copyOp.getSrc(), axisInfo);
    return getNumberOfLoadInstructions(copyOp.getSrc().getType(), ...);
  } else if (auto bufferOp = dyn_cast<amdgpu::BufferLoadToLocalOp>(op)) {
    // 새로 추가된 buffer_load_to_local 처리
    auto ptrType = cast<RankedTensorType>(LLVM::AMD::getPointerTypeWithShape(
        bufferOp.getPtr(), bufferOp.getOffsets()));
    int contig = LLVM::AMD::getVectorSize(bufferOp.getPtr(),
                                          bufferOp.getOffsets(), axisInfo);
    return getNumberOfLoadInstructions(ptrType, bufferOp.getDest().getType(),
                                       bufferOp.getMask(), contig, axisInfo);
  }
}

왜 이게 좋은가

1. 정확한 wait 카운트: buffer_load_to_local이 카운트에 포함되어 불필요한 보수적 wait이 제거됩니다.
2. 인터페이스 일반화: getNumberOfLoadInstructions가 contig를 매개변수로 받도록 변경되어 다양한 로드 유형을 통합적으로 처리합니다.
3. Gluon 호환: buffer ops를 직접 생성하는 Gluon 커널에서도 올바른 wait 카운트가 적용됩니다.

정리

비동기 메모리 연산의 wait 카운트는 GPU 성능에 직접적인 영향을 미칩니다. 새로운 명령어 유형이 추가될 때 기존 분석 패스를 반드시 업데이트해야 한다는 교훈을 주는 PR입니다.

참고 자료

• triton-lang/triton#8621

---

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