[CodeGen][NewPM] Port phi-node-elimination to new pass manager (#98867)

Summary:
- Add `PHIEliminationPass `.
- Support new pass manager in `MachineBasicBlock:: SplitCriticalEdge `

Test Plan: 

Reviewers: 

Subscribers: 

Tasks: 

Tags: 


Differential Revision: https://phabricator.intern.facebook.com/D60251632

Author: paperchalice

llvm/include/llvm/CodeGen/MachineBasicBlock.h

@@ -43,6 +43,8 @@ class raw_ostream;
 class LiveIntervals;
 class TargetRegisterClass;
 class TargetRegisterInfo;
+template <typename IRUnitT, typename... ExtraArgTs> class AnalysisManager;
+using MachineFunctionAnalysisManager = AnalysisManager<MachineFunction>;
 
 // This structure uniquely identifies a basic block section.
 // Possible values are
@@ -968,7 +970,16 @@ class MachineBasicBlock
   /// MachineLoopInfo, as applicable.
   MachineBasicBlock *
   SplitCriticalEdge(MachineBasicBlock *Succ, Pass &P,
-                    std::vector<SparseBitVector<>> *LiveInSets = nullptr);
+                    std::vector<SparseBitVector<>> *LiveInSets = nullptr) {
+    return SplitCriticalEdge(Succ, &P, nullptr, LiveInSets);
+  }
+
+  MachineBasicBlock *
+  SplitCriticalEdge(MachineBasicBlock *Succ,
+                    MachineFunctionAnalysisManager &MFAM,
+                    std::vector<SparseBitVector<>> *LiveInSets = nullptr) {
+    return SplitCriticalEdge(Succ, nullptr, &MFAM, LiveInSets);
+  }
 
   /// Check if the edge between this block and the given successor \p
   /// Succ, can be split. If this returns true a subsequent call to
@@ -1243,6 +1254,12 @@ class MachineBasicBlock
   /// unless you know what you're doing, because it doesn't update Pred's
   /// successors list. Use Pred->removeSuccessor instead.
   void removePredecessor(MachineBasicBlock *Pred);
+
+  // Helper method for new pass manager migration.
+  MachineBasicBlock *
+  SplitCriticalEdge(MachineBasicBlock *Succ, Pass *P,
+                    MachineFunctionAnalysisManager *MFAM,
+                    std::vector<SparseBitVector<>> *LiveInSets);
 };
 
 raw_ostream& operator<<(raw_ostream &OS, const MachineBasicBlock &MBB);

llvm/include/llvm/CodeGen/PHIElimination.h

@@ -0,0 +1,24 @@
+//===- llvm/CodeGen/PHIElimination.h ----------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_CODEGEN_PHIELIMINATION_H
+#define LLVM_CODEGEN_PHIELIMINATION_H
+
+#include "llvm/CodeGen/MachinePassManager.h"
+
+namespace llvm {
+
+class PHIEliminationPass : public PassInfoMixin<PHIEliminationPass> {
+public:
+  PreservedAnalyses run(MachineFunction &MF,
+                        MachineFunctionAnalysisManager &MFAM);
+};
+
+} // namespace llvm
+
+#endif // LLVM_CODEGEN_PHIELIMINATION_H

llvm/include/llvm/Passes/CodeGenPassBuilder.h

@@ -43,6 +43,7 @@
 #include "llvm/CodeGen/MachineFunctionAnalysis.h"
 #include "llvm/CodeGen/MachineModuleInfo.h"
 #include "llvm/CodeGen/MachinePassManager.h"
+#include "llvm/CodeGen/PHIElimination.h"
 #include "llvm/CodeGen/PreISelIntrinsicLowering.h"
 #include "llvm/CodeGen/RegAllocFast.h"
 #include "llvm/CodeGen/ReplaceWithVeclib.h"

llvm/include/llvm/Passes/MachinePassRegistry.def

@@ -132,6 +132,7 @@ MACHINE_FUNCTION_PASS("dead-mi-elimination", DeadMachineInstructionElimPass())
 MACHINE_FUNCTION_PASS("finalize-isel", FinalizeISelPass())
 MACHINE_FUNCTION_PASS("localstackalloc", LocalStackSlotAllocationPass())
 MACHINE_FUNCTION_PASS("no-op-machine-function", NoOpMachineFunctionPass())
+MACHINE_FUNCTION_PASS("phi-node-elimination", PHIEliminationPass())
 MACHINE_FUNCTION_PASS("print", PrintMIRPass())
 MACHINE_FUNCTION_PASS("print<live-intervals>", LiveIntervalsPrinterPass(dbgs()))
 MACHINE_FUNCTION_PASS("print<live-vars>", LiveVariablesPrinterPass(dbgs()))
@@ -231,7 +232,6 @@ DUMMY_MACHINE_FUNCTION_PASS("mirfs-discriminators", MIRAddFSDiscriminatorsPass)
 DUMMY_MACHINE_FUNCTION_PASS("opt-phis", OptimizePHIsPass)
 DUMMY_MACHINE_FUNCTION_PASS("patchable-function", PatchableFunctionPass)
 DUMMY_MACHINE_FUNCTION_PASS("peephole-opt", PeepholeOptimizerPass)
-DUMMY_MACHINE_FUNCTION_PASS("phi-node-elimination", PHIEliminationPass)
 DUMMY_MACHINE_FUNCTION_PASS("post-RA-sched", PostRASchedulerPass)
 DUMMY_MACHINE_FUNCTION_PASS("postmisched", PostMachineSchedulerPass)
 DUMMY_MACHINE_FUNCTION_PASS("postra-machine-sink", PostRAMachineSinkingPass)

llvm/lib/CodeGen/MachineBasicBlock.cpp

@@ -1135,9 +1135,19 @@ class SlotIndexUpdateDelegate : public MachineFunction::Delegate {
   }
 };
 
+#define GET_RESULT(RESULT, GETTER, INFIX)                                      \
+  [MF, P, MFAM]() {                                                            \
+    if (P) {                                                                   \
+      auto *Wrapper = P->getAnalysisIfAvailable<RESULT##INFIX##WrapperPass>(); \
+      return Wrapper ? &Wrapper->GETTER() : nullptr;                           \
+    }                                                                          \
+    return MFAM->getCachedResult<RESULT##Analysis>(*MF);                       \
+  }()
+
 MachineBasicBlock *MachineBasicBlock::SplitCriticalEdge(
-    MachineBasicBlock *Succ, Pass &P,
+    MachineBasicBlock *Succ, Pass *P, MachineFunctionAnalysisManager *MFAM,
     std::vector<SparseBitVector<>> *LiveInSets) {
+  assert((P || MFAM) && "Need a way to get analysis results!");
   if (!canSplitCriticalEdge(Succ))
     return nullptr;
 
@@ -1161,10 +1171,8 @@ MachineBasicBlock *MachineBasicBlock::SplitCriticalEdge(
                     << " -- " << printMBBReference(*NMBB) << " -- "
                     << printMBBReference(*Succ) << '\n');
 
-  auto *LISWrapper = P.getAnalysisIfAvailable<LiveIntervalsWrapperPass>();
-  LiveIntervals *LIS = LISWrapper ? &LISWrapper->getLIS() : nullptr;
-  auto *SIWrapper = P.getAnalysisIfAvailable<SlotIndexesWrapperPass>();
-  SlotIndexes *Indexes = SIWrapper ? &SIWrapper->getSI() : nullptr;
+  LiveIntervals *LIS = GET_RESULT(LiveIntervals, getLIS, );
+  SlotIndexes *Indexes = GET_RESULT(SlotIndexes, getSI, );
   if (LIS)
     LIS->insertMBBInMaps(NMBB);
   else if (Indexes)
@@ -1173,8 +1181,7 @@ MachineBasicBlock *MachineBasicBlock::SplitCriticalEdge(
   // On some targets like Mips, branches may kill virtual registers. Make sure
   // that LiveVariables is properly updated after updateTerminator replaces the
   // terminators.
-  auto *LVWrapper = P.getAnalysisIfAvailable<LiveVariablesWrapperPass>();
-  LiveVariables *LV = LVWrapper ? &LVWrapper->getLV() : nullptr;
+  LiveVariables *LV = GET_RESULT(LiveVariables, getLV, );
 
   // Collect a list of virtual registers killed by the terminators.
   SmallVector<Register, 4> KilledRegs;
@@ -1339,12 +1346,10 @@ MachineBasicBlock *MachineBasicBlock::SplitCriticalEdge(
     LIS->repairIntervalsInRange(this, getFirstTerminator(), end(), UsedRegs);
   }
 
-  if (auto *MDTWrapper =
-          P.getAnalysisIfAvailable<MachineDominatorTreeWrapperPass>())
-    MDTWrapper->getDomTree().recordSplitCriticalEdge(this, Succ, NMBB);
+  if (auto *MDT = GET_RESULT(MachineDominatorTree, getDomTree, ))
+    MDT->recordSplitCriticalEdge(this, Succ, NMBB);
 
-  auto *MLIWrapper = P.getAnalysisIfAvailable<MachineLoopInfoWrapperPass>();
-  if (MachineLoopInfo *MLI = MLIWrapper ? &MLIWrapper->getLI() : nullptr)
+  if (MachineLoopInfo *MLI = GET_RESULT(MachineLoop, getLI, Info))
     if (MachineLoop *TIL = MLI->getLoopFor(this)) {
       // If one or the other blocks were not in a loop, the new block is not
       // either, and thus LI doesn't need to be updated.

llvm/lib/CodeGen/PHIElimination.cpp

@@ -12,6 +12,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/CodeGen/PHIElimination.h"
 #include "PHIEliminationUtils.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/SmallPtrSet.h"
@@ -64,22 +65,13 @@ static cl::opt<bool> NoPhiElimLiveOutEarlyExit(
 
 namespace {
 
-class PHIElimination : public MachineFunctionPass {
+class PHIEliminationImpl {
   MachineRegisterInfo *MRI = nullptr; // Machine register information
   LiveVariables *LV = nullptr;
   LiveIntervals *LIS = nullptr;
+  MachineLoopInfo *MLI = nullptr;
+  MachineDominatorTree *MDT = nullptr;
 
-public:
-  static char ID; // Pass identification, replacement for typeid
-
-  PHIElimination() : MachineFunctionPass(ID) {
-    initializePHIEliminationPass(*PassRegistry::getPassRegistry());
-  }
-
-  bool runOnMachineFunction(MachineFunction &MF) override;
-  void getAnalysisUsage(AnalysisUsage &AU) const override;
-
-private:
   /// EliminatePHINodes - Eliminate phi nodes by inserting copy instructions
   /// in predecessor basic blocks.
   bool EliminatePHINodes(MachineFunction &MF, MachineBasicBlock &MBB);
@@ -118,10 +110,71 @@ class PHIElimination : public MachineFunctionPass {
   using LoweredPHIMap =
       DenseMap<MachineInstr *, unsigned, MachineInstrExpressionTrait>;
   LoweredPHIMap LoweredPHIs;
+
+  MachineFunctionPass *P = nullptr;
+  MachineFunctionAnalysisManager *MFAM = nullptr;
+
+public:
+  PHIEliminationImpl(MachineFunctionPass *P) : P(P) {
+    auto *LVWrapper = P->getAnalysisIfAvailable<LiveVariablesWrapperPass>();
+    auto *LISWrapper = P->getAnalysisIfAvailable<LiveIntervalsWrapperPass>();
+    auto *MLIWrapper = P->getAnalysisIfAvailable<MachineLoopInfoWrapperPass>();
+    auto *MDTWrapper =
+        P->getAnalysisIfAvailable<MachineDominatorTreeWrapperPass>();
+    LV = LVWrapper ? &LVWrapper->getLV() : nullptr;
+    LIS = LISWrapper ? &LISWrapper->getLIS() : nullptr;
+    MLI = MLIWrapper ? &MLIWrapper->getLI() : nullptr;
+    MDT = MDTWrapper ? &MDTWrapper->getDomTree() : nullptr;
+  }
+
+  PHIEliminationImpl(MachineFunction &MF, MachineFunctionAnalysisManager &AM)
+      : LV(AM.getCachedResult<LiveVariablesAnalysis>(MF)),
+        LIS(AM.getCachedResult<LiveIntervalsAnalysis>(MF)),
+        MLI(AM.getCachedResult<MachineLoopAnalysis>(MF)),
+        MDT(AM.getCachedResult<MachineDominatorTreeAnalysis>(MF)), MFAM(&AM) {}
+
+  bool run(MachineFunction &MF);
+};
+
+class PHIElimination : public MachineFunctionPass {
+public:
+  static char ID; // Pass identification, replacement for typeid
+
+  PHIElimination() : MachineFunctionPass(ID) {
+    initializePHIEliminationPass(*PassRegistry::getPassRegistry());
+  }
+
+  bool runOnMachineFunction(MachineFunction &MF) override {
+    PHIEliminationImpl Impl(this);
+    return Impl.run(MF);
+  }
+
+  MachineFunctionProperties getSetProperties() const override {
+    return MachineFunctionProperties().set(
+        MachineFunctionProperties::Property::NoPHIs);
+  }
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override;
 };
 
 } // end anonymous namespace
 
+PreservedAnalyses
+PHIEliminationPass::run(MachineFunction &MF,
+                        MachineFunctionAnalysisManager &MFAM) {
+  PHIEliminationImpl Impl(MF, MFAM);
+  bool Changed = Impl.run(MF);
+  if (!Changed)
+    return PreservedAnalyses::all();
+  auto PA = getMachineFunctionPassPreservedAnalyses();
+  PA.preserve<LiveIntervalsAnalysis>();
+  PA.preserve<LiveVariablesAnalysis>();
+  PA.preserve<SlotIndexesAnalysis>();
+  PA.preserve<MachineDominatorTreeAnalysis>();
+  PA.preserve<MachineLoopAnalysis>();
+  return PA;
+}
+
 STATISTIC(NumLowered, "Number of phis lowered");
 STATISTIC(NumCriticalEdgesSplit, "Number of critical edges split");
 STATISTIC(NumReused, "Number of reused lowered phis");
@@ -147,12 +200,8 @@ void PHIElimination::getAnalysisUsage(AnalysisUsage &AU) const {
   MachineFunctionPass::getAnalysisUsage(AU);
 }
 
-bool PHIElimination::runOnMachineFunction(MachineFunction &MF) {
+bool PHIEliminationImpl::run(MachineFunction &MF) {
   MRI = &MF.getRegInfo();
-  auto *LVWrapper = getAnalysisIfAvailable<LiveVariablesWrapperPass>();
-  LV = LVWrapper ? &LVWrapper->getLV() : nullptr;
-  auto *LISWrapper = getAnalysisIfAvailable<LiveIntervalsWrapperPass>();
-  LIS = LISWrapper ? &LISWrapper->getLIS() : nullptr;
 
   bool Changed = false;
 
@@ -187,9 +236,6 @@ bool PHIElimination::runOnMachineFunction(MachineFunction &MF) {
       }
     }
 
-    MachineLoopInfoWrapperPass *MLIWrapper =
-        getAnalysisIfAvailable<MachineLoopInfoWrapperPass>();
-    MachineLoopInfo *MLI = MLIWrapper ? &MLIWrapper->getLI() : nullptr;
     for (auto &MBB : MF)
       Changed |= SplitPHIEdges(MF, MBB, MLI, (LV ? &LiveInSets : nullptr));
   }
@@ -223,9 +269,8 @@ bool PHIElimination::runOnMachineFunction(MachineFunction &MF) {
   }
 
   // TODO: we should use the incremental DomTree updater here.
-  if (Changed)
-    if (auto *MDT = getAnalysisIfAvailable<MachineDominatorTreeWrapperPass>())
-      MDT->getDomTree().getBase().recalculate(MF);
+  if (Changed && MDT)
+    MDT->getBase().recalculate(MF);
 
   LoweredPHIs.clear();
   ImpDefs.clear();
@@ -238,8 +283,8 @@ bool PHIElimination::runOnMachineFunction(MachineFunction &MF) {
 
 /// EliminatePHINodes - Eliminate phi nodes by inserting copy instructions in
 /// predecessor basic blocks.
-bool PHIElimination::EliminatePHINodes(MachineFunction &MF,
-                                       MachineBasicBlock &MBB) {
+bool PHIEliminationImpl::EliminatePHINodes(MachineFunction &MF,
+                                           MachineBasicBlock &MBB) {
   if (MBB.empty() || !MBB.front().isPHI())
     return false; // Quick exit for basic blocks without PHIs.
 
@@ -286,9 +331,9 @@ static bool allPhiOperandsUndefined(const MachineInstr &MPhi,
   return true;
 }
 /// LowerPHINode - Lower the PHI node at the top of the specified block.
-void PHIElimination::LowerPHINode(MachineBasicBlock &MBB,
-                                  MachineBasicBlock::iterator LastPHIIt,
-                                  bool AllEdgesCritical) {
+void PHIEliminationImpl::LowerPHINode(MachineBasicBlock &MBB,
+                                      MachineBasicBlock::iterator LastPHIIt,
+                                      bool AllEdgesCritical) {
   ++NumLowered;
 
   MachineBasicBlock::iterator AfterPHIsIt = std::next(LastPHIIt);
@@ -689,7 +734,7 @@ void PHIElimination::LowerPHINode(MachineBasicBlock &MBB,
 /// particular, we want to map the number of uses of a virtual register which is
 /// used in a PHI node. We map that to the BB the vreg is coming from. This is
 /// used later to determine when the vreg is killed in the BB.
-void PHIElimination::analyzePHINodes(const MachineFunction &MF) {
+void PHIEliminationImpl::analyzePHINodes(const MachineFunction &MF) {
   for (const auto &MBB : MF) {
     for (const auto &BBI : MBB) {
       if (!BBI.isPHI())
@@ -705,9 +750,9 @@ void PHIElimination::analyzePHINodes(const MachineFunction &MF) {
   }
 }
 
-bool PHIElimination::SplitPHIEdges(MachineFunction &MF, MachineBasicBlock &MBB,
-                                   MachineLoopInfo *MLI,
-                                   std::vector<SparseBitVector<>> *LiveInSets) {
+bool PHIEliminationImpl::SplitPHIEdges(
+    MachineFunction &MF, MachineBasicBlock &MBB, MachineLoopInfo *MLI,
+    std::vector<SparseBitVector<>> *LiveInSets) {
   if (MBB.empty() || !MBB.front().isPHI() || MBB.isEHPad())
     return false; // Quick exit for basic blocks without PHIs.
 
@@ -774,7 +819,8 @@ bool PHIElimination::SplitPHIEdges(MachineFunction &MF, MachineBasicBlock &MBB,
       }
       if (!ShouldSplit && !SplitAllCriticalEdges)
         continue;
-      if (!PreMBB->SplitCriticalEdge(&MBB, *this, LiveInSets)) {
+      if (!(P ? PreMBB->SplitCriticalEdge(&MBB, *P, LiveInSets)
+              : PreMBB->SplitCriticalEdge(&MBB, *MFAM, LiveInSets))) {
         LLVM_DEBUG(dbgs() << "Failed to split critical edge.\n");
         continue;
       }
@@ -785,7 +831,7 @@ bool PHIElimination::SplitPHIEdges(MachineFunction &MF, MachineBasicBlock &MBB,
   return Changed;
 }
 
-bool PHIElimination::isLiveIn(Register Reg, const MachineBasicBlock *MBB) {
+bool PHIEliminationImpl::isLiveIn(Register Reg, const MachineBasicBlock *MBB) {
   assert((LV || LIS) &&
          "isLiveIn() requires either LiveVariables or LiveIntervals");
   if (LIS)
@@ -794,8 +840,8 @@ bool PHIElimination::isLiveIn(Register Reg, const MachineBasicBlock *MBB) {
     return LV->isLiveIn(Reg, *MBB);
 }
 
-bool PHIElimination::isLiveOutPastPHIs(Register Reg,
-                                       const MachineBasicBlock *MBB) {
+bool PHIEliminationImpl::isLiveOutPastPHIs(Register Reg,
+                                           const MachineBasicBlock *MBB) {
   assert((LV || LIS) &&
          "isLiveOutPastPHIs() requires either LiveVariables or LiveIntervals");
   // LiveVariables considers uses in PHIs to be in the predecessor basic block,

llvm/lib/Passes/PassBuilder.cpp

@@ -105,6 +105,7 @@
 #include "llvm/CodeGen/MachinePostDominators.h"
 #include "llvm/CodeGen/MachineRegisterInfo.h"
 #include "llvm/CodeGen/MachineVerifier.h"
+#include "llvm/CodeGen/PHIElimination.h"
 #include "llvm/CodeGen/PreISelIntrinsicLowering.h"
 #include "llvm/CodeGen/RegAllocFast.h"
 #include "llvm/CodeGen/SafeStack.h"

llvm/test/CodeGen/AArch64/PHIElimination-crash.mir

@@ -1,6 +1,9 @@
 # RUN: llc -mtriple=aarch64-linux-gnu -verify-machineinstrs -o /dev/null %s \
 # RUN:   -run-pass=livevars,phi-node-elimination,twoaddressinstruction \
 # RUN:   -no-phi-elim-live-out-early-exit=1 -phi-elim-split-all-critical-edges=1
+# RUN: llc -mtriple=aarch64-linux-gnu -verify-machineinstrs -o /dev/null %s \
+# RUN:   --passes='require<live-vars>,phi-node-elimination,two-address-instruction' \
+# RUN:   -no-phi-elim-live-out-early-exit=1 -phi-elim-split-all-critical-edges=1
 
 # Used to result in
 #

llvm/test/CodeGen/AArch64/PHIElimination-debugloc.mir

@@ -2,6 +2,10 @@
 # RUN:   -run-pass=livevars,phi-node-elimination,twoaddressinstruction \
 # RUN:   -no-phi-elim-live-out-early-exit=1 -phi-elim-split-all-critical-edges=1 \
 # RUN: | FileCheck %s
+# RUN: llc -mtriple=aarch64-linux-gnu -verify-machineinstrs -o - %s \
+# RUN:   --passes='require<live-vars>,phi-node-elimination,two-address-instruction' \
+# RUN:   -no-phi-elim-live-out-early-exit=1 -phi-elim-split-all-critical-edges=1 \
+# RUN: | FileCheck %s
 
 --- |
   define void @test() !dbg !7 {