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

Reverts c992690.

The problem is that if there is a sequence "{delete A->B} {delete A->B}
{insert A->B}" the net result is "{delete A->B}", which is not what we
want.

Duplicate successors may happen in cases like switch statements (as
shown in the unit test).

The second problem was that in `invoke` cases, some edges we speculate may get deleted don't, but are also not reachable from the inlined call site's basic block. We just need to check which edges are actually not present anymore.

The fix is to sanitize the list of deletes, just like we do for inserts.

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,20 @@ 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.
+  // Some BBs have (at this point) duplicate edges. Remove duplicates, otherwise
+  // the DT updater will not apply changes correctly.
+  DenseSet<const BasicBlock *> Inserted;
+  for (auto *Succ : successors(&CallSiteBB))
+    if (Inserted.insert(Succ).second)
+      DomTreeUpdates.emplace_back(DominatorTree::UpdateKind::Delete,
+                                  const_cast<BasicBlock *>(&CallSiteBB),
+                                  const_cast<BasicBlock *>(Succ));
+  // Reuse Inserted (which has some allocated capacity at this point) below, if
+  // we have an invoke.
+  Inserted.clear();
   // 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 +350,12 @@ 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))
+      if (Inserted.insert(Succ).second)
+        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 +376,33 @@ FunctionPropertiesUpdater::FunctionPropertiesUpdater(
     FPI.updateForBB(*BB, -1);
 }
 
+DominatorTree &FunctionPropertiesUpdater::getUpdatedDominatorTree(
+    FunctionAnalysisManager &FAM) const {
+  auto &DT =
+      FAM.getResult<DominatorTreeAnalysis>(const_cast<Function &>(Caller));
+
+  SmallVector<DominatorTree::UpdateType, 2> FinalDomTreeUpdates;
+
+  DenseSet<const BasicBlock *> Inserted;
+  for (auto *Succ : successors(&CallSiteBB))
+    if (Inserted.insert(Succ).second)
+      FinalDomTreeUpdates.push_back({DominatorTree::UpdateKind::Insert,
+                                     const_cast<BasicBlock *>(&CallSiteBB),
+                                     const_cast<BasicBlock *>(Succ)});
+
+  // Perform the deletes last, so that any new nodes connected to nodes
+  // participating in the edge deletion are known to the DT.
+  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 +430,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,11 +473,17 @@ 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,
                                               const FunctionPropertiesInfo &FPI,
                                               FunctionAnalysisManager &FAM) {
+  if (!FAM.getResult<DominatorTreeAnalysis>(F).verify(
+          DominatorTree::VerificationLevel::Full))
+    return false;
   DominatorTree DT(F);
   LoopInfo LI(DT);
   auto Fresh = FunctionPropertiesInfo::getFunctionPropertiesInfo(F, DT, LI);

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);
   }

llvm/unittests/Analysis/FunctionPropertiesAnalysisTest.cpp

@@ -999,4 +999,119 @@ declare <4 x ptr> @f4()
   EXPECT_EQ(DetailedF1Properties.CallReturnsVectorPointerCount, 1);
   EnableDetailedFunctionProperties.setValue(false);
 }
+
+TEST_F(FunctionPropertiesAnalysisTest, ReAddEdges) {
+  LLVMContext C;
+  std::unique_ptr<Module> M = makeLLVMModule(C, R"IR(
+define hidden void @f1(ptr noundef %destatep, i32 noundef %offset, i8 noundef zeroext %byte1) {
+entry:
+  %cmp = icmp eq i8 %byte1, 0
+  br i1 %cmp, label %if.then, label %if.else
+
+if.then:                                          ; preds = %entry
+  call fastcc void @f2(ptr noundef %destatep, i32 noundef 37, i32 noundef 600)
+  %and = and i32 %offset, 3
+  switch i32 %and, label %default.unreachable [
+    i32 0, label %sw.bb
+    i32 1, label %sw.bb1
+    i32 2, label %sw.bb1
+    i32 3, label %if.end
+  ]
+
+sw.bb:                                            ; preds = %if.then
+  call fastcc void @f2(ptr noundef %destatep, i32 noundef 57, i32 noundef 600)
+  br label %if.end
+
+sw.bb1:                                           ; preds = %if.then, %if.then
+  call fastcc void @f2(ptr noundef %destatep, i32 noundef 56, i32 noundef 600) #34
+  br label %if.end
+
+default.unreachable:                              ; preds = %if.then
+  unreachable
+
+if.else:                                          ; preds = %entry
+  call fastcc void @f2(ptr noundef %destatep, i32 noundef 56, i32 noundef 600)
+  br label %if.end
+
+if.end:                                           ; preds = %sw.bb, %sw.bb1, %if.then, %if.else
+  ret void
+}
+
+define internal fastcc void @f2(ptr nocapture noundef %destatep, i32 noundef %r_enc, i32 noundef %whack) {
+entry:
+  %enc_prob = getelementptr inbounds nuw i8, ptr %destatep, i32 512
+  %arrayidx = getelementptr inbounds [67 x i32], ptr %enc_prob, i32 0, i32 %r_enc
+  %0 = load i32, ptr %arrayidx, align 4
+  %sub = sub nsw i32 %0, %whack
+  store i32 %sub, ptr %arrayidx, align 4
+  ret void
+}
+  )IR");
+  auto *F1 = M->getFunction("f1");
+  auto *F2 = M->getFunction("f2");
+  auto *CB = [&]() -> CallBase * {
+    for (auto &BB : *F1)
+      for (auto &I : BB)
+        if (auto *CB = dyn_cast<CallBase>(&I);
+            CB && CB->getCalledFunction() && CB->getCalledFunction() == F2)
+          return CB;
+    return nullptr;
+  }();
+  ASSERT_NE(CB, nullptr);
+  auto FPI = buildFPI(*F1);
+  FunctionPropertiesUpdater FPU(FPI, *CB);
+  InlineFunctionInfo IFI;
+  auto IR = llvm::InlineFunction(*CB, IFI);
+  EXPECT_TRUE(IR.isSuccess());
+  invalidate(*F1);
+  EXPECT_TRUE(FPU.finishAndTest(FAM));
+}
+
+TEST_F(FunctionPropertiesAnalysisTest, InvokeLandingCanStillBeReached) {
+  LLVMContext C;
+  // %lpad is reachable from a block not involved in the inlining decision. We
+  // make sure that's not the entry - otherwise the DT will be recomputed from
+  // scratch. The idea here is that the edge known to the inliner to potentially
+  // disappear - %lpad->%ehcleanup -should survive because it is still reachable
+  // from %middle.
+  std::unique_ptr<Module> M = makeLLVMModule(C,
+                                             R"IR(
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-pc-linux-gnu"
+
+define i64 @f1(i32 noundef %value) {
+entry:
+  br label %middle
+middle:
+  %c = icmp eq i32 %value, 0
+  br i1 %c, label %invoke, label %lpad
+invoke:
+  invoke fastcc void @f2() to label %cont unwind label %lpad
+cont:
+  br label %exit
+lpad:
+  %lp = landingpad i32 cleanup
+  br label %ehcleanup
+ehcleanup:
+  resume i32 0
+exit:
+  ret i64 1
+}
+define void @f2() {
+  ret void
+}
+)IR");
+
+  Function *F1 = M->getFunction("f1");
+  CallBase *CB = findCall(*F1);
+  EXPECT_NE(CB, nullptr);
+
+  auto FPI = buildFPI(*F1);
+  FunctionPropertiesUpdater FPU(FPI, *CB);
+  InlineFunctionInfo IFI;
+  auto IR = llvm::InlineFunction(*CB, IFI);
+  EXPECT_TRUE(IR.isSuccess());
+  invalidate(*F1);
+  EXPECT_TRUE(FPU.finishAndTest(FAM));
+}
 } // end anonymous namespace