[MachineLateInstrsCleanup] Handle multiple kills for a preceding definition.

llvm/lib/CodeGen/MachineLateInstrsCleanup.cpp

@@ -48,17 +48,17 @@ class MachineLateInstrsCleanup {
       return MI && MI->isIdenticalTo(*ArgMI);
     }
   };
-
+  typedef SmallDenseMap<Register, TinyPtrVector<MachineInstr *>> Reg2MIVecMap;
   std::vector<Reg2MIMap> RegDefs;
-  std::vector<Reg2MIMap> RegKills;
+  std::vector<Reg2MIVecMap> RegKills;
 
   // Walk through the instructions in MBB and remove any redundant
   // instructions.
   bool processBlock(MachineBasicBlock *MBB);
 
   void removeRedundantDef(MachineInstr *MI);
   void clearKillsForDef(Register Reg, MachineBasicBlock *MBB,
-                        BitVector &VisitedPreds);
+                        BitVector &VisitedPreds, MachineInstr *ToRemoveMI);
 
 public:
   bool run(MachineFunction &MF);
@@ -135,33 +135,39 @@ bool MachineLateInstrsCleanup::run(MachineFunction &MF) {
 // definition.
 void MachineLateInstrsCleanup::clearKillsForDef(Register Reg,
                                                 MachineBasicBlock *MBB,
-                                                BitVector &VisitedPreds) {
+                                                BitVector &VisitedPreds,
+                                                MachineInstr *ToRemoveMI) {
   VisitedPreds.set(MBB->getNumber());
 
-  // Kill flag in MBB
-  if (MachineInstr *KillMI = RegKills[MBB->getNumber()].lookup(Reg)) {
-    KillMI->clearRegisterKills(Reg, TRI);
-    return;
-  }
+  // Clear kill flag(s) in MBB, that have been seen after the preceding
+  // definition. If Reg or one of its subregs was killed, it would actually
+  // be ok to stop after removing that (and any other) kill-flag, but it
+  // doesn't seem noticeably faster while it would be a bit more complicated.
+  Reg2MIVecMap &MBBKills = RegKills[MBB->getNumber()];
+  if (auto Kills = MBBKills.find(Reg); Kills != MBBKills.end())
+    for (auto *KillMI : Kills->second)
+      KillMI->clearRegisterKills(Reg, TRI);
 
-  // Def in MBB (missing kill flag)
-  if (MachineInstr *DefMI = RegDefs[MBB->getNumber()].lookup(Reg))
-    if (DefMI->getParent() == MBB)
-      return;
+  // Definition in current MBB: done.
+  Reg2MIMap &MBBDefs = RegDefs[MBB->getNumber()];
+  MachineInstr *DefMI = MBBDefs[Reg];
+  assert(DefMI->isIdenticalTo(*ToRemoveMI) && "Previous def not identical?");
+  if (DefMI->getParent() == MBB)
+    return;
 
   // If an earlier def is not in MBB, continue in predecessors.
   if (!MBB->isLiveIn(Reg))
     MBB->addLiveIn(Reg);
   assert(!MBB->pred_empty() && "Predecessor def not found!");
   for (MachineBasicBlock *Pred : MBB->predecessors())
     if (!VisitedPreds.test(Pred->getNumber()))
-      clearKillsForDef(Reg, Pred, VisitedPreds);
+      clearKillsForDef(Reg, Pred, VisitedPreds, ToRemoveMI);
 }
 
 void MachineLateInstrsCleanup::removeRedundantDef(MachineInstr *MI) {
   Register Reg = MI->getOperand(0).getReg();
   BitVector VisitedPreds(MI->getMF()->getNumBlockIDs());
-  clearKillsForDef(Reg, MI->getParent(), VisitedPreds);
+  clearKillsForDef(Reg, MI->getParent(), VisitedPreds, MI);
   MI->eraseFromParent();
   ++NumRemoved;
 }
@@ -197,7 +203,7 @@ static bool isCandidate(const MachineInstr *MI, Register &DefedReg,
 bool MachineLateInstrsCleanup::processBlock(MachineBasicBlock *MBB) {
   bool Changed = false;
   Reg2MIMap &MBBDefs = RegDefs[MBB->getNumber()];
-  Reg2MIMap &MBBKills = RegKills[MBB->getNumber()];
+  Reg2MIVecMap &MBBKills = RegKills[MBB->getNumber()];
 
   // Find reusable definitions in the predecessor(s).
   if (!MBB->pred_empty() && !MBB->isEHPad() &&
@@ -247,8 +253,8 @@ bool MachineLateInstrsCleanup::processBlock(MachineBasicBlock *MBB) {
         MBBDefs.erase(Reg);
         MBBKills.erase(Reg);
       } else if (MI.findRegisterUseOperandIdx(Reg, TRI, true /*isKill*/) != -1)
-        // Keep track of register kills.
-        MBBKills[Reg] = &MI;
+        // Keep track of all instructions that fully or partially kills Reg.
+        MBBKills[Reg].push_back(&MI);
     }
 
     // Record this MI for potential later reuse.

llvm/test/CodeGen/SystemZ/machine-latecleanup-kills.mir

@@ -0,0 +1,76 @@
+# NOTE: Assertions have been autogenerated by utils/update_mir_test_checks.py UTC_ARGS: --version 5
+# RUN: llc -mtriple=s390x-linux-gnu -run-pass=machine-latecleanup %s -o - -mcpu=z16 \
+# RUN:   -verify-machineinstrs 2>&1 | FileCheck %s
+
+# Kill flag of $r0q (super-reg) needs to be removed, and $r0l needs to be added as live-in.
+---
+name:            fun0
+tracksRegLiveness: true
+body:             |
+  ; CHECK-LABEL: name: fun0
+  ; CHECK: bb.0:
+  ; CHECK-NEXT:   successors: %bb.1(0x80000000)
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   renamable $r0l = LHIMux -1
+  ; CHECK-NEXT:   J %bb.1
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT: bb.1:
+  ; CHECK-NEXT:   liveins: $r0l
+  ; CHECK-NEXT: {{  $}}
+  ; CHECK-NEXT:   renamable $r1d = LGFI 0
+  ; CHECK-NEXT:   ST128 renamable $r0q, $r15d, 168, $noreg
+  ; CHECK-NEXT:   ST killed renamable $r0l, $r15d, 160, $noreg
+  ; CHECK-NEXT:   Return
+  bb.0:
+    renamable $r0l = LHIMux -1
+    J %bb.1
+
+  bb.1:
+    renamable $r1d = LGFI 0
+    ST128 killed renamable $r0q, $r15d, 168, $noreg
+    renamable $r0l = LHIMux -1
+    ST killed renamable $r0l, $r15d, 160, $noreg
+    Return
+...
+
+# Kill flags of both $r1d and $r0q (super-reg) need to be removed.
+---
+name:            fun1
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    ; CHECK-LABEL: name: fun1
+    ; CHECK: renamable $r1d = LLILL 1
+    ; CHECK-NEXT: STG renamable $r1d, killed $r15d, 8, $noreg
+    ; CHECK-NEXT: renamable $r0d = LLILL 0
+    ; CHECK-NEXT: ST128 renamable $r0q, $r15d, 0, $noreg
+    ; CHECK-NEXT: STG killed renamable $r1d, killed $r15d, 8, $noreg
+    ; CHECK-NEXT: Return
+    renamable $r1d = LLILL 1
+    STG killed renamable $r1d, killed $r15d, 8, $noreg
+    renamable $r0d = LLILL 0
+    ST128 killed renamable $r0q, $r15d, 0, $noreg
+    renamable $r1d = LLILL 1
+    STG killed renamable $r1d, killed $r15d, 8, $noreg
+    Return
+...
+
+# Kill flags of both $r1l (subreg) and $r1d need to be removed.
+---
+name:            fun2
+tracksRegLiveness: true
+body:             |
+  bb.0:
+    ; CHECK-LABEL: name: fun2
+    ; CHECK: renamable $r1d = LLILL 1
+    ; CHECK-NEXT: ST renamable $r1l, $r15d, 0, $noreg
+    ; CHECK-NEXT: STG renamable $r1d, $r15d, 8, $noreg
+    ; CHECK-NEXT: STG killed renamable $r1d, $r15d, 16, $noreg
+    ; CHECK-NEXT: Return
+    renamable $r1d = LLILL 1
+    ST killed renamable $r1l, $r15d, 0, $noreg
+    STG killed renamable $r1d, $r15d, 8, $noreg
+    renamable $r1d = LLILL 1
+    STG killed renamable $r1d, $r15d, 16, $noreg
+    Return
+...