본문으로 건너뛰기

[triton] AMD Async Load에 ROCDL Op 사용으로 전환

PR 링크: triton-lang/triton#9410 상태: Merged | 변경: +50 / -30

들어가며

이 PR은 AMD GPU의 async load(global to LDS) 및 cluster async load에서 문자열 기반 LLVM intrinsic 호출을 ROCDL dialect의 타입 안전한 op으로 교체합니다. 기능 변경은 없는 NFC(Non-Functional Change)입니다.

핵심 코드 분석

Before - 문자열 기반 intrinsic 호출

if (multicastMask) {
  std::string intrinsic =
      "llvm.amdgcn.cluster.load.async.to.lds.b" + std::to_string(vecBits);
  LLVM::createLLVMIntrinsicCallOp(
      rewriter, loc, intrinsic, {},
      {srcPtr, shmemAddr, b.i32_val(0), b.i32_val(cacheModifiers),
       multicastMask});
}

After - ROCDL op 사용

switch (vecBits) {
case 32:
  if (multicastMask)
    ROCDL::ClusterLoadAsyncToLDSB32Op::create(
        rewriter, loc, srcPtr, shmemAddr, 0, cacheModifiers,
        multicastMask, nullptr, nullptr, nullptr);
  else
    ROCDL::GlobalLoadAsyncToLDSB32Op::create(
        rewriter, loc, srcPtr, shmemAddr, 0, cacheModifiers,
        nullptr, nullptr, nullptr);
  break;
// case 64, 128 도 동일 패턴
}

왜 이게 좋은가

  1. 타입 안전성: 컴파일 타임에 인자 타입과 개수가 검증됩니다.
  2. 자동 완성: IDE에서 ROCDL op의 파라미터를 바로 확인할 수 있습니다.
  3. 유지보수성: intrinsic 이름이 변경되어도 한 곳만 수정하면 됩니다.

정리

문자열 기반 LLVM intrinsic 호출을 ROCDL op으로 바꾼 코드 품질 개선입니다. 기능은 동일하지만 타입 안전성이 강화되어 향후 버그 가능성이 줄었습니다.

참고 자료


이 글은 AI(Claude)의 도움을 받아 작성되었으며, 원본 PR의 코드 변경 사항을 기반으로 분석한 내용입니다.

댓글

관련 포스트

PR Analysis 의 다른글