[CodeGen] Use optimized domtree for MachineFunction (#102107)

The dominator tree gained an optimization to use block numbers instead
of a DenseMap to store blocks. Given that machine basic blocks already
have numbers, expose these via appropriate GraphTraits. For debugging,
block number epochs are added to MachineFunction -- this greatly helps
in finding uses of block numbers after RenumberBlocks().

In a few cases where dominator trees are preserved across renumberings,
the dominator tree is updated to use the new numbers.

Author: aengelke

llvm/include/llvm/CodeGen/MachineBasicBlock.h

@@ -1297,6 +1297,11 @@ template <> struct GraphTraits<MachineBasicBlock *> {
   static NodeRef getEntryNode(MachineBasicBlock *BB) { return BB; }
   static ChildIteratorType child_begin(NodeRef N) { return N->succ_begin(); }
   static ChildIteratorType child_end(NodeRef N) { return N->succ_end(); }
+
+  static unsigned getNumber(MachineBasicBlock *BB) {
+    assert(BB->getNumber() >= 0 && "negative block number");
+    return BB->getNumber();
+  }
 };
 
 template <> struct GraphTraits<const MachineBasicBlock *> {
@@ -1306,6 +1311,11 @@ template <> struct GraphTraits<const MachineBasicBlock *> {
   static NodeRef getEntryNode(const MachineBasicBlock *BB) { return BB; }
   static ChildIteratorType child_begin(NodeRef N) { return N->succ_begin(); }
   static ChildIteratorType child_end(NodeRef N) { return N->succ_end(); }
+
+  static unsigned getNumber(const MachineBasicBlock *BB) {
+    assert(BB->getNumber() >= 0 && "negative block number");
+    return BB->getNumber();
+  }
 };
 
 // Provide specializations of GraphTraits to be able to treat a
@@ -1324,6 +1334,11 @@ template <> struct GraphTraits<Inverse<MachineBasicBlock*>> {
 
   static ChildIteratorType child_begin(NodeRef N) { return N->pred_begin(); }
   static ChildIteratorType child_end(NodeRef N) { return N->pred_end(); }
+
+  static unsigned getNumber(MachineBasicBlock *BB) {
+    assert(BB->getNumber() >= 0 && "negative block number");
+    return BB->getNumber();
+  }
 };
 
 template <> struct GraphTraits<Inverse<const MachineBasicBlock*>> {
@@ -1336,6 +1351,11 @@ template <> struct GraphTraits<Inverse<const MachineBasicBlock*>> {
 
   static ChildIteratorType child_begin(NodeRef N) { return N->pred_begin(); }
   static ChildIteratorType child_end(NodeRef N) { return N->pred_end(); }
+
+  static unsigned getNumber(const MachineBasicBlock *BB) {
+    assert(BB->getNumber() >= 0 && "negative block number");
+    return BB->getNumber();
+  }
 };
 
 // These accessors are handy for sharing templated code between IR and MIR.

llvm/include/llvm/CodeGen/MachineFunction.h

@@ -293,6 +293,10 @@ class LLVM_EXTERNAL_VISIBILITY MachineFunction {
   // numbered and this vector keeps track of the mapping from ID's to MBB's.
   std::vector<MachineBasicBlock*> MBBNumbering;
 
+  // MBBNumbering epoch, incremented after renumbering to detect use of old
+  // block numbers.
+  unsigned MBBNumberingEpoch = 0;
+
   // Pool-allocate MachineFunction-lifetime and IR objects.
   BumpPtrAllocator Allocator;
 
@@ -856,6 +860,11 @@ class LLVM_EXTERNAL_VISIBILITY MachineFunction {
   /// getNumBlockIDs - Return the number of MBB ID's allocated.
   unsigned getNumBlockIDs() const { return (unsigned)MBBNumbering.size(); }
 
+  /// Return the numbering "epoch" of block numbers, incremented after each
+  /// numbering. Intended for asserting that no renumbering was performed when
+  /// used by, e.g., preserved analyses.
+  unsigned getBlockNumberEpoch() const { return MBBNumberingEpoch; }
+
   /// RenumberBlocks - This discards all of the MachineBasicBlock numbers and
   /// recomputes them.  This guarantees that the MBB numbers are sequential,
   /// dense, and match the ordering of the blocks within the function.  If a
@@ -1404,6 +1413,13 @@ template <> struct GraphTraits<MachineFunction*> :
   }
 
   static unsigned       size       (MachineFunction *F) { return F->size(); }
+
+  static unsigned getMaxNumber(MachineFunction *F) {
+    return F->getNumBlockIDs();
+  }
+  static unsigned getNumberEpoch(MachineFunction *F) {
+    return F->getBlockNumberEpoch();
+  }
 };
 template <> struct GraphTraits<const MachineFunction*> :
   public GraphTraits<const MachineBasicBlock*> {
@@ -1423,6 +1439,13 @@ template <> struct GraphTraits<const MachineFunction*> :
   static unsigned       size       (const MachineFunction *F)  {
     return F->size();
   }
+
+  static unsigned getMaxNumber(const MachineFunction *F) {
+    return F->getNumBlockIDs();
+  }
+  static unsigned getNumberEpoch(const MachineFunction *F) {
+    return F->getBlockNumberEpoch();
+  }
 };
 
 // Provide specializations of GraphTraits to be able to treat a function as a
@@ -1435,12 +1458,26 @@ template <> struct GraphTraits<Inverse<MachineFunction*>> :
   static NodeRef getEntryNode(Inverse<MachineFunction *> G) {
     return &G.Graph->front();
   }
+
+  static unsigned getMaxNumber(MachineFunction *F) {
+    return F->getNumBlockIDs();
+  }
+  static unsigned getNumberEpoch(MachineFunction *F) {
+    return F->getBlockNumberEpoch();
+  }
 };
 template <> struct GraphTraits<Inverse<const MachineFunction*>> :
   public GraphTraits<Inverse<const MachineBasicBlock*>> {
   static NodeRef getEntryNode(Inverse<const MachineFunction *> G) {
     return &G.Graph->front();
   }
+
+  static unsigned getMaxNumber(const MachineFunction *F) {
+    return F->getNumBlockIDs();
+  }
+  static unsigned getNumberEpoch(const MachineFunction *F) {
+    return F->getBlockNumberEpoch();
+  }
 };
 
 void verifyMachineFunction(const std::string &Banner,

llvm/lib/CodeGen/BasicBlockSections.cpp

@@ -72,8 +72,10 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/CodeGen/BasicBlockSectionUtils.h"
 #include "llvm/CodeGen/BasicBlockSectionsProfileReader.h"
+#include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
+#include "llvm/CodeGen/MachinePostDominators.h"
 #include "llvm/CodeGen/Passes.h"
 #include "llvm/CodeGen/TargetInstrInfo.h"
 #include "llvm/InitializePasses.h"
@@ -393,12 +395,21 @@ bool BasicBlockSections::runOnMachineFunction(MachineFunction &MF) {
   auto R1 = handleBBSections(MF);
   // Handle basic block address map after basic block sections are finalized.
   auto R2 = handleBBAddrMap(MF);
+
+  // We renumber blocks, so update the dominator tree we want to preserve.
+  if (auto *WP = getAnalysisIfAvailable<MachineDominatorTreeWrapperPass>())
+    WP->getDomTree().updateBlockNumbers();
+  if (auto *WP = getAnalysisIfAvailable<MachinePostDominatorTreeWrapperPass>())
+    WP->getPostDomTree().updateBlockNumbers();
+
   return R1 || R2;
 }
 
 void BasicBlockSections::getAnalysisUsage(AnalysisUsage &AU) const {
   AU.setPreservesAll();
   AU.addRequired<BasicBlockSectionsProfileReaderWrapperPass>();
+  AU.addUsedIfAvailable<MachineDominatorTreeWrapperPass>();
+  AU.addUsedIfAvailable<MachinePostDominatorTreeWrapperPass>();
   MachineFunctionPass::getAnalysisUsage(AU);
 }
 

llvm/lib/CodeGen/MIRSampleProfile.cpp

@@ -366,13 +366,18 @@ bool MIRProfileLoaderPass::runOnMachineFunction(MachineFunction &MF) {
   LLVM_DEBUG(dbgs() << "MIRProfileLoader pass working on Func: "
                     << MF.getFunction().getName() << "\n");
   MBFI = &getAnalysis<MachineBlockFrequencyInfoWrapperPass>().getMBFI();
+  auto *MDT = &getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree();
+  auto *MPDT =
+      &getAnalysis<MachinePostDominatorTreeWrapperPass>().getPostDomTree();
+
+  MF.RenumberBlocks();
+  MDT->updateBlockNumbers();
+  MPDT->updateBlockNumbers();
+
   MIRSampleLoader->setInitVals(
-      &getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree(),
-      &getAnalysis<MachinePostDominatorTreeWrapperPass>().getPostDomTree(),
-      &getAnalysis<MachineLoopInfoWrapperPass>().getLI(), MBFI,
+      MDT, MPDT, &getAnalysis<MachineLoopInfoWrapperPass>().getLI(), MBFI,
       &getAnalysis<MachineOptimizationRemarkEmitterPass>().getORE());
 
-  MF.RenumberBlocks();
   if (ViewBFIBefore && ViewBlockLayoutWithBFI != GVDT_None &&
       (ViewBlockFreqFuncName.empty() ||
        MF.getFunction().getName() == ViewBlockFreqFuncName)) {

llvm/lib/CodeGen/MachineBlockPlacement.cpp

@@ -3649,6 +3649,7 @@ void MachineBlockPlacement::assignBlockOrder(
     const std::vector<const MachineBasicBlock *> &NewBlockOrder) {
   assert(F->size() == NewBlockOrder.size() && "Incorrect size of block order");
   F->RenumberBlocks();
+  MPDT->updateBlockNumbers();
 
   bool HasChanges = false;
   for (size_t I = 0; I < NewBlockOrder.size(); I++) {

llvm/lib/CodeGen/MachineFunction.cpp

@@ -375,6 +375,7 @@ void MachineFunction::RenumberBlocks(MachineBasicBlock *MBB) {
   // numbering, shrink MBBNumbering now.
   assert(BlockNo <= MBBNumbering.size() && "Mismatch!");
   MBBNumbering.resize(BlockNo);
+  MBBNumberingEpoch++;
 }
 
 /// This method iterates over the basic blocks and assigns their IsBeginSection

llvm/lib/CodeGen/UnreachableBlockElim.cpp

@@ -195,6 +195,8 @@ bool UnreachableMachineBlockElim::runOnMachineFunction(MachineFunction &F) {
   }
 
   F.RenumberBlocks();
+  if (MDT)
+    MDT->updateBlockNumbers();
 
   return (!DeadBlocks.empty() || ModifiedPHI);
 }

llvm/lib/Target/ARM/ARMConstantIslandPass.cpp

@@ -414,6 +414,7 @@ bool ARMConstantIslands::runOnMachineFunction(MachineFunction &mf) {
   // Renumber all of the machine basic blocks in the function, guaranteeing that
   // the numbers agree with the position of the block in the function.
   MF->RenumberBlocks();
+  DT->updateBlockNumbers();
 
   // Try to reorder and otherwise adjust the block layout to make good use
   // of the TB[BH] instructions.
@@ -425,6 +426,7 @@ bool ARMConstantIslands::runOnMachineFunction(MachineFunction &mf) {
     T2JumpTables.clear();
     // Blocks may have shifted around. Keep the numbering up to date.
     MF->RenumberBlocks();
+    DT->updateBlockNumbers();
   }
 
   // Align any non-fallthrough blocks
@@ -670,8 +672,10 @@ void ARMConstantIslands::doInitialJumpTablePlacement(
   }
 
   // If we did anything then we need to renumber the subsequent blocks.
-  if (LastCorrectlyNumberedBB)
+  if (LastCorrectlyNumberedBB) {
     MF->RenumberBlocks(LastCorrectlyNumberedBB);
+    DT->updateBlockNumbers();
+  }
 }
 
 /// BBHasFallthrough - Return true if the specified basic block can fallthrough
@@ -972,6 +976,7 @@ static bool CompareMBBNumbers(const MachineBasicBlock *LHS,
 void ARMConstantIslands::updateForInsertedWaterBlock(MachineBasicBlock *NewBB) {
   // Renumber the MBB's to keep them consecutive.
   NewBB->getParent()->RenumberBlocks(NewBB);
+  DT->updateBlockNumbers();
 
   // Insert an entry into BBInfo to align it properly with the (newly
   // renumbered) block numbers.
@@ -1034,6 +1039,7 @@ MachineBasicBlock *ARMConstantIslands::splitBlockBeforeInstr(MachineInstr *MI) {
   // This is almost the same as updateForInsertedWaterBlock, except that
   // the Water goes after OrigBB, not NewBB.
   MF->RenumberBlocks(NewBB);
+  DT->updateBlockNumbers();
 
   // Insert an entry into BBInfo to align it properly with the (newly
   // renumbered) block numbers.
@@ -2485,6 +2491,7 @@ MachineBasicBlock *ARMConstantIslands::adjustJTTargetBlockForward(
     BB->updateTerminator(OldNext != MF->end() ? &*OldNext : nullptr);
     // Update numbering to account for the block being moved.
     MF->RenumberBlocks();
+    DT->updateBlockNumbers();
     ++NumJTMoved;
     return nullptr;
   }
@@ -2513,6 +2520,7 @@ MachineBasicBlock *ARMConstantIslands::adjustJTTargetBlockForward(
 
   // Update internal data structures to account for the newly inserted MBB.
   MF->RenumberBlocks(NewBB);
+  DT->updateBlockNumbers();
 
   // Update the CFG.
   NewBB->addSuccessor(BB);

llvm/lib/Target/CSKY/CSKYConstantIslandPass.cpp

@@ -28,7 +28,6 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineConstantPool.h"
-#include "llvm/CodeGen/MachineDominators.h"
 #include "llvm/CodeGen/MachineFrameInfo.h"
 #include "llvm/CodeGen/MachineFunction.h"
 #include "llvm/CodeGen/MachineFunctionPass.h"
@@ -217,11 +216,6 @@ class CSKYConstantIslands : public MachineFunctionPass {
 
   bool runOnMachineFunction(MachineFunction &F) override;
 
-  void getAnalysisUsage(AnalysisUsage &AU) const override {
-    AU.addRequired<MachineDominatorTreeWrapperPass>();
-    MachineFunctionPass::getAnalysisUsage(AU);
-  }
-
   MachineFunctionProperties getRequiredProperties() const override {
     return MachineFunctionProperties().set(
         MachineFunctionProperties::Property::NoVRegs);

llvm/lib/Target/WebAssembly/WebAssemblyCFGSort.cpp

@@ -186,10 +186,11 @@ struct Entry {
 /// Explore them.
 static void sortBlocks(MachineFunction &MF, const MachineLoopInfo &MLI,
                        const WebAssemblyExceptionInfo &WEI,
-                       const MachineDominatorTree &MDT) {
+                       MachineDominatorTree &MDT) {
   // Remember original layout ordering, so we can update terminators after
   // reordering to point to the original layout successor.
   MF.RenumberBlocks();
+  MDT.updateBlockNumbers();
 
   // Prepare for a topological sort: Record the number of predecessors each
   // block has, ignoring loop backedges.
@@ -330,6 +331,7 @@ static void sortBlocks(MachineFunction &MF, const MachineLoopInfo &MLI,
   }
   assert(Entries.empty() && "Active sort region list not finished");
   MF.RenumberBlocks();
+  MDT.updateBlockNumbers();
 
 #ifndef NDEBUG
   SmallSetVector<const SortRegion *, 8> OnStack;

llvm/lib/Target/WebAssembly/WebAssemblyCFGStackify.cpp

@@ -45,6 +45,8 @@ STATISTIC(NumCatchUnwindMismatches, "Number of catch unwind mismatches found");
 
 namespace {
 class WebAssemblyCFGStackify final : public MachineFunctionPass {
+  MachineDominatorTree *MDT;
+
   StringRef getPassName() const override { return "WebAssembly CFG Stackify"; }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
@@ -252,7 +254,6 @@ void WebAssemblyCFGStackify::unregisterScope(MachineInstr *Begin) {
 void WebAssemblyCFGStackify::placeBlockMarker(MachineBasicBlock &MBB) {
   assert(!MBB.isEHPad());
   MachineFunction &MF = *MBB.getParent();
-  auto &MDT = getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree();
   const auto &TII = *MF.getSubtarget<WebAssemblySubtarget>().getInstrInfo();
   const auto &MFI = *MF.getInfo<WebAssemblyFunctionInfo>();
 
@@ -264,7 +265,7 @@ void WebAssemblyCFGStackify::placeBlockMarker(MachineBasicBlock &MBB) {
   int MBBNumber = MBB.getNumber();
   for (MachineBasicBlock *Pred : MBB.predecessors()) {
     if (Pred->getNumber() < MBBNumber) {
-      Header = Header ? MDT.findNearestCommonDominator(Header, Pred) : Pred;
+      Header = Header ? MDT->findNearestCommonDominator(Header, Pred) : Pred;
       if (explicitlyBranchesTo(Pred, &MBB))
         IsBranchedTo = true;
     }
@@ -1439,6 +1440,7 @@ void WebAssemblyCFGStackify::recalculateScopeTops(MachineFunction &MF) {
   // Renumber BBs and recalculate ScopeTop info because new BBs might have been
   // created and inserted during fixing unwind mismatches.
   MF.RenumberBlocks();
+  MDT->updateBlockNumbers();
   ScopeTops.clear();
   ScopeTops.resize(MF.getNumBlockIDs());
   for (auto &MBB : reverse(MF)) {
@@ -1741,6 +1743,7 @@ bool WebAssemblyCFGStackify::runOnMachineFunction(MachineFunction &MF) {
                        "********** Function: "
                     << MF.getName() << '\n');
   const MCAsmInfo *MCAI = MF.getTarget().getMCAsmInfo();
+  MDT = &getAnalysis<MachineDominatorTreeWrapperPass>().getDomTree();
 
   releaseMemory();