Teach DTU to recalculate DT/PDT automatically when EntryBB is changed

[NutshellySima]

Summary

Previously, after the EntryBB is changed, due to restrictions of the DominatorTree, it is needed to call recalculate() explicitly.
This patch teaches DTU to detect the change of EntryBB when calling applyUpdates/insertEdge*/deleteEdge*/*deleteBB/getDT/getPDT/flush and recalculate available trees if it is needed, which makes DTU more friendly to use.

[brzycki]

This is much cleaner, I like this very much. Originally adding the method recalculate() was a compromise. With this change do we even need the public recalculate()method? Does it ever make sense to make a new entry BB that has zero successors?

[kuhar]

Seems like a nice improvement!

llvm/include/llvm/Analysis/DomTreeUpdater.h

245

Maybe NeedToRecalculate?

259

checkIfRootChanged?

264

internalRecalculate

llvm/include/llvm/Support/GenericDomTree.h

252

Why do we need this? maybe it would be better to expose/use some other function? Tree root can be accessed with the getRoots function.

[kuhar]

I believe that a real fix would be to always create a virtual root in DomTree and support fast rerooting, but that requires much more implementation and testing effort.

[NutshellySima]

In D58187#1396464, @brzycki wrote:

This is much cleaner, I like this very much. Originally adding the method recalculate() was a compromise. With this change do we even need the public recalculate()method? Does it ever make sense to make a new entry BB that has zero successors?

recalculate() can be deprecated. But there are still edge cases, for example, when we only insert a new block to be the entry block or move a block up to be the entry block, it is still needed to call flush() to detect this change. (But I don't see usage like the above scenario I mentioned currently.)

[brzycki]

recalculate() can be deprecated. But there are still edge cases, for example, when we only insert a new block to be the entry block or move a block up to be the entry block, it is still needed to call flush() to detect this change. (But I don't see usage like the above scenario I mentioned currently.)

If there are still edge cases then we need to keep it as a last-ditch method of resetting the DTU state without creating a new class instance. It might be a good idea to add documentation stating the interface isn't meant to be used except as a last-ditch fix that cannot be addressed by any other means.

[mkazantsev]

llvm/include/llvm/Analysis/DomTreeUpdater.h

257

Should we do it for PDT when exit block is changed?

[kuhar]

llvm/include/llvm/Analysis/DomTreeUpdater.h

257

PDT's root never changes -- we use a virtual exit node to support multiple exits. DT should be handled in the same way IMO, but that's a lot of work.

llvm/include/llvm/Analysis/DomTreeUpdater.h

 
 namespace llvm {
 class DomTreeUpdater {
 public:
   enum class UpdateStrategy : unsigned char { Eager = 0, Lazy = 1 };
 
   explicit DomTreeUpdater(UpdateStrategy Strategy_) : Strategy(Strategy_) {}
   DomTreeUpdater(DominatorTree &DT_, UpdateStrategy Strategy_)
-      : DT(&DT_), Strategy(Strategy_) {}
+      : DT(&DT_), Strategy(Strategy_) {
+    EntryBlock = &DT->Parent->getEntryBlock();
+  }
   DomTreeUpdater(DominatorTree *DT_, UpdateStrategy Strategy_)
-      : DT(DT_), Strategy(Strategy_) {}
+      : DT(DT_), Strategy(Strategy_) {
+    if (DT)
+      EntryBlock = &DT->Parent->getEntryBlock();
+  }
   DomTreeUpdater(PostDominatorTree &PDT_, UpdateStrategy Strategy_)
       : PDT(&PDT_), Strategy(Strategy_) {}
   DomTreeUpdater(PostDominatorTree *PDT_, UpdateStrategy Strategy_)
       : PDT(PDT_), Strategy(Strategy_) {}
   DomTreeUpdater(DominatorTree &DT_, PostDominatorTree &PDT_,
                  UpdateStrategy Strategy_)
-      : DT(&DT_), PDT(&PDT_), Strategy(Strategy_) {}
+      : DT(&DT_), PDT(&PDT_), Strategy(Strategy_) {
+    EntryBlock = &DT->Parent->getEntryBlock();
+  }
   DomTreeUpdater(DominatorTree *DT_, PostDominatorTree *PDT_,
                  UpdateStrategy Strategy_)
-      : DT(DT_), PDT(PDT_), Strategy(Strategy_) {}
+      : DT(DT_), PDT(PDT_), Strategy(Strategy_) {
+    if (DT)
+      EntryBlock = &DT->Parent->getEntryBlock();
+  }
 
   ~DomTreeUpdater() { flush(); }
 
   /// Returns true if the current strategy is Lazy.
   bool isLazy() const { return Strategy == UpdateStrategy::Lazy; };
 
   /// Returns true if the current strategy is Eager.
   bool isEager() const { return Strategy == UpdateStrategy::Eager; };
public:
   /// analyze the current state of the CFG to determine whether an update
   /// is valid.
   /// 2. It is illegal to submit any update that has already been submitted,
   /// i.e. you are supposed not to insert an existent edge or delete a
   /// nonexistent edge.
   void applyUpdates(ArrayRef<DominatorTree::UpdateType> Updates,
                     bool ForceRemoveDuplicates = false);
 
-  /// Notify DTU that the entry block was replaced.
   /// Recalculate all available trees and flush all BasicBlocks
-  /// awaiting deletion immediately.
+  /// awaiting deletion immediately. It is only used when you insert
+  /// a new BasicBlock to be the entry block without making other changes.
   void recalculate(Function &F);
 
   /// \deprecated { Submit an edge insertion to all available trees. The Eager
   /// Strategy flushes this update immediately while the Lazy Strategy queues
   /// the update. An internal function checks if the edge exists in the CFG in
   /// DEBUG mode. CAUTION! This function has to be called *after* making the
   /// update on the actual CFG. It is illegal to submit any update that has
   /// already been applied. }
public:
 
   /// Flush PostDomTree updates and return PostDomTree.
   /// It flushes Deleted BBs if both trees are up-to-date.
   /// It must only be called when it has a PostDomTree.
   PostDominatorTree &getPostDomTree();
 
   /// Apply all pending updates to available trees and flush all BasicBlocks
   /// awaiting deletion.
-
   void flush();
 
   ///@}
 
   /// Debug method to help view the internal state of this class.
   LLVM_DUMP_METHOD void dump() const;
 
 private:
   class CallBackOnDeletion final : public CallbackVH {
   public:
     CallBackOnDeletion(BasicBlock *V,
                        std::function<void(BasicBlock *)> Callback)
         : CallbackVH(V), DelBB(V), Callback_(Callback) {}
 
   private:
     BasicBlock *DelBB = nullptr;
     std::function<void(BasicBlock *)> Callback_;
 
     void deleted() override {
       Callback_(DelBB);
       CallbackVH::deleted();
     }
   };
 
+  BasicBlock *EntryBlock = nullptr;
+  bool NeedRecalculate = false;
   SmallVector<DominatorTree::UpdateType, 16> PendUpdates;
   size_t PendDTUpdateIndex = 0;
   size_t PendPDTUpdateIndex = 0;
   DominatorTree *DT = nullptr;
   PostDominatorTree *PDT = nullptr;
   const UpdateStrategy Strategy;
   SmallPtrSet<BasicBlock *, 8> DeletedBBs;
   std::vector<CallBackOnDeletion> Callbacks;
   bool IsRecalculatingDomTree = false;
   bool IsRecalculatingPostDomTree = false;
 
+  /// If the entry block is changed, recalculate available trees.
+  /// Returns true if available trees are recalculated.
+  bool checkEntryBlockChange();
+
+  /// Notify DTU that the entry block was replaced.
+  /// Recalculate all available trees and flush all BasicBlocks
+  /// awaiting deletion immediately.
+  void internalrecalculate();
+
   /// First remove all the instructions of DelBB and then make sure DelBB has a
   /// valid terminator instruction which is necessary to have when DelBB still
   /// has to be inside of its parent Function while awaiting deletion under Lazy
   /// UpdateStrategy to prevent other routines from asserting the state of the
   /// IR is inconsistent. Assert if DelBB is nullptr or has predecessors.
   void validateDeleteBB(BasicBlock *DelBB);
 
   /// Returns true if at least one BasicBlock is deleted.

llvm/include/llvm/Support/GenericDomTree.h

protected:
   DomTreeNodeMapType DomTreeNodes;
   DomTreeNodeBase<NodeT> *RootNode;
   ParentPtr Parent = nullptr;
 
   mutable bool DFSInfoValid = false;
   mutable unsigned int SlowQueries = 0;
 
   friend struct DomTreeBuilder::SemiNCAInfo<DominatorTreeBase>;
+  friend class DomTreeUpdater;
 
- public:
   DominatorTreeBase() {}
 
   DominatorTreeBase(DominatorTreeBase &&Arg)
       : Roots(std::move(Arg.Roots)),
         DomTreeNodes(std::move(Arg.DomTreeNodes)),
         RootNode(Arg.RootNode),
         Parent(Arg.Parent),
         DFSInfoValid(Arg.DFSInfoValid),

llvm/lib/Analysis/DomTreeUpdater.cpp

if (hasPendingDomTreeUpdates()) {
     const auto E = PendUpdates.end();
     assert(I < E && "Iterator range invalid; there should be DomTree updates.");
     DT->applyUpdates(ArrayRef<DominatorTree::UpdateType>(I, E));
     PendDTUpdateIndex = PendUpdates.size();
   }
 }
 
 void DomTreeUpdater::flush() {
+  if (isLazy() && NeedRecalculate) {
+    internalrecalculate();
+    return;
+  }
+
+  if (checkEntryBlockChange())
+    return;
+
   applyDomTreeUpdates();
   applyPostDomTreeUpdates();
   dropOutOfDateUpdates();
 }
 
 void DomTreeUpdater::applyPostDomTreeUpdates() {
   // No pending PostDomTreeUpdates.
   if (Strategy != UpdateStrategy::Lazy || !PDT)
for (auto *BB : DeletedBBs) {
     eraseDelBBNode(BB);
     delete BB;
   }
   DeletedBBs.clear();
   Callbacks.clear();
   return true;
 }
 
-void DomTreeUpdater::recalculate(Function &F) {
+bool DomTreeUpdater::checkEntryBlockChange() {
+  // There is no DT.
+  if (!DT)
+    return false;
+
+  BasicBlock *CurrentEntryBlock = &DT->Parent->getEntryBlock();
+  if (CurrentEntryBlock != EntryBlock) {
+    internalrecalculate();
+    return true;
+  }
+
+  return false;
+}
+
+void DomTreeUpdater::recalculate(Function &F) { internalrecalculate(); }
+
+void DomTreeUpdater::internalrecalculate() {
+
+  NeedRecalculate = false;
+
+  if (DT)
+    EntryBlock = &DT->Parent->getEntryBlock();
 
   if (Strategy == UpdateStrategy::Eager) {
     if (DT)
-      DT->recalculate(F);
+      DT->recalculate(*DT->Parent);
     if (PDT)
-      PDT->recalculate(F);
+      PDT->recalculate(*PDT->Parent);
     return;
   }
 
   // There is little performance gain if we pend the recalculation under
   // Lazy UpdateStrategy so we recalculate available trees immediately.
 
   // Prevent forceFlushDeletedBB() from erasing DomTree or PostDomTree nodes.
   IsRecalculatingDomTree = IsRecalculatingPostDomTree = true;
 
   // Because all trees are going to be up-to-date after recalculation,
   // flush awaiting deleted BasicBlocks.
   forceFlushDeletedBB();
   if (DT)
-    DT->recalculate(F);
+    DT->recalculate(*DT->Parent);
   if (PDT)
-    PDT->recalculate(F);
+    PDT->recalculate(*PDT->Parent);
 
   // Resume forceFlushDeletedBB() to erase DomTree or PostDomTree nodes.
   IsRecalculatingDomTree = IsRecalculatingPostDomTree = false;
   PendDTUpdateIndex = PendPDTUpdateIndex = PendUpdates.size();
   dropOutOfDateUpdates();
 }
 
 bool DomTreeUpdater::hasPendingUpdates() const {

 
 // The DT and PDT require the nodes related to updates
 // are not deleted when update functions are called.
 // So BasicBlock deletions must be pended when the
 // UpdateStrategy is Lazy. When the UpdateStrategy is
 // Eager, the BasicBlock will be deleted immediately.
 void DomTreeUpdater::deleteBB(BasicBlock *DelBB) {
   validateDeleteBB(DelBB);
+  bool IsEntryBlock = DelBB == EntryBlock;
   if (Strategy == UpdateStrategy::Lazy) {
     DeletedBBs.insert(DelBB);
+    if (IsEntryBlock)
+      NeedRecalculate = true;
     return;
   }
 
   DelBB->removeFromParent();
   eraseDelBBNode(DelBB);
   delete DelBB;
+  if (IsEntryBlock)
+    internalrecalculate();
 }
 
 void DomTreeUpdater::callbackDeleteBB(
     BasicBlock *DelBB, std::function<void(BasicBlock *)> Callback) {
   validateDeleteBB(DelBB);
+  bool IsEntryBlock = DelBB == EntryBlock;
   if (Strategy == UpdateStrategy::Lazy) {
     Callbacks.push_back(CallBackOnDeletion(DelBB, Callback));
     DeletedBBs.insert(DelBB);
+    if (IsEntryBlock)
+      NeedRecalculate = true;
     return;
   }
 
   DelBB->removeFromParent();
   eraseDelBBNode(DelBB);
   Callback(DelBB);
   delete DelBB;
+  if (IsEntryBlock)
+    internalrecalculate();
 }
 
 void DomTreeUpdater::eraseDelBBNode(BasicBlock *DelBB) {
   if (DT && !IsRecalculatingDomTree)
     if (DT->getNode(DelBB))
       DT->eraseNode(DelBB);
 
   if (PDT && !IsRecalculatingPostDomTree)
void DomTreeUpdater::validateDeleteBB(BasicBlock *DelBB) {
   }
   // Make sure DelBB has a valid terminator instruction. As long as DelBB is a
   // Child of Function F it must contain valid IR.
   new UnreachableInst(DelBB->getContext(), DelBB);
 }
 
 void DomTreeUpdater::applyUpdates(ArrayRef<DominatorTree::UpdateType> Updates,
                                   bool ForceRemoveDuplicates) {
+  if (isEager() && checkEntryBlockChange())
+    return;
+
   if (!DT && !PDT)
     return;
 
   if (Strategy == UpdateStrategy::Lazy || ForceRemoveDuplicates) {
     SmallVector<DominatorTree::UpdateType, 8> Seen;
     for (const auto U : Updates)
       // For Lazy UpdateStrategy, avoid duplicates to applyLazyUpdate() to save
       // on analysis.
void DomTreeUpdater::applyUpdates(ArrayRef<DominatorTree::UpdateType> Updates,
   if (DT)
     DT->applyUpdates(Updates);
   if (PDT)
     PDT->applyUpdates(Updates);
 }
 
 DominatorTree &DomTreeUpdater::getDomTree() {
   assert(DT && "Invalid acquisition of a null DomTree");
+  if (isLazy() && NeedRecalculate) {
+    internalrecalculate();
+    return *DT;
+  }
+
+  if (checkEntryBlockChange())
+    return *DT;
+
   applyDomTreeUpdates();
   dropOutOfDateUpdates();
   return *DT;
 }
 
 PostDominatorTree &DomTreeUpdater::getPostDomTree() {
   assert(PDT && "Invalid acquisition of a null PostDomTree");
+  if (isLazy() && NeedRecalculate) {
+    internalrecalculate();
+    return *PDT;
+  }
+
+  if (checkEntryBlockChange())
+    return *PDT;
+
   applyPostDomTreeUpdates();
   dropOutOfDateUpdates();
   return *PDT;
 }
 
 void DomTreeUpdater::insertEdge(BasicBlock *From, BasicBlock *To) {
 
 #ifndef NDEBUG
   assert(isUpdateValid({DominatorTree::Insert, From, To}) &&
          "Inserted edge does not appear in the CFG");
 #endif
 
+  if (isEager() && checkEntryBlockChange())
+    return;
+
   if (!DT && !PDT)
     return;
 
   // Won't affect DomTree and PostDomTree; discard update.
   if (From == To)
     return;
 
   if (Strategy == UpdateStrategy::Eager) {
     if (DT)
       DT->insertEdge(From, To);
     if (PDT)
       PDT->insertEdge(From, To);
     return;
   }
 
   applyLazyUpdate(DominatorTree::Insert, From, To);
 }
 
 void DomTreeUpdater::insertEdgeRelaxed(BasicBlock *From, BasicBlock *To) {
+  if (isEager() && checkEntryBlockChange())
+    return;
+
   if (From == To)
     return;
 
   if (!DT && !PDT)
     return;
 
   if (!isUpdateValid({DominatorTree::Insert, From, To}))
     return;

 
 void DomTreeUpdater::deleteEdge(BasicBlock *From, BasicBlock *To) {
 
 #ifndef NDEBUG
   assert(isUpdateValid({DominatorTree::Delete, From, To}) &&
          "Deleted edge still exists in the CFG!");
 #endif
 
+  if (isEager() && checkEntryBlockChange())
+    return;
+
   if (!DT && !PDT)
     return;
 
   // Won't affect DomTree and PostDomTree; discard update.
   if (From == To)
     return;
 
   if (Strategy == UpdateStrategy::Eager) {
     if (DT)
       DT->deleteEdge(From, To);
     if (PDT)
       PDT->deleteEdge(From, To);
     return;
   }
 
   applyLazyUpdate(DominatorTree::Delete, From, To);
 }
 
 void DomTreeUpdater::deleteEdgeRelaxed(BasicBlock *From, BasicBlock *To) {
+  if (isEager() && checkEntryBlockChange())
+    return;
+
   if (From == To)
     return;
 
   if (!DT && !PDT)
     return;
 
   if (!isUpdateValid({DominatorTree::Delete, From, To}))
     return;

llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp

static bool eliminateRecursiveTailCall(
   // OK! We can transform this tail call.  If this is the first one found,
   // create the new entry block, allowing us to branch back to the old entry.
   if (!OldEntry) {
     OldEntry = &F->getEntryBlock();
     BasicBlock *NewEntry = BasicBlock::Create(F->getContext(), "", F, OldEntry);
     NewEntry->takeName(OldEntry);
     OldEntry->setName("tailrecurse");
     BranchInst *BI = BranchInst::Create(OldEntry, NewEntry);
+    DTU.insertEdge(NewEntry, OldEntry);
     BI->setDebugLoc(CI->getDebugLoc());
 
     // If this tail call is marked 'tail' and if there are any allocas in the
     // entry block, move them up to the new entry block.
     TailCallsAreMarkedTail = CI->isTailCall();
     if (TailCallsAreMarkedTail)
       // Move all fixed sized allocas from OldEntry to NewEntry.
       for (BasicBlock::iterator OEBI = OldEntry->begin(), E = OldEntry->end(),
if (!OldEntry) {
     for (Function::arg_iterator I = F->arg_begin(), E = F->arg_end();
          I != E; ++I) {
       PHINode *PN = PHINode::Create(I->getType(), 2,
                                     I->getName() + ".tr", InsertPos);
       I->replaceAllUsesWith(PN); // Everyone use the PHI node now!
       PN->addIncoming(&*I, NewEntry);
       ArgumentPHIs.push_back(PN);
     }
-    // The entry block was changed from OldEntry to NewEntry.
-    // The forward DominatorTree needs to be recalculated when the EntryBB is
-    // changed. In this corner-case we recalculate the entire tree.
-    DTU.recalculate(*NewEntry->getParent());
   }
 
   // If this function has self recursive calls in the tail position where some
   // are marked tail and some are not, only transform one flavor or another.  We
   // have to choose whether we move allocas in the entry block to the new entry
   // block or not, so we can't make a good choice for both.  NOTE: We could do
   // slightly better here in the case that the function has no entry block
   // allocas.

llvm/lib/Transforms/Utils/Local.cpp

void llvm::MergeBasicBlockIntoOnlyPred(BasicBlock *DestBB,
 
   if (DTU) {
     assert(PredBB->getInstList().size() == 1 &&
            isa<UnreachableInst>(PredBB->getTerminator()) &&
            "The successor list of PredBB isn't empty before "
            "applying corresponding DTU updates.");
     DTU->applyUpdates(Updates, /*ForceRemoveDuplicates*/ true);
     DTU->deleteBB(PredBB);
-    // Recalculation of DomTree is needed when updating a forward DomTree and
-    // the Entry BB is replaced.
-    if (ReplaceEntryBB && DTU->hasDomTree()) {
-      // The entry block was removed and there is no external interface for
-      // the dominator tree to be notified of this change. In this corner-case
-      // we recalculate the entire tree.
-      DTU->recalculate(*(DestBB->getParent()));
-    }
+    if (ReplaceEntryBB)
+      DTU->flush();
   }
 
   else {
     PredBB->eraseFromParent(); // Nuke BB if DTU is nullptr.
   }
 }
 
 /// CanMergeValues - Return true if we can choose one of these values to use

llvm/unittests/Analysis/DomTreeUpdaterTest.cpp

BasicBlock *NewEntry =
       BasicBlock::Create(F->getContext(), "new_entry", F, BB0);
   BranchInst::Create(BB0, NewEntry);
   EXPECT_EQ(F->begin()->getName(), NewEntry->getName());
   EXPECT_TRUE(&F->getEntryBlock() == NewEntry);
 
   DTU.insertEdgeRelaxed(NewEntry, BB0);
 
   // Changing the Entry BB requires a full recalculation of DomTree.
-  DTU.recalculate(*F);
+  DTU.flush();
   ASSERT_TRUE(DT.verify());
   ASSERT_TRUE(PDT.verify());
 
   // CFG Change: remove new_edge -> bb0 and redirect to new_edge -> bb1.
   EXPECT_EQ(NewEntry->getTerminator()->getNumSuccessors(), 1u);
   NewEntry->getTerminator()->eraseFromParent();
   BranchInst::Create(BB1, NewEntry);
   EXPECT_EQ(BB0->getTerminator()->getNumSuccessors(), 1u);
TEST(DomTreeUpdater, LazyUpdateReplaceEntryBB) {
   EXPECT_TRUE(&F->getEntryBlock() == NewEntry);
 
   // Insert the new edge between new_entry -> bb0. Without this the
   // recalculate() call below will not actually recalculate the DT as there
   // are no changes pending and no blocks deleted.
   DTU.insertEdge(NewEntry, BB0);
 
   // Changing the Entry BB requires a full recalculation.
-  DTU.recalculate(*F);
+  DTU.flush();
   ASSERT_TRUE(DTU.getDomTree().verify());
   ASSERT_TRUE(DTU.getPostDomTree().verify());
 
   // CFG Change: remove new_edge -> bb0 and redirect to new_edge -> bb1.
   EXPECT_EQ(NewEntry->getTerminator()->getNumSuccessors(), 1u);
   NewEntry->getTerminator()->eraseFromParent();
   BranchInst::Create(BB1, NewEntry);
   EXPECT_EQ(BB0->getTerminator()->getNumSuccessors(), 1u);