[mlgo] Update FunctionPropertyCache after invalidating analyses

[mtrofin]

Summary

The update depends on LoopInfo, so we need that refreshed first, not
after.

[kazu]

LGTM. Thanks Mircea!

llvm/include/llvm/Analysis/FunctionPropertiesAnalysis.h

 /// 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);
+  void finish(const LoopInfo &LI) const;
 
 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

public:
   }
   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;
 
   std::unique_ptr<InlineAdvice> getMandatoryAdvice(CallBase &CB,
                                                    bool Advice) override;
 
   virtual std::unique_ptr<MLInlineAdvice> getMandatoryAdviceImpl(CallBase &CB);
 
   virtual std::unique_ptr<MLInlineAdvice>
   getAdviceFromModel(CallBase &CB, OptimizationRemarkEmitter &ORE);
 
   // Get the initial 'level' of the function, or 0 if the function has been
   // introduced afterwards.
   // TODO: should we keep this updated?
   unsigned getInitialFunctionLevel(const Function &F) const;
 
   std::unique_ptr<MLModelRunner> ModelRunner;
 
 private:
   int64_t getModuleIRSize() const;
 
-  void print(raw_ostream &OS) const override {
-    OS << "[MLInlineAdvisor] Nodes: " << NodeCount << " Edges: " << EdgeCount
-       << "\n";
-  }
+  void print(raw_ostream &OS) const override;
 
   mutable DenseMap<const Function *, FunctionPropertiesInfo> FPICache;
 
   LazyCallGraph &CG;
 
   int64_t NodeCount = 0;
   int64_t EdgeCount = 0;
   int64_t EdgesOfLastSeenNodes = 0;
public:
   void recordUnattemptedInliningImpl() override;
 
   Function *getCaller() const { return Caller; }
   Function *getCallee() const { return Callee; }
 
   const int64_t CallerIRSize;
   const int64_t CalleeIRSize;
   const int64_t CallerAndCalleeEdges;
+  void updateCachedCallerFPI(const LoopInfo &LI) const;
 
 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
 
 #endif // LLVM_ANALYSIS_MLINLINEADVISOR_H

llvm/lib/Analysis/FunctionPropertiesAnalysis.cpp

FunctionPropertiesUpdater::FunctionPropertiesUpdater(
   // 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) {
+void FunctionPropertiesUpdater::finish(const LoopInfo &LI) const {
   DenseSet<const BasicBlock *> ReIncluded;
   std::deque<const BasicBlock *> Worklist;
 
   if (&CallSiteBB != &*Caller.begin()) {
     FPI.reIncludeBB(*Caller.begin(), LI);
     ReIncluded.insert(&*Caller.begin());
   }
 

llvm/lib/Analysis/MLInlineAdvisor.cpp

 #define DEBUG_TYPE "inline-ml"
 
 static cl::opt<float> SizeIncreaseThreshold(
     "ml-advisor-size-increase-threshold", cl::Hidden,
     cl::desc("Maximum factor by which expected native size may increase before "
              "blocking any further inlining."),
     cl::init(2.0));
 
+static cl::opt<bool> KeepFPICache(
+    "ml-advisor-keep-fpi-cache", cl::Hidden,
+    cl::desc(
+        "For test - keep the ML Inline advisor's FunctionPropertiesInfo cache"),
+    cl::init(false));
+
 // clang-format off
 const std::array<TensorSpec, NumberOfFeatures> llvm::FeatureMap{
 #define POPULATE_NAMES(_, NAME) TensorSpec::createSpec<int64_t>(NAME, {1} ),
 // InlineCost features - these must come first
   INLINE_COST_FEATURE_ITERATOR(POPULATE_NAMES)
 #undef POPULATE_NAMES
 
 // Non-cost features
void MLInlineAdvisor::onPassEntry() {
   }
 
   EdgeCount -= EdgesOfLastSeenNodes;
   EdgesOfLastSeenNodes = 0;
 }
 
 void MLInlineAdvisor::onPassExit(LazyCallGraph::SCC *LastSCC) {
   // No need to keep this around - function passes will invalidate it.
+  if (!KeepFPICache)
     FPICache.clear();
   if (!LastSCC || ForceStop)
     return;
   // Keep track of the nodes and edges we last saw. Then, in onPassEntry,
   // we update the node count and edge count from the subset of these nodes that
   // survived.
   assert(NodesInLastSCC.empty());
   assert(NodeCount >= LastSCC->size());
   EdgesOfLastSeenNodes = 0;

 // number of functions and the number of static calls. We also keep track of the
 // total IR size in this module, to stop misbehaving policies at a certain bloat
 // factor (SizeIncreaseThreshold)
 void MLInlineAdvisor::onSuccessfulInlining(const MLInlineAdvice &Advice,
                                            bool CalleeWasDeleted) {
   assert(!ForceStop);
   Function *Caller = Advice.getCaller();
   Function *Callee = Advice.getCallee();
-
   // The caller features aren't valid anymore.
   {
-    PreservedAnalyses PA = PreservedAnalyses::all();
-    PA.abandon<FunctionPropertiesAnalysis>();
+    PreservedAnalyses PA = PreservedAnalyses::none();
     FAM.invalidate(*Caller, PA);
   }
+  Advice.updateCachedCallerFPI(FAM.getResult<LoopAnalysis>(*Caller));
   int64_t IRSizeAfter =
       getIRSize(*Caller) + (CalleeWasDeleted ? 0 : Advice.CalleeIRSize);
   CurrentIRSize += IRSizeAfter - (Advice.CallerIRSize + Advice.CalleeIRSize);
   if (CurrentIRSize > SizeIncreaseThreshold * InitialIRSize)
     ForceStop = true;
 
   // We can delta-update module-wide features. We know the inlining only changed
   // the caller, and maybe the callee (by deleting the latter).
std::unique_ptr<InlineAdvice> MLInlineAdvisor::getMandatoryAdvice(CallBase &CB,
 
   // If this is a "never inline" case, there won't be any changes to internal
   // state we need to track, so we can just return the base InlineAdvice, which
   // will do nothing interesting.
   // Same if we are forced to stop - we don't track anymore.
   return std::make_unique<InlineAdvice>(this, CB, getCallerORE(CB), Advice);
 }
 
-const LoopInfo &MLInlineAdvisor::getLoopInfo(Function &F) const {
-  return FAM.getResult<LoopAnalysis>(F);
-}
-
 std::unique_ptr<MLInlineAdvice>
 MLInlineAdvisor::getMandatoryAdviceImpl(CallBase &CB) {
   return std::make_unique<MLInlineAdvice>(this, CB, getCallerORE(CB), true);
 }
 
+void MLInlineAdvisor::print(raw_ostream &OS) const {
+  OS << "[MLInlineAdvisor] Nodes: " << NodeCount << " Edges: " << EdgeCount
+     << "\n";
+  OS << "[MLInlineAdvisor] FPI:\n";
+  for (auto I : FPICache) {
+    OS << I.getFirst()->getName() << ":\n";
+    I.getSecond().print(OS);
+    OS << "\n";
+  }
+  OS << "\n";
+}
+
 MLInlineAdvice::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
void MLInlineAdvice::reportContextForRemark(
   using namespace ore;
   OR << NV("Callee", Callee->getName());
   for (size_t I = 0; I < NumberOfFeatures; ++I)
     OR << NV(FeatureMap[I].name(),
              *getAdvisor()->getModelRunner().getTensor<int64_t>(I));
   OR << NV("ShouldInline", isInliningRecommended());
 }
 
-void MLInlineAdvice::updateCachedCallerFPI() {
-  FPU->finish(getAdvisor()->getLoopInfo(*Caller));
+void MLInlineAdvice::updateCachedCallerFPI(const LoopInfo &LI) const {
+  FPU->finish(LI);
 }
 
 void MLInlineAdvice::recordInliningImpl() {
-  updateCachedCallerFPI();
   ORE.emit([&]() {
     OptimizationRemark R(DEBUG_TYPE, "InliningSuccess", DLoc, Block);
     reportContextForRemark(R);
     return R;
   });
   getAdvisor()->onSuccessfulInlining(*this, /*CalleeWasDeleted*/ false);
 }
 
 void MLInlineAdvice::recordInliningWithCalleeDeletedImpl() {
-  updateCachedCallerFPI();
   ORE.emit([&]() {
     OptimizationRemark R(DEBUG_TYPE, "InliningSuccessWithCalleeDeleted", DLoc,
                          Block);
     reportContextForRemark(R);
     return R;
   });
   getAdvisor()->onSuccessfulInlining(*this, /*CalleeWasDeleted*/ true);
 }

llvm/test/Transforms/Inline/ML/fpi-update.ll

+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-grtev4-linux-gnu"
+
+; TODO: we could instantiate the MLInliner with a non-model generated evaluator
+; and drop the requirement
+; REQUIRES: llvm_inliner_model_autogenerated
+
+; RUN: opt -enable-ml-inliner=release -passes='scc-oz-module-inliner,print<inline-advisor>' \
+; RUN:     -keep-inline-advisor-for-printing -max-devirt-iterations=0 \
+; RUN:     -mandatory-inlining-first=0 -S < %s 2>&1 | FileCheck %s 
+
+define void @caller(i32 %i) #1 {
+  call void @callee(i32 %i)
+  ret void
+}
+
+define void @callee(i32 %i) #0 {
+entry:
+  br label %loop
+loop:
+  %cond = icmp slt i32 %i, 0
+  br i1 %cond, label %loop, label %exit
+exit:
+  ret void
+}
+
+attributes #0 = { alwaysinline }
+attributes #1 = { noinline optnone }
+
+; CHECK: [MLInlineAdvisor] FPI:
+; CHECK: caller:
+; CHECK: MaxLoopDepth: 1