RegisterCoalescer: Handle implicit-def of a super register when rematerializing

[arsenm]

Summary

Permit an implicit-def of a virtual register when rematerializing if
it defines a super register of a subregister def. The
rematerialization pre-legality check should really have been checking
the implicit operands, but that should be fixed separately.

[arsenm]

b1295dd5c923c6828775406b4063de1531fd4782

llvm/lib/CodeGen/RegisterCoalescer.cpp

bool RegisterCoalescer::reMaterializeTrivialDef(const CoalescerPair &CP,
   // If both SrcIdx and DstIdx are set, correct rematerialization would widen
   // the register substantially (beyond both source and dest size). This is bad
   // for performance since it can cascade through a function, introducing many
   // extra spills and fills (e.g. ARM can easily end up copying QQQQPR registers
   // around after a few subreg copies).
   if (SrcIdx && DstIdx)
     return false;
 
+  const unsigned DefSubIdx = DefMI->getOperand(0).getSubReg();
   const TargetRegisterClass *DefRC = TII->getRegClass(MCID, 0, TRI, *MF);
   if (!DefMI->isImplicitDef()) {
     if (DstReg.isPhysical()) {
       Register NewDstReg = DstReg;
 
       unsigned NewDstIdx = TRI->composeSubRegIndices(CP.getSrcIdx(),
                                               DefMI->getOperand(0).getSubReg());
       if (NewDstIdx)
bool RegisterCoalescer::reMaterializeTrivialDef(const CoalescerPair &CP,
   ImplicitOps.reserve(CopyMI->getNumOperands() -
                       CopyMI->getDesc().getNumOperands());
   for (unsigned I = CopyMI->getDesc().getNumOperands(),
                 E = CopyMI->getNumOperands();
        I != E; ++I) {
     MachineOperand &MO = CopyMI->getOperand(I);
     if (MO.isReg()) {
       assert(MO.isImplicit() && "No explicit operands after implicit operands.");
-      assert(MO.getReg().isPhysical() && "unexpected implicit virtual register def");
+      assert((MO.getReg().isPhysical() ||
+              (MO.getSubReg() == 0 && MO.getReg() == DstOperand.getReg())) &&
+             "unexpected implicit virtual register def");
       ImplicitOps.push_back(MO);
     }
   }
 
   CopyMI->eraseFromParent();
   ErasedInstrs.insert(CopyMI);
 
   // NewMI may have dead implicit defs (E.g. EFLAGS for MOV<bits>r0 on X86).
   // We need to remember these so we can add intervals once we insert
   // NewMI into SlotIndexes.
+  //
+  // We also expect to have tied implicit-defs of super registers originating
+  // from SUBREG_TO_REG, such as:
+  // $edi = MOV32r0 implicit-def dead $eflags, implicit-def $rdi
+  // undef %0.sub_32bit = MOV32r0 implicit-def dead $eflags, implicit-def %0
+
   SmallVector<MCRegister, 4> NewMIImplDefs;
   for (unsigned i = NewMI.getDesc().getNumOperands(),
                 e = NewMI.getNumOperands();
        i != e; ++i) {
     MachineOperand &MO = NewMI.getOperand(i);
     if (MO.isReg() && MO.isDef()) {
-      assert(MO.isImplicit() && MO.isDead() && MO.getReg().isPhysical());
+      assert(MO.isImplicit());
+      if (MO.getReg().isPhysical()) {
+        assert(MO.isImplicit() && MO.getReg().isPhysical() &&
+               (MO.isDead() ||
+                (DefSubIdx && (TRI->getSubReg(MO.getReg(), DefSubIdx) ==
+                               MCRegister(NewMI.getOperand(0).getReg())))));
         NewMIImplDefs.push_back(MO.getReg().asMCReg());
+      } else {
+        assert(MO.getReg() == NewMI.getOperand(0).getReg() &&
+               MO.getSubReg() == 0);
+        // We're only expecting another def of the main output, so the range
+        // should get updated with the regular output range.
+        //
+        // FIXME: The range updating below probably needs updating to look at
+        // the super register if subranges are tracked.
+        assert(!MRI->shouldTrackSubRegLiveness(DstReg) &&
+               "subrange update for implicit-def of super register may not be "
+               "properly handled");
+      }
     }
   }
 
   if (DstReg.isVirtual()) {
     unsigned NewIdx = NewMI.getOperand(0).getSubReg();
 
     if (DefRC != nullptr) {
       if (NewIdx)

llvm/test/CodeGen/X86/rematerialize-sub-super-reg.mir

bb.2:
     %t3:gr64 = ADD64ri8 %t3, 10, implicit-def $eflags
 
   bb.3:
     $rax = COPY %t3
     RET 0, $rax
 
 ...
 
-# FIXME: This currently asserts
-# ---
-# name:            rematerialize_subregister_into_superreg_def_with_impdef_physreg
-# tracksRegLiveness: true
-# body:             |
-#   bb.0:
-#     undef %t1.sub_32bit:gr64_with_sub_8bit = MOV32ri -11, implicit-def %t1
-#     CMP64ri8 %t1, 1, implicit-def $eflags
-#     JCC_1 %bb.2, 4, implicit killed $eflags
-#     JMP_1 %bb.1
-
-#   bb.1:
-#     %t2:gr64 = COPY %t1
-#     $rax = COPY %t2
-#     CMP64ri8 %t2, 1, implicit-def $eflags
-#     JCC_1 %bb.1, 4, implicit killed $eflags
-#     RET 0, $rax
-
-#   bb.2:
-#     %t3:gr64 = COPY %t1
-#     %t3:gr64 = ADD64ri8 %t3, 10, implicit-def $eflags
-
-#   bb.3:
-#     $rax = COPY %t3
-#     RET 0, $rax
+# Handle that rematerializing an instruction with an implicit def of a
+# virtual super register into a physical register works.
+#
+# FIXME: Resulting rematerializing has a redundant implicit-def
+---
+name:            rematerialize_subregister_into_superreg_def_with_impdef_physreg
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: name: rematerialize_subregister_into_superreg_def_with_impdef_physreg
+  ; CHECK: bb.0:
+  ; CHECK-NEXT:   successors: %bb.2(0x40000000), %bb.1(0x40000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   undef %t3.sub_32bit:gr64_with_sub_8bit = MOV32ri -11, implicit-def %t3
+  ; CHECK-NEXT:   CMP64ri8 %t3, 1, implicit-def $eflags
+  ; CHECK-NEXT:   JCC_1 %bb.2, 4, implicit killed $eflags
+  ; CHECK-NEXT:   JMP_1 %bb.1
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1:
+  ; CHECK-NEXT:   successors: %bb.1(0x80000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   dead $eax = MOV32ri -11, implicit-def $rax, implicit-def $rax
+  ; CHECK-NEXT:   CMP64ri8 %t3, 1, implicit-def $eflags
+  ; CHECK-NEXT:   JCC_1 %bb.1, 4, implicit killed $eflags
+  ; CHECK-NEXT:   RET 0, $rax
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.2:
+  ; CHECK-NEXT:   successors: %bb.3(0x80000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   %t3:gr64_with_sub_8bit = ADD64ri8 %t3, 10, implicit-def $eflags
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.3:
+  ; CHECK-NEXT:   $rax = COPY %t3
+  ; CHECK-NEXT:   RET 0, $rax
+  bb.0:
+    undef %t1.sub_32bit:gr64_with_sub_8bit = MOV32ri -11, implicit-def %t1
+    CMP64ri8 %t1, 1, implicit-def $eflags
+    JCC_1 %bb.2, 4, implicit killed $eflags
+    JMP_1 %bb.1
+
+  bb.1:
+    %t2:gr64 = COPY %t1
+    $rax = COPY %t2
+    CMP64ri8 %t2, 1, implicit-def $eflags
+    JCC_1 %bb.1, 4, implicit killed $eflags
+    RET 0, $rax
+
+  bb.2:
+    %t3:gr64 = COPY %t1
+    %t3:gr64 = ADD64ri8 %t3, 10, implicit-def $eflags
 
-# ...
+  bb.3:
+    $rax = COPY %t3
+    RET 0, $rax
+
+...