[triton] CLCTryCancel이 Async Proxy를 사용하도록 수정

PR 링크: triton-lang/triton#9464 상태: Merged | 변경: +5 / -1

들어가며

NVIDIA GPU의 CLC(Compute Linking and Caching) 명령어 중 CLCTryCancel은 진행 중인 비동기 복사를 취소하는 연산입니다. 이 연산은 shared memory에 데이터를 쓰는 async proxy write이지만, 기존에는 proxy fence 삽입 pass에서 인식되지 않아 fence가 누락될 수 있었습니다.

핵심 코드 분석

Before

bool isAsyncProxyWrite(Operation *op) {
  return isa<triton::nvidia_gpu::AsyncTMACopyGlobalToLocalOp,
             triton::nvidia_gpu::AsyncTMAGatherOp>(op);
  // CLCTryCancelOp 누락
}

After

bool isAsyncProxyWrite(Operation *op) {
  return isa<triton::nvidia_gpu::AsyncTMACopyGlobalToLocalOp,
             triton::nvidia_gpu::AsyncTMAGatherOp,
             triton::nvidia_gpu::CLCTryCancelOp>(op);  // 추가
}

Value getSmemDest(Operation *op) {
  // ...
  if (auto clcTryCancelOp = dyn_cast<triton::nvidia_gpu::CLCTryCancelOp>(op)) {
    return clcTryCancelOp.getResult();  // SMEM 대상 반환 추가
  }
}

왜 이게 좋은가

1. 정확한 fence 삽입: CLCTryCancel 후 shared memory 데이터를 사용할 때 proxy fence가 올바르게 삽입됩니다.
2. 메모리 안전성: fence 누락으로 인한 데이터 경쟁 버그를 방지합니다.
3. 작은 변경, 큰 영향: 5줄 추가만으로 잠재적인 correctness 버그를 해결합니다.

정리

비동기 메모리 연산의 proxy fence 누락 버그를 5줄 수정으로 해결한 PR입니다. 새로운 async 연산을 추가할 때 fence 관련 목록도 함께 업데이트해야 한다는 교훈을 줍니다.

참고 자료

• PR #9464

---

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