[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 ↗

(On Diff #178816)

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 ↗

(On Diff #178816)

Done. Thanks for the suggestion!

[skatkov]

Ping, any progress on this?

[wmi]

friendly ping. Please take another look.

[qcolombet]

LGTM

llvm/trunk/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());

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

+# RUN: llc -mtriple=x86_64-- -run-pass=simple-register-coalescing -run-pass=regallocbasic -run-pass=virtregrewriter -late-remat-update-threshold=0 %s -o - | 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 wants to check the live interval is up-to-date after
+# register coalescing.
+#
+# To prevent the test from taking effect only in assert enabled mode, we want
+# to achieve the test goal without dumping regalloc trace. We add strong hint
+# to allocate both %t1 and %t2 to $rax register. If the %t1's live interval is
+# not shrinked properly after register coalescing, %t1 and %t2 will not be
+# both allocated to $rax because of inference, and we utilize the fact to
+# achieve the test goal. But note that the assumption only holds when we use
+# regallocbasic instead of greedy because greedy can update the live interval
+# in the process of splitting.
+#
+# CHECK-LABEL: name: foo
+# CHECK: bb.0.entry:
+# CHECK: $rax = MOV64ri32 -11
+# CHECK: bb.1:
+# CHECK: $rax = MOV64ri32 -11
+# CHECK: $rax = ADD64ri8 killed renamable $rax, 5
+# CHECK: CMP64ri8 renamable $rax
+# CHECK: RET 0, $rax
+# CHECK: bb.2:
+# CHECK: $rax = ADD64ri8 killed renamable $rax, 10
+# CHECK: bb.3:
+# CHECK: RET 0, $rax
+---
+name:            foo
+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
+    %t2:gr64 = ADD64ri8 %t2, 5, implicit-def $eflags
+    $rax = COPY %t2
+    CMP64ri8 %t2, 1, implicit-def $eflags
+    JE_1 %bb.1, implicit killed $eflags
+    RET 0, $rax
+
+  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
+
+...