[triton] Concurrency Sanitizer를 Vendor Target Hooks로 리팩터링

PR 링크: triton-lang/triton#9658 상태: Merged | 변경: +456 / -295

들어가며

Triton의 Concurrency Sanitizer(ConSan)는 GPU 커널의 공유 메모리 경쟁 조건을 탐지합니다. 기존에는 NVIDIA GPU의 연산(TMA, wgmma 등)이 sanitizer 코드에 직접 하드코딩되어 있어 다른 벤더를 추가하기 어려웠습니다. 이 PR은 벤더별 로직을 ConSanTargetHooks 인터페이스로 분리합니다.

핵심 코드 분석

Before - 벤더별 하드코딩

// GlobalSanitizer.cpp 내부에 NVIDIA 연산이 직접 참조됨
if (auto tma = dyn_cast<ttng::AsyncTMACopyGlobalToLocalOp>(op)) { ... }
if (auto mma = dyn_cast<ttng::TCGen5MMAOp>(op)) { ... }

After - Target Hooks 인터페이스

class ConSanTargetHooks {
public:
  virtual ~ConSanTargetHooks() = default;
  virtual bool isTMAOp(Operation *op) const = 0;
  virtual std::optional<BarrierInitInfo>
  getBarrierInitInfo(Operation *op) const = 0;
  virtual std::optional<WaitOpInfo> getWaitOpInfo(Operation *op) const = 0;
  virtual std::optional<MemEffectsOpInfo>
  getMemEffectsOpInfo(Operation *op) const { /* default impl */ }
};

// 팩토리 패턴으로 벤더별 hooks 등록
void registerConSanHooks(llvm::StringRef key, ConSanHooksFactory factory);
std::unique_ptr<ConSanTargetHooks> createConSanHooks(llvm::StringRef key);

NVIDIA hooks 등록

// RegisterTritonDialects.h
mlir::triton::nvidia_gpu::registerConSanNVIDIAHooks();

왜 이게 좋은가

1. 벤더 독립적 설계: AMD, Intel 등의 GPU도 ConSanTargetHooks를 구현하면 ConSan을 활용할 수 있습니다.
2. 의존성 역전: sanitizer 코어 코드가 벤더 dialect에 의존하지 않고 인터페이스에만 의존합니다.
3. 팩토리 패턴: 런타임에 target에 맞는 hooks를 동적으로 선택하여 컴파일 의존성을 줄입니다.

정리

ConSan의 벤더별 로직을 추상 인터페이스로 분리한 전형적인 의존성 역전 리팩터링입니다. 멀티벤더 GPU 지원의 기반을 마련하면서도 기존 NVIDIA 기능은 그대로 유지합니다.

참고 자료

• PR #9658

---

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