[MachineLateInstrsCleanup] Handle multiple kills for a preceding definition. (#119132)

When removing a redundant definition in order to reuse an earlier
identical one it is necessary to remove any earlier kill flag as well.

Previously, the assumption has been that any register that kills the
defined Reg is enough to handle for this purpose, but this is actually
not quite enough. A kill of a super-register does not necessarily imply
that all of its subregs (including Reg) is defined at that point: a
partial definition of a register is legal. This means Reg may have been
killed earlier and is not live at that point.

This patch changes the tracking of kill flags to allow for multiple
flags to be removed: instead of remembering just the single / latest
kill flag, a vector is now used to track and remove them all.
TinyPtrVector seems ideal for this as there are only very rarely more
than one kill flag, and it doesn't seem to give much difference in
compile time.

The kill flags handling here is making this pass much more complicated
than it would have to be. This pass does not depend on kill flags for
its own use, so an interesting alternative to all this handling would be
to just remove them all. If there actually is a serious user, maybe that pass
could instead recompute them.

Also adding an assertion which is unrelated to kill flags, but it seems
to make sense (according to liberal assertion policy), to verify that
the preceding definition is in fact identical in clearKillsForDef().

Fixes #117783

Author: JonPsson1

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
+...