[triton] AMD GFX1250에서 Buffer Atomic 연산 활성화

PR 링크: triton-lang/triton#9744 상태: Merged | 변경: +113 / -7

들어가며

Buffer atomic 연산은 GPU의 global memory에 대한 원자적 읽기-수정-쓰기(RMW)를 수행합니다. 기존에는 CDNA3/CDNA4/RDNA4에서만 지원되었는데, 이 PR은 AMD의 GFX1250 아키텍처에서도 buffer atomic을 활성화하고, 해당 아키텍처에 맞는 cache policy를 적용합니다.

핵심 코드 분석

Before:

bool TargetInfo::supportsBufferAtomicRMW() const {
  return llvm::is_contained(
      {ISAFamily::CDNA3, ISAFamily::CDNA4, ISAFamily::RDNA4}, getISAFamily());
}

After:

bool TargetInfo::supportsBufferAtomicRMW() const {
  return llvm::is_contained({ISAFamily::CDNA3, ISAFamily::CDNA4,
                             ISAFamily::RDNA4, ISAFamily::GFX1250},
                            getISAFamily());
}

GFX1250 전용 cache policy 설정:

int32_t TargetInfo::getBufferAtomicCachePolicy(bool hasUsers) const {
  const int sc0Bit = 0b1;          // TH_ATOMIC_RETURN (cpol bit 0)
  const int scopeDevBit = 0b10000; // SCOPE_DEV = 2 << 3 (cpol bits [4:3])
  int32_t aux = 0;
  if (hasUsers)
    aux |= sc0Bit;
  if (getISAFamily() == ISAFamily::GFX1250)
    aux |= scopeDevBit;  // device-wide visibility
  return aux;
}

Resource descriptor에서도 GFX1250의 GFX12+ v2i64 형식을 올바르게 처리합니다.

왜 이게 좋은가

GFX1250은 GFX12 기반이므로 기존 GFX9 계열과 cache policy 인코딩이 다릅니다. SCOPE_DEV 비트(cpol bits [4:3])를 설정하여 device-wide 가시성을 보장하는 것이 핵심입니다. 반환값이 필요한 경우에만 SC0(TH_ATOMIC_RETURN)을 설정하여 불필요한 캐시 무효화를 방지합니다. 또한 packed bf16 fadd(BUFFER_ATOMIC_PK_ADD_BF16)를 지원하여 mixed-precision 학습에서의 gradient accumulation 성능을 향상시킵니다.

정리

GFX1250을 buffer atomic RMW/CAS 지원 목록에 추가하고, SCOPE_DEV cache policy와 packed bf16 fadd를 포함한 GFX12+ 전용 설정을 구현했습니다.

참고 자료

• triton-lang/triton#9744

이 분석은 AI가 실제 코드 diff를 기반으로 작성했습니다.