[triton] AMD ConvertWarpPipeline에서 AsyncWaitOp 인식 및 Barrier 정렬 수정

PR 링크: triton-lang/triton#9593 상태: Merged | 변경: +83 / -2

들어가며

AMD의 warp pipeline은 루프 본문의 연산을 여러 스테이지로 나누어 파이프라인 병렬 처리합니다. ConvertPipelinedForPattern은 스테이지 사이의 barrier를 인식하여 동기화 포인트를 삽입하는데, amdg.async_wait 연산이 누락되어 있었습니다.

핵심 코드 분석

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 추가로 cluster/bar 어긋남

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) 제거

테스트

tt.func @async_wait_between_stages(%n: index, %ptr: !tt.ptr<f32>) {
  scf.for %i = %c0 to %n step %c1 {
    scf.execute_region { ... } {triton.warp_pipeline.stage = "stage1"}
    amdg.async_wait {num_inst = 0 : i32}  // 이제 barrier로 인식
    scf.execute_region { ... } {triton.warp_pipeline.stage = "stage2"}
  } {triton.warp_pipeline.pipelined_for}
}

왜 이게 좋은가

1. 완전한 barrier 인식: async wait를 포함한 모든 동기화 연산이 파이프라인에서 올바르게 처리됩니다.
2. 배열 정렬 수정: bars와 clusterBlocks 배열의 1:1 대응이 보장됩니다.
3. E2E 테스트: Python 수준의 warp pipeline 컴파일 테스트도 추가되었습니다.

정리

Warp pipeline의 barrier 인식 목록에 amdg.async_wait를 추가하고, bars 배열 정렬 버그를 수정한 PR입니다. 파이프라인 pass에서 새로운 동기화 연산의 누락은 미묘한 correctness 버그를 유발하므로 주의가 필요합니다.

참고 자료

• PR #9593

---

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