[CodeGen] Extract copy-paste on PHI MachineInstr income removal. (llvm#158634)

Author: ivafanas

llvm/include/llvm/CodeGen/MachineBasicBlock.h

@@ -1287,6 +1287,15 @@ class MachineBasicBlock
   // Helper function for MIRPrinter.
   LLVM_ABI bool canPredictBranchProbabilities() const;
 
+  /// Iterate over block PHI instructions and remove all incoming values for
+  /// PredMBB.
+  ///
+  /// Method does not erase PHI instructions even if they have single income or
+  /// do not have incoming values ar all. It is a caller responsibility to make
+  /// decision how to process PHI instructions after incoming values removal.
+  LLVM_ABI void
+  removePHIsIncomingValuesForPredecessor(const MachineBasicBlock &PredMBB);
+
 private:
   /// Return probability iterator corresponding to the I successor iterator.
   probability_iterator getProbabilityIterator(succ_iterator I);

llvm/include/llvm/CodeGen/MachineInstr.h

@@ -2000,6 +2000,15 @@ class MachineInstr
   /// and point them to \p Reg instead.
   LLVM_ABI void changeDebugValuesDefReg(Register Reg);
 
+  /// Remove all incoming values of Phi instruction for the given block.
+  ///
+  /// Return deleted operands count.
+  ///
+  /// Method does not erase PHI instruction even if it has single income or does
+  /// not have incoming values at all. It is a caller responsibility to make
+  /// decision how to process PHI instruction after incoming values removed.
+  LLVM_ABI unsigned removePHIIncomingValueFor(const MachineBasicBlock &MBB);
+
   /// Sets all register debug operands in this debug value instruction to be
   /// undef.
   void setDebugValueUndef() {

llvm/lib/CodeGen/MachineBasicBlock.cpp

@@ -1802,6 +1802,12 @@ bool MachineBasicBlock::sizeWithoutDebugLargerThan(unsigned Limit) const {
   return false;
 }
 
+void MachineBasicBlock::removePHIsIncomingValuesForPredecessor(
+    const MachineBasicBlock &PredMBB) {
+  for (MachineInstr &Phi : phis())
+    Phi.removePHIIncomingValueFor(PredMBB);
+}
+
 const MBBSectionID MBBSectionID::ColdSectionID(MBBSectionID::SectionType::Cold);
 const MBBSectionID
     MBBSectionID::ExceptionSectionID(MBBSectionID::SectionType::Exception);

llvm/lib/CodeGen/MachineInstr.cpp

@@ -2747,3 +2747,18 @@ bool MachineInstr::mayFoldInlineAsmRegOp(unsigned OpId) const {
     return F.getRegMayBeFolded();
   return false;
 }
+
+unsigned MachineInstr::removePHIIncomingValueFor(const MachineBasicBlock &MBB) {
+  assert(isPHI());
+
+  // Phi might have multiple entries for MBB. Need to remove them all.
+  unsigned RemovedCount = 0;
+  for (unsigned N = getNumOperands(); N > 2; N -= 2) {
+    if (getOperand(N - 1).getMBB() == &MBB) {
+      removeOperand(N - 1);
+      removeOperand(N - 2);
+      RemovedCount += 2;
+    }
+  }
+  return RemovedCount;
+}

llvm/lib/CodeGen/ModuloSchedule.cpp

@@ -10,6 +10,7 @@
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Analysis/MemoryLocation.h"
 #include "llvm/CodeGen/LiveIntervals.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
 #include "llvm/CodeGen/MachineLoopInfo.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
@@ -859,20 +860,6 @@ void ModuloScheduleExpander::splitLifetimes(MachineBasicBlock *KernelBB,
   }
 }
 
-/// Remove the incoming block from the Phis in a basic block.
-static void removePhis(MachineBasicBlock *BB, MachineBasicBlock *Incoming) {
-  for (MachineInstr &MI : *BB) {
-    if (!MI.isPHI())
-      break;
-    for (unsigned i = 1, e = MI.getNumOperands(); i != e; i += 2)
-      if (MI.getOperand(i + 1).getMBB() == Incoming) {
-        MI.removeOperand(i + 1);
-        MI.removeOperand(i);
-        break;
-      }
-  }
-}
-
 /// Create branches from each prolog basic block to the appropriate epilog
 /// block.  These edges are needed if the loop ends before reaching the
 /// kernel.
@@ -906,7 +893,7 @@ void ModuloScheduleExpander::addBranches(MachineBasicBlock &PreheaderBB,
       Prolog->removeSuccessor(LastPro);
       LastEpi->removeSuccessor(Epilog);
       numAdded = TII->insertBranch(*Prolog, Epilog, nullptr, Cond, DebugLoc());
-      removePhis(Epilog, LastEpi);
+      Epilog->removePHIsIncomingValuesForPredecessor(*LastEpi);
       // Remove the blocks that are no longer referenced.
       if (LastPro != LastEpi) {
         for (auto &MI : *LastEpi)
@@ -924,7 +911,7 @@ void ModuloScheduleExpander::addBranches(MachineBasicBlock &PreheaderBB,
       LastPro->eraseFromParent();
     } else {
       numAdded = TII->insertBranch(*Prolog, LastPro, nullptr, Cond, DebugLoc());
-      removePhis(Epilog, Prolog);
+      Epilog->removePHIsIncomingValuesForPredecessor(*Prolog);
     }
     LastPro = Prolog;
     LastEpi = Epilog;

llvm/lib/CodeGen/TailDuplicator.cpp

@@ -375,13 +375,7 @@ void TailDuplicator::processPHI(
   if (!Remove)
     return;
 
-  // MI might have multiple entries for PredBB. Need to remove them all.
-  for (unsigned N = MI->getNumOperands(); N > 2; N -= 2) {
-    if (MI->getOperand(N - 1).getMBB() == PredBB) {
-      MI->removeOperand(N - 1);
-      MI->removeOperand(N - 2);
-    }
-  }
+  MI->removePHIIncomingValueFor(*PredBB);
 
   if (MI->getNumOperands() == 1 && !TailBB->hasAddressTaken())
     MI->eraseFromParent();

llvm/lib/CodeGen/UnreachableBlockElim.cpp

@@ -22,6 +22,7 @@
 #include "llvm/CodeGen/UnreachableBlockElim.h"
 #include "llvm/ADT/DepthFirstIterator.h"
 #include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
 #include "llvm/CodeGen/MachineInstrBuilder.h"
@@ -155,18 +156,7 @@ bool UnreachableMachineBlockElim::run(MachineFunction &F) {
       if (MDT && MDT->getNode(&BB)) MDT->eraseNode(&BB);
 
       while (!BB.succ_empty()) {
-        MachineBasicBlock* succ = *BB.succ_begin();
-
-        for (MachineInstr &Phi : succ->phis()) {
-          for (unsigned i = Phi.getNumOperands() - 1; i >= 2; i -= 2) {
-            if (Phi.getOperand(i).isMBB() &&
-                Phi.getOperand(i).getMBB() == &BB) {
-              Phi.removeOperand(i);
-              Phi.removeOperand(i - 1);
-            }
-          }
-        }
-
+        (*BB.succ_begin())->removePHIsIncomingValuesForPredecessor(BB);
         BB.removeSuccessor(BB.succ_begin());
       }
     }