[mlgo] Incrementally update FunctionPropertiesInfo during inlining

Re-computing FunctionPropertiesInfo after each inlining may be very time
consuming: in certain cases, e.g. large caller with lots of callsites,
and when the overall IR doesn't increase (thus not tripping a size bloat
threshold).

This patch addresses this by incrementally updating
FunctionPropertiesInfo.

Differential Revision: https://reviews.llvm.org/D125841

Author: mtrofin

llvm/include/llvm/Analysis/FunctionPropertiesAnalysis.h

@@ -14,17 +14,33 @@
 #ifndef LLVM_ANALYSIS_FUNCTIONPROPERTIESANALYSIS_H
 #define LLVM_ANALYSIS_FUNCTIONPROPERTIESANALYSIS_H
 
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/iterator_range.h"
+#include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/PassManager.h"
 
 namespace llvm {
 class Function;
 class LoopInfo;
 
 class FunctionPropertiesInfo {
+  friend class FunctionPropertiesUpdater;
+  void updateForBB(const BasicBlock &BB, int64_t Direction);
+  void updateAggregateStats(const Function &F, const LoopInfo &LI);
+  void reIncludeBB(const BasicBlock &BB, const LoopInfo &LI);
+
 public:
   static FunctionPropertiesInfo getFunctionPropertiesInfo(const Function &F,
                                                           const LoopInfo &LI);
 
+  bool operator==(const FunctionPropertiesInfo &FPI) const {
+    return std::memcmp(this, &FPI, sizeof(FunctionPropertiesInfo)) == 0;
+  }
+
+  bool operator!=(const FunctionPropertiesInfo &FPI) const {
+    return !(*this == FPI);
+  }
+
   void print(raw_ostream &OS) const;
 
   /// Number of basic blocks
@@ -57,6 +73,9 @@ class FunctionPropertiesInfo {
 
   // Number of Top Level Loops in the Function
   int64_t TopLevelLoopCount = 0;
+
+  // All non-debug instructions
+  int64_t TotalInstructionCount = 0;
 };
 
 // Analysis pass
@@ -82,5 +101,24 @@ class FunctionPropertiesPrinterPass
   PreservedAnalyses run(Function &F, FunctionAnalysisManager &AM);
 };
 
+/// Correctly update FunctionPropertiesInfo post-inlining. A
+/// FunctionPropertiesUpdater keeps the state necessary for tracking the changes
+/// llvm::InlineFunction makes. The idea is that inlining will at most modify
+/// a few BBs of the Caller (maybe the entry BB and definitely the callsite BB)
+/// and potentially affect exception handling BBs in the case of invoke
+/// inlining.
+class FunctionPropertiesUpdater {
+public:
+  FunctionPropertiesUpdater(FunctionPropertiesInfo &FPI, const CallBase &CB);
+
+  void finish(const LoopInfo &LI);
+
+private:
+  FunctionPropertiesInfo &FPI;
+  const BasicBlock &CallSiteBB;
+  const Function &Caller;
+
+  DenseSet<const BasicBlock *> Successors;
+};
 } // namespace llvm
 #endif // LLVM_ANALYSIS_FUNCTIONPROPERTIESANALYSIS_H

llvm/include/llvm/Analysis/MLInlineAdvisor.h

@@ -9,6 +9,7 @@
 #ifndef LLVM_ANALYSIS_MLINLINEADVISOR_H
 #define LLVM_ANALYSIS_MLINLINEADVISOR_H
 
+#include "llvm/Analysis/FunctionPropertiesAnalysis.h"
 #include "llvm/Analysis/InlineAdvisor.h"
 #include "llvm/Analysis/LazyCallGraph.h"
 #include "llvm/Analysis/MLModelRunner.h"
@@ -33,13 +34,17 @@ class MLInlineAdvisor : public InlineAdvisor {
   void onPassEntry() override;
   void onPassExit(LazyCallGraph::SCC *SCC) override;
 
-  int64_t getIRSize(const Function &F) const { return F.getInstructionCount(); }
+  int64_t getIRSize(Function &F) const {
+    return getCachedFPI(F).TotalInstructionCount;
+  }
   void onSuccessfulInlining(const MLInlineAdvice &Advice,
                             bool CalleeWasDeleted);
 
   bool isForcedToStop() const { return ForceStop; }
   int64_t getLocalCalls(Function &F);
   const MLModelRunner &getModelRunner() const { return *ModelRunner.get(); }
+  FunctionPropertiesInfo &getCachedFPI(Function &) const;
+  const LoopInfo &getLoopInfo(Function &F) const;
 
 protected:
   std::unique_ptr<InlineAdvice> getAdviceImpl(CallBase &CB) override;
@@ -67,6 +72,8 @@ class MLInlineAdvisor : public InlineAdvisor {
        << "\n";
   }
 
+  mutable DenseMap<const Function *, FunctionPropertiesInfo> FPICache;
+
   LazyCallGraph &CG;
 
   int64_t NodeCount = 0;
@@ -86,16 +93,7 @@ class MLInlineAdvisor : public InlineAdvisor {
 class MLInlineAdvice : public InlineAdvice {
 public:
   MLInlineAdvice(MLInlineAdvisor *Advisor, CallBase &CB,
-                 OptimizationRemarkEmitter &ORE, bool Recommendation)
-      : InlineAdvice(Advisor, CB, ORE, Recommendation),
-        CallerIRSize(Advisor->isForcedToStop() ? 0
-                                               : Advisor->getIRSize(*Caller)),
-        CalleeIRSize(Advisor->isForcedToStop() ? 0
-                                               : Advisor->getIRSize(*Callee)),
-        CallerAndCalleeEdges(Advisor->isForcedToStop()
-                                 ? 0
-                                 : (Advisor->getLocalCalls(*Caller) +
-                                    Advisor->getLocalCalls(*Callee))) {}
+                 OptimizationRemarkEmitter &ORE, bool Recommendation);
   virtual ~MLInlineAdvice() = default;
 
   void recordInliningImpl() override;
@@ -112,10 +110,14 @@ class MLInlineAdvice : public InlineAdvice {
 
 private:
   void reportContextForRemark(DiagnosticInfoOptimizationBase &OR);
-
+  void updateCachedCallerFPI();
   MLInlineAdvisor *getAdvisor() const {
     return static_cast<MLInlineAdvisor *>(Advisor);
   };
+  // Make a copy of the FPI of the caller right before inlining. If inlining
+  // fails, we can just update the cache with that value.
+  const FunctionPropertiesInfo PreInlineCallerFPI;
+  Optional<FunctionPropertiesUpdater> FPU;
 };
 
 } // namespace llvm

llvm/lib/Analysis/FunctionPropertiesAnalysis.cpp

@@ -12,49 +12,75 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/FunctionPropertiesAnalysis.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/Analysis/LoopInfo.h"
+#include "llvm/IR/CFG.h"
 #include "llvm/IR/Instructions.h"
+#include <deque>
 
 using namespace llvm;
 
-FunctionPropertiesInfo
-FunctionPropertiesInfo::getFunctionPropertiesInfo(const Function &F,
-                                                  const LoopInfo &LI) {
-
-  FunctionPropertiesInfo FPI;
+namespace {
+int64_t getNrBlocksFromCond(const BasicBlock &BB) {
+  int64_t Ret = 0;
+  if (const auto *BI = dyn_cast<BranchInst>(BB.getTerminator())) {
+    if (BI->isConditional())
+      Ret += BI->getNumSuccessors();
+  } else if (const auto *SI = dyn_cast<SwitchInst>(BB.getTerminator())) {
+    Ret += (SI->getNumCases() + (nullptr != SI->getDefaultDest()));
+  }
+  return Ret;
+}
 
-  FPI.Uses = ((!F.hasLocalLinkage()) ? 1 : 0) + F.getNumUses();
+int64_t getUses(const Function &F) {
+  return ((!F.hasLocalLinkage()) ? 1 : 0) + F.getNumUses();
+}
+} // namespace
 
-  for (const auto &BB : F) {
-    ++FPI.BasicBlockCount;
+void FunctionPropertiesInfo::reIncludeBB(const BasicBlock &BB,
+                                         const LoopInfo &LI) {
+  updateForBB(BB, +1);
+  MaxLoopDepth =
+      std::max(MaxLoopDepth, static_cast<int64_t>(LI.getLoopDepth(&BB)));
+}
 
-    if (const auto *BI = dyn_cast<BranchInst>(BB.getTerminator())) {
-      if (BI->isConditional())
-        FPI.BlocksReachedFromConditionalInstruction += BI->getNumSuccessors();
-    } else if (const auto *SI = dyn_cast<SwitchInst>(BB.getTerminator())) {
-      FPI.BlocksReachedFromConditionalInstruction +=
-          (SI->getNumCases() + (nullptr != SI->getDefaultDest()));
+void FunctionPropertiesInfo::updateForBB(const BasicBlock &BB,
+                                         int64_t Direction) {
+  assert(Direction == 1 || Direction == -1);
+  BasicBlockCount += Direction;
+  BlocksReachedFromConditionalInstruction +=
+      (Direction * getNrBlocksFromCond(BB));
+  for (const auto &I : BB) {
+    if (auto *CS = dyn_cast<CallBase>(&I)) {
+      const auto *Callee = CS->getCalledFunction();
+      if (Callee && !Callee->isIntrinsic() && !Callee->isDeclaration())
+        DirectCallsToDefinedFunctions += Direction;
     }
-
-    for (const auto &I : BB) {
-      if (auto *CS = dyn_cast<CallBase>(&I)) {
-        const auto *Callee = CS->getCalledFunction();
-        if (Callee && !Callee->isIntrinsic() && !Callee->isDeclaration())
-          ++FPI.DirectCallsToDefinedFunctions;
-      }
-      if (I.getOpcode() == Instruction::Load) {
-        ++FPI.LoadInstCount;
-      } else if (I.getOpcode() == Instruction::Store) {
-        ++FPI.StoreInstCount;
-      }
+    if (I.getOpcode() == Instruction::Load) {
+      LoadInstCount += Direction;
+    } else if (I.getOpcode() == Instruction::Store) {
+      StoreInstCount += Direction;
     }
-    // Loop Depth of the Basic Block
-    int64_t LoopDepth;
-    LoopDepth = LI.getLoopDepth(&BB);
-    if (FPI.MaxLoopDepth < LoopDepth)
-      FPI.MaxLoopDepth = LoopDepth;
   }
-  FPI.TopLevelLoopCount += llvm::size(LI);
+  TotalInstructionCount += Direction * BB.sizeWithoutDebug();
+}
+
+void FunctionPropertiesInfo::updateAggregateStats(const Function &F,
+                                                  const LoopInfo &LI) {
+
+  Uses = getUses(F);
+  TopLevelLoopCount = llvm::size(LI);
+}
+
+FunctionPropertiesInfo
+FunctionPropertiesInfo::getFunctionPropertiesInfo(const Function &F,
+                                                  const LoopInfo &LI) {
+
+  FunctionPropertiesInfo FPI;
+  for (const auto &BB : F)
+    if (!pred_empty(&BB) || BB.isEntryBlock())
+      FPI.reIncludeBB(BB, LI);
+  FPI.updateAggregateStats(F, LI);
   return FPI;
 }
 
@@ -68,7 +94,8 @@ void FunctionPropertiesInfo::print(raw_ostream &OS) const {
      << "LoadInstCount: " << LoadInstCount << "\n"
      << "StoreInstCount: " << StoreInstCount << "\n"
      << "MaxLoopDepth: " << MaxLoopDepth << "\n"
-     << "TopLevelLoopCount: " << TopLevelLoopCount << "\n\n";
+     << "TopLevelLoopCount: " << TopLevelLoopCount << "\n"
+     << "TotalInstructionCount: " << TotalInstructionCount << "\n\n";
 }
 
 AnalysisKey FunctionPropertiesAnalysis::Key;
@@ -87,3 +114,60 @@ FunctionPropertiesPrinterPass::run(Function &F, FunctionAnalysisManager &AM) {
   AM.getResult<FunctionPropertiesAnalysis>(F).print(OS);
   return PreservedAnalyses::all();
 }
+
+FunctionPropertiesUpdater::FunctionPropertiesUpdater(
+    FunctionPropertiesInfo &FPI, const CallBase &CB)
+    : FPI(FPI), CallSiteBB(*CB.getParent()), Caller(*CallSiteBB.getParent()) {
+
+  // For BBs that are likely to change, we subtract from feature totals their
+  // contribution. Some features, like max loop counts or depths, are left
+  // invalid, as they will be updated post-inlining.
+  SmallPtrSet<const BasicBlock *, 4> LikelyToChangeBBs;
+  // The CB BB will change - it'll either be split or the callee's body (single
+  // BB) will be pasted in.
+  LikelyToChangeBBs.insert(&CallSiteBB);
+
+  // The caller's entry BB may change due to new alloca instructions.
+  LikelyToChangeBBs.insert(&*Caller.begin());
+
+  // The successors may become unreachable in the case of `invoke` inlining.
+  // We track successors separately, too, because they form a boundary, together
+  // with the CB BB ('Entry') between which the inlined callee will be pasted.
+  Successors.insert(succ_begin(&CallSiteBB), succ_end(&CallSiteBB));
+  for (const auto *BB : Successors)
+    LikelyToChangeBBs.insert(BB);
+
+  // Commit the change. While some of the BBs accounted for above may play dual
+  // role - e.g. caller's entry BB may be the same as the callsite BB - set
+  // insertion semantics make sure we account them once. This needs to be
+  // followed in `finish`, too.
+  for (const auto *BB : LikelyToChangeBBs)
+    FPI.updateForBB(*BB, -1);
+}
+
+void FunctionPropertiesUpdater::finish(const LoopInfo &LI) {
+  DenseSet<const BasicBlock *> ReIncluded;
+  std::deque<const BasicBlock *> Worklist;
+
+  if (&CallSiteBB != &*Caller.begin()) {
+    FPI.reIncludeBB(*Caller.begin(), LI);
+    ReIncluded.insert(&*Caller.begin());
+  }
+
+  // Update feature values from the BBs that were copied from the callee, or
+  // might have been modified because of inlining. The latter have been
+  // subtracted in the FunctionPropertiesUpdater ctor.
+  Worklist.push_back(&CallSiteBB);
+  while (!Worklist.empty()) {
+    const auto *BB = Worklist.front();
+    Worklist.pop_front();
+    if (!ReIncluded.insert(BB).second)
+      continue;
+    FPI.reIncludeBB(*BB, LI);
+    if (!Successors.contains(BB))
+      for (const auto *Succ : successors(BB))
+        Worklist.push_back(Succ);
+  }
+  FPI.updateAggregateStats(Caller, LI);
+  assert(FPI == FunctionPropertiesInfo::getFunctionPropertiesInfo(Caller, LI));
+}