[DomTreeUpdater] Move critical edge splitting code to updater

Author: paperchalice

llvm/include/llvm/Analysis/GenericDomTreeUpdater.h

@@ -105,6 +105,12 @@ class GenericDomTreeUpdater {
     return PendUpdates.size() != PendPDTUpdateIndex;
   }
 
+  bool hasPendingCriticalEdges() const {
+    if (!PDT && !DT)
+      return false;
+    return !CriticalEdgesToSplit.empty();
+  }
+
   ///@{
   /// \name Mutation APIs
   ///
@@ -146,8 +152,25 @@ class GenericDomTreeUpdater {
   /// 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.
+  /// 3. This kind updates are incompatible with critical edge splitting
+  /// updates, call this method will apply all critical edge updates in
+  /// lazy mode. It is not recommended to interleave applyUpdates and
+  /// applyUpdatesForCriticalEdgeSplitting.
   void applyUpdates(ArrayRef<typename DomTreeT::UpdateType> Updates);
 
+  /// Apply updates that the critical edge (FromBB, ToBB) has been
+  /// split with NewBB.
+  ///
+  /// \note Do not use this method with regular edges.
+  ///
+  /// \note This kind updates are incompatible with generic updates,
+  /// call this method will submit all generic updates in lazy mode.
+  /// It is not recommended to interleave applyUpdates and
+  /// applyUpdatesForCriticalEdgeSplitting.
+  void applyUpdatesForCriticalEdgeSplitting(BasicBlockT *FromBB,
+                                            BasicBlockT *ToBB,
+                                            BasicBlockT *NewBB);
+
   /// Submit updates to all available trees. It will also
   /// 1. discard duplicated updates,
   /// 2. remove invalid updates. (Invalid updates means deletion of an edge that
@@ -197,6 +220,7 @@ class GenericDomTreeUpdater {
     applyDomTreeUpdates();
     applyPostDomTreeUpdates();
     dropOutOfDateUpdates();
+    applySplitCriticalEdges();
   }
 
   ///@}
@@ -243,6 +267,35 @@ class GenericDomTreeUpdater {
   /// Drop all updates applied by all available trees and delete BasicBlocks if
   /// all available trees are up-to-date.
   void dropOutOfDateUpdates();
+
+private:
+  /// Helper structure used to hold all the basic blocks
+  /// involved in the split of a critical edge.
+  struct CriticalEdge {
+    BasicBlockT *FromBB;
+    BasicBlockT *ToBB;
+    BasicBlockT *NewBB;
+  };
+
+  /// Pile up all the critical edges to be split.
+  /// The splitting of a critical edge is local and thus, it is possible
+  /// to apply several of those changes at the same time.
+  SmallVector<CriticalEdge, 32> CriticalEdgesToSplit;
+
+  /// Remember all the basic blocks that are inserted during
+  /// edge splitting.
+  /// Invariant: NewBBs == all the basic blocks contained in the NewBB
+  /// field of all the elements of CriticalEdgesToSplit.
+  /// I.e., forall elt in CriticalEdgesToSplit, it exists BB in NewBBs
+  /// such as BB == elt.NewBB.
+  SmallSet<BasicBlockT *, 32> NewBBs;
+
+  /// 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.
+  ///
+  /// \post CriticalEdgesToSplit.empty().
+  void applySplitCriticalEdges();
 };
 
 } // namespace llvm

llvm/include/llvm/Analysis/GenericDomTreeUpdaterImpl.h

@@ -16,6 +16,7 @@
 #ifndef LLVM_ANALYSIS_GENERICDOMTREEUPDATERIMPL_H
 #define LLVM_ANALYSIS_GENERICDOMTREEUPDATERIMPL_H
 
+#include "llvm/ADT/SmallBitVector.h"
 #include "llvm/Analysis/GenericDomTreeUpdater.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
@@ -61,6 +62,9 @@ void GenericDomTreeUpdater<DerivedT, DomTreeT, PostDomTreeT>::applyUpdates(
     return;
 
   if (Strategy == UpdateStrategy::Lazy) {
+    if (!CriticalEdgesToSplit.empty())
+      applySplitCriticalEdges();
+
     PendUpdates.reserve(PendUpdates.size() + Updates.size());
     for (const auto &U : Updates)
       if (!isSelfDominance(U))
@@ -348,6 +352,123 @@ void GenericDomTreeUpdater<DerivedT, DomTreeT,
   PendPDTUpdateIndex -= dropIndex;
 }
 
+template <typename DerivedT, typename DomTreeT, typename PostDomTreeT>
+void GenericDomTreeUpdater<DerivedT, DomTreeT, PostDomTreeT>::
+    applyUpdatesForCriticalEdgeSplitting(BasicBlockT *FromBB, BasicBlockT *ToBB,
+                                         BasicBlockT *NewBB) {
+  if (!DT && !PDT)
+    return;
+  CriticalEdgesToSplit.push_back({FromBB, ToBB, NewBB});
+  bool Inserted = NewBBs.insert(NewBB).second;
+  (void)Inserted;
+  assert(Inserted &&
+         "A basic block inserted via edge splitting cannot appear twice");
+  if (Strategy == UpdateStrategy::Lazy) {
+    applyDomTreeUpdates();
+    applyPostDomTreeUpdates();
+  }
+  if (Strategy == UpdateStrategy::Eager)
+    applySplitCriticalEdges();
+}
+
+template <typename DerivedT, typename DomTreeT, typename PostDomTreeT>
+void GenericDomTreeUpdater<DerivedT, DomTreeT,
+                           PostDomTreeT>::applySplitCriticalEdges() {
+  // Bail out early if there is nothing to do.
+  if (CriticalEdgesToSplit.empty())
+    return;
+
+  // For each element in CriticalEdgesToSplit, remember whether or not element
+  // is the new immediate domminator of its successor. The mapping is done by
+  // index, i.e., the information for the ith element of CriticalEdgesToSplit is
+  // the ith element of IsNewIDom.
+  SmallBitVector IsNewIDom(CriticalEdgesToSplit.size(), true);
+  SmallBitVector IsNewIPDom(CriticalEdgesToSplit.size(), true);
+  size_t Idx = 0;
+
+  // Collect all the dominance properties info, before invalidating
+  // the underlying DT.
+  for (CriticalEdge &Edge : CriticalEdgesToSplit) {
+    // Update dominator information.
+    if (DT) {
+      BasicBlockT *Succ = Edge.ToBB;
+      auto *SuccDTNode = DT->getNode(Succ);
+
+      for (BasicBlockT *PredBB : predecessors(Succ)) {
+        if (PredBB == Edge.NewBB)
+          continue;
+        // If we are in this situation:
+        // FromBB1        FromBB2
+        //    +              +
+        //   + +            + +
+        //  +   +          +   +
+        // ...  Split1  Split2 ...
+        //           +   +
+        //            + +
+        //             +
+        //            Succ
+        // Instead of checking the domiance property with Split2, we check it
+        // with FromBB2 since Split2 is still unknown of the underlying DT
+        // structure.
+        if (NewBBs.count(PredBB)) {
+          assert(pred_size(PredBB) == 1 && "A basic block resulting from a "
+                                           "critical edge split has more "
+                                           "than one predecessor!");
+          PredBB = *pred_begin(PredBB);
+        }
+        if (!DT->dominates(SuccDTNode, DT->getNode(PredBB))) {
+          IsNewIDom[Idx] = false;
+          break;
+        }
+      }
+    }
+
+    // Same as DT version but from another direction.
+    if (PDT) {
+      BasicBlockT *Pred = Edge.FromBB;
+      auto *PredDTNode = PDT->getNode(Pred);
+      for (BasicBlockT *SuccBB : successors(Pred)) {
+        if (SuccBB == Edge.NewBB)
+          continue;
+        if (NewBBs.count(SuccBB)) {
+          assert(succ_size(SuccBB) == 1 && "A basic block resulting from a "
+                                           "critical edge split has more "
+                                           "than one predecessor!");
+          SuccBB = *succ_begin(SuccBB);
+        }
+        if (!PDT->dominates(PredDTNode, PDT->getNode(SuccBB))) {
+          IsNewIPDom[Idx] = false;
+          break;
+        }
+      }
+    }
+    ++Idx;
+  }
+
+  // Now, update DT with the collected dominance properties info.
+  Idx = 0;
+  for (CriticalEdge &Edge : CriticalEdgesToSplit) {
+    if (DT) {
+      // We know FromBB dominates NewBB.
+      auto *NewDTNode = DT->addNewBlock(Edge.NewBB, Edge.FromBB);
+
+      // If all the other predecessors of "Succ" are dominated by "Succ" itself
+      // then the new block is the new immediate dominator of "Succ". Otherwise,
+      // the new block doesn't dominate anything.
+      if (IsNewIDom[Idx])
+        DT->changeImmediateDominator(DT->getNode(Edge.ToBB), NewDTNode);
+    }
+    if (PDT) {
+      auto *NewPDTNode = PDT->addNewBlock(Edge.NewBB, Edge.ToBB);
+      if (IsNewIPDom[Idx])
+        PDT->changeImmediateDominator(PDT->getNode(Edge.FromBB), NewPDTNode);
+    }
+    ++Idx;
+  }
+  NewBBs.clear();
+  CriticalEdgesToSplit.clear();
+}
+
 } // namespace llvm
 
 #endif // LLVM_ANALYSIS_GENERICDOMTREEUPDATERIMPL_H

llvm/include/llvm/CodeGen/MachineBasicBlock.h

@@ -32,6 +32,7 @@
 namespace llvm {
 
 class BasicBlock;
+class MachineDomTreeUpdater;
 class MachineFunction;
 class MCSymbol;
 class ModuleSlotTracker;
@@ -970,15 +971,17 @@ class MachineBasicBlock
   /// MachineLoopInfo, as applicable.
   MachineBasicBlock *
   SplitCriticalEdge(MachineBasicBlock *Succ, Pass &P,
-                    std::vector<SparseBitVector<>> *LiveInSets = nullptr) {
-    return SplitCriticalEdge(Succ, &P, nullptr, LiveInSets);
+                    std::vector<SparseBitVector<>> *LiveInSets = nullptr,
+                    MachineDomTreeUpdater *MDTU = nullptr) {
+    return SplitCriticalEdge(Succ, &P, nullptr, LiveInSets, MDTU);
   }
 
   MachineBasicBlock *
   SplitCriticalEdge(MachineBasicBlock *Succ,
                     MachineFunctionAnalysisManager &MFAM,
-                    std::vector<SparseBitVector<>> *LiveInSets = nullptr) {
-    return SplitCriticalEdge(Succ, nullptr, &MFAM, LiveInSets);
+                    std::vector<SparseBitVector<>> *LiveInSets = nullptr,
+                    MachineDomTreeUpdater *MDTU = nullptr) {
+    return SplitCriticalEdge(Succ, nullptr, &MFAM, LiveInSets, MDTU);
   }
 
   /// Check if the edge between this block and the given successor \p
@@ -1256,10 +1259,9 @@ class MachineBasicBlock
   void removePredecessor(MachineBasicBlock *Pred);
 
   // Helper method for new pass manager migration.
-  MachineBasicBlock *
-  SplitCriticalEdge(MachineBasicBlock *Succ, Pass *P,
-                    MachineFunctionAnalysisManager *MFAM,
-                    std::vector<SparseBitVector<>> *LiveInSets);
+  MachineBasicBlock *SplitCriticalEdge(
+      MachineBasicBlock *Succ, Pass *P, MachineFunctionAnalysisManager *MFAM,
+      std::vector<SparseBitVector<>> *LiveInSets, MachineDomTreeUpdater *MDTU);
 };
 
 raw_ostream& operator<<(raw_ostream &OS, const MachineBasicBlock &MBB);
@@ -1341,6 +1343,12 @@ inline auto successors(const MachineBasicBlock *BB) { return BB->successors(); }
 inline auto predecessors(const MachineBasicBlock *BB) {
   return BB->predecessors();
 }
+inline auto succ_size(const MachineBasicBlock *BB) { return BB->succ_size(); }
+inline auto pred_size(const MachineBasicBlock *BB) { return BB->pred_size(); }
+inline auto succ_begin(const MachineBasicBlock *BB) { return BB->succ_begin(); }
+inline auto pred_begin(const MachineBasicBlock *BB) { return BB->pred_begin(); }
+inline auto succ_end(const MachineBasicBlock *BB) { return BB->succ_end(); }
+inline auto pred_end(const MachineBasicBlock *BB) { return BB->pred_end(); }
 
 /// MachineInstrSpan provides an interface to get an iteration range
 /// containing the instruction it was initialized with, along with all