[Triton] AMD amdgpu.async_wait Op 도입으로 비동기 트랜잭션 의미론 명확화

PR 링크: triton-lang/triton#8575 상태: Merged | 변경: +98 / -43

들어가며

Triton의 ttg.async_wait은 outstanding ttg.commit_group 수를 카운트한다. 그러나 AMD GPU에서 LLVM으로 lowering할 때는 outstanding async intrinsic/assembly 명령어 수가 필요하다. 기존에는 UpdateAsyncWaitCnt pass가 ttg.async_wait의 num attribute를 in-place로 수정했지만, 이는 의미론이 변경되는 것을 IR에서 명시적으로 표현하지 못했다. 이 PR은 amdgpu.async_wait Op을 새로 도입하여 이 의미론 변경을 명확히 한다.

핵심 코드 분석

Before: 같은 Op에서 의미가 바뀜

// TTGIR 수준: num = outstanding commit groups
ttg.async_wait {num = 0 : i32}

// UpdateAsyncWaitCnt 후: 같은 Op이지만 num이 하드웨어 명령어 수로 변경됨
ttg.async_wait {num = 2 : i32}

After: 별도 Op으로 의미론 분리

// TTGIR 수준: commit group 기반
ttg.async_wait %token {num = 0 : i32}

// UpdateAsyncWaitCount 후: 새 Op으로 변환, num_inst는 하드웨어 명령어 수
// CHECK: amdgpu.async_wait %token {num_inst = 2}
amdgpu.async_wait %token {num_inst = 2 : i32}

membar 호환성

amdgpu.async_wait 후에 LocalBarrier를 자동 삽입하여, membar pass가 기존 ttg.async_wait과 동일하게 동작하도록 한다.

// amdgpu.async_wait 후 자동으로 barrier 삽입
amdgpu.async_wait {num_inst = 0 : i32}
// rocdl.s.waitcnt ...
// rocdl.s.barrier  <- LocalBarrier에 의해 생성

왜 이게 좋은가

1. IR 의미론 명확성: commit group 기반 ttg.async_wait과 하드웨어 명령어 기반 amdgpu.async_wait이 IR에서 명시적으로 구분된다.
2. Gluon 호환: UpdateAsyncWaitCount를 TTGIR->LLVM 단계로 이동하여 Gluon 커널에서도 자동 적용된다.
3. membar 투명성: LocalBarrier 추가로 기존 membar pass와의 호환성을 유지한다.
4. 디버깅 용이성: 변환 전후의 Op이 다르므로, IR 덤프 시 어느 단계에 있는지 명확히 알 수 있다.

정리

이 PR은 AMD GPU 전용 amdgpu.async_wait Op을 도입하여, ttg.async_wait(commit group 의미론)과 하드웨어 명령어 수 기반 대기를 IR 수준에서 명확히 분리했다. UpdateAsyncWaitCount pass를 LLVM lowering 단계로 이동하여 Gluon 커널에서도 자동 적용되며, LocalBarrier 삽입으로 기존 membar pass와의 호환성을 유지한다.

참고 자료

• triton-lang/triton#8575

---

이 글은 AI를 활용하여 PR의 핵심 변경사항을 분석하고 정리한 것입니다. 실제 코드의 맥락은 원본 PR을 참고해 주세요.