Refactor the ProfileSummaryInfo to use doInitialization and doFinalization to handle Module update.

Summary: This refactors the change in r282616

Reviewers: davidxl, eraman, mehdi_amini

Subscribers: mehdi_amini, davide, llvm-commits

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

git-svn-id: https://llvm.org/svn/llvm-project/llvm/trunk@282630 91177308-0d34-0410-b5e6-96231b3b80d8

include/llvm/Analysis/ProfileSummaryInfo.h

@@ -40,7 +40,7 @@ class ProfileSummary;
 // units. This would require making this depend on BFI.
 class ProfileSummaryInfo {
 private:
-  Module *M;
+  Module &M;
   std::unique_ptr<ProfileSummary> Summary;
   void computeSummary();
   void computeThresholds();
@@ -48,7 +48,7 @@ private:
   Optional<uint64_t> HotCountThreshold, ColdCountThreshold;
 
 public:
-  ProfileSummaryInfo(Module *M) : M(M) {}
+  ProfileSummaryInfo(Module &M) : M(M) {}
   ProfileSummaryInfo(ProfileSummaryInfo &&Arg)
       : M(Arg.M), Summary(std::move(Arg.Summary)) {}
   /// \brief Returns true if \p F is a hot function.
@@ -59,8 +59,6 @@ public:
   bool isHotCount(uint64_t C);
   /// \brief Returns true if count \p C is considered cold.
   bool isColdCount(uint64_t C);
-  /// \brief Checks if \p NewM is up-to-date, if not, invalidate Summary.
-  void resetModule(Module *NewM);
 };
 
 /// An analysis pass based on legacy pass manager to deliver ProfileSummaryInfo.
@@ -71,7 +69,12 @@ public:
   static char ID;
   ProfileSummaryInfoWrapperPass();
 
-  ProfileSummaryInfo *getPSI(Module &M);
+  ProfileSummaryInfo *getPSI() {
+    return &*PSI;
+  }
+
+  bool doInitialization(Module &M) override;
+  bool doFinalization(Module &M) override;
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.setPreservesAll();
   }

lib/Analysis/ModuleSummaryAnalysis.cpp

@@ -228,7 +228,7 @@ ModuleSummaryIndexWrapperPass::ModuleSummaryIndexWrapperPass()
 }
 
 bool ModuleSummaryIndexWrapperPass::runOnModule(Module &M) {
-  auto &PSI = *getAnalysis<ProfileSummaryInfoWrapperPass>().getPSI(M);
+  auto &PSI = *getAnalysis<ProfileSummaryInfoWrapperPass>().getPSI();
   Index = buildModuleSummaryIndex(
       M,
       [this](const Function &F) {

lib/Analysis/ProfileSummaryInfo.cpp

@@ -57,7 +57,7 @@ static uint64_t getMinCountForPercentile(SummaryEntryVector &DS,
 void ProfileSummaryInfo::computeSummary() {
   if (Summary)
     return;
-  auto *SummaryMD = M->getProfileSummary();
+  auto *SummaryMD = M.getProfileSummary();
   if (!SummaryMD)
     return;
   Summary.reset(ProfileSummary::getFromMD(SummaryMD));
@@ -113,13 +113,6 @@ void ProfileSummaryInfo::computeThresholds() {
       getMinCountForPercentile(DetailedSummary, ProfileSummaryCutoffCold);
 }
 
-void ProfileSummaryInfo::resetModule(Module *NewM) {
-  if (NewM == M)
-    return;
-  M = NewM;
-  Summary.reset(nullptr);
-}
-
 bool ProfileSummaryInfo::isHotCount(uint64_t C) {
   if (!HotCountThreshold)
     computeThresholds();
@@ -132,14 +125,6 @@ bool ProfileSummaryInfo::isColdCount(uint64_t C) {
   return ColdCountThreshold && C <= ColdCountThreshold.getValue();
 }
 
-ProfileSummaryInfo *ProfileSummaryInfoWrapperPass::getPSI(Module &M) {
-  if (!PSI)
-    PSI.reset(new ProfileSummaryInfo(&M));
-  else
-    PSI->resetModule(&M);
-  return PSI.get();
-}
-
 INITIALIZE_PASS(ProfileSummaryInfoWrapperPass, "profile-summary-info",
                 "Profile summary info", false, true)
 
@@ -148,10 +133,20 @@ ProfileSummaryInfoWrapperPass::ProfileSummaryInfoWrapperPass()
   initializeProfileSummaryInfoWrapperPassPass(*PassRegistry::getPassRegistry());
 }
 
+bool ProfileSummaryInfoWrapperPass::doInitialization(Module &M) {
+  PSI.reset(new ProfileSummaryInfo(M));
+  return false;
+}
+
+bool ProfileSummaryInfoWrapperPass::doFinalization(Module &M) {
+  PSI.reset();
+  return false;
+}
+
 char ProfileSummaryAnalysis::PassID;
 ProfileSummaryInfo ProfileSummaryAnalysis::run(Module &M,
                                                ModuleAnalysisManager &) {
-  return ProfileSummaryInfo(&M);
+  return ProfileSummaryInfo(M);
 }
 
 // FIXME: This only tests isHotFunction and isColdFunction and not the

lib/LTO/ThinLTOCodeGenerator.cpp

@@ -378,7 +378,7 @@ ProcessThinLTOModule(Module &TheModule, ModuleSummaryIndex &Index,
     SmallVector<char, 128> OutputBuffer;
     {
       raw_svector_ostream OS(OutputBuffer);
-      ProfileSummaryInfo PSI(&TheModule);
+      ProfileSummaryInfo PSI(TheModule);
       auto Index = buildModuleSummaryIndex(TheModule, nullptr, nullptr);
       WriteBitcodeToFile(&TheModule, OS, true, &Index);
     }

lib/Transforms/IPO/Inliner.cpp

@@ -634,7 +634,7 @@ inlineCallsImpl(CallGraphSCC &SCC, CallGraph &CG,
 bool Inliner::inlineCalls(CallGraphSCC &SCC) {
   CallGraph &CG = getAnalysis<CallGraphWrapperPass>().getCallGraph();
   ACT = &getAnalysis<AssumptionCacheTracker>();
-  PSI = getAnalysis<ProfileSummaryInfoWrapperPass>().getPSI(CG.getModule());
+  PSI = getAnalysis<ProfileSummaryInfoWrapperPass>().getPSI();
   auto &TLI = getAnalysis<TargetLibraryInfoWrapperPass>().getTLI();
   // We compute dedicated AA results for each function in the SCC as needed. We
   // use a lambda referencing external objects so that they live long enough to