[triton] AMD membarFilter에 bufferID 고려 추가

PR 링크: triton-lang/triton#9265 상태: Merged | 변경: +82 / -20

들어가며

Triton의 membar 분석은 shared memory 접근 간에 barrier가 필요한지 판단합니다. AMD 백엔드의 membarFilter는 async load와 sync된 local load 간에 barrier를 건너뛸 수 있는데, 기존에는 두 연산이 동일한 buffer를 사용하는지 확인하지 않아 false negative이 발생할 수 있었습니다.

핵심 코드 분석

1. MembarFilterFn 시그니처 변경

Before:

using MembarFilterFn = std::function<bool(Operation *, Operation *)>;

After:

using MembarFilterFn =
    std::function<bool(Operation *, Operation *, Allocation *)>;

Allocation 포인터를 추가로 받아, filter 함수에서 buffer 할당 정보에 접근할 수 있게 되었습니다.

2. Buffer ID 기반 memdesc 추출

auto getMemdescValue = [](Operation *op) -> Value {
    return llvm::TypeSwitch<Operation *, Value>(op)
        .Case<triton::amdgpu::BufferLoadToLocalOp>(
            [](auto op) { return op.getDest(); })
        .Case<triton::gpu::AsyncCopyGlobalToLocalOp>(
            [](auto op) { return op.getResult(); })
        .Case<triton::gpu::LocalLoadOp>([](auto op) { return op.getSrc(); })
        .Default([](Operation *) { return Value(); });
};

각 연산 유형에서 memdesc value를 추출하여, 두 연산이 같은 shared memory buffer를 참조하는지 확인합니다.

3. 재사용된 allocation 감지 테스트

// dealloc 후 동일 offset에 재할당된 경우 barrier 필요
%alloc1 = ttg.local_alloc {allocation.offset = 0 : i32}
// ... load into alloc1, local_load from alloc1 ...
ttg.local_dealloc %alloc1
%alloc2 = ttg.local_alloc {allocation.offset = 0 : i32}
// CHECK: ttg.barrier local
// CHECK-NEXT: amdg.buffer_load_to_local
%async2 = amdg.buffer_load_to_local %B[%offset] into %slice2_0

왜 이게 좋은가

• 정확성 보장: Buffer ID 확인으로 재사용된 allocation 간 누락된 barrier를 감지합니다.
• 성능 유지: 같은 buffer를 공유하는 async load/sync load 간의 불필요한 barrier는 여전히 건너뜁니다.
• 범용 인터페이스: Allocation *를 filter에 전달하여 다른 백엔드에서도 유사한 최적화가 가능합니다.

정리

AMD membar filter에 buffer ID 검증을 추가하여, 같은 물리적 메모리를 다른 논리적 buffer로 재사용하는 경우의 barrier 누락 버그를 수정한 PR입니다.

참고 자료

• triton-lang/triton#9265

---

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