2 files changed, 13 insertions, 18 deletions
diff --git a/llvm/lib/Analysis/ModuleSummaryAnalysis.cpp b/llvm/lib/Analysis/ModuleSummaryAnalysis.cpp
index c736c6aa169..5a9d9a6604f 100644
--- a/llvm/lib/Analysis/ModuleSummaryAnalysis.cpp
+++ b/llvm/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) {
diff --git a/llvm/lib/Analysis/ProfileSummaryInfo.cpp b/llvm/lib/Analysis/ProfileSummaryInfo.cpp
index ecfa5c2a3fc..df1414d6811 100644
--- a/llvm/lib/Analysis/ProfileSummaryInfo.cpp
+++ b/llvm/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