[triton] AMD 백엔드에 Concurrency Sanitizer(ConSan) 지원 추가

PR 링크: triton-lang/triton#9692 상태: Merged | 변경: +3219 / -104

들어가며

ConSan(Concurrency Sanitizer)은 GPU 커널에서 데이터 레이스, barrier 오용, 동기화 누락 등의 동시성 버그를 감지합니다. 기존에는 NVIDIA CUDA 백엔드에서만 지원되었는데, 이 PR은 AMD HIP 백엔드에서도 ConSan을 사용할 수 있도록 대규모 인프라를 구축합니다.

핵심 코드 분석

MBarrierOpInterface 도입

기존에는 NVIDIA 전용 dialect에 직접 의존했으나, 벤더 중립적인 인터페이스를 추가:

def MBarrierOpInterface : OpInterface<"MBarrierOpInterface"> {
  let methods = [
    InterfaceMethod<
      "Return the barrier memdesc operand, or null if absent.",
      "::mlir::Value", "getBarrierMemDesc">
  ];
}

Target Hook 확장

AMD 전용 commit kind와 ordered commit 처리:

virtual SmallVector<CommitKindDesc> getOutstandingWriteCommitKinds() const {
    return {{CommitKind::AsyncCp, "async_copy_global_to_shared"}};
}

// TDM ops complete in-order within a warp
virtual bool isOrderedCommitKind(CommitKind::Kind kind) const {
    return false;  // AMD override에서 true 반환 가능
}

Capture Count 추정

WarpSpecialize의 shared memory 할당을 위해 ConSan이 추가할 capture 수를 미리 계산:

inline int estimateConSanCaptureCount(int numActiveMemTypes,
                                      bool hasBarriers, int numCommitKinds) {
    int perMemType = kCapturesPerMemType * numActiveMemTypes;
    int barrierCaptures = hasBarriers ?
        kBarrierBaseCaptures + kBarrierTrackingCapturesPerMemType * numActiveMemTypes : 0;
    return perMemType + barrierCaptures + kFixedCaptures + numCommitKinds;
}

왜 이게 좋은가

벤더 중립적인 MBarrierOpInterface를 도입하여 ConSan이 NVIDIA/AMD 어느 쪽의 barrier 연산이든 동일한 인터페이스로 분석할 수 있습니다. isOrderedCommitKind로 TDM 같은 in-order completion을 가진 연산에서의 false positive를 방지합니다. Barrier state의 비트 필드를 10비트로 확장(bits [1..10] initial count, bits [11..20] current count)하여 더 많은 warp를 지원합니다.

정리

AMD 백엔드에 ConSan을 지원하기 위해 MBarrierOpInterface, AMD 전용 target hook, capture count 추정, 그리고 barrier state 비트 확장을 구현했습니다.

참고 자료

• triton-lang/triton#9692

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