[HotColdSplitting] Outline more than once per function

[vsk]

Summary

Algorithm: Identify maximal cold regions and put them in a worklist. If
a candidate region overlaps with another, discard it. While the worklist
is full, remove a single-entry sub-region from the worklist and attempt
to outline it. By the non-overlap property, this should not invalidate
parts of the domtree pertaining to other outlining regions.

Testing: LNT results on X86 are clean. With test-suite + externals, llvm
outlines 134KB pre-patch, and 383KB post-patch (+ ~2.8x). The file
483.xalancbmk/src/Constants.cpp stands out as an extreme case where llvm
outlines over 100 times in some functions (mostly EH paths). There was
not a significant performance impact pre vs. post-patch.

[hiraditya]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

348

I think, this can go before the previous check.

489

Add optimize for size attribute?

[vsk]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

348

I think that would cause splitting to fail when a predecessor outside of the cold region is the entry block. Am I missing something?

489

Good point.

[hiraditya]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

480

Creating scope '{}' may not be required.

517

Maybe print the region which didn't get outlined, or at least the root-block.

[hiraditya]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

348

When SinkPostDom is true and PredBB is the entry, we can just return.

[tejohnson]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

522

The last argument should be changed from a constant 1 to a running count of outlined regions. Otherwise all outlined functions will have the same name.

[sebpop]

Maybe you can add the testcase from my previous patch: https://reviews.llvm.org/D53588

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

378

s/anu/any/

517

I think that would generate too much debug output. Maybe completely remove the LLVM_DEBUG stmt.

[vsk]

• While doing performance testing I found a miscompile in ./SingleSource/Regression/C++/EH/Regression-C++-class_hierarchy. I'll file a PR with more details by next week. It looks like it could be an existing bug that surfaces due to more aggressive outlining.
• Added tests (including the one from @sebpop's earlier patch).

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

348

I see what you mean.

480

It's not, I just find it easier to read multiple statements when there are surrounding curly braces.

517

I think some form of debug statement here would help to tune the outlining threshold. Could we keep it for now? I'll remove some of the other redundant debug statements to keep the output reasonable.

522

Thanks for pointing this out.

[hiraditya]

Can we push this patch? This is not enabled by default so we can continue development in subsequent patches.

[vsk]

In D53887#1286544, @hiraditya wrote:

Can we push this patch? This is not enabled by default so we can continue development in subsequent patches.

I'd like to form a plan for addressing llvm.org/PR39545 first, to avoid inadvertently leaving the EH regression in tree.

Edit: Starting to think that the underlying bug in PR39545 is faulty lowering of llvm.eh.typeid.for in outlined thunks. At a conceptual level it shouldn't be hard to fix. I really think we should tackle that first as looks like a pretty bad miscompile affecting all in-tree clients of CodeExtractor.

[junbuml]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

159

This function set MinSize, so I think the function name should be something like markMinSize().

402

We should have additional check before adding SuccBB into Blocks (maybe successors of SuccBB). If a dom-frontier have a phi taking different incoming values from multiple cold blocks, it will assert. In the example below, %if.end takes different incoming values from %coldbb and %coldbb2.

define void @foo(i32 %cond) {
entry:
  %tobool = icmp eq i32 %cond, 0
  br i1 %tobool, label %if.end, label %coldbb

coldbb:
  call void (...) @sink()
  br i1 undef, label %if.end, label %coldbb2

coldbb2:
  br label %if.end

if.end:
  %p = phi i32 [0, %entry], [1, %coldbb], [3, %coldbb2]
  ret void
}

declare void @sink(...) cold

503

I don't think it's good idea to add MinSize in this case. It is possible that a hot function or main() itself can be successfully ended up with exit(0).

[hiraditya]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

503

The static analysis here is pretty conservative, unless the runtime-profile information is broken, it seems unlikely that this pass will mark a hot function as cold.

[junbuml]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

503

I agree that this pass will unlikely mark hot blocks as cold. But, in line 509~511, it mark MinSize on the function, if all blocks are post-domed by a single cold block (e.g., mark-the-whole-func-cold.ll). In the example below, we may not want to mark MinSize on main().

void main() {
  // .. 
  // hot loop .. 
  // 
  exit (0);
}

[junbuml]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

274

Considering the case below, I believed we should not mark the MinSize here as well.

void foo(bool c) {
  if(c) {
    ....
    return;
  }

   // hot code
   exit (0);
}

[vsk]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

159

That's true, but as a follow-up, I'd like to have this helper add Attribute::Cold, and move the function to the end of the text segment. Would it be all right to keep the aspirational name?

274

I think this is illustrative of a more general problem: I don't think we should treat NoReturn functions as cold at all. On Darwin, functions like longjmp are NoReturn but warm. Our kernel uses continuation functions (conceptually similar to longjmp) to reduce stack usage: these are also NoReturn but hot.

IMO, noreturn functions which are actually cold should be marked as such (abort, etc., but not exit). That's what we do on Darwin.

402

As noted in llvm.org/PR39545, I'll send out a separate review to address this.

[vsk]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

274

If the concern is that it's hard to add attributes to functions in libc, we could special-case them as needed. But I think treating all noreturn functions as cold when they might not be will cause problems, or at least prevent us from attaching MinSize, etc.

[junbuml]

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

274

IMHO, I doubt if unlikelyExecuted() is conservative enough without profile data. Except a specific cold mark(Attribute::Cold), all other checks may not perfectly clear to say the blocks are really cold.

[vsk]

Removing the WIP label, as I'm confident in the results now.

• Fix an iterator invalidation bug in takeSingleEntrySubRegion.
• I've addressed the immediate miscompile reported in llvm.org/PR39545, and plan on improving handling of eh.typeid.for in a follow-up.
• @junbuml has concerns about applying MinSize to functions which call exit. IMHO the right fix is to not treat noreturn calls as cold (also as a follow-up).

[junbuml]

IMHO the right fix is to not treat noreturn calls as cold (also as a follow-up).

I'm not sure if handling noreturn is a right fix. A block containing exit(0) will have "unreachable", so it must be still considered as a cold block even after removing noreturn from unlikelyExecuted().

[vsk]

In D53887#1290571, @junbuml wrote:

IMHO the right fix is to not treat noreturn calls as cold (also as a follow-up).

I'm not sure if handling noreturn is a right fix. A block containing exit(0) will have "unreachable", so it must be still considered as a cold block even after removing noreturn from unlikelyExecuted().

I had in mind:

if (blockEndsInUnreachable(BB)) {
  // Calls to noreturn functions are followed by an unreachable inst, but
  // the call itself may be warm (e.g. longjmp, or exit).
  if (auto *CI =
          dyn_cast_or_null<CallInst>(BB.getTerminator()->getPrevNode()))
    if (CI->hasFnAttr(Attribute::NoReturn))
      return false;
  return true;
}

[junbuml]

I guess it would be good to give some stress test on this pass to see if there is any hidden bug by relaxing conditions in unlikelyExecuted(). As we treat more blocks as cold without being limited on unlikelyExecuted(), we maybe able to expose hidden issues with it.

[vsk]

In D53887#1290584, @junbuml wrote:

I guess it would be good to give some stress test on this pass to see if there is any hidden bug by relaxing conditions in unlikelyExecuted(). As we treat more blocks as cold without being limited on unlikelyExecuted(), we maybe able to expose hidden issues with it.

That's a great idea. I'll do that now with the test-suite + externals. FWIW, I also built all of iOS with this pass enabled, and the only compiler crash I found was https://bugs.llvm.org/show_bug.cgi?id=39564. I'm still investigating possible miscompiles.

[vsk]

In D53887#1290587, @vsk wrote:

In D53887#1290584, @junbuml wrote:

I guess it would be good to give some stress test on this pass to see if there is any hidden bug by relaxing conditions in unlikelyExecuted(). As we treat more blocks as cold without being limited on unlikelyExecuted(), we maybe able to expose hidden issues with it.

That's a great idea. I'll do that now with the test-suite + externals. FWIW, I also built all of iOS with this pass enabled, and the only compiler crash I found was https://bugs.llvm.org/show_bug.cgi?id=39564. I'm still investigating possible miscompiles.

So, I defined unlikelyExecuted(BB) = true, and found a bug while running through the test-suite. The issue is that two do two DFS's (one on predecessor blocks, one on successor blocks). The second DFS can mark a block already marked by the first DFS. In practice this isn't a problem, because CodeExtractor maintains a set of blocks. But it's weird to have duplicate blocks in the extraction list, and it's a simple issue to fix.

Other than that, I found no miscompiles. The test suite passed cleanly.

Edit: the unlikelyExecuted(BB) = true experiment is superseded by the one described below.

In D53887#1290587, @vsk wrote:

In D53887#1290584, @junbuml wrote:

I guess it would be good to give some stress test on this pass to see if there is any hidden bug by relaxing conditions in unlikelyExecuted(). As we treat more blocks as cold without being limited on unlikelyExecuted(), we maybe able to expose hidden issues with it.

That's a great idea. I'll do that now with the test-suite + externals. FWIW, I also built all of iOS with this pass enabled, and the only compiler crash I found was https://bugs.llvm.org/show_bug.cgi?id=39564. I'm still investigating possible miscompiles.

So, I defined bool Cold = !pred_empty(BB), and this resulted in 29MB of text being outlined in test suite (out of 43MB total). I found a small issue in D54189, but no actual miscompilations. All tests passed and matched the reference output.

[vsk]

Ping, are there any outstanding concerns about this one? It'd be nice to have this in-tree, as I have a few follow-ups based on it.

[vsk]

Friendly ping. I've rebased this on top of r348205, which fixes the assertion failure pointed out in llvm.org/PR39564.

I've stress-tested this by:

• Building LNT+externals with hot/cold splitting enabled. I forced outlining to occur whenever a block has more than 1 predecessor, so long as it wouldn't result in the entire function being outlined. All output validation tests still passed.
• Running check-llvm in a stage2 build with hot/cold splitting enabled in the same way described above, but with stack coloring disabled due to llvm.org/PR39671.

[hiraditya]

LGTM, if there are outstanding comments from other reviewers we can address them in subsequent patches.

llvm/lib/Transforms/IPO/HotColdSplitting.cpp

 //===- HotColdSplitting.cpp -- Outline Cold Regions -------------*- C++ -*-===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // Outline cold regions to a separate function.
 // TODO: Update BFI and BPI
 // TODO: Add all the outlined functions to a separate section.
 //
 //===----------------------------------------------------------------------===//
 
+#include "llvm/ADT/PostOrderIterator.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/AliasAnalysis.h"
 #include "llvm/Analysis/BlockFrequencyInfo.h"
 #include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/CFG.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/PostDominators.h"
static bool unlikelyExecuted(const BasicBlock &BB) {
   return false;
 }
 
 /// Check whether it's safe to outline \p BB.
 static bool mayExtractBlock(const BasicBlock &BB) {
   return !BB.hasAddressTaken();
 }
 
-/// Check whether \p BB is profitable to outline (i.e. its code size cost meets
-/// the threshold set in \p MinOutliningThreshold).
-static bool isProfitableToOutline(const BasicBlock &BB,
+/// Check whether \p Region is profitable to outline.
+static bool isProfitableToOutline(const BlockSequence &Region,
                                   TargetTransformInfo &TTI) {
+  if (Region.size() > 1)
+    return true;
+
   int Cost = 0;
+  const BasicBlock &BB = *Region[0];
   for (const Instruction &I : BB) {
     if (isa<DbgInfoIntrinsic>(&I) || &I == BB.getTerminator())
       continue;
 
     Cost += TTI.getInstructionCost(&I, TargetTransformInfo::TCK_CodeSize);
 
     if (Cost >= (MinOutliningThreshold * TargetTransformInfo::TCC_Basic))
       return true;
   }
   return false;
 }
 
-/// Identify the maximal region of cold blocks which includes \p SinkBB.
-///
-/// Include all blocks post-dominated by \p SinkBB, \p SinkBB itself, and all
-/// blocks dominated by \p SinkBB. Exclude all other blocks, and blocks which
-/// cannot be outlined.
-///
-/// Return an empty sequence if the cold region is too small to outline, or if
-/// the cold region has no warm predecessors.
-static BlockSequence findMaximalColdRegion(BasicBlock &SinkBB,
-                                           TargetTransformInfo &TTI,
-                                           DominatorTree &DT,
-                                           PostDomTree &PDT) {
-  // The maximal cold region.
-  BlockSequence ColdRegion = {};
-
-  // The ancestor farthest-away from SinkBB, and also post-dominated by it.
-  BasicBlock *MaxAncestor = &SinkBB;
-  unsigned MaxAncestorHeight = 0;
-
-  // Visit SinkBB's ancestors using inverse DFS.
-  auto PredIt = ++idf_begin(&SinkBB);
-  auto PredEnd = idf_end(&SinkBB);
-  while (PredIt != PredEnd) {
-    BasicBlock &PredBB = **PredIt;
-    bool SinkPostDom = PDT.dominates(&SinkBB, &PredBB);
-
-    // If SinkBB does not post-dominate a predecessor, do not mark the
-    // predecessor (or any of its predecessors) cold.
-    if (!SinkPostDom || !mayExtractBlock(PredBB)) {
-      PredIt.skipChildren();
-      continue;
-    }
-
-    // Keep track of the post-dominated ancestor farthest away from the sink.
-    unsigned AncestorHeight = PredIt.getPathLength();
-    if (AncestorHeight > MaxAncestorHeight) {
-      MaxAncestor = &PredBB;
-      MaxAncestorHeight = AncestorHeight;
-    }
-
-    ColdRegion.push_back(&PredBB);
-    ++PredIt;
-  }
-
-  // CodeExtractor requires that all blocks to be extracted must be dominated
-  // by the first block to be extracted.
-  //
-  // To avoid spurious or repeated outlining, require that the max ancestor
-  // has a predecessor. By construction this predecessor is not in the cold
-  // region, i.e. its existence implies we don't outline the whole function.
-  //
-  // TODO: If MaxAncestor has no predecessors, we may be able to outline the
-  // second largest cold region that has a predecessor.
-  if (pred_empty(MaxAncestor) ||
-      MaxAncestor->getSinglePredecessor() == MaxAncestor)
-    return {};
-
-  // Filter out predecessors not dominated by the max ancestor.
-  //
-  // TODO: Blocks not dominated by the max ancestor could be extracted as
-  // other cold regions. Marking outlined calls as noreturn when appropriate
-  // and outlining more than once per function could achieve most of the win.
-  auto EraseIt = remove_if(ColdRegion, [&](BasicBlock *PredBB) {
-    return PredBB != MaxAncestor && !DT.dominates(MaxAncestor, PredBB);
-  });
-  ColdRegion.erase(EraseIt, ColdRegion.end());
-
-  // Add SinkBB to the cold region.
-  ColdRegion.push_back(&SinkBB);
-
-  // Ensure that the first extracted block is the max ancestor.
-  if (ColdRegion[0] != MaxAncestor) {
-    auto AncestorIt = find(ColdRegion, MaxAncestor);
-    *AncestorIt = ColdRegion[0];
-    ColdRegion[0] = MaxAncestor;
-  }
-
-  // Find all successors of SinkBB dominated by SinkBB using DFS.
-  auto SuccIt = ++df_begin(&SinkBB);
-  auto SuccEnd = df_end(&SinkBB);
-  while (SuccIt != SuccEnd) {
-    BasicBlock &SuccBB = **SuccIt;
-    bool SinkDom = DT.dominates(&SinkBB, &SuccBB);
-
-    // If SinkBB does not dominate a successor, do not mark the successor (or
-    // any of its successors) cold.
-    if (!SinkDom || !mayExtractBlock(SuccBB)) {
-      SuccIt.skipChildren();
-      continue;
-    }
-
-    ColdRegion.push_back(&SuccBB);
-    ++SuccIt;
-  }
-
-  if (ColdRegion.size() == 1 && !isProfitableToOutline(*ColdRegion[0], TTI))
-    return {};
-
-  return ColdRegion;
-}
-
-/// Get the largest cold region in \p F.
-static BlockSequence getLargestColdRegion(Function &F, ProfileSummaryInfo &PSI,
-                                          BlockFrequencyInfo *BFI,
-                                          TargetTransformInfo &TTI,
-                                          DominatorTree &DT, PostDomTree &PDT) {
-  // Keep track of the largest cold region.
-  BlockSequence LargestColdRegion = {};
-
-  for (BasicBlock &BB : F) {
-    // Identify cold blocks.
-    if (!mayExtractBlock(BB))
-      continue;
-    bool Cold =
-        PSI.isColdBB(&BB, BFI) || (EnableStaticAnalyis && unlikelyExecuted(BB));
-    if (!Cold)
-      continue;
-
-    LLVM_DEBUG({
-      dbgs() << "Found cold block:\n";
-      BB.dump();
-    });
-
-    // Find a maximal cold region we can outline.
-    BlockSequence ColdRegion = findMaximalColdRegion(BB, TTI, DT, PDT);
-    if (ColdRegion.empty()) {
-      LLVM_DEBUG(dbgs() << "  Skipping (block not profitable to extract)\n");
-      continue;
-    }
-
-    ++NumColdRegionsFound;
-
-    LLVM_DEBUG({
-      llvm::dbgs() << "Identified cold region with " << ColdRegion.size()
-                   << " blocks:\n";
-      for (BasicBlock *BB : ColdRegion)
-        BB->dump();
-    });
-
-    // TODO: Outline more than one region.
-    if (ColdRegion.size() > LargestColdRegion.size())
-      LargestColdRegion = std::move(ColdRegion);
-  }
-
-  return LargestColdRegion;
-}
-
 class HotColdSplitting {
 public:
   HotColdSplitting(ProfileSummaryInfo *ProfSI,
                    function_ref<BlockFrequencyInfo *(Function &)> GBFI,
                    function_ref<TargetTransformInfo &(Function &)> GTTI,
                    std::function<OptimizationRemarkEmitter &(Function &)> *GORE)
       : PSI(ProfSI), GetBFI(GBFI), GetTTI(GTTI), GetORE(GORE) {}
   bool run(Module &M);
 
 private:
   bool shouldOutlineFrom(const Function &F) const;
+  bool outlineColdRegions(Function &F, ProfileSummaryInfo &PSI,
+                          BlockFrequencyInfo *BFI, TargetTransformInfo &TTI,
+                          DominatorTree &DT, PostDomTree &PDT,
+                          OptimizationRemarkEmitter &ORE);
   Function *extractColdRegion(const BlockSequence &Region, DominatorTree &DT,
                               BlockFrequencyInfo *BFI, TargetTransformInfo &TTI,
                               OptimizationRemarkEmitter &ORE, unsigned Count);
   SmallPtrSet<const Function *, 2> OutlinedFunctions;
   ProfileSummaryInfo *PSI;
   function_ref<BlockFrequencyInfo *(Function &)> GetBFI;
   function_ref<TargetTransformInfo &(Function &)> GetTTI;
   std::function<OptimizationRemarkEmitter &(Function &)> *GetORE;
if (TTI.useColdCCForColdCall(*OutF)) {
       CS.setCallingConv(CallingConv::Cold);
     }
     CI->setIsNoInline();
 
     // Try to make the outlined code as small as possible on the assumption
     // that it's cold.
     assert(!OutF->hasFnAttribute(Attribute::OptimizeNone) &&
            "An outlined function should never be marked optnone");
     OutF->addFnAttr(Attribute::MinSize);
 
     LLVM_DEBUG(llvm::dbgs() << "Outlined Region: " << *OutF);
     ORE.emit([&]() {
       return OptimizationRemark(DEBUG_TYPE, "HotColdSplit",
                                 &*Region[0]->begin())
              << ore::NV("Original", OrigF) << " split cold code into "
              << ore::NV("Split", OutF);
     });
     return OutF;
   }
 
   ORE.emit([&]() {
     return OptimizationRemarkMissed(DEBUG_TYPE, "ExtractFailed",
                                     &*Region[0]->begin())
            << "Failed to extract region at block "
            << ore::NV("Block", Region.front());
   });
   return nullptr;
 }
 
+/// A pair of (basic block, score).
+using BlockTy = std::pair<BasicBlock *, unsigned>;
+
+/// A maximal outlining region. This contains all blocks post-dominated by a
+/// sink block, the sink block itself, and all blocks dominated by the sink.
+class OutliningRegion {
+  /// A list of (block, score) pairs. A block's score is non-zero iff it's a
+  /// viable sub-region entry point. Blocks with higher scores are better entry
+  /// points (i.e. they are more distant ancestors of the sink block).
+  SmallVector<BlockTy, 0> Blocks = {};
+
+  /// The suggested entry point into the region. If the region has multiple
+  /// entry points, all blocks within the region may not be reachable from this
+  /// entry point.
+  BasicBlock *SuggestedEntryPoint = nullptr;
+
+  /// Whether the entire function is cold.
+  bool EntireFunctionCold = false;
+
+  /// Whether or not \p BB could be the entry point of an extracted region.
+  static bool isViableEntryPoint(BasicBlock &BB) { return !BB.isEHPad(); }
+
+  /// If \p BB is a viable entry point, return \p Score. Return 0 otherwise.
+  static unsigned getEntryPointScore(BasicBlock &BB, unsigned Score) {
+    return isViableEntryPoint(BB) ? Score : 0;
+  }
+
+public:
+  static OutliningRegion create(BasicBlock &SinkBB, const DominatorTree &DT,
+                                const PostDomTree &PDT) {
+    OutliningRegion ColdRegion;
+
+    // The ancestor farthest-away from SinkBB, and also post-dominated by it.
+    unsigned SinkScore = getEntryPointScore(SinkBB, 1);
+    ColdRegion.SuggestedEntryPoint = (SinkScore > 0) ? &SinkBB : nullptr;
+    unsigned BestScore = SinkScore;
+
+    // Visit SinkBB's ancestors using inverse DFS.
+    auto PredIt = ++idf_begin(&SinkBB);
+    auto PredEnd = idf_end(&SinkBB);
+    while (PredIt != PredEnd) {
+      BasicBlock &PredBB = **PredIt;
+      bool SinkPostDom = PDT.dominates(&SinkBB, &PredBB);
+
+      // If SinkBB does not post-dominate a predecessor, do not mark the
+      // predecessor (or any of its predecessors) cold.
+      if (!SinkPostDom || !mayExtractBlock(PredBB)) {
+        PredIt.skipChildren();
+        continue;
+      }
+
+      // If the predecessor is cold and has no predecessors, the entire
+      // function must be cold.
+      if (pred_empty(&PredBB)) {
+        ColdRegion.EntireFunctionCold = true;
+        return ColdRegion;
+      }
+
+      // Keep track of the post-dominated ancestor farthest away from the sink.
+      // Add a large constant to its score as it makes for a better entry point
+      // (compared to sink-successor blocks).
+      unsigned PredScore =
+          getEntryPointScore(PredBB, PredIt.getPathLength() + INT_MAX);
+      if (PredScore > BestScore) {
+        ColdRegion.SuggestedEntryPoint = &PredBB;
+        BestScore = PredScore;
+      }
+
+      ColdRegion.Blocks.push_back({&PredBB, PredScore});
+      ++PredIt;
+    }
+
+    // Add SinkBB to the cold region.
+    ColdRegion.Blocks.push_back({&SinkBB, SinkScore});
+
+    // Find all successors of SinkBB dominated by SinkBB using DFS.
+    auto SuccIt = ++df_begin(&SinkBB);
+    auto SuccEnd = df_end(&SinkBB);
+    while (SuccIt != SuccEnd) {
+      BasicBlock &SuccBB = **SuccIt;
+      bool SinkDom = DT.dominates(&SinkBB, &SuccBB);
+
+      // If SinkBB does not dominate a successor, do not mark the successor (or
+      // anu of its successors) cold.
+      if (!SinkDom || !mayExtractBlock(SuccBB)) {
+        SuccIt.skipChildren();
+        continue;
+      }
+
+      unsigned SuccScore = getEntryPointScore(SuccBB, 1);
+      if (SuccScore > BestScore) {
+        ColdRegion.SuggestedEntryPoint = &SuccBB;
+        BestScore = SuccScore;
+      }
+
+      ColdRegion.Blocks.push_back({&SuccBB, SuccScore});
+      ++SuccIt;
+    }
+
+    return ColdRegion;
+  }
+
+  /// Whether this region has nothing to extract.
+  bool empty() const { return !SuggestedEntryPoint; }
+
+  /// The blocks in this region.
+  ArrayRef<std::pair<BasicBlock *, unsigned>> blocks() const { return Blocks; }
+
+  /// Whether the entire function containing this region is cold.
+  bool isEntireFunctionCold() const { return EntireFunctionCold; }
+
+  /// Remove a sub-region from this region and return it as a block sequence.
+  BlockSequence takeSingleEntrySubRegion(DominatorTree &DT) {
+    assert(!empty() && !isEntireFunctionCold() && "Nothing to extract");
+
+    // Remove blocks dominated by the suggested entry point from this region.
+    // During the removal, identify the next best entry point into the region.
+    BasicBlock *NextEntryPoint = nullptr;
+    unsigned NextScore = 0;
+    auto RegionStartIt = remove_if(Blocks, [&](const BlockTy &Block) {
+      BasicBlock *BB = Block.first;
+      unsigned Score = Block.second;
+      bool InSubRegion =
+          BB == SuggestedEntryPoint || DT.dominates(SuggestedEntryPoint, BB);
+      if (!InSubRegion && Score > NextScore) {
+        NextEntryPoint = BB;
+        NextScore = Score;
+      }
+      return InSubRegion;
+    });
+
+    // Construct the sub-region to extract from the set of removed blocks,
+    // then erase it from the containing region.
+    BlockSequence SubRegion;
+    for (const BlockTy &Block : make_range(RegionStartIt, Blocks.end()))
+      SubRegion.push_back(Block.first);
+    Blocks.erase(RegionStartIt, Blocks.end());
+
+    // Ensure that the first extracted block is the suggested entry point.
+    if (SubRegion[0] != SuggestedEntryPoint) {
+      auto *AncestorIt = find(SubRegion, SuggestedEntryPoint);
+      *AncestorIt = SubRegion[0];
+      SubRegion[0] = SuggestedEntryPoint;
+    }
+
+    // Update the suggested entry point.
+    SuggestedEntryPoint = NextEntryPoint;
+
+    return SubRegion;
+  }
+};
+
+bool HotColdSplitting::outlineColdRegions(Function &F, ProfileSummaryInfo &PSI,
+                                          BlockFrequencyInfo *BFI,
+                                          TargetTransformInfo &TTI,
+                                          DominatorTree &DT, PostDomTree &PDT,
+                                          OptimizationRemarkEmitter &ORE) {
+  bool Changed = false;
+
+  // The set of cold blocks.
+  SmallPtrSet<BasicBlock *, 4> ColdBlocks;
+
+  // The worklist of non-intersecting regions left to outline.
+  SmallVector<OutliningRegion, 2> OutliningWorklist;
+
+  // Set up an RPO traversal. Experimentally, this performs better (outlines
+  // more) than a PO traversal, because we prevent region overlap by keeping
+  // the first region to contain a block.
+  ReversePostOrderTraversal<Function *> RPOT(&F);
+
+  // Find all cold regions.
+  for (BasicBlock *BB : RPOT) {
+    // Skip blocks which can't be outlined.
+    if (!mayExtractBlock(*BB))
+      continue;
+
+    // This block is already part of some outlining region.
+    if (ColdBlocks.count(BB))
+      continue;
+
+    bool Cold =
+        PSI.isColdBB(BB, BFI) || (EnableStaticAnalyis && unlikelyExecuted(*BB));
+    if (!Cold)
+      continue;
+
+    LLVM_DEBUG({
+      dbgs() << "Found cold block:\n";
+      BB->dump();
+    });
+
+    auto Region = OutliningRegion::create(*BB, DT, PDT);
+    if (Region.empty())
+      continue;
+
+    if (Region.isEntireFunctionCold()) {
+      // TODO: Move this function into a cold section.
+      LLVM_DEBUG(dbgs() << "Entire function is cold\n");
+      return false;
+    }
+
+    // If this outlining regions intersects with another, drop the new region.
+    //
+    // TODO: It's theoretically possible to outline more by only keeping the
+    // largest region which contains a block, but the extra bookkeeping to do
+    // this is tricky/expensive.
+    bool RegionsOverlap = any_of(Region.blocks(), [&](const BlockTy &Block) {
+      return !ColdBlocks.insert(Block.first).second;
+    });
+    if (RegionsOverlap)
+      continue;
+
+    OutliningWorklist.push_back(std::move(Region));
+    ++NumColdRegionsFound;
+  }
+
+  // Outline single-entry cold regions, splitting up larger regions as needed.
+  while (!OutliningWorklist.empty()) {
+    OutliningRegion Region = OutliningWorklist.pop_back_val();
+    assert(!Region.empty() && "Empty outlining region in worklist");
+    do {
+      BlockSequence SubRegion = Region.takeSingleEntrySubRegion(DT);
+      if (!isProfitableToOutline(SubRegion, TTI)) {
+        LLVM_DEBUG(dbgs() << "Skipping outlining; not profitable to outline\n");
+        continue;
+      }
+
+      Function *Outlined =
+          extractColdRegion(SubRegion, DT, BFI, TTI, ORE, /*Count=*/1);
+      if (Outlined) {
+        OutlinedFunctions.insert(Outlined);
+        Changed = true;
+      }
+    } while (!Region.empty());
+  }
+
+  return Changed;
+}
+
 bool HotColdSplitting::run(Module &M) {
   bool Changed = false;
+  OutlinedFunctions.clear();
   for (auto &F : M) {
     if (!shouldOutlineFrom(F)) {
-      LLVM_DEBUG(llvm::dbgs() << "Not outlining in " << F.getName() << "\n");
+      LLVM_DEBUG(llvm::dbgs() << "Skipping " << F.getName() << "\n");
       continue;
     }
-
     LLVM_DEBUG(llvm::dbgs() << "Outlining in " << F.getName() << "\n");
     DominatorTree DT(F);
     PostDomTree PDT(F);
     PDT.recalculate(F);
     BlockFrequencyInfo *BFI = GetBFI(F);
     TargetTransformInfo &TTI = GetTTI(F);
-
-    BlockSequence ColdRegion = getLargestColdRegion(F, *PSI, BFI, TTI, DT, PDT);
-    if (ColdRegion.empty())
-      continue;
-
     OptimizationRemarkEmitter &ORE = (*GetORE)(F);
-    Function *Outlined =
-        extractColdRegion(ColdRegion, DT, BFI, TTI, ORE, /*Count=*/1);
-    if (Outlined) {
-      OutlinedFunctions.insert(Outlined);
-      Changed = true;
-    }
+    Changed |= outlineColdRegions(F, *PSI, BFI, TTI, DT, PDT, ORE);
   }
   return Changed;
 }
 
 bool HotColdSplittingLegacyPass::runOnModule(Module &M) {
   if (skipModule(M))
     return false;
   ProfileSummaryInfo *PSI =