[RegisterCoalescer] dst register's live interval needs to be updated when merging a src register in ToBeUpdated set

[wmi]

Summary

This is to fix the bug in PR40061 related with https://reviews.llvm.org/rL339035.

In https://reviews.llvm.org/rL339035, live interval of source pseudo register in rematerialized copy is saved in ToBeUpdated set and its update is postponed.

In PR40061, %t2 = %t1 is rematerialized and %t1 is added into toBeUpdated set to postpone its live interval update. After the rematerialization, the live interval of %t1 is larger than necessary. Then %t1 is merged into %t3 and %t1 gets removed. After the merge, %t3 contains live interval larger than necessary. Because %t3 is not in toBeUpdated set, its live interval is not updated after register coalescing and it will break some assumption in regalloc.

The patch requires the live interval of destination register in a merge to be updated if the source register is in ToBeUpdated.

[wmi]

A minor update to the comment.

[skatkov]

I've verified that this patch fixes the original problem.

[qcolombet]

Looks mostly good to me.

test/CodeGen/X86/late-remat-update-2.mir

1

Could you come up with checks that don't require asserts?

[wmi]

Address Quentin's comment.

[wmi]

test/CodeGen/X86/late-remat-update-2.mir

1

Done. Thanks for the suggestion!

[skatkov]

Ping, any progress on this?

[wmi]

friendly ping. Please take another look.

[qcolombet]

LGTM

lib/CodeGen/RegisterCoalescer.cpp

for (LiveInterval::SubRange &S : LI.subranges()) {
       if ((S.LaneMask & ShrinkMask).none())
         continue;
       LLVM_DEBUG(dbgs() << "Shrink LaneUses (Lane " << PrintLaneMask(S.LaneMask)
                         << ")\n");
       LIS->shrinkToUses(S, LI.reg);
     }
     LI.removeEmptySubRanges();
   }
+
+  // CP.getSrcReg()'s live interval has been merged into CP.getDstReg's live
+  // interval. Since CP.getSrcReg() is in ToBeUpdated set and its live interval
+  // is not up-to-date, need to update the merged live interval here.
+  if (ToBeUpdated.count(CP.getSrcReg()))
+    ShrinkMainRange = true;
+
   if (ShrinkMainRange) {
     LiveInterval &LI = LIS->getInterval(CP.getDstReg());
     shrinkToUses(&LI);
   }
 
   // SrcReg is guaranteed to be the register whose live interval that is
   // being merged.
   LIS->removeInterval(CP.getSrcReg());

test/CodeGen/X86/late-remat-update-2.mir

+# REQUIRES: asserts
+# RUN: llc -mtriple=x86_64-- -run-pass=simple-register-coalescing -late-remat-update-threshold=0 -debug-only=regalloc %s -o /dev/null 2>&1 | FileCheck %s
+#
+# PR40061: %t2 = %t1 is rematerialized and %t1 is added into toBeUpdated set
+# to postpone its live interval update. After the rematerialization, the live
+# interval of %t1 is larger than necessary. Then %t1 is merged into %t3 and %t1
+# gets removed. After the merge, %t3 contains live interval larger than
+# necessary. Because %t3 is not in toBeUpdated set so its live interval is not
+# updated after register coalescing, and it will break some assumption in
+# regalloc. This test checks the live interval is up-to-date after register
+# coalescing.
+#
+# Check the interval of %t3 after %t1 is merged into %t3 in register coalescing.
+# Make sure the interval of %t3 doesn't cover the entry of bb.1 because %t1 is
+# rematerialized in bb.1.
+# CHECK: INTERVALS
+# CHECK: %2 [16r,80B:1)[240B,272r:1)[272r,304r:0)  0@272r 1@16r
+#
+# CHECK: MACHINEINSTRS
+# Check %t1 has been merged with %t3
+# CHECK: 0B bb.0.entry:
+# CHECK: 16B    %t3:gr64 = MOV64ri32 -11
+# CHECK: 32B    CMP64ri8 %t3:gr64, 1
+# Check %t1 here has been rematerialized.
+# CHECK: 80B  bb.1:
+# CHECK: 96B    %t2:gr64 = MOV64ri32 -11
+
+---
+name:            _Z3fooi
+body:             |
+  bb.0.entry:
+    successors: %bb.1(0x15555555), %bb.2(0x6aaaaaab)
+
+    %t1:gr64 = MOV64ri32 -11
+    CMP64ri8 %t1, 1, implicit-def $eflags
+    JE_1 %bb.2, implicit killed $eflags
+    JMP_1 %bb.1
+
+  bb.1:
+    successors: %bb.1(0x80000000)
+
+    %t2:gr64 = COPY %t1
+    ADJCALLSTACKDOWN64 0, 0, 0, implicit-def dead $rsp, implicit-def dead $eflags, implicit-def dead $ssp, implicit $rsp, implicit $ssp
+    %t2:gr64 = ADD64ri8 %t2, 5, implicit-def $eflags
+    $rdi = COPY %t2
+    CALL64pcrel32 @_Z3fooi, csr_64, implicit $rsp, implicit $ssp, implicit killed $rdi, implicit-def $rsp, implicit-def $ssp
+    ADJCALLSTACKUP64 0, 0, implicit-def dead $rsp, implicit-def dead $eflags, implicit-def dead $ssp, implicit $rsp, implicit $ssp
+    CMP64ri8 %t2, 1, implicit-def $eflags
+    JE_1 %bb.1, implicit killed $eflags
+    RET -1
+
+  bb.2:
+    successors: %bb.3(0x80000000)
+    %t3:gr64 = COPY %t1
+    %t3:gr64 = ADD64ri8 %t3, 10, implicit-def $eflags
+
+  bb.3:
+    $rax = COPY %t3
+    RET 0, $rax
+
+...