[triton] Global Sanitizer에 TMA 및 cp.async 연산 부분 지원 추가

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

들어가며

Triton의 Global Sanitizer(GSan)는 GPU 커널의 메모리 접근 오류를 런타임에 탐지하는 도구입니다. 기존에는 일반 load/store만 추적했지만, TMA(Tensor Memory Accelerator)와 cp.async 같은 비동기 메모리 연산은 지원하지 않았습니다. 이 PR은 tensor descriptor에서 base pointer/shape/stride를 디코딩하는 기능을 추가하여 이러한 연산의 메모리 접근도 검증할 수 있게 합니다.

핵심 코드 분석

Tensor Descriptor 디코딩 LLVM 변환

// tensor descriptor 구조체에서 shape/stride 추출
static SmallVector<Value>
decodeTensorDescShape(ConversionPatternRewriter &rewriter, Location loc,
                      Value descStruct, unsigned rank) {
  SmallVector<Value> shape;
  for (unsigned dim = 0; dim < rank; ++dim) {
    unsigned packedIdx = rank - 1 - dim;
    Value dimMinusOne = extractTensorDescWord(
        rewriter, loc, descStruct, kTensorMapShapeWordBase + packedIdx);
    shape.push_back(b.add(dimMinusOne, b.i64_val(1)));  // shape는 -1로 저장됨
  }
  return shape;
}

Before - TMA 연산 미추적

// GlobalSanitizer.cpp에서 TMA 연산 미처리
// 일반 load/store만 instrumentation

After - TMA 정보 추출 연산 추가

// 새로운 ExperimentalGSanTensorDescInfoOp
def TTI_ExperimentalGSanTensorDescInfoOp
    : TTI_Op<"experimental_gsan_tensordesc_info", [Pure]> {
  let summary = "Decode GSan descriptor metadata from a native tensor descriptor";
  let arguments = (ins TT_AnyTensorDescType:$desc);
  let results = (outs Variadic<AnyType>:$result);
}

왜 이게 좋은가

1. 비동기 메모리 접근 검증: TMA와 cp.async의 out-of-bounds 접근을 런타임에 탐지할 수 있습니다.
2. Descriptor 투명성: tensor descriptor의 내부 구조(base ptr, shape, stride)를 IR 수준에서 접근할 수 있어 디버깅이 용이해집니다.
3. 점진적 확장: "Partially support"라는 제목처럼 완전한 지원이 아닌 점진적 확장으로 안정성을 확보합니다.

정리

GSan의 커버리지를 TMA/cp.async까지 확장한 중요한 개선입니다. 특히 tensor descriptor에서 shape/stride를 디코딩하는 로직은 sanitizer뿐 아니라 다른 분석 도구에서도 활용 가능합니다.

참고 자료

• PR #9699

---

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