[AMDGPU] Preserve chain when selecting llvm.amdgcn.pops.exiting.wave.id

[jayfoad]

Without this SelectionDAG could fail assertions when using the intrinsic
in a non-entry BB.

[llvmbot]

@llvm/pr-subscribers-backend-amdgpu

Author: Jay Foad (jayfoad)

Changes

Without this SelectionDAG could fail assertions when using the intrinsic
in a non-entry BB.

---

Full diff: https://github.com/llvm/llvm-project/pull/96167.diff

2 Files Affected:

• (modified) llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp (+2-1)
• (modified) llvm/test/CodeGen/AMDGPU/llvm.amdgcn.pops.exiting.wave.id.ll (+39)

diff --git a/llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp b/llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp
index d60c62a30e7ed..5b616d8cc77d1 100644
--- a/llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp
+++ b/llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp
@@ -2554,7 +2554,8 @@ void AMDGPUDAGToDAGISel::SelectPOPSExitingWaveID(SDNode *N) {
   // intrinsic is IntrReadMem/IntrWriteMem but the instruction is not marked
   // mayLoad/mayStore and tablegen complains about the mismatch.
   SDValue Reg = CurDAG->getRegister(AMDGPU::SRC_POPS_EXITING_WAVE_ID, MVT::i32);
-  CurDAG->SelectNodeTo(N, AMDGPU::S_MOV_B32, N->getVTList(), Reg);
+  SDValue Chain = N->getOperand(0);
+  CurDAG->SelectNodeTo(N, AMDGPU::S_MOV_B32, N->getVTList(), {Reg, Chain});
 }
 
 static unsigned gwsIntrinToOpcode(unsigned IntrID) {
diff --git a/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.pops.exiting.wave.id.ll b/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.pops.exiting.wave.id.ll
index 4927c2ffcdf30..de9e40727b8ed 100644
--- a/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.pops.exiting.wave.id.ll
+++ b/llvm/test/CodeGen/AMDGPU/llvm.amdgcn.pops.exiting.wave.id.ll
@@ -32,3 +32,42 @@ define amdgpu_ps void @test(ptr addrspace(1) inreg %ptr) {
   store i32 %id, ptr addrspace(1) %ptr
   ret void
 }
+
+define amdgpu_ps void @test_loop() {
+; SDAG-LABEL: test_loop:
+; SDAG:       ; %bb.0:
+; SDAG-NEXT:    s_mov_b32 s0, src_pops_exiting_wave_id
+; SDAG-NEXT:  .LBB1_1: ; %loop
+; SDAG-NEXT:    ; =>This Inner Loop Header: Depth=1
+; SDAG-NEXT:    s_cmp_eq_u32 s0, 0
+; SDAG-NEXT:    s_cbranch_scc1 .LBB1_1
+; SDAG-NEXT:  ; %bb.2: ; %exit
+; SDAG-NEXT:    s_endpgm
+;
+; GFX9-GISEL-LABEL: test_loop:
+; GFX9-GISEL:       ; %bb.0:
+; GFX9-GISEL-NEXT:    s_mov_b32 s0, src_pops_exiting_wave_id
+; GFX9-GISEL-NEXT:  .LBB1_1: ; %loop
+; GFX9-GISEL-NEXT:    ; =>This Inner Loop Header: Depth=1
+; GFX9-GISEL-NEXT:    s_cmp_eq_u32 s0, 0
+; GFX9-GISEL-NEXT:    s_cbranch_scc1 .LBB1_1
+; GFX9-GISEL-NEXT:  ; %bb.2: ; %exit
+; GFX9-GISEL-NEXT:    s_endpgm
+;
+; GFX10-GISEL-LABEL: test_loop:
+; GFX10-GISEL:       ; %bb.0:
+; GFX10-GISEL-NEXT:    s_mov_b32 s0, src_pops_exiting_wave_id
+; GFX10-GISEL-NEXT:  .LBB1_1: ; %loop
+; GFX10-GISEL-NEXT:    ; =>This Inner Loop Header: Depth=1
+; GFX10-GISEL-NEXT:    s_cmp_eq_u32 s0, 0
+; GFX10-GISEL-NEXT:    s_cbranch_scc1 .LBB1_1
+; GFX10-GISEL-NEXT:  ; %bb.2: ; %exit
+; GFX10-GISEL-NEXT:    s_endpgm
+  br label %loop
+loop:
+  %id = call i32 @llvm.amdgcn.pops.exiting.wave.id()
+  %cond = icmp eq i32 %id, 0
+  br i1 %cond, label %loop, label %exit
+exit:
+  ret void
+}