[CodeGen][NewPM] Split MachineDominatorTree into a concrete analysis result (#94571)

Prepare for new pass manager version of `MachineDominatorTreeAnalysis`.
We may need a machine dominator tree version of `DomTreeUpdater` to
handle `SplitCriticalEdge` in some CodeGen passes.

Author: paperchalice

llvm/include/llvm/CodeGen/MachineDominators.h

@@ -24,6 +24,7 @@
 #include "llvm/Support/GenericDomTreeConstruction.h"
 #include <cassert>
 #include <memory>
+#include <optional>
 
 namespace llvm {
 class AnalysisUsage;
@@ -39,16 +40,39 @@ inline void DominatorTreeBase<MachineBasicBlock, false>::addRoot(
 
 extern template class DomTreeNodeBase<MachineBasicBlock>;
 extern template class DominatorTreeBase<MachineBasicBlock, false>; // DomTree
-extern template class DominatorTreeBase<MachineBasicBlock, true>; // PostDomTree
 
-using MachineDomTree = DomTreeBase<MachineBasicBlock>;
 using MachineDomTreeNode = DomTreeNodeBase<MachineBasicBlock>;
 
+namespace DomTreeBuilder {
+using MBBDomTree = DomTreeBase<MachineBasicBlock>;
+using MBBUpdates = ArrayRef<llvm::cfg::Update<MachineBasicBlock *>>;
+using MBBDomTreeGraphDiff = GraphDiff<MachineBasicBlock *, false>;
+
+extern template void Calculate<MBBDomTree>(MBBDomTree &DT);
+extern template void CalculateWithUpdates<MBBDomTree>(MBBDomTree &DT,
+                                                      MBBUpdates U);
+
+extern template void InsertEdge<MBBDomTree>(MBBDomTree &DT,
+                                            MachineBasicBlock *From,
+                                            MachineBasicBlock *To);
+
+extern template void DeleteEdge<MBBDomTree>(MBBDomTree &DT,
+                                            MachineBasicBlock *From,
+                                            MachineBasicBlock *To);
+
+extern template void ApplyUpdates<MBBDomTree>(MBBDomTree &DT,
+                                              MBBDomTreeGraphDiff &,
+                                              MBBDomTreeGraphDiff *);
+
+extern template bool Verify<MBBDomTree>(const MBBDomTree &DT,
+                                        MBBDomTree::VerificationLevel VL);
+} // namespace DomTreeBuilder
+
 //===-------------------------------------
 /// DominatorTree Class - Concrete subclass of DominatorTreeBase that is used to
 /// compute a normal dominator tree.
 ///
-class MachineDominatorTree : public MachineFunctionPass {
+class MachineDominatorTree : public DomTreeBase<MachineBasicBlock> {
   /// Helper structure used to hold all the basic blocks
   /// involved in the split of a critical edge.
   struct CriticalEdge {
@@ -70,70 +94,64 @@ class MachineDominatorTree : public MachineFunctionPass {
   /// such as BB == elt.NewBB.
   mutable SmallSet<MachineBasicBlock *, 32> NewBBs;
 
-  /// The DominatorTreeBase that is used to compute a normal dominator tree.
-  std::unique_ptr<MachineDomTree> DT;
-
   /// Apply all the recorded critical edges to the DT.
   /// This updates the underlying DT information in a way that uses
   /// the fast query path of DT as much as possible.
+  /// FIXME: This method should not be a const member!
   ///
   /// \post CriticalEdgesToSplit.empty().
   void applySplitCriticalEdges() const;
 
 public:
-  static char ID; // Pass ID, replacement for typeid
+  using Base = DomTreeBase<MachineBasicBlock>;
 
-  MachineDominatorTree();
-  explicit MachineDominatorTree(MachineFunction &MF) : MachineFunctionPass(ID) {
-    calculate(MF);
-  }
+  MachineDominatorTree() = default;
+  explicit MachineDominatorTree(MachineFunction &MF) { calculate(MF); }
 
-  MachineDomTree &getBase() {
-    if (!DT)
-      DT.reset(new MachineDomTree());
+  // FIXME: If there is an updater for MachineDominatorTree,
+  // migrate to this updater and remove these wrappers.
+
+  MachineDominatorTree &getBase() {
     applySplitCriticalEdges();
-    return *DT;
+    return *this;
   }
 
-  void getAnalysisUsage(AnalysisUsage &AU) const override;
-
   MachineBasicBlock *getRoot() const {
     applySplitCriticalEdges();
-    return DT->getRoot();
+    return Base::getRoot();
   }
 
   MachineDomTreeNode *getRootNode() const {
     applySplitCriticalEdges();
-    return DT->getRootNode();
+    return const_cast<MachineDomTreeNode *>(Base::getRootNode());
   }
 
-  bool runOnMachineFunction(MachineFunction &F) override;
-
   void calculate(MachineFunction &F);
 
   bool dominates(const MachineDomTreeNode *A,
                  const MachineDomTreeNode *B) const {
     applySplitCriticalEdges();
-    return DT->dominates(A, B);
+    return Base::dominates(A, B);
   }
 
   void getDescendants(MachineBasicBlock *A,
                       SmallVectorImpl<MachineBasicBlock *> &Result) {
     applySplitCriticalEdges();
-    DT->getDescendants(A, Result);
+    Base::getDescendants(A, Result);
   }
 
   bool dominates(const MachineBasicBlock *A, const MachineBasicBlock *B) const {
     applySplitCriticalEdges();
-    return DT->dominates(A, B);
+    return Base::dominates(A, B);
   }
 
   // dominates - Return true if A dominates B. This performs the
   // special checks necessary if A and B are in the same basic block.
   bool dominates(const MachineInstr *A, const MachineInstr *B) const {
     applySplitCriticalEdges();
     const MachineBasicBlock *BBA = A->getParent(), *BBB = B->getParent();
-    if (BBA != BBB) return DT->dominates(BBA, BBB);
+    if (BBA != BBB)
+      return Base::dominates(BBA, BBB);
 
     // Loop through the basic block until we find A or B.
     MachineBasicBlock::const_iterator I = BBA->begin();
@@ -146,34 +164,34 @@ class MachineDominatorTree : public MachineFunctionPass {
   bool properlyDominates(const MachineDomTreeNode *A,
                          const MachineDomTreeNode *B) const {
     applySplitCriticalEdges();
-    return DT->properlyDominates(A, B);
+    return Base::properlyDominates(A, B);
   }
 
   bool properlyDominates(const MachineBasicBlock *A,
                          const MachineBasicBlock *B) const {
     applySplitCriticalEdges();
-    return DT->properlyDominates(A, B);
+    return Base::properlyDominates(A, B);
   }
 
   /// findNearestCommonDominator - Find nearest common dominator basic block
   /// for basic block A and B. If there is no such block then return NULL.
   MachineBasicBlock *findNearestCommonDominator(MachineBasicBlock *A,
                                                 MachineBasicBlock *B) {
     applySplitCriticalEdges();
-    return DT->findNearestCommonDominator(A, B);
+    return Base::findNearestCommonDominator(A, B);
   }
 
   MachineDomTreeNode *operator[](MachineBasicBlock *BB) const {
     applySplitCriticalEdges();
-    return DT->getNode(BB);
+    return Base::getNode(BB);
   }
 
   /// getNode - return the (Post)DominatorTree node for the specified basic
   /// block.  This is the same as using operator[] on this class.
   ///
   MachineDomTreeNode *getNode(MachineBasicBlock *BB) const {
     applySplitCriticalEdges();
-    return DT->getNode(BB);
+    return Base::getNode(BB);
   }
 
   /// addNewBlock - Add a new node to the dominator tree information.  This
@@ -182,7 +200,7 @@ class MachineDominatorTree : public MachineFunctionPass {
   MachineDomTreeNode *addNewBlock(MachineBasicBlock *BB,
                                   MachineBasicBlock *DomBB) {
     applySplitCriticalEdges();
-    return DT->addNewBlock(BB, DomBB);
+    return Base::addNewBlock(BB, DomBB);
   }
 
   /// changeImmediateDominator - This method is used to update the dominator
@@ -191,43 +209,37 @@ class MachineDominatorTree : public MachineFunctionPass {
   void changeImmediateDominator(MachineBasicBlock *N,
                                 MachineBasicBlock *NewIDom) {
     applySplitCriticalEdges();
-    DT->changeImmediateDominator(N, NewIDom);
+    Base::changeImmediateDominator(N, NewIDom);
   }
 
   void changeImmediateDominator(MachineDomTreeNode *N,
                                 MachineDomTreeNode *NewIDom) {
     applySplitCriticalEdges();
-    DT->changeImmediateDominator(N, NewIDom);
+    Base::changeImmediateDominator(N, NewIDom);
   }
 
   /// eraseNode - Removes a node from  the dominator tree. Block must not
   /// dominate any other blocks. Removes node from its immediate dominator's
   /// children list. Deletes dominator node associated with basic block BB.
   void eraseNode(MachineBasicBlock *BB) {
     applySplitCriticalEdges();
-    DT->eraseNode(BB);
+    Base::eraseNode(BB);
   }
 
   /// splitBlock - BB is split and now it has one successor. Update dominator
   /// tree to reflect this change.
   void splitBlock(MachineBasicBlock* NewBB) {
     applySplitCriticalEdges();
-    DT->splitBlock(NewBB);
+    Base::splitBlock(NewBB);
   }
 
   /// isReachableFromEntry - Return true if A is dominated by the entry
   /// block of the function containing it.
   bool isReachableFromEntry(const MachineBasicBlock *A) {
     applySplitCriticalEdges();
-    return DT->isReachableFromEntry(A);
+    return Base::isReachableFromEntry(A);
   }
 
-  void releaseMemory() override;
-
-  void verifyAnalysis() const override;
-
-  void print(raw_ostream &OS, const Module*) const override;
-
   /// Record that the critical edge (FromBB, ToBB) has been
   /// split with NewBB.
   /// This is best to use this method instead of directly update the
@@ -251,6 +263,34 @@ class MachineDominatorTree : public MachineFunctionPass {
   }
 };
 
+/// \brief Analysis pass which computes a \c MachineDominatorTree.
+class MachineDominatorTreeWrapperPass : public MachineFunctionPass {
+  // MachineFunctionPass may verify the analysis result without running pass,
+  // e.g. when `F.hasAvailableExternallyLinkage` is true.
+  std::optional<MachineDominatorTree> DT;
+
+public:
+  static char ID;
+
+  MachineDominatorTreeWrapperPass();
+
+  MachineDominatorTree &getDomTree() { return *DT; }
+  const MachineDominatorTree &getDomTree() const { return *DT; }
+
+  bool runOnMachineFunction(MachineFunction &MF) override;
+
+  void verifyAnalysis() const override;
+
+  void getAnalysisUsage(AnalysisUsage &AU) const override {
+    AU.setPreservesAll();
+    MachineFunctionPass::getAnalysisUsage(AU);
+  }
+
+  void releaseMemory() override;
+
+  void print(raw_ostream &OS, const Module *M = nullptr) const override;
+};
+
 //===-------------------------------------
 /// DominatorTree GraphTraits specialization so the DominatorTree can be
 /// iterable by generic graph iterators.

llvm/include/llvm/CodeGen/MachineUniformityAnalysis.h

@@ -30,7 +30,7 @@ using MachineUniformityInfo = GenericUniformityInfo<MachineSSAContext>;
 /// everything is uniform.
 MachineUniformityInfo computeMachineUniformityInfo(
     MachineFunction &F, const MachineCycleInfo &cycleInfo,
-    const MachineDomTree &domTree, bool HasBranchDivergence);
+    const MachineDominatorTree &domTree, bool HasBranchDivergence);
 
 /// Legacy analysis pass which computes a \ref MachineUniformityInfo.
 class MachineUniformityAnalysisPass : public MachineFunctionPass {

llvm/include/llvm/InitializePasses.h

@@ -190,7 +190,7 @@ void initializeMachineCopyPropagationPass(PassRegistry&);
 void initializeMachineCycleInfoPrinterPassPass(PassRegistry &);
 void initializeMachineCycleInfoWrapperPassPass(PassRegistry &);
 void initializeMachineDominanceFrontierPass(PassRegistry&);
-void initializeMachineDominatorTreePass(PassRegistry&);
+void initializeMachineDominatorTreeWrapperPassPass(PassRegistry &);
 void initializeMachineFunctionPrinterPassPass(PassRegistry&);
 void initializeMachineFunctionSplitterPass(PassRegistry &);
 void initializeMachineLateInstrsCleanupPass(PassRegistry&);

llvm/lib/CodeGen/AsmPrinter/AsmPrinter.cpp

@@ -1728,7 +1728,8 @@ void AsmPrinter::emitFunctionBody() {
 
   if (isVerbose()) {
     // Get MachineDominatorTree or compute it on the fly if it's unavailable
-    MDT = getAnalysisIfAvailable<MachineDominatorTree>();
+    auto MDTWrapper = getAnalysisIfAvailable<MachineDominatorTreeWrapperPass>();
+    MDT = MDTWrapper ? &MDTWrapper->getDomTree() : nullptr;
     if (!MDT) {
       OwnedMDT = std::make_unique<MachineDominatorTree>();
       OwnedMDT->getBase().recalculate(*MF);

llvm/lib/CodeGen/CodeGen.cpp

@@ -80,7 +80,7 @@ void llvm::initializeCodeGen(PassRegistry &Registry) {
   initializeMachineCopyPropagationPass(Registry);
   initializeMachineCycleInfoPrinterPassPass(Registry);
   initializeMachineCycleInfoWrapperPassPass(Registry);
-  initializeMachineDominatorTreePass(Registry);
+  initializeMachineDominatorTreeWrapperPassPass(Registry);
   initializeMachineFunctionPrinterPassPass(Registry);
   initializeMachineLateInstrsCleanupPass(Registry);
   initializeMachineLICMPass(Registry);

llvm/lib/CodeGen/EarlyIfConversion.cpp

@@ -790,15 +790,15 @@ char &llvm::EarlyIfConverterID = EarlyIfConverter::ID;
 INITIALIZE_PASS_BEGIN(EarlyIfConverter, DEBUG_TYPE,
                       "Early If Converter", false, false)
 INITIALIZE_PASS_DEPENDENCY(MachineBranchProbabilityInfo)
-INITIALIZE_PASS_DEPENDENCY(MachineDominatorTree)
+INITIALIZE_PASS_DEPENDENCY(MachineDominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(MachineTraceMetrics)
 INITIALIZE_PASS_END(EarlyIfConverter, DEBUG_TYPE,
                     "Early If Converter", false, false)
 
 void EarlyIfConverter::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.addRequired<MachineBranchProbabilityInfo>();
-  AU.addRequired<MachineDominatorTree>();
-  AU.addPreserved<MachineDominatorTree>();
+  AU.addRequired<MachineDominatorTreeWrapperPass>();
+  AU.addPreserved<MachineDominatorTreeWrapperPass>();
   AU.addRequired<MachineLoopInfo>();
   AU.addPreserved<MachineLoopInfo>();
   AU.addRequired<MachineTraceMetrics>();
@@ -1089,7 +1089,7 @@ bool EarlyIfConverter::runOnMachineFunction(MachineFunction &MF) {
   TRI = STI.getRegisterInfo();
   SchedModel = STI.getSchedModel();
   MRI = &MF.getRegInfo();
-  DomTree = &getAnalysis<MachineDominatorTree>();
+  DomTree = &getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree();
   Loops = &getAnalysis<MachineLoopInfo>();
   Traces = &getAnalysis<MachineTraceMetrics>();
   MinInstr = nullptr;
@@ -1144,15 +1144,15 @@ char &llvm::EarlyIfPredicatorID = EarlyIfPredicator::ID;
 
 INITIALIZE_PASS_BEGIN(EarlyIfPredicator, DEBUG_TYPE, "Early If Predicator",
                       false, false)
-INITIALIZE_PASS_DEPENDENCY(MachineDominatorTree)
+INITIALIZE_PASS_DEPENDENCY(MachineDominatorTreeWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(MachineBranchProbabilityInfo)
 INITIALIZE_PASS_END(EarlyIfPredicator, DEBUG_TYPE, "Early If Predicator", false,
                     false)
 
 void EarlyIfPredicator::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.addRequired<MachineBranchProbabilityInfo>();
-  AU.addRequired<MachineDominatorTree>();
-  AU.addPreserved<MachineDominatorTree>();
+  AU.addRequired<MachineDominatorTreeWrapperPass>();
+  AU.addPreserved<MachineDominatorTreeWrapperPass>();
   AU.addRequired<MachineLoopInfo>();
   AU.addPreserved<MachineLoopInfo>();
   MachineFunctionPass::getAnalysisUsage(AU);
@@ -1223,7 +1223,7 @@ bool EarlyIfPredicator::runOnMachineFunction(MachineFunction &MF) {
   TRI = STI.getRegisterInfo();
   MRI = &MF.getRegInfo();
   SchedModel.init(&STI);
-  DomTree = &getAnalysis<MachineDominatorTree>();
+  DomTree = &getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree();
   Loops = &getAnalysis<MachineLoopInfo>();
   MBPI = &getAnalysis<MachineBranchProbabilityInfo>();