[triton] FPSan에서 Warp Specialization + TMem 사용 시 크래시 수정

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

들어가며

FPSan(Floating-Point Sanitizer)은 Triton 커널의 부동소수점 연산 오류를 감지하는 도구입니다. Warp Specialization과 tensor memory를 동시에 사용할 때, FPSan이 생성하는 scratch pointer가 partition region 외부에 정의되어 있어 접근할 수 없는 크래시가 발생했습니다.

핵심 코드 분석

Before (scope 확인 없이 직접 사용):

ScratchInfo info{ptr, tensorTy};
scratchMap[memdesc][scope] = info;
return info;

After (scope 가용성 확인 + capture로 전달):

Value remapToScope(Value value, PatternRewriter &rewriter, Region *scope,
                   Location loc) {
    if (!scope || isValueAvailableInScope(value, scope))
        return value;

    auto partitions = dyn_cast_or_null<ttg::WarpSpecializePartitionsOp>(
        parentOp);
    if (!partitions) return value;

    // 기존 capture에서 찾거나 새로 추가
    unsigned captureIdx = partitions.getNumOperands();
    for (auto [i, capture] : llvm::enumerate(partitions.getExplicitCaptures())) {
        if (capture == value) { captureIdx = i; break; }
    }
    if (captureIdx == partitions.getNumOperands()) {
        partitions->insertOperands(captureIdx, value);
        for (Region &region : partitions.getPartitionRegions())
            region.addArgument(value.getType(), loc);
    }
    return scope->getArgument(captureIdx);
}

isValueAvailableInScope 유틸리티 함수로 값이 현재 scope에서 접근 가능한지 확인합니다:

static bool isValueAvailableInScope(Value value, Region *scope) {
    if (auto arg = dyn_cast<BlockArgument>(value)) {
        Region *argRegion = arg.getOwner()->getParent();
        return argRegion == scope || scope->isAncestor(argRegion);
    }
    // DefiningOp 체크도 유사
}

왜 이게 좋은가

WarpSpecialize는 각 파티션을 별도의 Region으로 분리하므로, 외부에서 정의된 값은 명시적 capture를 통해서만 접근할 수 있습니다. FPSan이 생성하는 scratch pointer는 global_scratch_alloc으로 함수 시작에서 할당되므로 파티션 Region 외부에 있습니다. remapToScope는 이를 자동으로 감지하고 capture list에 추가하여 MLIR의 scope 규칙을 위반하지 않도록 합니다. 이는 ptxas 13.x 컴파일러 버그 workaround(-Ofc mid)도 함께 적용합니다.

정리

FPSan의 TmemScratchManager에서 생성되는 pointer를 WarpSpecialize 파티션의 capture list에 자동으로 추가하는 remapToScope 메서드를 도입하여 scope 위반 크래시를 수정했습니다.

참고 자료

• triton-lang/triton#9415

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