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

[Mircea Trofin via llvm-commits]

https://github.com/mtrofin updated https://github.com/llvm/llvm-project/pull/104867

>From c5e48e05356a6532d2e8ec5f8b344ef81ed2fc07 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Mon, 19 Aug 2024 14:37:23 -0700
Subject: [PATCH] [nfc] Incrementally update DominatorTreeAnalysis in
 FunctionPropertiesAnalysis

---
 .../Analysis/FunctionPropertiesAnalysis.h     |  6 ++
 .../Analysis/FunctionPropertiesAnalysis.cpp   | 58 ++++++++++++++++++-
 llvm/lib/Analysis/MLInlineAdvisor.cpp         |  1 -
 3 files changed, 62 insertions(+), 3 deletions(-)

diff --git a/llvm/include/llvm/Analysis/FunctionPropertiesAnalysis.h b/llvm/include/llvm/Analysis/FunctionPropertiesAnalysis.h
index ee447d3e4ebb6a..af72f6e0f90b11 100644
--- a/llvm/include/llvm/Analysis/FunctionPropertiesAnalysis.h
+++ b/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
diff --git a/llvm/lib/Analysis/FunctionPropertiesAnalysis.cpp b/llvm/lib/Analysis/FunctionPropertiesAnalysis.cpp
index 6d6ec6c7b1cc76..479cfc58ab38f5 100644
--- a/llvm/lib/Analysis/FunctionPropertiesAnalysis.cpp
+++ b/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,
diff --git a/llvm/lib/Analysis/MLInlineAdvisor.cpp b/llvm/lib/Analysis/MLInlineAdvisor.cpp
index b59aa4810005bc..8bb5efcf1b2ecb 100644
--- a/llvm/lib/Analysis/MLInlineAdvisor.cpp
+++ b/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);
   }