[llvm] [AMDGPU] Preserve chain when selecting llvm.amdgcn.pops.exiting.wave.id (PR #96167)

[Jay Foad via llvm-commits]

https://github.com/jayfoad created https://github.com/llvm/llvm-project/pull/96167

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


>From 14f481a456f817c2de35330d63a0306342179ecd Mon Sep 17 00:00:00 2001
From: Jay Foad <jay.foad at amd.com>
Date: Thu, 20 Jun 2024 12:10:30 +0100
Subject: [PATCH] [AMDGPU] Preserve chain when selecting
 llvm.amdgcn.pops.exiting.wave.id

Without this SelectionDAG could fail assertions when using the intrinsic
in a non-entry BB.
---
 llvm/lib/Target/AMDGPU/AMDGPUISelDAGToDAG.cpp |  3 +-
 .../llvm.amdgcn.pops.exiting.wave.id.ll       | 39 +++++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)

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
+}