[nfc][mlgo] Incrementally update DominatorTreeAnalysis in FunctionPropertiesAnalysis (#104867)

We need the dominator tree analysis for loop info analysis, which we need to get features like most nested loop and number of top level loops. Invalidating and recomputing these from scratch after each successful inlining can sometimes lead to lengthy compile times. We don't need to recompute from scratch, though, since we have some boundary information about where the changes to the CFG happen; moreover, for dom tree, the API supports incrementally updating the analysis result.

This change addresses the dom tree part. The loop info is still recomputed from scratch. This does reduce the compile time quite significantly already, though (~5x in a specific case)

The loop info change might be more involved and would follow in a subsequent PR.

Author: mtrofin

llvm/include/llvm/Analysis/FunctionPropertiesAnalysis.h

@@ -15,6 +15,7 @@
 #define LLVM_ANALYSIS_FUNCTIONPROPERTIESANALYSIS_H
 
 #include "llvm/ADT/DenseSet.h"
+#include "llvm/IR/Dominators.h"
 #include "llvm/IR/PassManager.h"
 
 namespace llvm {
@@ -186,7 +187,12 @@ class FunctionPropertiesUpdater {
   static bool isUpdateValid(Function &F, const FunctionPropertiesInfo &FPI,
                             FunctionAnalysisManager &FAM);
 
+  DominatorTree &getUpdatedDominatorTree(FunctionAnalysisManager &FAM) const;
+
   DenseSet<const BasicBlock *> Successors;
+
+  // Edges we might potentially need to remove from the dominator tree.
+  SmallVector<DominatorTree::UpdateType, 2> DomTreeUpdates;
 };
 } // namespace llvm
 #endif // LLVM_ANALYSIS_FUNCTIONPROPERTIESANALYSIS_H

llvm/lib/Analysis/FunctionPropertiesAnalysis.cpp

@@ -326,6 +326,14 @@ FunctionPropertiesUpdater::FunctionPropertiesUpdater(
   // with the CB BB ('Entry') between which the inlined callee will be pasted.
   Successors.insert(succ_begin(&CallSiteBB), succ_end(&CallSiteBB));
 
+  // the outcome of the inlining may be that some edges get lost (DCEd BBs
+  // because inlining brought some constant, for example). We don't know which
+  // edges will be removed, so we list all of them as potentially removable.
+  for (auto *Succ : successors(&CallSiteBB))
+    DomTreeUpdates.emplace_back(DominatorTree::UpdateKind::Delete,
+                                const_cast<BasicBlock *>(&CallSiteBB),
+                                const_cast<BasicBlock *>(Succ));
+
   // Inlining only handles invoke and calls. If this is an invoke, and inlining
   // it pulls another invoke, the original landing pad may get split, so as to
   // share its content with other potential users. So the edge up to which we
@@ -336,6 +344,11 @@ FunctionPropertiesUpdater::FunctionPropertiesUpdater(
   if (const auto *II = dyn_cast<InvokeInst>(&CB)) {
     const auto *UnwindDest = II->getUnwindDest();
     Successors.insert(succ_begin(UnwindDest), succ_end(UnwindDest));
+    // Same idea as above, we pretend we lose all these edges.
+    for (auto *Succ : successors(UnwindDest))
+      DomTreeUpdates.emplace_back(DominatorTree::UpdateKind::Delete,
+                                  const_cast<BasicBlock *>(UnwindDest),
+                                  const_cast<BasicBlock *>(Succ));
   }
 
   // Exclude the CallSiteBB, if it happens to be its own successor (1-BB loop).
@@ -356,6 +369,45 @@ FunctionPropertiesUpdater::FunctionPropertiesUpdater(
     FPI.updateForBB(*BB, -1);
 }
 
+DominatorTree &FunctionPropertiesUpdater::getUpdatedDominatorTree(
+    FunctionAnalysisManager &FAM) const {
+  auto &DT =
+      FAM.getResult<DominatorTreeAnalysis>(const_cast<Function &>(Caller));
+
+  SetVector<const BasicBlock *> NewSucc;
+  NewSucc.insert(succ_begin(&CallSiteBB), succ_end(&CallSiteBB));
+
+  // tell the DomTree about the new edges
+  std::deque<const BasicBlock *> Worklist;
+  Worklist.push_back(&CallSiteBB);
+
+  // Build the list of edges to actually remove. Those are those edges in the
+  // DomTreeUpdates that cannot be found in the CFG anymore.
+  SmallVector<DominatorTree::UpdateType, 2> FinalDomTreeUpdates;
+  while (!Worklist.empty()) {
+    auto *BB = Worklist.front();
+    Worklist.pop_front();
+    assert(DT.getNode(BB));
+
+    for (auto *Succ : NewSucc) {
+      if (!DT.getNode(Succ))
+        Worklist.push_back(Succ);
+      FinalDomTreeUpdates.push_back({DominatorTree::UpdateKind::Insert,
+                                     const_cast<BasicBlock *>(BB),
+                                     const_cast<BasicBlock *>(Succ)});
+    }
+  }
+  for (auto &Upd : DomTreeUpdates)
+    if (!llvm::is_contained(successors(Upd.getFrom()), Upd.getTo()))
+      FinalDomTreeUpdates.push_back(Upd);
+
+  DT.applyUpdates(FinalDomTreeUpdates);
+#ifdef EXPENSIVE_CHECKS
+  assert(DT.verify(DominatorTree::VerificationLevel::Full));
+#endif
+  return DT;
+}
+
 void FunctionPropertiesUpdater::finish(FunctionAnalysisManager &FAM) const {
   // Update feature values from the BBs that were copied from the callee, or
   // might have been modified because of inlining. The latter have been
@@ -383,8 +435,7 @@ void FunctionPropertiesUpdater::finish(FunctionAnalysisManager &FAM) const {
   // remove E.
   SetVector<const BasicBlock *> Reinclude;
   SetVector<const BasicBlock *> Unreachable;
-  const auto &DT =
-      FAM.getResult<DominatorTreeAnalysis>(const_cast<Function &>(Caller));
+  auto &DT = getUpdatedDominatorTree(FAM);
 
   if (&CallSiteBB != &*Caller.begin())
     Reinclude.insert(&*Caller.begin());
@@ -427,6 +478,9 @@ void FunctionPropertiesUpdater::finish(FunctionAnalysisManager &FAM) const {
 
   const auto &LI = FAM.getResult<LoopAnalysis>(const_cast<Function &>(Caller));
   FPI.updateAggregateStats(Caller, LI);
+#ifdef EXPENSIVE_CHECKS
+  assert(isUpdateValid(Caller, FPI, FAM));
+#endif
 }
 
 bool FunctionPropertiesUpdater::isUpdateValid(Function &F,

llvm/lib/Analysis/MLInlineAdvisor.cpp

@@ -288,7 +288,6 @@ void MLInlineAdvisor::onSuccessfulInlining(const MLInlineAdvice &Advice,
   {
     PreservedAnalyses PA = PreservedAnalyses::all();
     PA.abandon<FunctionPropertiesAnalysis>();
-    PA.abandon<DominatorTreeAnalysis>();
     PA.abandon<LoopAnalysis>();
     FAM.invalidate(*Caller, PA);
   }