1 files changed, 32 insertions, 9 deletions

diff --git a/llvm/lib/Transforms/IPO/HotColdSplitting.cpp b/llvm/lib/Transforms/IPO/HotColdSplitting.cpp
index 710a9e72569..0ff902d735c 100644
--- a/llvm/lib/Transforms/IPO/HotColdSplitting.cpp
+++ b/llvm/lib/Transforms/IPO/HotColdSplitting.cpp
@@ -168,8 +168,11 @@ public:
   HotColdSplitting(ProfileSummaryInfo *ProfSI,
                    function_ref<BlockFrequencyInfo *(Function &)> GBFI,
                    function_ref<TargetTransformInfo &(Function &)> GTTI,
+                   function_ref<DominatorTree *(Function &)> GetDT,
+                   function_ref<PostDominatorTree *(Function &)> GetPDT,
                    std::function<OptimizationRemarkEmitter &(Function &)> *GORE)
-      : PSI(ProfSI), GetBFI(GBFI), GetTTI(GTTI), GetORE(GORE) {}
+      : PSI(ProfSI), GetBFI(GBFI), GetTTI(GTTI), GetDT(GetDT), GetPDT(GetPDT),
+        GetORE(GORE) {}
   bool run(Module &M);
 
 private:
@@ -182,6 +185,8 @@ private:
   ProfileSummaryInfo *PSI;
   function_ref<BlockFrequencyInfo *(Function &)> GetBFI;
   function_ref<TargetTransformInfo &(Function &)> GetTTI;
+  function_ref<DominatorTree *(Function &)> GetDT;
+  function_ref<PostDominatorTree *(Function &)> GetPDT;
   std::function<OptimizationRemarkEmitter &(Function &)> *GetORE;
 };
 
@@ -195,6 +200,8 @@ public:
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<AssumptionCacheTracker>();
     AU.addRequired<BlockFrequencyInfoWrapperPass>();
+    AU.addRequired<DominatorTreeWrapperPass>();
+    AU.addRequired<PostDominatorTreeWrapperPass>();
     AU.addRequired<ProfileSummaryInfoWrapperPass>();
     AU.addRequired<TargetTransformInfoWrapperPass>();
   }
@@ -462,12 +469,11 @@ bool HotColdSplitting::outlineColdRegions(Function &F, bool HasProfileSummary) {
   ReversePostOrderTraversal<Function *> RPOT(&F);
 
   // Calculate domtrees lazily. This reduces compile-time significantly.
-  std::unique_ptr<DominatorTree> DT;
-  std::unique_ptr<PostDominatorTree> PDT;
+  DominatorTree *DT = nullptr;
+  PostDominatorTree *PDT = nullptr;
 
   // Calculate BFI lazily (it's only used to query ProfileSummaryInfo). This
-  // reduces compile-time significantly. TODO: When we *do* use BFI, we should
-  // be able to salvage its domtrees instead of recomputing them.
+  // reduces compile-time significantly.
   BlockFrequencyInfo *BFI = nullptr;
   if (HasProfileSummary)
     BFI = GetBFI(F);
@@ -492,9 +498,9 @@ bool HotColdSplitting::outlineColdRegions(Function &F, bool HasProfileSummary) {
     });
 
     if (!DT)
-      DT = make_unique<DominatorTree>(F);
+      DT = GetDT(F);
     if (!PDT)
-      PDT = make_unique<PostDominatorTree>(F);
+      PDT = GetPDT(F);
 
     auto Region = OutliningRegion::create(*BB, *DT, *PDT);
     if (Region.empty())
@@ -595,6 +601,13 @@ bool HotColdSplittingLegacyPass::runOnModule(Module &M) {
   auto GBFI = [this](Function &F) {
     return &this->getAnalysis<BlockFrequencyInfoWrapperPass>(F).getBFI();
   };
+  auto GetDT = [this](Function &F) {
+    return &this->getAnalysis<DominatorTreeWrapperPass>(F).getDomTree();
+  };
+  auto GetPDT = [this](Function &F) {
+    return &this->getAnalysis<PostDominatorTreeWrapperPass>(F).getPostDomTree();
+  };
+
   std::unique_ptr<OptimizationRemarkEmitter> ORE;
   std::function<OptimizationRemarkEmitter &(Function &)> GetORE =
       [&ORE](Function &F) -> OptimizationRemarkEmitter & {
@@ -602,7 +615,7 @@ bool HotColdSplittingLegacyPass::runOnModule(Module &M) {
     return *ORE.get();
   };
 
-  return HotColdSplitting(PSI, GBFI, GTTI, &GetORE).run(M);
+  return HotColdSplitting(PSI, GBFI, GTTI, GetDT, GetPDT, &GetORE).run(M);
 }
 
 PreservedAnalyses
@@ -623,6 +636,14 @@ HotColdSplittingPass::run(Module &M, ModuleAnalysisManager &AM) {
     return FAM.getResult<TargetIRAnalysis>(F);
   };
 
+  auto GetDT = [&FAM](Function &F) {
+    return &FAM.getResult<DominatorTreeAnalysis>(F);
+  };
+
+  auto GetPDT = [&FAM](Function &F) {
+    return &FAM.getResult<PostDominatorTreeAnalysis>(F);
+  };
+
   std::unique_ptr<OptimizationRemarkEmitter> ORE;
   std::function<OptimizationRemarkEmitter &(Function &)> GetORE =
       [&ORE](Function &F) -> OptimizationRemarkEmitter & {
@@ -632,7 +653,7 @@ HotColdSplittingPass::run(Module &M, ModuleAnalysisManager &AM) {
 
   ProfileSummaryInfo *PSI = &AM.getResult<ProfileSummaryAnalysis>(M);
 
-  if (HotColdSplitting(PSI, GBFI, GTTI, &GetORE).run(M))
+  if (HotColdSplitting(PSI, GBFI, GTTI, GetDT, GetPDT, &GetORE).run(M))
     return PreservedAnalyses::none();
   return PreservedAnalyses::all();
 }
@@ -642,6 +663,8 @@ INITIALIZE_PASS_BEGIN(HotColdSplittingLegacyPass, "hotcoldsplit",
                       "Hot Cold Splitting", false, false)
 INITIALIZE_PASS_DEPENDENCY(ProfileSummaryInfoWrapperPass)
 INITIALIZE_PASS_DEPENDENCY(BlockFrequencyInfoWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)
+INITIALIZE_PASS_DEPENDENCY(PostDominatorTreeWrapperPass)
 INITIALIZE_PASS_END(HotColdSplittingLegacyPass, "hotcoldsplit",
                     "Hot Cold Splitting", false, false)