From 10c2f6bb1bdb84c6305329c97d3daa5923d7bdb7 Mon Sep 17 00:00:00 2001
From: Matt Arsenault <Matthew.Arsenault@amd.com>
Date: Sat, 21 Sep 2024 10:18:54 +0400
Subject: [PATCH] RegisterCoalescer: Set undef on full register uses when
 coalescing implicit_def

Previously this would delete the IMPLICIT_DEF and not introduce the undef
flag on the use operand.

Fixes sub-issue found while reducing #109294
---
 llvm/lib/CodeGen/RegisterCoalescer.cpp        |  5 ++-
 ...reg-use-implicit-def-erase-issue109249.mir | 43 +++++++++++++++++++
 2 files changed, 47 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/CodeGen/AMDGPU/register-coalescer--set-undef-full-reg-use-implicit-def-erase-issue109249.mir

diff --git a/llvm/lib/CodeGen/RegisterCoalescer.cpp b/llvm/lib/CodeGen/RegisterCoalescer.cpp
index 5de873fd41578..20ad6445344d8 100644
--- a/llvm/lib/CodeGen/RegisterCoalescer.cpp
+++ b/llvm/lib/CodeGen/RegisterCoalescer.cpp
@@ -1842,9 +1842,12 @@ void RegisterCoalescer::updateRegDefsUses(Register SrcReg, Register DstReg,
 
   if (DstInt && DstInt->hasSubRanges() && DstReg != SrcReg) {
     for (MachineOperand &MO : MRI->reg_operands(DstReg)) {
+      if (MO.isUndef())
+        continue;
       unsigned SubReg = MO.getSubReg();
-      if (SubReg == 0 || MO.isUndef())
+      if (SubReg == 0 && MO.isDef())
         continue;
+
       MachineInstr &MI = *MO.getParent();
       if (MI.isDebugInstr())
         continue;
diff --git a/llvm/test/CodeGen/AMDGPU/register-coalescer--set-undef-full-reg-use-implicit-def-erase-issue109249.mir b/llvm/test/CodeGen/AMDGPU/register-coalescer--set-undef-full-reg-use-implicit-def-erase-issue109249.mir
new file mode 100644
index 0000000000000..3f5df7a7b868a
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/register-coalescer--set-undef-full-reg-use-implicit-def-erase-issue109249.mir
@@ -0,0 +1,43 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 5
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -mcpu=gfx900 -run-pass=register-coalescer -verify-coalescing -o - %s | FileCheck %s
+
+# Make sure that the undef flag is set on %0 after the IMPLICIT_DEF is
+# deleted when coalescing %0 with %1
+
+---
+name:            test
+tracksRegLiveness: true
+machineFunctionInfo:
+  stackPtrOffsetReg: '$sgpr32'
+body:             |
+  bb.0:
+    ; CHECK-LABEL: name: test
+    ; CHECK: [[S_BUFFER_LOAD_DWORDX2_IMM:%[0-9]+]]:sreg_64_xexec = S_BUFFER_LOAD_DWORDX2_IMM undef %0, 36, 0 :: (dereferenceable invariant load (s64))
+    ; CHECK-NEXT: undef [[S_ADD_U32_:%[0-9]+]].sub1:sgpr_128 = S_ADD_U32 [[S_BUFFER_LOAD_DWORDX2_IMM]].sub0, 32, implicit-def dead $scc
+    ; CHECK-NEXT: SI_RETURN implicit [[S_ADD_U32_]].sub1
+    %0:sgpr_128 = IMPLICIT_DEF
+    %1:sreg_64_xexec = S_BUFFER_LOAD_DWORDX2_IMM %0, 36, 0 :: (dereferenceable invariant load (s64))
+    %2:sreg_32 = S_ADD_U32 %1.sub0, 32, implicit-def dead $scc
+    %0.sub1:sgpr_128 = COPY killed %2
+    SI_RETURN implicit %0.sub1
+
+...
+
+---
+name:            test_w_undef_dead
+tracksRegLiveness: true
+machineFunctionInfo:
+  stackPtrOffsetReg: '$sgpr32'
+body:             |
+  bb.0:
+    ; CHECK-LABEL: name: test_w_undef_dead
+    ; CHECK: dead [[S_BUFFER_LOAD_DWORDX2_IMM:%[0-9]+]]:sreg_64_xexec = S_BUFFER_LOAD_DWORDX2_IMM undef %0, 36, 0 :: (dereferenceable invariant load (s64))
+    ; CHECK-NEXT: undef [[S_ADD_U32_:%[0-9]+]].sub1:sgpr_128 = S_ADD_U32 undef [[S_BUFFER_LOAD_DWORDX2_IMM]].sub0, 32, implicit-def dead $scc
+    ; CHECK-NEXT: SI_RETURN implicit [[S_ADD_U32_]].sub1
+    %0:sgpr_128 = IMPLICIT_DEF
+    dead %1:sreg_64_xexec = S_BUFFER_LOAD_DWORDX2_IMM %0, 36, 0 :: (dereferenceable invariant load (s64))
+    %2:sreg_32 = S_ADD_U32 undef %1.sub0, 32, implicit-def dead $scc
+    %0.sub1:sgpr_128 = COPY killed %2
+    SI_RETURN implicit %0.sub1
+
+...