[mlgo][inline] Improve global state tracking

[mtrofin]

Summary

The global state refers to the number of the nodes currently in the module, and the number of direct calls between nodes, across the module.

Node counts are not a problem; edge counts are because we want strictly the kind of edges that affect inlining (direct calls), and that is not easily obtainable without iteration over the whole module.

This patch avoids relying on analysis invalidation because it turned out to be too aggressive in some cases. It leverages the fact that Node objects are stable - they do not get deleted while cgscc passes are run over the module; and cgscc pass manager invariants.

[mtrofin]

an alternative that is more robust and arguably easier to understand, as it makes no assumption about scc contents after we leave the inliner: any function could change, and we would capture that. The case of node addition, however, is handled the same way.

[phosek]

This looks reasonable although I'm not familiar with this part of the codebase so I'd like for @aeubanks to take a look as well.

llvm/include/llvm/Analysis/MLInlineAdvisor.h

80

llvm/lib/Analysis/MLInlineAdvisor.cpp

177

llvm/test/Transforms/Inline/ML/state-tracking.ll

25

Can you add a newline here?

[mtrofin]

feedback

[aeubanks]

This patch avoids relying on analysis invalidation because it turned out to be too aggressive in some cases.

Can you explain this a bit more?

We shouldn't need to completely recalculate everything when adding a new function. Any new functions have to be children of the current SCC, so we should be able to find them by iterating through the outgoing edges from the current SCC and seeing if we encounter any nodes we haven't seen before.

llvm/include/llvm/Analysis/InlineAdvisor.h

12

do you need this? this is the legacy PM one

llvm/lib/Analysis/MLInlineAdvisor.cpp

146

if ShuttingDown is ever true, that we means the object has already been destructed and calling any method on it is bad.
In general having things done at a distance in the destructor of an analysis seems iffy. Keeping everything in the pass's run method should be doable here, we just need to keep track of the nodes in the most recently visited SCC.

415

newline?

llvm/test/Transforms/Inline/ML/state-tracking-coro.ll

2

I'm not sure I'd want to depend on exact coroutine semantics for these tests especially since it seems like they keep changing
but if we don't do that then I don't think there's any other way to test new functions being added to the call graph
maybe it's ok if these tests aren't run normally and you're fine with keeping these up to date if coroutine semantics/pipelines change

[mtrofin]

I think there are 2 questions, is that correct? If so:

In D115847#3228512, @aeubanks wrote:

This patch avoids relying on analysis invalidation because it turned out to be too aggressive in some cases.

Can you explain this a bit more?

This refers to the behavior before. Invalidation of the whole analysis (the MLInlineAdvisor one) was still happening too frequently.

We shouldn't need to completely recalculate everything when adding a new function. Any new functions have to be children of the current SCC, so we should be able to find them by iterating through the outgoing edges from the current SCC and seeing if we encounter any nodes we haven't seen before.

I think you're referring here to the behavior in the patch, correct (i.e. not to the quoted sentence in the patch description)?

I would prefer not assuming that SCC passes adding functions relate them to the SCC they are currently processing. I realize that's currently the case, but this assumption made me uncomfortable: I am not sure why this is something that a pass would necessarily respect in the future.

I think there is a partial update approach that is only dependent on the fact that Nodes are only added for the duration of a module-wide cgscc traversal: we remember a "watermark" of the last node and iterate from that watermark up next time we re-enter the inliner. That's a change that can be done incrementally though, from the current change.

llvm/lib/Analysis/MLInlineAdvisor.cpp

146

See the dtor - this just avoids spending time doing inserts due to the abandon-ing of the NotifyOnChangeFunctionAnalysis in the ~MLInlineAdvisor. The object isn't destructed just yet.

[mtrofin]

feedback

llvm/test/Transforms/Inline/ML/state-tracking-coro.ll

2

ack - this is just taking a dep on the coro pass generating a few fcts, I'm hoping that won't change, but if it does and it becomes burdensome, we can cross that bridge then, I think?

[aeubanks]

In D115847#3228603, @mtrofin wrote:

I think there are 2 questions, is that correct? If so:

In D115847#3228512, @aeubanks wrote:

This patch avoids relying on analysis invalidation because it turned out to be too aggressive in some cases.

Can you explain this a bit more?

This refers to the behavior before. Invalidation of the whole analysis (the MLInlineAdvisor one) was still happening too frequently.

Oh I thought this was in reference to FunctionPropertiesAnalysis being invalidated too much.

We shouldn't need to completely recalculate everything when adding a new function. Any new functions have to be children of the current SCC, so we should be able to find them by iterating through the outgoing edges from the current SCC and seeing if we encounter any nodes we haven't seen before.

I think you're referring here to the behavior in the patch, correct (i.e. not to the quoted sentence in the patch description)?

I would prefer not assuming that SCC passes adding functions relate them to the SCC they are currently processing. I realize that's currently the case, but this assumption made me uncomfortable: I am not sure why this is something that a pass would necessarily respect in the future.

An SCC pass should not be adding functions that don't relate to some function in the current SCC. That wouldn't fit the CGSCC pass model and would break LCG. So I think looking around all the nodes in the current SCC to discover new functions is fine.

I think there is a partial update approach that is only dependent on the fact that Nodes are only added for the duration of a module-wide cgscc traversal: we remember a "watermark" of the last node and iterate from that watermark up next time we re-enter the inliner. That's a change that can be done incrementally though, from the current change.

later changes sgtm.

llvm/lib/Analysis/MLInlineAdvisor.cpp

146

ah I missed that

llvm/test/Transforms/Inline/ML/state-tracking-scc-splits.ll

5

not related to this patch but you could add this as a lit.local.cfg in the ML directory so you don't have to put it on every test

[mtrofin]

I think there are 2 questions, is that correct? If so:

In D115847#3228512, @aeubanks wrote:

This patch avoids relying on analysis invalidation because it turned out to be too aggressive in some cases.

Can you explain this a bit more?

This refers to the behavior before. Invalidation of the whole analysis (the MLInlineAdvisor one) was still happening too frequently.

We shouldn't need to completely recalculate everything when adding a new function. Any new functions have to be children of the current SCC, so we should be able to find them by iterating through the outgoing edges from the current SCC and seeing if we encounter any nodes we haven't seen before.

I think you're referring here to the behavior in the patch, correct (i.e. not to the quoted sentence in the patch description)?

I would prefer not assuming that SCC passes adding functions relate them to the SCC they are currently processing. I realize that's currently the case, but this assumption made me uncomfortable: I am not sure why this is something that a pass would necessarily respect in the future.

I think there is a partial update approach that is only dependent on the fact that Nodes are only added for the duration of a module-wide cgscc traversal: we remember a "watermark" of the last node and iterate from that watermark up next time we re-enter the inliner. That's a change that can be done incrementally though, from the current change.

In D115847#3232350, @aeubanks wrote:

In D115847#3228603, @mtrofin wrote:

I think there are 2 questions, is that correct? If so:

In D115847#3228512, @aeubanks wrote:

This patch avoids relying on analysis invalidation because it turned out to be too aggressive in some cases.

Can you explain this a bit more?

This refers to the behavior before. Invalidation of the whole analysis (the MLInlineAdvisor one) was still happening too frequently.

Oh I thought this was in reference to FunctionPropertiesAnalysis being invalidated too much.

We shouldn't need to completely recalculate everything when adding a new function. Any new functions have to be children of the current SCC, so we should be able to find them by iterating through the outgoing edges from the current SCC and seeing if we encounter any nodes we haven't seen before.

I think you're referring here to the behavior in the patch, correct (i.e. not to the quoted sentence in the patch description)?

I would prefer not assuming that SCC passes adding functions relate them to the SCC they are currently processing. I realize that's currently the case, but this assumption made me uncomfortable: I am not sure why this is something that a pass would necessarily respect in the future.

An SCC pass should not be adding functions that don't relate to some function in the current SCC. That wouldn't fit the CGSCC pass model and would break LCG. So I think looking around all the nodes in the current SCC to discover new functions is fine.

Ah, ok, so should such a pass try to get introduced, some test(s) would fail, I assume. OK. Let me try and make sure I understand how this check would work:

• I keep track of the set of Nodes I've seen last, and the total set of Nodes
• when onPassEntry is called again, I check that the set of Nodes adjacent (through what? Ref or Call edges? I guess both) to the Nodes in the set above are in the total set of Nodes, otherwise I found new nodes

... (kind of obvious what follows)

Is this correct? (thanks!)

I think there is a partial update approach that is only dependent on the fact that Nodes are only added for the duration of a module-wide cgscc traversal: we remember a "watermark" of the last node and iterate from that watermark up next time we re-enter the inliner. That's a change that can be done incrementally though, from the current change.

later changes sgtm.

[mtrofin]

llvm/test/Transforms/Inline/ML/state-tracking-scc-splits.ll

5

some tests have REQUIRES: have_tf_api, some have_tf_aot, and some both. ('api' is for training, 'aot' is for deployment compiler, i.e. where we embed a pre-trained model as a .o)

[aeubanks]

In D115847#3232394, @mtrofin wrote:

I think there are 2 questions, is that correct? If so:

In D115847#3228512, @aeubanks wrote:

This patch avoids relying on analysis invalidation because it turned out to be too aggressive in some cases.

Can you explain this a bit more?

This refers to the behavior before. Invalidation of the whole analysis (the MLInlineAdvisor one) was still happening too frequently.

We shouldn't need to completely recalculate everything when adding a new function. Any new functions have to be children of the current SCC, so we should be able to find them by iterating through the outgoing edges from the current SCC and seeing if we encounter any nodes we haven't seen before.

I think you're referring here to the behavior in the patch, correct (i.e. not to the quoted sentence in the patch description)?

I would prefer not assuming that SCC passes adding functions relate them to the SCC they are currently processing. I realize that's currently the case, but this assumption made me uncomfortable: I am not sure why this is something that a pass would necessarily respect in the future.

I think there is a partial update approach that is only dependent on the fact that Nodes are only added for the duration of a module-wide cgscc traversal: we remember a "watermark" of the last node and iterate from that watermark up next time we re-enter the inliner. That's a change that can be done incrementally though, from the current change.

In D115847#3232350, @aeubanks wrote:

In D115847#3228603, @mtrofin wrote:

I think there are 2 questions, is that correct? If so:

In D115847#3228512, @aeubanks wrote:

This patch avoids relying on analysis invalidation because it turned out to be too aggressive in some cases.

Can you explain this a bit more?

This refers to the behavior before. Invalidation of the whole analysis (the MLInlineAdvisor one) was still happening too frequently.

Oh I thought this was in reference to FunctionPropertiesAnalysis being invalidated too much.

We shouldn't need to completely recalculate everything when adding a new function. Any new functions have to be children of the current SCC, so we should be able to find them by iterating through the outgoing edges from the current SCC and seeing if we encounter any nodes we haven't seen before.

I think you're referring here to the behavior in the patch, correct (i.e. not to the quoted sentence in the patch description)?

I would prefer not assuming that SCC passes adding functions relate them to the SCC they are currently processing. I realize that's currently the case, but this assumption made me uncomfortable: I am not sure why this is something that a pass would necessarily respect in the future.

An SCC pass should not be adding functions that don't relate to some function in the current SCC. That wouldn't fit the CGSCC pass model and would break LCG. So I think looking around all the nodes in the current SCC to discover new functions is fine.

Ah, ok, so should such a pass try to get introduced, some test(s) would fail, I assume. OK. Let me try and make sure I understand how this check would work:

• I keep track of the set of Nodes I've seen last, and the total set of Nodes
• when onPassEntry is called again, I check that the set of Nodes adjacent (through what? Ref or Call edges? I guess both) to the Nodes in the set above are in the total set of Nodes, otherwise I found new nodes

yeah it could be a ref or a call edge, and theoretically we could add functions that are only referenced from another added function (e.g. A becomes A'->A1->A2), although LCG doesn't currently support that

... (kind of obvious what follows)

Is this correct? (thanks!)

yup sounds right to me

I think there is a partial update approach that is only dependent on the fact that Nodes are only added for the duration of a module-wide cgscc traversal: we remember a "watermark" of the last node and iterate from that watermark up next time we re-enter the inliner. That's a change that can be done incrementally though, from the current change.

later changes sgtm.

[mtrofin]

In D115847#3232408, @aeubanks wrote:

In D115847#3232394, @mtrofin wrote:

I think there are 2 questions, is that correct? If so:

In D115847#3228512, @aeubanks wrote:

This patch avoids relying on analysis invalidation because it turned out to be too aggressive in some cases.

Can you explain this a bit more?

This refers to the behavior before. Invalidation of the whole analysis (the MLInlineAdvisor one) was still happening too frequently.

We shouldn't need to completely recalculate everything when adding a new function. Any new functions have to be children of the current SCC, so we should be able to find them by iterating through the outgoing edges from the current SCC and seeing if we encounter any nodes we haven't seen before.

I think you're referring here to the behavior in the patch, correct (i.e. not to the quoted sentence in the patch description)?

I would prefer not assuming that SCC passes adding functions relate them to the SCC they are currently processing. I realize that's currently the case, but this assumption made me uncomfortable: I am not sure why this is something that a pass would necessarily respect in the future.

I think there is a partial update approach that is only dependent on the fact that Nodes are only added for the duration of a module-wide cgscc traversal: we remember a "watermark" of the last node and iterate from that watermark up next time we re-enter the inliner. That's a change that can be done incrementally though, from the current change.

In D115847#3232350, @aeubanks wrote:

In D115847#3228603, @mtrofin wrote:

I think there are 2 questions, is that correct? If so:

In D115847#3228512, @aeubanks wrote:

This patch avoids relying on analysis invalidation because it turned out to be too aggressive in some cases.

Can you explain this a bit more?

This refers to the behavior before. Invalidation of the whole analysis (the MLInlineAdvisor one) was still happening too frequently.

Oh I thought this was in reference to FunctionPropertiesAnalysis being invalidated too much.

We shouldn't need to completely recalculate everything when adding a new function. Any new functions have to be children of the current SCC, so we should be able to find them by iterating through the outgoing edges from the current SCC and seeing if we encounter any nodes we haven't seen before.

I think you're referring here to the behavior in the patch, correct (i.e. not to the quoted sentence in the patch description)?

I would prefer not assuming that SCC passes adding functions relate them to the SCC they are currently processing. I realize that's currently the case, but this assumption made me uncomfortable: I am not sure why this is something that a pass would necessarily respect in the future.

An SCC pass should not be adding functions that don't relate to some function in the current SCC. That wouldn't fit the CGSCC pass model and would break LCG. So I think looking around all the nodes in the current SCC to discover new functions is fine.

Ah, ok, so should such a pass try to get introduced, some test(s) would fail, I assume. OK. Let me try and make sure I understand how this check would work:

• I keep track of the set of Nodes I've seen last, and the total set of Nodes
• when onPassEntry is called again, I check that the set of Nodes adjacent (through what? Ref or Call edges? I guess both) to the Nodes in the set above are in the total set of Nodes, otherwise I found new nodes

yeah it could be a ref or a call edge, and theoretically we could add functions that are only referenced from another added function (e.g. A becomes A'->A1->A2), although LCG doesn't currently support that

... (kind of obvious what follows)

Is this correct? (thanks!)

yup sounds right to me

Oh OK - then I can go back to the first version of this patch, where I didn't need the function pass part (but need this Node accounting)

I think there is a partial update approach that is only dependent on the fact that Nodes are only added for the duration of a module-wide cgscc traversal: we remember a "watermark" of the last node and iterate from that watermark up next time we re-enter the inliner. That's a change that can be done incrementally though, from the current change.

later changes sgtm.

[mtrofin]

Improved handing of "new nodes" - also went back to the first variant of this patch.

[aeubanks]

llvm/include/llvm/Analysis/MLInlineAdvisor.h

12

not used?

llvm/lib/Analysis/MLInlineAdvisor.cpp

147

then the pipeline is restarted on the merged SCC

153

or adjacent to other new nodes. no existing passes do this, but it's possible
this isn't properly handled below, but maybe a TODO is good enough for now?

158

does this come up? the CGSCC infra only visits function definitions

164

ditto

[mtrofin]

feedback

llvm/lib/Analysis/MLInlineAdvisor.cpp

153

fixed.

158

N->isDead could happen if the function for N died since. The second part - not sure, I could imagine a pass converting an implementation to an intrinsic?

164

here I think you're right, this should be an assert.

[aeubanks]

looks good except for the one comment

llvm/lib/Analysis/MLInlineAdvisor.cpp

158

a pass wouldn't replace the node with an intrinsic, every node is required to have a function definition
so I'd remove the check for a function declaration

[mtrofin]

feedback

llvm/lib/Analysis/MLInlineAdvisor.cpp

158

made into an assert.

llvm/include/llvm/Analysis/InlineAdvisor.h

 //===- InlineAdvisor.h - Inlining decision making abstraction -*- C++ ---*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 #ifndef LLVM_ANALYSIS_INLINEADVISOR_H
 #define LLVM_ANALYSIS_INLINEADVISOR_H
 
 #include "llvm/Analysis/InlineCost.h"
+#include "llvm/Analysis/LazyCallGraph.h"
 #include "llvm/Analysis/Utils/ImportedFunctionsInliningStatistics.h"
 #include "llvm/Config/llvm-config.h"
 #include "llvm/IR/PassManager.h"
 #include <memory>
 #include <unordered_set>
 
 namespace llvm {
 class BasicBlock;
public:
 
   /// This must be called when the Inliner pass is entered, to allow the
   /// InlineAdvisor update internal state, as result of function passes run
   /// between Inliner pass runs (for the same module).
   virtual void onPassEntry() {}
 
   /// This must be called when the Inliner pass is exited, as function passes
   /// may be run subsequently. This allows an implementation of InlineAdvisor
-  /// to prepare for a partial update.
-  virtual void onPassExit() {}
+  /// to prepare for a partial update, based on the optional SCC.
+  virtual void onPassExit(LazyCallGraph::SCC *SCC = nullptr) {}
 
-  /// Called when the module is invalidated. We let the advisor implementation
-  /// decide what to refresh - in the case of the development mode
-  /// implementation, for example, we wouldn't want to delete the whole object
-  /// and need to re-load the model evaluator.
-  virtual void onModuleInvalidated() {}
+  /// Support for printer pass
+  virtual void print(raw_ostream &OS) const {
+    OS << "Unimplemented InlineAdvisor print\n";
+  }
 
 protected:
   InlineAdvisor(Module &M, FunctionAnalysisManager &FAM);
   virtual std::unique_ptr<InlineAdvice> getAdviceImpl(CallBase &CB) = 0;
   virtual std::unique_ptr<InlineAdvice> getMandatoryAdvice(CallBase &CB,
                                                            bool Advice);
 
   Module &M;

 class InlineAdvisorAnalysis : public AnalysisInfoMixin<InlineAdvisorAnalysis> {
 public:
   static AnalysisKey Key;
   InlineAdvisorAnalysis() = default;
   struct Result {
     Result(Module &M, ModuleAnalysisManager &MAM) : M(M), MAM(MAM) {}
     bool invalidate(Module &, const PreservedAnalyses &PA,
                     ModuleAnalysisManager::Invalidator &) {
-      if (Advisor && !PA.areAllPreserved())
-        Advisor->onModuleInvalidated();
       // Check whether the analysis has been explicitly invalidated. Otherwise,
       // it's stateless and remains preserved.
       auto PAC = PA.getChecker<InlineAdvisorAnalysis>();
       return !PAC.preservedWhenStateless();
     }
     bool tryCreate(InlineParams Params, InliningAdvisorMode Mode,
                    const ReplayInlinerSettings &ReplaySettings);
     InlineAdvisor *getAdvisor() const { return Advisor.get(); }
 
   private:
     Module &M;
     ModuleAnalysisManager &MAM;
     std::unique_ptr<InlineAdvisor> Advisor;
   };
 
   Result run(Module &M, ModuleAnalysisManager &MAM) { return Result(M, MAM); }
 };
 
+/// Printer pass for the FunctionPropertiesAnalysis results.
+class InlineAdvisorAnalysisPrinterPass
+    : public PassInfoMixin<InlineAdvisorAnalysisPrinterPass> {
+  raw_ostream &OS;
+
+public:
+  explicit InlineAdvisorAnalysisPrinterPass(raw_ostream &OS) : OS(OS) {}
+
+  PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
+};
+
 #ifdef LLVM_HAVE_TF_AOT
 std::unique_ptr<InlineAdvisor>
 getReleaseModeAdvisor(Module &M, ModuleAnalysisManager &MAM);
 #endif
 
 #ifdef LLVM_HAVE_TF_API
 std::unique_ptr<InlineAdvisor>
 getDevelopmentModeAdvisor(Module &M, ModuleAnalysisManager &MAM,

llvm/include/llvm/Analysis/MLInlineAdvisor.h

 //===- MLInlineAdvisor.h - ML - based InlineAdvisor factories ---*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef LLVM_ANALYSIS_MLINLINEADVISOR_H
 #define LLVM_ANALYSIS_MLINLINEADVISOR_H
 
+#include "llvm/ADT/StringSet.h"
 #include "llvm/Analysis/InlineAdvisor.h"
 #include "llvm/Analysis/LazyCallGraph.h"
 #include "llvm/Analysis/MLModelRunner.h"
 #include "llvm/IR/PassManager.h"
 
 #include <memory>
 #include <unordered_map>
 
 namespace llvm {
 class Module;
 class MLInlineAdvice;
 
 class MLInlineAdvisor : public InlineAdvisor {
 public:
   MLInlineAdvisor(Module &M, ModuleAnalysisManager &MAM,
                   std::unique_ptr<MLModelRunner> ModelRunner);
 
   virtual ~MLInlineAdvisor() = default;
 
   void onPassEntry() override;
+  void onPassExit(LazyCallGraph::SCC *SCC) override;
 
   int64_t getIRSize(const Function &F) const { return F.getInstructionCount(); }
   void onSuccessfulInlining(const MLInlineAdvice &Advice,
                             bool CalleeWasDeleted);
 
   bool isForcedToStop() const { return ForceStop; }
   int64_t getLocalCalls(Function &F);
   const MLModelRunner &getModelRunner() const { return *ModelRunner.get(); }
-  void onModuleInvalidated() override { Invalid = true; }
 
 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;
 
-  bool Invalid = true;
+  void print(raw_ostream &OS) const override {
+    OS << "[MLInlineAdvisor] Nodes: " << NodeCount << " Edges: " << EdgeCount
+       << "\n";
+  }
+
   LazyCallGraph &CG;
 
   int64_t NodeCount = 0;
   int64_t EdgeCount = 0;
+  int64_t EdgesOfLastSeenNodes = 0;
+
   std::map<const LazyCallGraph::Node *, unsigned> FunctionLevels;
   const int32_t InitialIRSize = 0;
   int32_t CurrentIRSize = 0;
-
+  DenseSet<const LazyCallGraph::Node *> NodesInLastSCC;
+  DenseSet<const LazyCallGraph::Node *> AllNodes;
   bool ForceStop = false;
 };
 
 /// InlineAdvice that tracks changes post inlining. For that reason, it only
 /// overrides the "successful inlining" extension points.
 class MLInlineAdvice : public InlineAdvice {
 public:
   MLInlineAdvice(MLInlineAdvisor *Advisor, CallBase &CB,
                  OptimizationRemarkEmitter &ORE, bool Recommendation)

llvm/lib/Analysis/InlineAdvisor.cpp

 #include "llvm/Analysis/InlineCost.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/Analysis/ReplayInlineAdvisor.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/DebugInfoMetadata.h"
 #include "llvm/IR/Instructions.h"
+#include "llvm/IR/PassManager.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace llvm;
 #define DEBUG_TYPE "inline"
 
 // This weirdly named statistic tracks the number of times that, when attempting
 // to inline a function A into B, we analyze the callers of B in order to see
bool Advice = CB.getCaller() != CB.getCalledFunction() &&
                 MandatoryInliningKind::Always ==
                     getMandatoryKind(CB, FAM, getCallerORE(CB));
   return getMandatoryAdvice(CB, Advice);
 }
 
 OptimizationRemarkEmitter &InlineAdvisor::getCallerORE(CallBase &CB) {
   return FAM.getResult<OptimizationRemarkEmitterAnalysis>(*CB.getCaller());
 }
+
+PreservedAnalyses
+InlineAdvisorAnalysisPrinterPass::run(Module &M, ModuleAnalysisManager &MAM) {
+  const auto *IA = MAM.getCachedResult<InlineAdvisorAnalysis>(M);
+  if (!IA)
+    OS << "No Inline Advisor\n";
+  else
+    IA->getAdvisor()->print(OS);
+  return PreservedAnalyses::all();
+}

llvm/lib/Analysis/MLInlineAdvisor.cpp

for (auto *CGNode : CGNodes) {
       }
     }
     for (auto *CGNode : CGNodes) {
       Function *F = CGNode->getFunction();
       if (F && !F->isDeclaration())
         FunctionLevels[&CG.get(*F)] = Level;
     }
   }
+  for (auto KVP : FunctionLevels) {
+    AllNodes.insert(KVP.first);
+    EdgeCount += getLocalCalls(KVP.first->getFunction());
+  }
+  NodeCount = AllNodes.size();
 }
 
 unsigned MLInlineAdvisor::getInitialFunctionLevel(const Function &F) const {
   return CG.lookup(F) ? FunctionLevels.at(CG.lookup(F)) : 0;
 }
 
 void MLInlineAdvisor::onPassEntry() {
   // Function passes executed between InlinerPass runs may have changed the
   // module-wide features.
-  if (!Invalid)
-    return;
-  NodeCount = 0;
-  EdgeCount = 0;
-  for (auto &F : M)
-    if (!F.isDeclaration()) {
-      ++NodeCount;
-      EdgeCount += getLocalCalls(F);
-    }
-  Invalid = false;
+  // The cgscc pass manager rules are such that:
+  // - if a pass leads to merging SCCs, then the pipeline is resumed
+  // - if a pass leads to splitting the SCC, then we continue with one of the
+  // splits
+  // This means that the NodesInLastSCC is a superset (not strict) of the nodes
+  // that subsequent passes would have processed
+  // - in addition, if new Nodes were created by a pass (e.g. CoroSplit),
+  // they'd be adjacent to Nodes in the last SCC. So we just need to check the
+  // boundary of Nodes in NodesInLastSCC for Nodes we haven't seen. We don't
+  // care about the nature of the Edge (call or ref).
+  NodeCount -= static_cast<int64_t>(NodesInLastSCC.size());
+  for (const auto *N : NodesInLastSCC) {
+    if (N->isDead() || N->getFunction().isDeclaration())
+      continue;
+    EdgeCount += getLocalCalls(N->getFunction());
+    for (const auto &E : *(*N)) {
+      const auto *AdjNode = &E.getNode();
+      if (AdjNode->isDead() || AdjNode->getFunction().isDeclaration())
+        continue;
+      auto I = AllNodes.insert(AdjNode);
+      if (I.second) {
+        ++NodeCount;
+        EdgeCount += getLocalCalls(AdjNode->getFunction());
+      }
+    }
+  }
+
+  NodesInLastSCC.clear();
+  EdgeCount -= EdgesOfLastSeenNodes;
+  EdgesOfLastSeenNodes = 0;
+}
+
+void MLInlineAdvisor::onPassExit(LazyCallGraph::SCC *LastSCC) {
+  if (!LastSCC)
+    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;
+  for (const auto &N : *LastSCC) {
+    assert(!N.isDead());
+    EdgesOfLastSeenNodes += getLocalCalls(N.getFunction());
+    NodesInLastSCC.insert(&N);
+  assert(EdgeCount >= EdgesOfLastSeenNodes);
 }
 
 int64_t MLInlineAdvisor::getLocalCalls(Function &F) {
   return FAM.getResult<FunctionPropertiesAnalysis>(F)
       .DirectCallsToDefinedFunctions;
 }
 
 // Update the internal state of the advisor, and force invalidate feature

 }
 void MLInlineAdvice::recordUnattemptedInliningImpl() {
   ORE.emit([&]() {
     OptimizationRemarkMissed R(DEBUG_TYPE, "IniningNotAttempted", DLoc, Block);
     reportContextForRemark(R);
     return R;
   });
 }
 #endif // defined(LLVM_HAVE_TF_AOT) || defined(LLVM_HAVE_TF_API)

llvm/lib/Passes/PassRegistry.def

 MODULE_PASS("function-import", FunctionImportPass())
 MODULE_PASS("function-specialization", FunctionSpecializationPass())
 MODULE_PASS("globaldce", GlobalDCEPass())
 MODULE_PASS("globalopt", GlobalOptPass())
 MODULE_PASS("globalsplit", GlobalSplitPass())
 MODULE_PASS("hotcoldsplit", HotColdSplittingPass())
 MODULE_PASS("inferattrs", InferFunctionAttrsPass())
 MODULE_PASS("inliner-wrapper", ModuleInlinerWrapperPass())
+MODULE_PASS("print<inline-advisor>", InlineAdvisorAnalysisPrinterPass(dbgs()))
 MODULE_PASS("inliner-wrapper-no-mandatory-first", ModuleInlinerWrapperPass(
   getInlineParams(),
   false))
 MODULE_PASS("insert-gcov-profiling", GCOVProfilerPass())
 MODULE_PASS("instrorderfile", InstrOrderFilePass())
 MODULE_PASS("instrprof", InstrProfiling())
 MODULE_PASS("internalize", InternalizePass())
 MODULE_PASS("invalidate<all>", InvalidateAllAnalysesPass())

llvm/lib/Transforms/IPO/Inliner.cpp

 /// prior to LLVM's code generator having support for stack coloring based on
 /// lifetime markers. It is now in the process of being removed. To experiment
 /// with disabling it and relying fully on lifetime marker based stack
 /// coloring, you can pass this flag to LLVM.
 static cl::opt<bool>
     DisableInlinedAllocaMerging("disable-inlined-alloca-merging",
                                 cl::init(false), cl::Hidden);
 
+/// A flag for test, so we can print the content of the advisor when running it
+/// as part of the default (e.g. -O3) pipeline.
+static cl::opt<bool> KeepAdvisorForPrinting("keep-inline-advisor-for-printing",
+                                            cl::init(false), cl::Hidden);
+
 extern cl::opt<InlinerFunctionImportStatsOpts> InlinerFunctionImportStats;
 
 static cl::opt<std::string> CGSCCInlineReplayFile(
     "cgscc-inline-replay", cl::init(""), cl::value_desc("filename"),
     cl::desc(
         "Optimization remarks file containing inline remarks to be replayed "
         "by cgscc inlining."),
     cl::Hidden);
PreservedAnalyses InlinerPass::run(LazyCallGraph::SCC &InitialC,
 
   FunctionAnalysisManager &FAM =
       AM.getResult<FunctionAnalysisManagerCGSCCProxy>(InitialC, CG)
           .getManager();
 
   InlineAdvisor &Advisor = getAdvisor(MAMProxy, FAM, M);
   Advisor.onPassEntry();
 
-  auto AdvisorOnExit = make_scope_exit([&] { Advisor.onPassExit(); });
+  auto AdvisorOnExit = make_scope_exit([&] { Advisor.onPassExit(&InitialC); });
 
   // We use a single common worklist for calls across the entire SCC. We
   // process these in-order and append new calls introduced during inlining to
   // the end. The PriorityInlineOrder is optional here, in which the smaller
   // callee would have a higher priority to inline.
   //
   // Note that this particular order of processing is actually critical to
   // avoid very bad behaviors. Consider *highly connected* call graphs where
MPM.addPass(createModuleToPostOrderCGSCCPassAdaptor(
         createDevirtSCCRepeatedPass(std::move(PM), MaxDevirtIterations)));
 
   MPM.addPass(std::move(AfterCGMPM));
   MPM.run(M, MAM);
 
   // Discard the InlineAdvisor, a subsequent inlining session should construct
   // its own.
   auto PA = PreservedAnalyses::all();
+  if (!KeepAdvisorForPrinting)
     PA.abandon<InlineAdvisorAnalysis>();
   return PA;
 }
 
 void InlinerPass::printPipeline(
     raw_ostream &OS, function_ref<StringRef(StringRef)> MapClassName2PassName) {
   static_cast<PassInfoMixin<InlinerPass> *>(this)->printPipeline(
       OS, MapClassName2PassName);
   if (OnlyMandatory)

llvm/test/Transforms/Inline/ML/state-tracking-coro.ll

+; Based on llvm/test/Transforms/Coroutines/coro-split-02.ll
+; Corosplit will keep f1 and add 3 more functions.
+; RUN: opt -passes='default<O1>,print<inline-advisor>' -training-log=/dev/null \
+; RUN:   -S -enable-ml-inliner=development -keep-inline-advisor-for-printing < %s 2>&1 | FileCheck %s
+; REQUIRES: have_tf_api
+;
+; CHECK: [MLInlineAdvisor] Nodes: 4 Edges: 0
+
+%"struct.std::coroutine_handle" = type { i8* }
+%"struct.std::coroutine_handle.0" = type { %"struct.std::coroutine_handle" }
+%"struct.lean_future<int>::Awaiter" = type { i32, %"struct.std::coroutine_handle.0" }
+
+declare i8* @malloc(i64)
+declare void @print(i32)
+
+define void @a() "coroutine.presplit"="1" {
+entry:
+  %ref.tmp7 = alloca %"struct.lean_future<int>::Awaiter", align 8
+  %testval = alloca i32
+  %id = call token @llvm.coro.id(i32 0, i8* null, i8* null, i8* null)
+  %alloc = call i8* @malloc(i64 16) #3
+  %vFrame = call noalias nonnull i8* @llvm.coro.begin(token %id, i8* %alloc)
+
+  %save = call token @llvm.coro.save(i8* null)
+  %Result.i19 = getelementptr inbounds %"struct.lean_future<int>::Awaiter", %"struct.lean_future<int>::Awaiter"* %ref.tmp7, i64 0, i32 0
+  %suspend = call i8 @llvm.coro.suspend(token %save, i1 false)
+  switch i8 %suspend, label %exit [
+    i8 0, label %await.ready
+    i8 1, label %exit
+  ]
+await.ready:
+  %StrayCoroSave = call token @llvm.coro.save(i8* null)
+  %val = load i32, i32* %Result.i19
+  %cast = bitcast i32* %testval to i8*
+  call void @llvm.lifetime.start.p0i8(i64 4, i8* %cast)
+  %test = load i32, i32* %testval
+  call void @print(i32 %test)
+  call void @llvm.lifetime.end.p0i8(i64 4, i8*  %cast)
+  call void @print(i32 %val)
+  br label %exit
+exit:
+  call i1 @llvm.coro.end(i8* null, i1 false)
+  ret void
+}
+
+declare token @llvm.coro.id(i32, i8* readnone, i8* nocapture readonly, i8*)
+declare i1 @llvm.coro.alloc(token) #3
+declare noalias nonnull i8* @"\01??2@YAPEAX_K@Z"(i64) local_unnamed_addr
+declare i64 @llvm.coro.size.i64() #5
+declare i8* @llvm.coro.begin(token, i8* writeonly) #3
+declare void @"\01?puts@@YAXZZ"(...)
+declare token @llvm.coro.save(i8*) #3
+declare i8* @llvm.coro.frame() #5
+declare i8 @llvm.coro.suspend(token, i1) #3
+declare void @"\01??3@YAXPEAX@Z"(i8*) local_unnamed_addr #10
+declare i8* @llvm.coro.free(token, i8* nocapture readonly) #2
+declare i1 @llvm.coro.end(i8*, i1) #3
+declare void @llvm.lifetime.start.p0i8(i64, i8* nocapture) #4
+declare void @llvm.lifetime.end.p0i8(i64, i8* nocapture) #4

llvm/test/Transforms/Inline/ML/state-tracking-scc-splits.ll

+; Based on llvm/test/Other/cgscc-iterate-function-mutation.ll
+; RUN: opt -passes='default<O3>,print<inline-advisor>' -training-log=/dev/null \
+; RUN:   -S -enable-ml-inliner=development -keep-inline-advisor-for-printing < %s 2>&1 | FileCheck %s
+; REQUIRES: have_tf_api
+;
+; CHECK: [MLInlineAdvisor] Nodes: 36 Edges: 0
+
+declare void @readnone() nofree nosync readnone
+declare void @unknown()
+declare void @reference_function_pointer(void()*) nofree nosync readnone
+
+; The @test1_* set of functions checks that when we mutate functions with
+; simplifycfg to delete call edges and this ends up splitting both the SCCs
+; and the RefSCCs that those functions are in, we re-run the CGSCC passes to
+; observe the refined call graph structure.
+
+define void @test1_a() {
+  call void @test1_b1()
+  call void @test1_b2()
+  call void @test1_b3()
+  call void @test1_b4()
+  ret void
+}
+
+define void @test1_b1() {
+  call void @readnone()
+  ret void
+}
+
+define void @test1_b2() {
+  call void @readnone()
+  br i1 false, label %dead, label %exit
+
+dead:
+  call void @test1_a()
+  br label %exit
+
+exit:
+  ret void
+}
+
+define void @test1_b3() {
+  call void @unknown()
+  br i1 false, label %dead, label %exit
+
+dead:
+  call void @test1_a()
+  br label %exit
+
+exit:
+  ret void
+}
+
+define void @test1_b4() {
+  call void @readnone()
+  br i1 false, label %dead, label %exit
+
+dead:
+  call void @test1_a()
+  br label %exit
+
+exit:
+  ret void
+}
+
+define void @test2_a() {
+  call void @test2_b1()
+  call void @test2_b2()
+  call void @test2_b3()
+  call void @test2_b4()
+  ret void
+}
+
+define void @test2_b1() {
+  call void @readnone()
+  ret void
+}
+
+define void @test2_b2() {
+  call void @reference_function_pointer(void()* @test2_a)
+  br i1 false, label %dead, label %exit
+
+dead:
+  call void @test2_a()
+  br label %exit
+
+exit:
+  ret void
+}
+
+define void @test2_b3() {
+  call void @reference_function_pointer(void()* @test2_a)
+  call void @unknown()
+  br i1 false, label %dead, label %exit
+
+dead:
+  call void @test2_a()
+  br label %exit
+
+exit:
+  ret void
+}
+
+define void @test2_b4() {
+  call void @reference_function_pointer(void()* @test2_a)
+  br i1 false, label %dead, label %exit
+
+dead:
+  call void @test2_a()
+  br label %exit
+
+exit:
+  ret void
+}
+
+define void @test3_a() {
+  call void @test3_b11()
+  call void @test3_b21()
+  call void @test3_b31()
+  call void @test3_b41()
+  ret void
+}
+
+define void @test3_b11() {
+  call void @test3_b12()
+  ret void
+}
+
+define void @test3_b12() {
+  call void @test3_b13()
+  ret void
+}
+
+define void @test3_b13() {
+  call void @readnone()
+  ret void
+}
+
+define void @test3_b21() {
+  call void @test3_b22()
+  ret void
+}
+
+define void @test3_b22() {
+  call void @test3_b23()
+  ret void
+}
+
+define void @test3_b23() {
+  call void @readnone()
+  br i1 false, label %dead, label %exit
+
+dead:
+  call void @test3_a()
+  br label %exit
+
+exit:
+  ret void
+}
+
+define void @test3_b31() {
+  call void @test3_b32()
+  ret void
+}
+
+define void @test3_b32() {
+  call void @test3_b33()
+  ret void
+}
+
+define void @test3_b33() {
+  call void @unknown()
+  br i1 false, label %dead, label %exit
+
+dead:
+  call void @test3_a()
+  br label %exit
+
+exit:
+  ret void
+}
+
+define void @test3_b41() {
+  call void @test3_b42()
+  ret void
+}
+
+define void @test3_b42() {
+  call void @test3_b43()
+  ret void
+}
+
+define void @test3_b43() {
+  call void @readnone()
+  br i1 false, label %dead, label %exit
+
+dead:
+  call void @test3_a()
+  br label %exit
+
+exit:
+  ret void
+}
+
+define void @test4_a() {
+  call void @test4_b11()
+  call void @test4_b21()
+  call void @test4_b31()
+  call void @test4_b41()
+  ret void
+}
+
+define void @test4_b11() {
+  call void @test4_b12()
+  ret void
+}
+
+define void @test4_b12() {
+  call void @test4_b13()
+  ret void
+}
+
+define void @test4_b13() {
+  call void @readnone()
+  ret void
+}
+
+define void @test4_b21() {
+  call void @test4_b22()
+  ret void
+}
+
+define void @test4_b22() {
+  call void @test4_b23()
+  ret void
+}
+
+define void @test4_b23() {
+  call void @reference_function_pointer(void()* @test4_a)
+  br i1 false, label %dead, label %exit
+
+dead:
+  call void @test4_a()
+  br label %exit
+
+exit:
+  ret void
+}
+
+define void @test4_b31() {
+  call void @test4_b32()
+  ret void
+}
+
+define void @test4_b32() {
+  call void @test4_b33()
+  ret void
+}
+
+define void @test4_b33() {
+  call void @reference_function_pointer(void()* @test4_a)
+  call void @unknown()
+  br i1 false, label %dead, label %exit
+
+dead:
+  call void @test4_a()
+  br label %exit
+
+exit:
+  ret void
+}
+
+define void @test4_b41() {
+  call void @test4_b42()
+  ret void
+}
+
+define void @test4_b42() {
+  call void @test4_b43()
+  ret void
+}
+
+define void @test4_b43() {
+  call void @reference_function_pointer(void()* @test4_a)
+  br i1 false, label %dead, label %exit
+
+dead:
+  call void @test4_a()
+  br label %exit
+
+exit:
+  ret void
+}

llvm/test/Transforms/Inline/ML/state-tracking.ll

+; RUN: opt -passes='default<O3>,print<inline-advisor>' -training-log=/dev/null \
+; RUN:   -S -enable-ml-inliner=development -keep-inline-advisor-for-printing < %s 2>&1 | FileCheck %s
+; REQUIRES: have_tf_api
+;
+; CHECK: [MLInlineAdvisor] Nodes: 3 Edges: 1
+
+declare i32 @f1()
+
+define i32 @f2() {
+    ret i32 1
+}
+
+define i32 @f3() noinline {
+    ret i32 2
+}
+
+define i32 @f4() {
+    %a = call i32 @f1()
+    %b = call i32 @f2()
+    %c = call i32 @f3()
+    %d = add i32 %a, %b
+    %e = add i32 %d, %c
+    ret i32 %e
+}
 No newline at end of file