[Triton] GSan(Global Sanitizer) warp specialized 커널에서의 deadlock 수정
들어가며
Triton의 GSan(Global Sanitizer)은 메모리 접근 오류를 런타임에 감지하는 도구다. Warp specialized 커널에 GSan을 적용하면 두 가지 deadlock이 발생했다: (1) GSan 함수의 높은 레지스터 요구로 인한 USETMAXREG.TRY_ALLOC deadlock, (2) __triton_gsan_init의 bar.sync가 worker warp에서 호출되지 않아 발생하는 barrier deadlock.
핵심 코드 분석
Deadlock 1: Dynamic register allocation
// After: warp_specialize op에 속성 추가하여 동적 레지스터 할당 비활성화
module.walk([&](ttg::WarpSpecializeOp op) {
op->setAttr(kDisableSetMaxRegisterAttr, builder.getUnitAttr());
});
// lowering 시 속성 확인
bool emitDynamicRegRealloc = llvm::none_of(wsOps, [](WarpSpecializeOp ws) {
return ws->hasAttr(kDisableSetMaxRegisterAttr);
});
if (emitDynamicRegRealloc && maxnreg)
createRegRealloc(b, maxnreg.getInt(), defRegs);
Deadlock 2: Init barrier
// Before: CUDA device code에서 bar.sync 호출 (모든 warp 참여 필요)
__device__ void initThread(GlobalState *globals, Location loc) {
// ... 초기화 ...
ctaBarrier(); // worker warp가 호출하지 않아 deadlock!
}
// After: LLVM IR lowering에서 barrier를 삽입 (warp specialization을 이해하는 위치)
b.call(runtimeFunc, ValueRange{gsanGlobalStatePtr, ...});
b.barrier(ttg::AddrSpace::Local); // 여기서 올바르게 삽입
왜 이게 좋은가
- 실용적 해결: 두 가지 독립적인 deadlock을 각각 적절한 수준에서 해결했다.
- 관심사 분리: barrier 동기화를 CUDA 런타임이 아닌 LLVM IR lowering에서 처리하여, warp specialization 문맥을 올바르게 반영한다.
- 테스트 추가: warp specialized 커널에서 GSan이 정상 동작하는 end-to-end 테스트가 포함되었다.
정리
+76/-6의 작은 변경이지만, GPU 도구 개발에서 warp specialization과의 상호작용이 얼마나 까다로운지를 보여주는 PR이다. 동기화 프리미티브는 모든 참여 스레드가 도달할 수 있는 위치에 있어야 한다.
참고 자료
이 글은 AI(Claude)의 도움을 받아 작성되었습니다. 코드 분석 내용은 실제 PR diff를 기반으로 합니다.
댓글