[triton] AMD Software Warp Pipeline에서 크래시 수정

PR 링크: triton-lang/triton#9631 상태: Merged | 변경: +10 / -6

들어가며

AMD GPU의 Warp Pipeline 변환 pass에서 amdg.async_wait 연산이 스테이지 간 barrier로 인식되지 않아 파이프라인 스케줄링이 실패하고 크래시가 발생했습니다. 이 PR은 해당 연산을 barrier 목록에 추가하고, barrier 카운트 정렬 버그도 함께 수정합니다.

핵심 코드 분석

Before

} else if (isa<ROCDL::BarrierOp, gpu::BarrierOp, triton::gpu::AsyncWaitOp,
               triton::amdgpu::AsyncTDMWait,
               triton::amdgpu::AsyncTDMIntrinsicWait>(op)) {
  // amdgpu::AsyncWaitOp 누락!
  existingBarrierMap[currCluster] = &op;
  bars.push_back(false);  // 잘못된 bars 추가

After

} else if (isa<ROCDL::BarrierOp, gpu::BarrierOp, triton::gpu::AsyncWaitOp,
               triton::amdgpu::AsyncWaitOp,  // 추가
               triton::amdgpu::AsyncTDMWait,
               triton::amdgpu::AsyncTDMIntrinsicWait>(op)) {
  existingBarrierMap[currCluster] = &op;
  // bars.push_back(false) 제거 - barrier는 cluster가 아님

왜 이게 좋은가

1. 크래시 수정: amdg.async_wait가 스테이지 사이에 올 때 더 이상 크래시하지 않습니다.
2. bars 정렬 수정: barrier를 bars 벡터에 잘못 추가하여 cluster와 barrier 배열이 어긋나던 문제를 해결했습니다.
3. 테스트 추가: async_wait_between_stages MLIR 테스트로 이 시나리오를 검증합니다.

정리

10줄 변경으로 warp pipeline pass의 크래시를 수정한 PR입니다. 새로운 async 연산을 추가할 때 pipeline pass의 barrier 인식 목록도 함께 업데이트해야 한다는 교훈을 줍니다.

참고 자료

• PR #9631

---

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