[llvm] [CodeGen][AMDGPU] Reg. Coalescer: Update live intervals of Src reg when replacing COPY for IMPLICIT_DEF. (PR #175986)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-backend-amdgpu

Author: Felipe Quezada (kezada94)

<details>
<summary>Changes</summary>

This change tries to fix live intervals when coalescing a COPY instr into an IMPLICIT_DEF. 
If the COPY is replaced by an IMPLICIT_DEF and the operand is a sub register without live subranges, removing it could result in incorrect live information for the complete virtual register, specifically when that COPY lies at the end index of the register live interval.
This PR checks if that condition is met, in such cases, it updates the live intervals of the Src register.
Although it's not a common case, adding this check should improve robustness.

Fixes issue: https://github.com/llvm/llvm-project/issues/175985

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


2 Files Affected:

- (modified) llvm/lib/CodeGen/RegisterCoalescer.cpp (+10-1) 
- (added) llvm/test/CodeGen/AMDGPU/regcoal-update-src-subrange-on-removed-lane-use.mir (+26) 


``````````diff
diff --git a/llvm/lib/CodeGen/RegisterCoalescer.cpp b/llvm/lib/CodeGen/RegisterCoalescer.cpp
index ffbdff4f45150..c2cf838378a40 100644
--- a/llvm/lib/CodeGen/RegisterCoalescer.cpp
+++ b/llvm/lib/CodeGen/RegisterCoalescer.cpp
@@ -1770,13 +1770,15 @@ MachineInstr *RegisterCoalescer::eliminateUndefCopy(MachineInstr *CopyMI) {
     return nullptr;
 
   SlotIndex Idx = LIS->getInstructionIndex(*CopyMI);
-  const LiveInterval &SrcLI = LIS->getInterval(SrcReg);
+  LiveInterval &SrcLI = LIS->getInterval(SrcReg);
+  bool subRegHasLiveSubrange = false;
   // CopyMI is undef iff SrcReg is not live before the instruction.
   if (SrcSubIdx != 0 && SrcLI.hasSubRanges()) {
     LaneBitmask SrcMask = TRI->getSubRegIndexLaneMask(SrcSubIdx);
     for (const LiveInterval::SubRange &SR : SrcLI.subranges()) {
       if ((SR.LaneMask & SrcMask).none())
         continue;
+      subRegHasLiveSubrange = true;
       if (SR.liveAt(Idx))
         return nullptr;
     }
@@ -1807,6 +1809,13 @@ MachineInstr *RegisterCoalescer::eliminateUndefCopy(MachineInstr *CopyMI) {
     }
 
     CopyMI->setDesc(TII->get(TargetOpcode::IMPLICIT_DEF));
+    // If there wasn't a live subrange for this Src subregister, and it's at the
+    // end of the register segment, we need to update the live intervals of the
+    // Src reg.
+    if (SrcSubIdx != 0 && !subRegHasLiveSubrange &&
+        SrcLI.endIndex().getBaseIndex() == Idx) {
+      LIS->shrinkToUses(&SrcLI);
+    }
     LLVM_DEBUG(dbgs() << "\tReplaced copy of <undef> value with an "
                          "implicit def\n");
     return CopyMI;
diff --git a/llvm/test/CodeGen/AMDGPU/regcoal-update-src-subrange-on-removed-lane-use.mir b/llvm/test/CodeGen/AMDGPU/regcoal-update-src-subrange-on-removed-lane-use.mir
new file mode 100644
index 0000000000000..61b757ab47978
--- /dev/null
+++ b/llvm/test/CodeGen/AMDGPU/regcoal-update-src-subrange-on-removed-lane-use.mir
@@ -0,0 +1,26 @@
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -run-pass=register-coalescer -debug-only=regalloc -o /dev/null %s 2>&1 | FileCheck %s --check-prefix=CHECK-DBG
+# RUN: llc -mtriple=amdgcn-amd-amdhsa -run-pass=register-coalescer -o - %s | FileCheck %s --check-prefix=CHECK
+
+# CHECK-DBG: ********** REGISTER COALESCER **********
+# CHECK-DBG: ********** Function: test
+# CHECK-DBG: ********** JOINING INTERVALS ***********
+# CHECK-DBG: ********** INTERVALS **********
+# CHECK-DBG: %0 [16r,32r:0) 0 at 16r  weight:0.000000e+00
+# CHECK-DBG: %1 [32r,32d:0) 0 at 32r  L0000000000000003 [32r,32d:0) 0 at 32r  weight:0.000000e+00
+# CHECK-DBG: %2 [48r,64r:0) 0 at 48r  weight:0.000000e+00
+# CHECK-DBG: %3 [64r,64d:0) 0 at 64r  weight:0.000000e+00
+---
+name: test
+tracksRegLiveness: true
+body: |
+  bb.0:
+    ; CHECK-LABEL: name: test
+    ; CHECK: [[DST:%[0-9]+]]:sgpr_32 = IMPLICIT_DEF
+
+    %0:vgpr_32 = IMPLICIT_DEF
+    undef %9.sub0:sgpr_64 = V_READFIRSTLANE_B32 %0:vgpr_32, implicit $exec
+    %12:sgpr_32 = COPY %9.sub1:sgpr_64, implicit $exec
+    %16:vgpr_32 = V_MOV_B32_e32 %12:sgpr_32, implicit $exec
+
+    S_ENDPGM 0
+...

``````````

</details>


https://github.com/llvm/llvm-project/pull/175986