[NFC][MLGO] Use LazyCallGraph::Node to track functions.

[mtrofin]

Summary

This avoids the InlineAdvisor carrying the responsibility of deleting
Function objects. We use LazyCallGraph::Node objects instead, which are
stable in memory for the duration of the Module-wide performance of CGSCC
passes started under the same ModuleToPostOrderCGSCCPassAdaptor (which
is the case here)

[mtrofin]

removed supurious include

[aeubanks]

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

408

AFAICT, this is only called in the constructor of DevelopmentModeMLInlineAdvisor where we shouldn't have any deleted functions. and isFunctionDeleted() is only referenced here. is isFunctionDeleted() still necessary?

[mtrofin]

no need for isFunctionDeleted

[mtrofin]

llvm/lib/Analysis/DevelopmentModeInlineAdvisor.cpp

408

yup - thanks! removed it.

[aeubanks]

lg, thanks!

llvm/include/llvm/Analysis/MLInlineAdvisor.h

12

shouldn't need this anymore

[nikic]

llvm/include/llvm/Analysis/MLInlineAdvisor.h

77

Where is DeletedFunctions used now?

[mtrofin]

llvm/include/llvm/Analysis/MLInlineAdvisor.h

12

Right. I'll move it to the .cpp because i'm still using it just to populate some initial state (I want to yank it though if I figure out how)

77

Nowhere - thanks for the catch!

[fhahn]

It looks like this commit may be causing sanitizer errors in the inliner. They started with this commit: https://lab.llvm.org/buildbot/#/builders/5/builds/17176

********************
Testing:  0.. 10.. 20.. 30.. 40.. 50.. 60.. 70..
FAIL: LLVM :: Transforms/Inline/nested-inline.ll (63479 of 81196)
******************** TEST 'LLVM :: Transforms/Inline/nested-inline.ll' FAILED ********************
Script:
--
: 'RUN: at line 1';   /b/sanitizer-x86_64-linux-fast/build/llvm_build_asan/bin/opt < /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/test/Transforms/Inline/nested-inline.ll -inline -S | /b/sanitizer-x86_64-linux-fast/build/llvm_build_asan/bin/FileCheck /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/test/Transforms/Inline/nested-inline.ll
: 'RUN: at line 2';   /b/sanitizer-x86_64-linux-fast/build/llvm_build_asan/bin/opt < /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/test/Transforms/Inline/nested-inline.ll -passes='cgscc(inline)' -S | /b/sanitizer-x86_64-linux-fast/build/llvm_build_asan/bin/FileCheck /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/test/Transforms/Inline/nested-inline.ll
: 'RUN: at line 3';   /b/sanitizer-x86_64-linux-fast/build/llvm_build_asan/bin/opt < /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/test/Transforms/Inline/nested-inline.ll -passes='module-inline' -S | /b/sanitizer-x86_64-linux-fast/build/llvm_build_asan/bin/FileCheck /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/test/Transforms/Inline/nested-inline.ll
--
Exit Code: 1
Command Output (stderr):
--
=================================================================
==52172==ERROR: LeakSanitizer: detected memory leaks
Direct leak of 41 byte(s) in 2 object(s) allocated from:
    #0 0x51ccc8d in operator new(unsigned long) /b/s/w/ir/cache/builder/src/third_party/llvm/compiler-rt/lib/asan/asan_new_delete.cpp:95:3
    #1 0xb01dfcc in Allocate /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/Support/AllocatorBase.h:85:12
    #2 0xb01dfcc in allocateWithKey<llvm::MallocAllocator> /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/StringMapEntry.h:50:32
    #3 0xb01dfcc in Create<llvm::MallocAllocator> /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/StringMapEntry.h:122:17
    #4 0xb01dfcc in llvm::Value::setNameImpl(llvm::Twine const&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/IR/Value.cpp:356:18
    #5 0xb01e3bd in llvm::Value::setName(llvm::Twine const&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/IR/Value.cpp:377:3
    #6 0xadeecd9 in GlobalValue /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/IR/GlobalValue.h:85:5
    #7 0xadeecd9 in GlobalObject /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/IR/GlobalObject.h:46:9
    #8 0xadeecd9 in llvm::Function::Function(llvm::FunctionType*, llvm::GlobalValue::LinkageTypes, unsigned int, llvm::Twine const&, llvm::Module*) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/IR/Function.cpp:383:7
    #9 0x99928d4 in Create /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/IR/Function.h:142:16
    #10 0x99928d4 in llvm::LLParser::parseFunctionHeader(llvm::Function*&, bool) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/AsmParser/LLParser.cpp:5644:8
    #11 0x9985a96 in llvm::LLParser::parseDefine() /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/AsmParser/LLParser.cpp:540:10
    #12 0x997dca8 in llvm::LLParser::parseTopLevelEntities() /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/AsmParser/LLParser.cpp:352:11
    #13 0x997d8eb in llvm::LLParser::Run(bool, llvm::function_ref<llvm::Optional<std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> > > (llvm::StringRef)>) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/AsmParser/LLParser.cpp:81:10
    #14 0x9963a0c in parseAssemblyInto(llvm::MemoryBufferRef, llvm::Module*, llvm::ModuleSummaryIndex*, llvm::SMDiagnostic&, llvm::SlotMapping*, bool, llvm::function_ref<llvm::Optional<std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> > > (llvm::StringRef)>) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/AsmParser/Parser.cpp:36:8
    #15 0x9963dd6 in parseAssemblyInto /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/AsmParser/Parser.cpp:43:10
    #16 0x9963dd6 in llvm::parseAssembly(llvm::MemoryBufferRef, llvm::SMDiagnostic&, llvm::LLVMContext&, llvm::SlotMapping*, llvm::function_ref<llvm::Optional<std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> > > (llvm::StringRef)>) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/AsmParser/Parser.cpp:54:7
    #17 0xb7db7f8 in llvm::parseIR(llvm::MemoryBufferRef, llvm::SMDiagnostic&, llvm::LLVMContext&, llvm::function_ref<llvm::Optional<std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> > > (llvm::StringRef)>) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/IRReader/IRReader.cpp:88:10
    #18 0xb7dcaee in llvm::parseIRFile(llvm::StringRef, llvm::SMDiagnostic&, llvm::LLVMContext&, llvm::function_ref<llvm::Optional<std::__1::basic_string<char, std::__1::char_traits<char>, std::__1::allocator<char> > > (llvm::StringRef)>) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/IRReader/IRReader.cpp:102:10
    #19 0x5227789 in main /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/tools/opt/opt.cpp:628:9
    #20 0x7f853d78209a in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x2409a) (BuildId: 18b9a9a8c523e5cfe5b5d946d605d09242f09798)
Indirect leak of 780 byte(s) in 1 object(s) allocated from:
    #0 0x51922a8 in calloc /b/s/w/ir/cache/builder/src/third_party/llvm/compiler-rt/lib/asan/asan_malloc_linux.cpp:77:3
    #1 0xc9aa451 in llvm::safe_calloc(unsigned long, unsigned long) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/Support/MemAlloc.h:40:18
    #2 0xc9ab9ad in llvm::StringMapImpl::RehashTable(unsigned int) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/Support/StringMap.cpp:219:59
    #3 0xb02ef01 in insert /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/StringMap.h:286:5
    #4 0xb02ef01 in llvm::ValueSymbolTable::reinsertValue(llvm::Value*) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/IR/ValueSymbolTable.cpp:76:12
    #5 0xab97270 in void llvm::SymbolTableListTraits<llvm::Instruction>::setSymTabObject<llvm::Function*>(llvm::Function**, llvm::Function*) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/IR/SymbolTableListTraitsImpl.h:62:16
    #6 0xade9c59 in llvm::SymbolTableListTraits<llvm::BasicBlock>::addNodeToList(llvm::BasicBlock*) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/IR/SymbolTableListTraitsImpl.h:71:6
    #7 0xcb87e64 in insert /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/ilist.h:229:11
    #8 0xcb87e64 in push_back /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/ADT/ilist.h:313:33
    #9 0xcb87e64 in llvm::CloneAndPruneIntoFromInst(llvm::Function*, llvm::Function const*, llvm::Instruction const*, llvm::ValueMap<llvm::Value const*, llvm::WeakTrackingVH, llvm::ValueMapConfig<llvm::Value const*, llvm::sys::SmartMutex<false> > >&, bool, llvm::SmallVectorImpl<llvm::ReturnInst*>&, char const*, llvm::ClonedCodeInfo*) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/Transforms/Utils/CloneFunction.cpp:548:34
    #10 0xcb94c99 in llvm::CloneAndPruneFunctionInto(llvm::Function*, llvm::Function const*, llvm::ValueMap<llvm::Value const*, llvm::WeakTrackingVH, llvm::ValueMapConfig<llvm::Value const*, llvm::sys::SmartMutex<false> > >&, bool, llvm::SmallVectorImpl<llvm::ReturnInst*>&, char const*, llvm::ClonedCodeInfo*) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/Transforms/Utils/CloneFunction.cpp:782:3
    #11 0xcc66eda in llvm::InlineFunction(llvm::CallBase&, llvm::InlineFunctionInfo&, llvm::AAResults*, bool, llvm::Function*) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/Transforms/Utils/InlineFunction.cpp:1949:5
    #12 0xd4f9e97 in llvm::ModuleInlinerPass::run(llvm::Module&, llvm::AnalysisManager<llvm::Module>&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/lib/Transforms/IPO/ModuleInliner.cpp:252:11
    #13 0xd43c4f1 in llvm::detail::PassModel<llvm::Module, llvm::ModuleInlinerPass, llvm::PreservedAnalyses, llvm::AnalysisManager<llvm::Module> >::run(llvm::Module&, llvm::AnalysisManager<llvm::Module>&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/IR/PassManagerInternal.h:88:17
    #14 0xafaf29a in llvm::PassManager<llvm::Module, llvm::AnalysisManager<llvm::Module> >::run(llvm::Module&, llvm::AnalysisManager<llvm::Module>&) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/include/llvm/IR/PassManager.h:525:21
    #15 0x51e2b4e in llvm::runPassPipeline(llvm::StringRef, llvm::Module&, llvm::TargetMachine*, llvm::TargetLibraryInfoImpl*, llvm::ToolOutputFile*, llvm::ToolOutputFile*, llvm::ToolOutputFile*, llvm::StringRef, llvm::ArrayRef<llvm::StringRef>, llvm::opt_tool::OutputKind, llvm::opt_tool::VerifierKind, bool, bool, bool, bool, bool) /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/tools/opt/NewPMDriver.cpp:487:7
    #16 0x522d67a in main /b/sanitizer-x86_64-linux-fast/build/llvm-project/llvm/tools/opt/opt.cpp:818:12
    #17 0x7f853d78209a in __libc_start_main (/lib/x86_64-linux-gnu/libc.so.6+0x2409a) (BuildId: 18b9a9a8c523e5cfe5b5d946d605d09242f09798)

[mtrofin]

Sorry about that, looking

[aeubanks]

In D116964#3237844, @mtrofin wrote:

Sorry about that, looking

fixed in fe827a93f69ddea80308d63a6b54b17106779354

caused by code duplication with the module inliner. is it still in use?

[mtrofin]

thanks, Arthur, for the quick fix!

llvm/include/llvm/Analysis/InlineAdvisor.h

public:
   }
 
   /// Exactly one of the record* APIs must be called. Implementers may extend
   /// behavior by implementing the corresponding record*Impl.
   ///
   /// Call after inlining succeeded, and did not result in deleting the callee.
   void recordInlining();
 
-  /// Call after inlining succeeded, and resulted in deleting the callee.
+  /// Call after inlining succeeded, and results in the callee being
+  /// delete-able, meaning, it has no more users, and will be cleaned up
+  /// subsequently.
   void recordInliningWithCalleeDeleted();
 
   /// Call after the decision for a call site was to not inline.
   void recordUnsuccessfulInlining(const InlineResult &Result) {
     markRecorded();
     recordUnsuccessfulInliningImpl(Result);
   }
 
protected:
   virtual std::unique_ptr<InlineAdvice> getAdviceImpl(CallBase &CB) = 0;
   virtual std::unique_ptr<InlineAdvice> getMandatoryAdvice(CallBase &CB,
                                                            bool Advice);
 
   Module &M;
   FunctionAnalysisManager &FAM;
   std::unique_ptr<ImportedFunctionsInliningStatistics> ImportedFunctionsStats;
 
-  /// We may want to defer deleting functions to after the inlining for a whole
-  /// module has finished. This allows us to reliably use function pointers as
-  /// unique identifiers, as an efficient implementation detail of the
-  /// InlineAdvisor. Otherwise, it is possible the memory allocator
-  /// re-allocate Function objects at the same address of a deleted Function;
-  /// and Functions are potentially created during the function passes called
-  /// after each SCC inlining (e.g. argument promotion does that).
-  void freeDeletedFunctions();
-
-  bool isFunctionDeleted(const Function *F) const {
-    return DeletedFunctions.count(F);
-  }
-
   enum class MandatoryInliningKind { NotMandatory, Always, Never };
 
   static MandatoryInliningKind getMandatoryKind(CallBase &CB,
                                                 FunctionAnalysisManager &FAM,
                                                 OptimizationRemarkEmitter &ORE);
 
   OptimizationRemarkEmitter &getCallerORE(CallBase &CB);
 
 private:
   friend class InlineAdvice;
-  void markFunctionAsDeleted(Function *F);
-  std::unordered_set<const Function *> DeletedFunctions;
 };
 
 /// The default (manual heuristics) implementation of the InlineAdvisor. This
 /// implementation does not need to keep state between inliner pass runs, and is
 /// reusable as-is for inliner pass test scenarios, as well as for regular use.
 class DefaultInlineAdvisor : public InlineAdvisor {
 public:
   DefaultInlineAdvisor(Module &M, FunctionAnalysisManager &FAM,
                        InlineParams Params)
       : InlineAdvisor(M, FAM), Params(Params) {}
 
 private:
   std::unique_ptr<InlineAdvice> getAdviceImpl(CallBase &CB) override;
 
-  void onPassExit() override { freeDeletedFunctions(); }
-
   InlineParams Params;
 };
 
 /// The InlineAdvisorAnalysis is a module pass because the InlineAdvisor
 /// needs to capture state right before inlining commences over a module.
 class InlineAdvisorAnalysis : public AnalysisInfoMixin<InlineAdvisorAnalysis> {
 public:
   static AnalysisKey Key;

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/Analysis/CallGraph.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);
 
-  CallGraph *callGraph() const { return CG.get(); }
   virtual ~MLInlineAdvisor() = default;
 
   void onPassEntry() 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;
+  bool isFunctionDeleted(const Function &F) const;
 
 private:
   int64_t getModuleIRSize() const;
 
   bool Invalid = true;
-  std::unique_ptr<CallGraph> CG;
+  LazyCallGraph &CG;
 
   int64_t NodeCount = 0;
   int64_t EdgeCount = 0;
-  std::map<const Function *, unsigned> FunctionLevels;
+  std::map<const LazyCallGraph::Node *, unsigned> FunctionLevels;
   const int32_t InitialIRSize = 0;
   int32_t CurrentIRSize = 0;
 
   bool ForceStop = false;
+  // Track deleted functions aggressively, because the inliner invalidates nodes
+  // infrequently.
+  DenseSet<const LazyCallGraph::Node *> DeletedFunctions;
 };
 
 /// 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/DevelopmentModeInlineAdvisor.cpp

return std::make_unique<LoggingMLInlineAdvice>(
       /*CB=*/CB, /*ORE=*/ORE, /*Recommendation=*/Recommendation,
       /*Logger=*/*Logger,
       /*CallerSizeEstimateBefore=*/getNativeSizeEstimate(*CB.getCaller()),
       /*CalleeSizeEstimateBefore=*/
       getNativeSizeEstimate(*CB.getCalledFunction()),
       /*DefaultDecision=*/DefaultAdvice);
 }
 
 size_t DevelopmentModeMLInlineAdvisor::getTotalSizeEstimate() {
   if (!InlineSizeEstimatorAnalysis::isEvaluatorRequested())
     return 0;
   size_t Ret = 0;
   for (auto &F : M) {
     if (F.isDeclaration())
       continue;
-    if (isFunctionDeleted(&F))
+    if (isFunctionDeleted(F))
       continue;
     Ret += *getNativeSizeEstimate(F);
   }
   return Ret;
 }
 
 std::unique_ptr<InlineAdvisor> llvm::getDevelopmentModeAdvisor(
     Module &M, ModuleAnalysisManager &MAM,

llvm/lib/Analysis/InlineAdvisor.cpp

 
 InlineAdvice::InlineAdvice(InlineAdvisor *Advisor, CallBase &CB,
                            OptimizationRemarkEmitter &ORE,
                            bool IsInliningRecommended)
     : Advisor(Advisor), Caller(CB.getCaller()), Callee(CB.getCalledFunction()),
       DLoc(CB.getDebugLoc()), Block(CB.getParent()), ORE(ORE),
       IsInliningRecommended(IsInliningRecommended) {}
 
-void InlineAdvisor::markFunctionAsDeleted(Function *F) {
-  assert((!DeletedFunctions.count(F)) &&
-         "Cannot put cause a function to become dead twice!");
-  DeletedFunctions.insert(F);
-}
-
-void InlineAdvisor::freeDeletedFunctions() {
-  for (auto *F : DeletedFunctions)
-    delete F;
-  DeletedFunctions.clear();
-}
-
 void InlineAdvice::recordInlineStatsIfNeeded() {
   if (Advisor->ImportedFunctionsStats)
     Advisor->ImportedFunctionsStats->recordInline(*Caller, *Callee);
 }
 
 void InlineAdvice::recordInlining() {
   markRecorded();
   recordInlineStatsIfNeeded();
   recordInliningImpl();
 }
 
 void InlineAdvice::recordInliningWithCalleeDeleted() {
   markRecorded();
   recordInlineStatsIfNeeded();
-  Advisor->markFunctionAsDeleted(Callee);
   recordInliningWithCalleeDeletedImpl();
 }
 
 AnalysisKey InlineAdvisorAnalysis::Key;
 
 bool InlineAdvisorAnalysis::Result::tryCreate(
     InlineParams Params, InliningAdvisorMode Mode,
     const ReplayInlinerSettings &ReplaySettings) {

 }
 
 InlineAdvisor::~InlineAdvisor() {
   if (ImportedFunctionsStats) {
     assert(InlinerFunctionImportStats != InlinerFunctionImportStatsOpts::No);
     ImportedFunctionsStats->dump(InlinerFunctionImportStats ==
                                  InlinerFunctionImportStatsOpts::Verbose);
   }
-
-  freeDeletedFunctions();
 }
 
 std::unique_ptr<InlineAdvice> InlineAdvisor::getMandatoryAdvice(CallBase &CB,
                                                                 bool Advice) {
   return std::make_unique<MandatoryInlineAdvice>(this, CB, getCallerORE(CB),
                                                  Advice);
 }
 

llvm/lib/Analysis/MLInlineAdvisor.cpp

 #include "llvm/Config/config.h"
 #if defined(LLVM_HAVE_TF_AOT) || defined(LLVM_HAVE_TF_API)
 
 #include <limits>
 #include <unordered_map>
 #include <unordered_set>
 
 #include "llvm/ADT/SCCIterator.h"
-#include "llvm/Analysis/CallGraph.h"
 #include "llvm/Analysis/FunctionPropertiesAnalysis.h"
 #include "llvm/Analysis/InlineCost.h"
+#include "llvm/Analysis/LazyCallGraph.h"
 #include "llvm/Analysis/MLInlineAdvisor.h"
 #include "llvm/Analysis/MLModelRunner.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/InstIterator.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/PassManager.h"
if (auto *CS = dyn_cast<CallBase>(&I))
     }
   return nullptr;
 }
 
 MLInlineAdvisor::MLInlineAdvisor(Module &M, ModuleAnalysisManager &MAM,
                                  std::unique_ptr<MLModelRunner> Runner)
     : InlineAdvisor(
           M, MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager()),
-      ModelRunner(std::move(Runner)), CG(new CallGraph(M)),
+      ModelRunner(std::move(Runner)),
+      CG(MAM.getResult<LazyCallGraphAnalysis>(M)),
       InitialIRSize(getModuleIRSize()), CurrentIRSize(InitialIRSize) {
   assert(ModelRunner);
 
   // Extract the 'call site height' feature - the position of a call site
   // relative to the farthest statically reachable SCC node. We don't mutate
   // this value while inlining happens. Empirically, this feature proved
   // critical in behavioral cloning - i.e. training a model to mimic the manual
   // heuristic's decisions - and, thus, equally important for training for
   // improvement.
-  for (auto I = scc_begin(CG.get()); !I.isAtEnd(); ++I) {
+  CallGraph CGraph(M);
+  for (auto I = scc_begin(&CGraph); !I.isAtEnd(); ++I) {
     const std::vector<CallGraphNode *> &CGNodes = *I;
     unsigned Level = 0;
     for (auto *CGNode : CGNodes) {
       Function *F = CGNode->getFunction();
       if (!F || F->isDeclaration())
         continue;
       for (auto &I : instructions(F)) {
         if (auto *CS = getInlinableCS(I)) {
           auto *Called = CS->getCalledFunction();
-          auto Pos = FunctionLevels.find(Called);
+          auto Pos = FunctionLevels.find(&CG.get(*Called));
           // In bottom up traversal, an inlinable callee is either in the
           // same SCC, or to a function in a visited SCC. So not finding its
           // level means we haven't visited it yet, meaning it's in this SCC.
           if (Pos == FunctionLevels.end())
             continue;
           Level = std::max(Level, Pos->second + 1);
         }
       }
     }
     for (auto *CGNode : CGNodes) {
       Function *F = CGNode->getFunction();
       if (F && !F->isDeclaration())
-        FunctionLevels[F] = Level;
+        FunctionLevels[&CG.get(*F)] = Level;
     }
   }
 }
 
+unsigned MLInlineAdvisor::getInitialFunctionLevel(const Function &F) const {
+  return CG.lookup(F) ? FunctionLevels.at(CG.lookup(F)) : 0;
+}
+
+bool MLInlineAdvisor::isFunctionDeleted(const Function &F) const {
+  const auto *N = CG.lookup(F);
+  assert(N);
+  return DeletedFunctions.find(N) != DeletedFunctions.end();
+}
+
 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)

 // 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();
+  if (CalleeWasDeleted) {
+    auto R = DeletedFunctions.insert(&CG.get(*Callee));
+    (void)R;
+    assert(R.second);
+  }
 
   // The caller features aren't valid anymore.
   {
     PreservedAnalyses PA = PreservedAnalyses::all();
     PA.abandon<FunctionPropertiesAnalysis>();
     FAM.invalidate(*Caller, PA);
   }
   int64_t IRSizeAfter =
NewCallerAndCalleeEdges +=
         FAM.getResult<FunctionPropertiesAnalysis>(*Callee)
             .DirectCallsToDefinedFunctions;
   EdgeCount += (NewCallerAndCalleeEdges - Advice.CallerAndCalleeEdges);
   assert(CurrentIRSize >= 0 && EdgeCount >= 0 && NodeCount >= 0);
 }
 
 int64_t MLInlineAdvisor::getModuleIRSize() const {
   int64_t Ret = 0;
-  for (auto &F : CG->getModule())
+  for (auto &F : M)
     if (!F.isDeclaration())
       Ret += getIRSize(F);
   return Ret;
 }
 
 std::unique_ptr<InlineAdvice> MLInlineAdvisor::getAdviceImpl(CallBase &CB) {
   auto &Caller = *CB.getCaller();
   auto &Callee = *CB.getCalledFunction();
std::unique_ptr<InlineAdvice> MLInlineAdvisor::getAdviceImpl(CallBase &CB) {
   }
 
   auto &CallerBefore = FAM.getResult<FunctionPropertiesAnalysis>(Caller);
   auto &CalleeBefore = FAM.getResult<FunctionPropertiesAnalysis>(Callee);
 
   *ModelRunner->getTensor<int64_t>(FeatureIndex::CalleeBasicBlockCount) =
       CalleeBefore.BasicBlockCount;
   *ModelRunner->getTensor<int64_t>(FeatureIndex::CallSiteHeight) =
-      FunctionLevels[&Caller];
+      getInitialFunctionLevel(Caller);
   *ModelRunner->getTensor<int64_t>(FeatureIndex::NodeCount) = NodeCount;
   *ModelRunner->getTensor<int64_t>(FeatureIndex::NrCtantParams) = NrCtantParams;
   *ModelRunner->getTensor<int64_t>(FeatureIndex::EdgeCount) = EdgeCount;
   *ModelRunner->getTensor<int64_t>(FeatureIndex::CallerUsers) =
       CallerBefore.Uses;
   *ModelRunner->getTensor<int64_t>(
       FeatureIndex::CallerConditionallyExecutedBlocks) =
       CallerBefore.BlocksReachedFromConditionalInstruction;

llvm/lib/Transforms/IPO/Inliner.cpp

for (Function *DeadF : DeadFunctions) {
     UR.InvalidatedSCCs.insert(&DeadC);
     UR.InvalidatedRefSCCs.insert(&DeadRC);
 
     // If the updated SCC was the one containing the deleted function, clear it.
     if (&DeadC == UR.UpdatedC)
       UR.UpdatedC = nullptr;
 
     // And delete the actual function from the module.
-    // The Advisor may use Function pointers to efficiently index various
-    // internal maps, e.g. for memoization. Function cleanup passes like
-    // argument promotion create new functions. It is possible for a new
-    // function to be allocated at the address of a deleted function. We could
-    // index using names, but that's inefficient. Alternatively, we let the
-    // Advisor free the functions when it sees fit.
-    DeadF->getBasicBlockList().clear();
-    M.getFunctionList().remove(DeadF);
+    M.getFunctionList().erase(DeadF);
 
     ++NumDeleted;
   }
 
   if (!Changed)
     return PreservedAnalyses::all();
 
   PreservedAnalyses PA;