Partial Inlining with multi-region outlining based on PGO information

[gyiu]

Summary

With PGO information, we can do more aggressive outlining of cold regions in the inline candidate function. This contrasts with the scheme of keeping only the 'early return' portion of the inline candidate and outlining the rest of the function as a single function call.

Support for outlining multiple regions of each function is added, as well as some basic heuristics to determine which regions are good to outline. Outline candidates limited to regions that are single-entry & single-exit. Also we don't account for live-ranges we may be killing across the region with a function. These are enhancements we can consider in another patch.

Fallback to the regular partial inlining scheme is retained when either i) no regions are identified for outlining in the function, or ii) the outlined function could not be inlined in any of its callers.

[davidxl]

lib/Transforms/IPO/PartialInlining.cpp

45

Make the name more meaningful? New does not tells us anything.

63

Make this an enum option?

66

Merge with the enum trace option

77

Indicate the purpose of the ratio, e.g, 'to decide xxx'

83

Why is option needed?

87

Why using BP to determine region coldness instead of PSI and profile count?

161

Add documentation.

417

Add a high level description of how the cost/savings are computed below.

1060

Is DT updated by the extractor?

1150

return { false, nullptr };

1173

Better check if the function has EntryCount

[gyiu]

lib/Transforms/IPO/PartialInlining.cpp

87

No particular reason, mainly for convenience. I figure I would still need to compare the profile count of a block to its sibling blocks to know its relative coldness. Is PSI and profile count more accurate?

1060

I don't believe so. You're thinking we should update the DT each time we extract a code region?

[gyiu]

lib/Transforms/IPO/PartialInlining.cpp

83

Just for added tune-ability. I wasn't sure if '100' is a reasonable minimum execution count to have.

[echristo]

One high level comment that I added inline as well:

What's the criteria for deleting the old code? When should we plan on doing this?

lib/Transforms/IPO/PartialInlining.cpp

168–169

When is this?

[gyiu]

lib/Transforms/IPO/PartialInlining.cpp

168–169

This is a leftover comment that's no longer valid. We'll need to keep both constructors. However, there's a refinement to how we track outlined functions that is now redundant (class members OutlinedFunctions vs. OutlinedFunc+OutliningCallBB), so I'll make that change now.

[gyiu]

• Updated code based on review comments.
• Also added check for candidate region outputs. If cold region has an output (ie. a live exit variable) then we could be blocking some code motion in the caller after it has been outlined.
• Also added code to mark outlined functions/callsites with coldCC.

[davidxl]

please also add test cases.

lib/Transforms/IPO/PartialInlining.cpp

87

Using PSI gives you information about global hotness of a block, so that the compiler only focus on regions that are globally hot or cold.

112

Make these two member variables.

360

This skips SamplePGO. Is it intended? If yes, only the second check should be enough.

369

Is this too restrictive? Single entry region does not mean the entry block can not have multiple predecessors.

438

should check if the the Succ block is actually cold enough.

452

if (!IsSingleExit(..)))

continue;

477

Why this limitation? This should be handled by the code extractor

1346

Most of the code here can be shared with the existing partial inliner. Refactor them to avoid code duplication?

[gyiu]

lib/Transforms/IPO/PartialInlining.cpp

87

Interesting. One concern I have, though, is whether we should check the predecessor block for 'isHotBB', or the cold successor block for 'isColdBB', or both. Checking if the predecessor is globally hot will still require some branch probability ratio to indicate a relatively cold successor block. Similarly with a globally cold successor, we will still need to check the predecessor block for relative hotness. Checking both could be good, though we will likely miss a lot of opportunities.

360

Yeah, I'm not confident that the sampling profile will give us proper coldness information, and may lead us to outline more than we should.

369

I guess I'm unsure what a dominator block with multiple predecessors actually means in this case, because the code looks for a cold successor block to a branch. If that cold successor has multiple predecessors, then doesn't that we can enter this region through multiple blocks, hence not be single entry at all?

438

You mean check the global coldness of the Succ Block? I guess this is where we need to weigh the pros and cons of global coldnress vs. relative coldness (similarly for hotness). Are you thinking we should check relative AND global coldness? Would this be reasonable in a situation where the callgraph is disjoint (eg. building a shared library) and we can't make reasonable comparisons of the counters between subgraphs?

477

You're right. I didn't realize that.

[gyiu]

Updated based on latest comments. Still missing testcase that I'm working on currently. Will update patch when ready.

[gyiu]

• Added testcase.
• Remove the requirement for a conditional branch as the terminator of the entry block, for multi-region outlining specifically.

[davidxl]

lib/Transforms/IPO/PartialInlining.cpp

361

Consider making this into a message with -Rpass-missed=partial-inining

429

This trace does not look useful. Consider removing it.

440

merge these into one output stream

446

Consider changing this into -Rpass-analysis=partial-inlining

1092

-Rpass-missed

1150

clang-format and similar other places.

1165–1197

Consider removing this trace. Should be handled by -print-before-xxx and -print-after-xxx

1413

You can achieve the same thing with -Rpass=partial-inlining. Consider remove this tracing output.

test/Transforms/CodeExtractor/PartialInlinePGORegion.ll

5 ↗

(On Diff #121831)

Perhaps add another case with multiple regions outlined?

[gyiu]

• Added new multi-outlined-region testcase.
• Removed unnecessary traces and moved useful ones to 'ifndef NDEBUG' macro.
• Modified/added ORE messages according to suggestions.
• General clean-up using clang-format.

[davidxl]

lgtm

include/llvm/Analysis/ProfileSummaryInfo.h

public:
   /// \brief Returns true if BasicBlock \p B is considered hot.
   bool isHotBB(const BasicBlock *B, BlockFrequencyInfo *BFI);
   /// \brief Returns true if BasicBlock \p B is considered cold.
   bool isColdBB(const BasicBlock *B, BlockFrequencyInfo *BFI);
   /// \brief Returns true if CallSite \p CS is considered hot.
   bool isHotCallSite(const CallSite &CS, BlockFrequencyInfo *BFI);
   /// \brief Returns true if Callsite \p CS is considered cold.
   bool isColdCallSite(const CallSite &CS, BlockFrequencyInfo *BFI);
+  /// \brief Returns HotCountThreshold if set.
+  uint64_t getHotCountThreshold() {
+    return HotCountThreshold ? HotCountThreshold.getValue() : 0;
+  }
+  /// \brief Returns ColdCountThreshold if set.
+  uint64_t getColdCountThreshold() {
+    return ColdCountThreshold ? ColdCountThreshold.getValue() : 0;
+  }
 };
 
 /// An analysis pass based on legacy pass manager to deliver ProfileSummaryInfo.
 class ProfileSummaryInfoWrapperPass : public ImmutablePass {
   std::unique_ptr<ProfileSummaryInfo> PSI;
 
 public:
   static char ID;

lib/Transforms/IPO/PartialInlining.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/IPO/PartialInlining.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/BlockFrequencyInfo.h"
 #include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/CodeMetrics.h"
+#include "llvm/Analysis/DominanceFrontier.h"
 #include "llvm/Analysis/InlineCost.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/OptimizationDiagnosticInfo.h"
+#include "llvm/Analysis/PostDominators.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
+#include "llvm/Analysis/RegionInfo.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/DiagnosticInfo.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Pass.h"
 #include "llvm/Transforms/IPO.h"
 #include "llvm/Transforms/Utils/Cloning.h"
 #include "llvm/Transforms/Utils/CodeExtractor.h"
 using namespace llvm;
 
 #define DEBUG_TYPE "partial-inlining"
 
 STATISTIC(NumPartialInlined,
           "Number of callsites functions partially inlined into.");
+STATISTIC(NewNumPartialInlined,
+           "Number of callsites functions (NEW) partially inlined into.");
+STATISTIC(NumColdRegionsFound,
+           "Number of cold single entry/exit regions found.");
+STATISTIC(NumColdRegionsOutlined,
+           "Number of cold single entry/exit regions outlined.");
 
 // Command line option to disable partial-inlining. The default is false:
 static cl::opt<bool>
     DisablePartialInlining("disable-partial-inlining", cl::init(false),
-                           cl::Hidden, cl::desc("Disable partial ininling"));
+                           cl::Hidden, cl::desc("Disable partial inlining"));
+// Command line option to disable multi-region partial-inlining. The default is
+// false:
+static cl::opt<bool> DisableMultiRegionPartialInline(
+    "disable-mr-partial-inlining", cl::init(false), cl::Hidden,
+    cl::desc("Disable multi-region partial inlining"));
+
+// Command line option to debug partial-inlining. The default is false:
+static cl::opt<bool>
+    TracePartialInlining("trace-partial-inlining", cl::init(false),
+                           cl::Hidden, cl::desc("Trace partial inlining"));
+static cl::opt<bool>
+    TracePartialInliningSparse("trace-partial-inlining-sparse", cl::init(false),
+                                                    cl::Hidden, cl::desc("Trace partial inlining sparse"));
 // This is an option used by testing:
 static cl::opt<bool> SkipCostAnalysis("skip-partial-inlining-cost-analysis",
                                       cl::init(false), cl::ZeroOrMore,
                                       cl::ReallyHidden,
                                       cl::desc("Skip Cost Analysis"));
+// Use for tuning outlining heuristics.
+static cl::opt<float> MinRegionSizeRatio(
+    "min-region-size-ratio", cl::init(0.1), cl::Hidden,
+    cl::desc("Minimum ratio comparing relative sizes of each "
+             "outline candidate and original function"));
+// Use for tuning outlining heuristics.
+static cl::opt<unsigned>
+    MinBlockCounterExecution("min-block-execution", cl::init(100), cl::Hidden,
+                             cl::desc("Minimum block executions to consider "
+                                      "its BranchProbabilityInfo valid"));
+// Use for tuning outlining heuristics.
+static cl::opt<float> ColdBranchRatio(
+    "cold-branch-ratio", cl::init(0.1), cl::Hidden,
+    cl::desc("Minimum BranchProbability to consider a region cold."));
 
 static cl::opt<unsigned> MaxNumInlineBlocks(
     "max-num-inline-blocks", cl::init(5), cl::Hidden,
     cl::desc("Max number of blocks to be partially inlined"));
 
 // Command line option to set the maximum number of partial inlining allowed
 // for the module. The default value of -1 means no limit.
 static cl::opt<int> MaxNumPartialInlining(
     "max-partial-inlining", cl::init(-1), cl::Hidden, cl::ZeroOrMore,
     cl::desc("Max number of partial inlining. The default is unlimited"));
 
 // Used only when PGO or user annotated branch data is absent. It is
 // the least value that is used to weigh the outline region. If BFI
 // produces larger value, the BFI value will be used.
 static cl::opt<int>
     OutlineRegionFreqPercent("outline-region-freq-percent", cl::init(75),
                              cl::Hidden, cl::ZeroOrMore,
                              cl::desc("Relative frequency of outline region to "
                                       "the entry block"));
 
 static cl::opt<unsigned> ExtraOutliningPenalty(
     "partial-inlining-extra-penalty", cl::init(0), cl::Hidden,
     cl::desc("A debug option to add additional penalty to the computed one."));
 
 namespace {
 
 struct FunctionOutliningInfo {
   FunctionOutliningInfo()
       : Entries(), ReturnBlock(nullptr), NonReturnBlock(nullptr),
         ReturnBlockPreds() {}
   // Returns the number of blocks to be inlined including all blocks
   // in Entries and one return block.
   unsigned GetNumInlinedBlocks() const { return Entries.size() + 1; }
 
   // A set of blocks including the function entry that guard
   // the region to be outlined.
   SmallVector<BasicBlock *, 4> Entries;
   // The return block that is not included in the outlined region.
   BasicBlock *ReturnBlock;
   // The dominating block of the region to be outlined.
   BasicBlock *NonReturnBlock;
   // The set of blocks in Entries that that are predecessors to ReturnBlock
   SmallVector<BasicBlock *, 4> ReturnBlockPreds;
 };
 
+struct FunctionOutliningMultiRegionInfo {
+  FunctionOutliningMultiRegionInfo()
+      : ORI() {}
+
+  // Container for outline regions
+  struct OutlineRegionInfo {
+    OutlineRegionInfo(SmallVector<BasicBlock *, 8> Region,
+                      BasicBlock *EntryBlock, BasicBlock *ExitBlock,
+                      BasicBlock *ReturnBlock)
+        : Region(Region), EntryBlock(EntryBlock), ExitBlock(ExitBlock),
+          ReturnBlock(ReturnBlock) {}
+    SmallVector<BasicBlock *, 8> Region;
+    BasicBlock *EntryBlock;
+    BasicBlock *ExitBlock;
+    BasicBlock *ReturnBlock;
+  };
+
+  SmallVector<OutlineRegionInfo, 4> ORI;
+};
+
 struct PartialInlinerImpl {
   PartialInlinerImpl(
       std::function<AssumptionCache &(Function &)> *GetAC,
       std::function<TargetTransformInfo &(Function &)> *GTTI,
       Optional<function_ref<BlockFrequencyInfo &(Function &)>> GBFI,
       ProfileSummaryInfo *ProfSI)
       : GetAssumptionCache(GetAC), GetTTI(GTTI), GetBFI(GBFI), PSI(ProfSI) {}
   bool run(Module &M);
-  Function *unswitchFunction(Function *F);
+  std::pair<bool, Function *> unswitchFunction(Function *F);
 
   // This class speculatively clones the the function to be partial inlined.
   // At the end of partial inlining, the remaining callsites to the cloned
   // function that are not partially inlined will be fixed up to reference
   // the original function, and the cloned function will be erased.
   struct FunctionCloner {
+    // Constructor with FunctionOutliningInfo should be removed once
+    // multi-region outlining work is complete.
     FunctionCloner(Function *F, FunctionOutliningInfo *OI);
+    FunctionCloner(Function *F, FunctionOutliningMultiRegionInfo *OMRI);
     ~FunctionCloner();
 
     // Prepare for function outlining: making sure there is only
     // one incoming edge from the extracted/outlined region to
     // the return block.
     void NormalizeReturnBlock();
 
-    // Do function outlining:
-    Function *doFunctionOutlining();
+    // Do function outlining for cold regions.
+    bool doMultiRegionFunctionOutlining();
+    // Do function outlining for region after early return block(s).
+    Function *doSingleRegionFunctionOutlining();
 
     Function *OrigFunc = nullptr;
     Function *ClonedFunc = nullptr;
+
+    typedef std::pair<Function *, BasicBlock *> FuncBodyCallerPair;
+    // Keep track of Outlined Functions and the basic block they're called from.
+    SmallVector<FuncBodyCallerPair, 4> OutlinedFunctions;
     Function *OutlinedFunc = nullptr;
     BasicBlock *OutliningCallBB = nullptr;
+
     // ClonedFunc is inlined in one of its callers after function
     // outlining.
     bool IsFunctionInlined = false;
     // The cost of the region to be outlined.
     int OutlinedRegionCost = 0;
+    // ClonedOI is specific to outlining non-early return blocks.
     std::unique_ptr<FunctionOutliningInfo> ClonedOI = nullptr;
+    // ClonedOMRI is specific to outlining cold regions.
+    std::unique_ptr<FunctionOutliningMultiRegionInfo> ClonedOMRI = nullptr;
     std::unique_ptr<BlockFrequencyInfo> ClonedFuncBFI = nullptr;
   };
 
 private:
   int NumPartialInlining = 0;
   std::function<AssumptionCache &(Function &)> *GetAssumptionCache;
   std::function<TargetTransformInfo &(Function &)> *GetTTI;
   Optional<function_ref<BlockFrequencyInfo &(Function &)>> GetBFI;
bool shouldPartialInline(CallSite CS, FunctionCloner &Cloner,
                            BlockFrequency WeightedOutliningRcost,
                            OptimizationRemarkEmitter &ORE);
 
   // Try to inline DuplicateFunction (cloned from F with call to
   // the OutlinedFunction into its callers. Return true
   // if there is any successful inlining.
   bool tryPartialInline(FunctionCloner &Cloner);
 
+  bool tryNewPartialInline(FunctionCloner &Cloner);
+
   // Compute the mapping from use site of DuplicationFunction to the enclosing
   // BB's profile count.
   void computeCallsiteToProfCountMap(Function *DuplicateFunction,
                                      DenseMap<User *, uint64_t> &SiteCountMap);
 
   bool IsLimitReached() {
     return (MaxNumPartialInlining != -1 &&
             NumPartialInlining >= MaxNumPartialInlining);
private:
   //   basic block Cloner.OutliningCallBB;
   std::tuple<int, int> computeOutliningCosts(FunctionCloner &Cloner);
   // Compute the 'InlineCost' of block BB. InlineCost is a proxy used to
   // approximate both the size and runtime cost (Note that in the current
   // inline cost analysis, there is no clear distinction there either).
   static int computeBBInlineCost(BasicBlock *BB);
 
   std::unique_ptr<FunctionOutliningInfo> computeOutliningInfo(Function *F);
+  std::unique_ptr<FunctionOutliningMultiRegionInfo> computeOutliningColdRegionsInfo(Function *F);
 
 };
 
 struct PartialInlinerLegacyPass : public ModulePass {
   static char ID; // Pass identification, replacement for typeid
   PartialInlinerLegacyPass() : ModulePass(ID) {
     initializePartialInlinerLegacyPassPass(*PassRegistry::getPassRegistry());
   }
std::function<TargetTransformInfo &(Function &)> GetTTI =
       return TTIWP->getTTI(F);
     };
 
     return PartialInlinerImpl(&GetAssumptionCache, &GetTTI, None, PSI).run(M);
   }
 };
 }
 
+struct PostDomTree : PostDomTreeBase<BasicBlock> {
+  PostDomTree(Function &F) { recalculate(F); }
+};
+
+std::unique_ptr<FunctionOutliningMultiRegionInfo>
+PartialInlinerImpl::computeOutliningColdRegionsInfo(Function *F) {
+  BasicBlock *EntryBlock = &F->front();
+  BranchInst *BR = dyn_cast<BranchInst>(EntryBlock->getTerminator());
+  if (!BR || BR->isUnconditional())
+    return std::unique_ptr<FunctionOutliningMultiRegionInfo>();
+
+  DominatorTree DT(*F);
+  PostDominatorTree PDT;
+  PDT.recalculate(*F);
+  DominanceFrontier DF;
+  DF.analyze(DT);
+  RegionInfo RI ;
+  RI.recalculate(*F, &DT, &PDT, &DF);
+  LoopInfo LI(DT);
+  BranchProbabilityInfo BPI(*F, LI);
+  if (TracePartialInlining) {
+    dbgs() << "BranchProbabilityInfo:\n";
+    BPI.print(dbgs());
+  }
+  std::unique_ptr<BlockFrequencyInfo> ScopedBFI;
+  BlockFrequencyInfo *BFI;
+  if (!GetBFI) {
+    ScopedBFI.reset(new BlockFrequencyInfo(*F, BPI, LI));
+    BFI = ScopedBFI.get();
+  } else
+    BFI = &(*GetBFI)(*F);
+
+  if (TracePartialInlining) {
+    dbgs() << "BlockFrequencyInfo:\n";
+    BFI->print(dbgs());
+    dbgs() << "DominatorTree:\n";
+    DT.print(dbgs());
+  }
+
+  auto IsReturnBlock = [](BasicBlock *BB) {
+    TerminatorInst *TI = BB->getTerminator();
+    return isa<ReturnInst>(TI);
+  };
+
+  // Return if we don't have profiling information.
+  if (!PSI->hasProfileSummary() || !PSI->hasInstrumentationProfile())
+    return std::unique_ptr<FunctionOutliningMultiRegionInfo>();
+
+  std::unique_ptr<FunctionOutliningMultiRegionInfo> OutliningInfo =
+      llvm::make_unique<FunctionOutliningMultiRegionInfo>();
+
+  auto IsSingleEntry = [](SmallVectorImpl<BasicBlock *> &BlockList) {
+    BasicBlock *Dom = BlockList.front();
+    return BlockList.size() > 1 &&
+           std::distance(pred_begin(Dom), pred_end(Dom)) == 1;
+  };
+
+  auto IsSingleExit =
+      [IsReturnBlock](SmallVectorImpl<BasicBlock *> &BlockList) -> BasicBlock * {
+    BasicBlock *ExitBlock = nullptr;
+    for (auto *Block : BlockList) {
+      if (IsReturnBlock(Block))
+        return nullptr;
+      for (auto SI = succ_begin(Block); SI != succ_end(Block); ++SI) {
+        if (!is_contained(BlockList, *SI)) {
+          if (TracePartialInlining)
+            dbgs() << "Found exit block = " << (*SI)->getName() << "\n";
+          if (ExitBlock)
+            return nullptr;
+          else
+            ExitBlock = Block;
+        }
+      }
+    }
+    return ExitBlock;
+  };
+
+  auto BBProfileCount = [BFI](BasicBlock *BB) {
+    return BFI->getBlockProfileCount(BB)
+               ? BFI->getBlockProfileCount(BB).getValue()
+               : 0;
+  };
+
+  // Use the same computeBBInlineCost function to compute the cost savings of
+  // the outlining the candidate region.
+  int OverallFunctionCost = 0;
+  for (auto &BB : *F)
+    OverallFunctionCost += computeBBInlineCost(&BB);
+
+  int MinOutlineRegionCost =
+      static_cast<int>(OverallFunctionCost * MinRegionSizeRatio);
+  BranchProbability MinBranchProbability(
+      static_cast<int>(ColdBranchRatio * MinBlockCounterExecution),
+      MinBlockCounterExecution);
+  bool ColdCandidateFound = false;
+  BasicBlock *CurrEntry = EntryBlock;
+  std::vector<BasicBlock *> DFS;
+  DenseMap<BasicBlock *, bool> VisitedMap;
+  DFS.push_back(CurrEntry);
+  VisitedMap[CurrEntry] = true;
+  // Use Depth First Search on the basic blocks to find CFG edges that are
+  // considered cold.
+  while (!DFS.empty()) {
+    auto *thisBB = DFS.back();
+    DFS.pop_back();
+    // Only consider regions with predecessor blocks that are considered
+    // not-cold (default: part of the top 99.99% of all block counters)
+    // AND greater than our minimum block execution count (default: 100).
+    if (PSI->isColdBB(thisBB, BFI) ||
+        BBProfileCount(thisBB) < MinBlockCounterExecution)
+      continue;
+    for (auto SI = succ_begin(thisBB); SI != succ_end(thisBB); ++SI) {
+      if (VisitedMap[*SI])
+        continue;
+      VisitedMap[*SI] = true;
+      DFS.push_back(*SI);
+      // If branch isn't cold, we skip to the next one.
+      BranchProbability SuccProb = BPI.getEdgeProbability(thisBB, *SI);
+      if (SuccProb > MinBranchProbability)
+        continue;
+      if (TracePartialInlining) {
+        dbgs() << "Found cold edge: " << thisBB->getName() << "->"
+               << (*SI)->getName() << "\n";
+        dbgs() << "Branch Probability = " << SuccProb << "\n";
+      }
+      SmallVector<BasicBlock *, 8> DominateVector;
+      DT.getDescendants(*SI, DominateVector);
+      // We can only outline single entry regions (for now).
+      if (!IsSingleEntry(DominateVector))
+        continue;
+      if (TracePartialInlining)
+        dbgs() << "Dominates more than one block. Single predecessor.\n";
+      // We can only outline single exit regions (for now).
+      if (auto *ExitBlock = IsSingleExit(DominateVector)) {
+        int OutlineRegionCost = 0;
+        for (auto *BB : DominateVector)
+          OutlineRegionCost += computeBBInlineCost(BB);
+
+        if (TracePartialInlining) {
+          dbgs() << "Single Exit Block = " << ExitBlock->getName() << "\n";
+          dbgs() << "OutlineRegionCost = " << OutlineRegionCost << "\n";
+          dbgs() << "OverallFunctionCost = " << OverallFunctionCost << "\n";
+        }
+
+        if (OutlineRegionCost < MinOutlineRegionCost) {
+          if (TracePartialInlining) {
+            dbgs() << "Inline cost-savings smaller than "
+                   << MinOutlineRegionCost << ". Skip.\n";
+          }
+          continue;
+        }
+        // For now, ignore blocks that belong to a SISE region that is a
+        // candidate for outlining.  In the future, we may want to look
+        // at inner regions because the outer region may have live-exit
+        // variables.
+        for (auto *BB : DominateVector)
+          VisitedMap[BB] = true;
+        // can't outline if region contains a return block
+        if (IsReturnBlock(ExitBlock))
+          continue;
+        // ReturnBlock here means the block after the outline call
+        BasicBlock *ReturnBlock = ExitBlock->getSingleSuccessor();
+        // assert(ReturnBlock && "ReturnBlock is NULL somehow!");
+        FunctionOutliningMultiRegionInfo::OutlineRegionInfo RegInfo(
+            DominateVector, DominateVector.front(), ExitBlock, ReturnBlock);
+        RegInfo.Region = DominateVector;
+        OutliningInfo->ORI.push_back(RegInfo);
+        if (TracePartialInlining) {
+          dbgs() << "Found Cold Candidate starting at block: "
+                 << DominateVector.front()->getName() << "\n";
+        }
+        ColdCandidateFound = true;
+        NumColdRegionsFound++;
+      }
+    }
+  }
+  if (ColdCandidateFound)
+    return OutliningInfo;
+  else
+    return std::unique_ptr<FunctionOutliningMultiRegionInfo>();
+}
+
 std::unique_ptr<FunctionOutliningInfo>
 PartialInlinerImpl::computeOutliningInfo(Function *F) {
   BasicBlock *EntryBlock = &F->front();
   BranchInst *BR = dyn_cast<BranchInst>(EntryBlock->getTerminator());
   if (!BR || BR->isUnconditional())
     return std::unique_ptr<FunctionOutliningInfo>();
 
   // Returns true if Succ is BB's successor

 
 BranchProbability
 PartialInlinerImpl::getOutliningCallBBRelativeFreq(FunctionCloner &Cloner) {
 
   auto EntryFreq =
       Cloner.ClonedFuncBFI->getBlockFreq(&Cloner.ClonedFunc->getEntryBlock());
   auto OutliningCallFreq =
       Cloner.ClonedFuncBFI->getBlockFreq(Cloner.OutliningCallBB);
-
-  auto OutlineRegionRelFreq =
-      BranchProbability::getBranchProbability(OutliningCallFreq.getFrequency(),
-                                              EntryFreq.getFrequency());
+  // FIXME Hackery needed because ClonedFuncBFI is based on the function BEFORE
+  // we outlined any regions, so we may encounter situations where the
+  // OutliningCallFreq is *slightly* bigger than the EntryFreq.
+  if (OutliningCallFreq.getFrequency() > EntryFreq.getFrequency()) {
+    OutliningCallFreq = EntryFreq;
+  }
+  auto OutlineRegionRelFreq = BranchProbability::getBranchProbability(
+      OutliningCallFreq.getFrequency(), EntryFreq.getFrequency());
 
   if (hasProfileData(Cloner.OrigFunc, Cloner.ClonedOI.get()))
     return OutlineRegionRelFreq;
 
   // When profile data is not available, we need to be conservative in
   // estimating the overall savings. Static branch prediction can usually
   // guess the branch direction right (taken/non-taken), but the guessed
   // branch probability is usually not biased enough. In case when the
for (BasicBlock *E : OI->ReturnBlockPreds) {
     BasicBlock *NewE = cast<BasicBlock>(VMap[E]);
     ClonedOI->ReturnBlockPreds.push_back(NewE);
   }
   // Go ahead and update all uses to the duplicate, so that we can just
   // use the inliner functionality when we're done hacking.
   F->replaceAllUsesWith(ClonedFunc);
 }
 
+PartialInlinerImpl::FunctionCloner::FunctionCloner(
+    Function *F, FunctionOutliningMultiRegionInfo *OI)
+    : OrigFunc(F) {
+  ClonedOMRI = llvm::make_unique<FunctionOutliningMultiRegionInfo>();
+
+  // Clone the function, so that we can hack away on it.
+  ValueToValueMapTy VMap;
+  ClonedFunc = CloneFunction(F, VMap);
+
+  // Go through all Outline Candidate Regions and update all BasicBlock
+  // information.
+  for (FunctionOutliningMultiRegionInfo::OutlineRegionInfo RegionInfo :
+       OI->ORI) {
+    SmallVector<BasicBlock *, 8> Region;
+    for (BasicBlock *BB : RegionInfo.Region) {
+      Region.push_back(cast<BasicBlock>(VMap[BB]));
+    }
+    BasicBlock *NewEntryBlock = cast<BasicBlock>(VMap[RegionInfo.EntryBlock]);
+    BasicBlock *NewExitBlock = cast<BasicBlock>(VMap[RegionInfo.ExitBlock]);
+    BasicBlock *NewReturnBlock = nullptr;
+    if (RegionInfo.ReturnBlock)
+      NewReturnBlock = cast<BasicBlock>(VMap[RegionInfo.ReturnBlock]);
+    FunctionOutliningMultiRegionInfo::OutlineRegionInfo MappedRegionInfo(
+        Region, NewEntryBlock, NewExitBlock, NewReturnBlock);
+    ClonedOMRI->ORI.push_back(MappedRegionInfo);
+  }
+  // Go ahead and update all uses to the duplicate, so that we can just
+  // use the inliner functionality when we're done hacking.
+  F->replaceAllUsesWith(ClonedFunc);
+}
+
 void PartialInlinerImpl::FunctionCloner::NormalizeReturnBlock() {
 
   auto getFirstPHI = [](BasicBlock *BB) {
     BasicBlock::iterator I = BB->begin();
     PHINode *FirstPhi = nullptr;
     while (I != BB->end()) {
       PHINode *Phi = dyn_cast<PHINode>(I);
       if (!Phi)
         break;
       if (!FirstPhi) {
         FirstPhi = Phi;
         break;
       }
     }
     return FirstPhi;
   };
 
+  // Shouldn't need to normalize PHIs if we're not outlining non-early return
+  // blocks.
+  if (!ClonedOI)
+    return;
+
   // Special hackery is needed with PHI nodes that have inputs from more than
   // one extracted block.  For simplicity, just split the PHIs into a two-level
   // sequence of PHIs, some of which will go in the extracted region, and some
   // of which will go outside.
   BasicBlock *PreReturn = ClonedOI->ReturnBlock;
   // only split block when necessary:
   PHINode *FirstPhi = getFirstPHI(PreReturn);
   unsigned NumPredsFromEntries = ClonedOI->ReturnBlockPreds.size();
while (I != PreReturn->end()) {
     // Keeping the trivial phi can introduce definition inside the outline
     // region which is live-out, causing necessary overhead (load, store
     // arg passing etc).
     if (auto *OldPhiVal = IsTrivialPhi(OldPhi)) {
       OldPhi->replaceAllUsesWith(OldPhiVal);
       DeadPhis.push_back(OldPhi);
     }
     ++I;
-    }
-    for (auto *DP : DeadPhis)
-      DP->eraseFromParent();
 
-    for (auto E : ClonedOI->ReturnBlockPreds) {
-      E->getTerminator()->replaceUsesOfWith(PreReturn, ClonedOI->ReturnBlock);
-    }
 }
 
-Function *PartialInlinerImpl::FunctionCloner::doFunctionOutlining() {
+bool PartialInlinerImpl::FunctionCloner::doMultiRegionFunctionOutlining() {
+
+  auto ComputeRegionCost = [](SmallVectorImpl<BasicBlock *> &Region) {
+    int Cost = 0;
+    for (BasicBlock* BB : Region)
+      Cost += computeBBInlineCost(BB);
+    return Cost;
+  };
+
+  assert(ClonedOMRI && "Expecting OutlineInfo for multi region outline");
+
+  if (ClonedOMRI->ORI.empty())
+    return false;
+
+  // The CodeExtractor needs a dominator tree.
+  DominatorTree DT;
+  DT.recalculate(*ClonedFunc);
+
+  // Manually calculate a BlockFrequencyInfo and BranchProbabilityInfo.
+  LoopInfo LI(DT);
+  BranchProbabilityInfo BPI(*ClonedFunc, LI);
+  ClonedFuncBFI.reset(new BlockFrequencyInfo(*ClonedFunc, BPI, LI));
+
+  for (FunctionOutliningMultiRegionInfo::OutlineRegionInfo RegionInfo : ClonedOMRI->ORI) {
+
+    OutlinedRegionCost += ComputeRegionCost(RegionInfo.Region);
+
+    Function *OutlinedFunc = CodeExtractor(RegionInfo.Region, &DT, /*AggregateArgs*/ false,
+                                 ClonedFuncBFI.get(), &BPI)
+                       .extractCodeRegion();
+
+    if (OutlinedFunc) {
+      BasicBlock *OutliningCallBB = PartialInlinerImpl::getOneCallSiteTo(OutlinedFunc)
+          .getInstruction()
+          ->getParent();
+      assert(OutliningCallBB->getParent() == ClonedFunc);
+      OutlinedFunctions.push_back(std::make_pair(OutlinedFunc,OutliningCallBB));
+      NumColdRegionsOutlined++;
+    }
+  }
+
+  return !OutlinedFunctions.empty();
+}
+
+Function * PartialInlinerImpl::FunctionCloner::doSingleRegionFunctionOutlining() {
   // Returns true if the block is to be partial inlined into the caller
   // (i.e. not to be extracted to the out of line function)
   auto ToBeInlined = [&, this](BasicBlock *BB) {
     return BB == ClonedOI->ReturnBlock ||
            (std::find(ClonedOI->Entries.begin(), ClonedOI->Entries.end(), BB) !=
             ClonedOI->Entries.end());
   };
 
+  assert(ClonedOI && "Expecting OutlineInfo for single region outline");
+  // The CodeExtractor needs a dominator tree.
+  DominatorTree DT;
+  DT.recalculate(*ClonedFunc);
+
+  // Manually calculate a BlockFrequencyInfo and BranchProbabilityInfo.
+  LoopInfo LI(DT);
+  BranchProbabilityInfo BPI(*ClonedFunc, LI);
+  ClonedFuncBFI.reset(new BlockFrequencyInfo(*ClonedFunc, BPI, LI));
+
   // Gather up the blocks that we're going to extract.
   std::vector<BasicBlock *> ToExtract;
   ToExtract.push_back(ClonedOI->NonReturnBlock);
   OutlinedRegionCost +=
       PartialInlinerImpl::computeBBInlineCost(ClonedOI->NonReturnBlock);
   for (BasicBlock &BB : *ClonedFunc)
     if (!ToBeInlined(&BB) && &BB != ClonedOI->NonReturnBlock) {
       ToExtract.push_back(&BB);
       // FIXME: the code extractor may hoist/sink more code
       // into the outlined function which may make the outlining
       // overhead (the difference of the outlined function cost
       // and OutliningRegionCost) look larger.
       OutlinedRegionCost += computeBBInlineCost(&BB);
     }
 
-  // The CodeExtractor needs a dominator tree.
-  DominatorTree DT;
-  DT.recalculate(*ClonedFunc);
-
-  // Manually calculate a BlockFrequencyInfo and BranchProbabilityInfo.
-  LoopInfo LI(DT);
-  BranchProbabilityInfo BPI(*ClonedFunc, LI);
-  ClonedFuncBFI.reset(new BlockFrequencyInfo(*ClonedFunc, BPI, LI));
-
   // Extract the body of the if.
   OutlinedFunc = CodeExtractor(ToExtract, &DT, /*AggregateArgs*/ false,
                                ClonedFuncBFI.get(), &BPI)
                      .extractCodeRegion();
 
   if (OutlinedFunc) {
     OutliningCallBB = PartialInlinerImpl::getOneCallSiteTo(OutlinedFunc)
         .getInstruction()
         ->getParent();
     assert(OutliningCallBB->getParent() == ClonedFunc);
   }
 
   return OutlinedFunc;
 }
 
 PartialInlinerImpl::FunctionCloner::~FunctionCloner() {
   // Ditch the duplicate, since we're done with it, and rewrite all remaining
   // users (function pointers, etc.) back to the original function.
   ClonedFunc->replaceAllUsesWith(OrigFunc);
   ClonedFunc->eraseFromParent();
     if (!IsFunctionInlined) {
+      if (!ClonedOMRI) {
       // Remove the function that is speculatively created if there is no
       // reference.
       if (OutlinedFunc)
         OutlinedFunc->eraseFromParent();
+    } else {
+      // Erase each outlined function.
+      for (auto FuncBBPair : OutlinedFunctions) {
+        Function *Func = FuncBBPair.first;
+        Func->eraseFromParent();
+      }
+    }
    }
 }
 
-Function *PartialInlinerImpl::unswitchFunction(Function *F) {
+std::pair<bool, Function *> PartialInlinerImpl::unswitchFunction(Function *F) {
 
   if (F->hasAddressTaken())
-    return nullptr;
+    return std::make_pair(false,nullptr);
 
   // Let inliner handle it
   if (F->hasFnAttribute(Attribute::AlwaysInline))
-    return nullptr;
+    return std::make_pair(false,nullptr);
 
   if (F->hasFnAttribute(Attribute::NoInline))
-    return nullptr;
+    return std::make_pair(false,nullptr);
 
   if (PSI->isFunctionEntryCold(F))
-    return nullptr;
+    return std::make_pair(false,nullptr);
 
   if (F->user_begin() == F->user_end())
-    return nullptr;
+    return std::make_pair(false,nullptr);
 
-  std::unique_ptr<FunctionOutliningInfo> OI = computeOutliningInfo(F);
+  if (TracePartialInlining) {
+    dbgs() << ">>>>>> Original Function >>>>>>\n";
+    F->print(dbgs());
+    dbgs() << "<<<<<< Original Function <<<<<<\n";
+  }
+
+  // Only try to outline cold regions if we have a profile summary, which
+  // implies we have profiling information.
+  if (PSI->hasProfileSummary() && !DisableMultiRegionPartialInline) {
+    std::unique_ptr<FunctionOutliningMultiRegionInfo> OMRI =
+        computeOutliningColdRegionsInfo(F);
+    if (OMRI) {
+      FunctionCloner Cloner(F, OMRI.get());
+
+      if (TracePartialInlining) {
+        dbgs() << "HotCountThreshold = " << PSI->getHotCountThreshold() << "\n";
+        dbgs() << "ColdCountThreshold = " << PSI->getColdCountThreshold()
+               << "\n";
+      }
+      bool DidOutline = Cloner.doMultiRegionFunctionOutlining();
+
+      if (DidOutline) {
+        if (TracePartialInlining) {
+          dbgs() << ">>>>>> Outlined (Cloned) Function >>>>>>\n";
+          Cloner.ClonedFunc->print(dbgs());
+          dbgs() << "<<<<<< Outlined (Cloned) Function <<<<<<\n";
+        }
+
+        if (tryNewPartialInline(Cloner))
+          return std::make_pair (true, nullptr);
+      }
+    }
+  }
 
+  // Fall-thru to regular partial inlining if we:
+  //    i) can't find any cold regions to outline, or
+  //   ii) can't inline the outlined function anywhere.
+  std::unique_ptr<FunctionOutliningInfo> OI = computeOutliningInfo(F);
   if (!OI)
-    return nullptr;
+    return std::make_pair(false, nullptr);
 
   FunctionCloner Cloner(F, OI.get());
   Cloner.NormalizeReturnBlock();
-  Function *OutlinedFunction = Cloner.doFunctionOutlining();
+
+  if (TracePartialInlining) {
+    dbgs() << "ClonedFunc: \n";
+    Cloner.ClonedFunc->print(dbgs());
+  }
+
+  Function *OutlinedFunction = Cloner.doSingleRegionFunctionOutlining();
+
+  if (OutlinedFunction) {
+    if (TracePartialInlining)
+      dbgs() << "Outlined at least one region. Try Inlining.\n";
+  } else {
+    if (TracePartialInlining)
+      dbgs() << "Didn't outline any functions.\n";
+    return std::make_pair (false, nullptr);
+  }
 
   bool AnyInline = tryPartialInline(Cloner);
 
   if (AnyInline)
-    return OutlinedFunction;
+    return std::make_pair (true, OutlinedFunction);
 
-  return nullptr;
+  return std::make_pair (false, nullptr);
 }
 
 bool PartialInlinerImpl::tryPartialInline(FunctionCloner &Cloner) {
   int NonWeightedRcost;
   int SizeCost;
 
   if (Cloner.OutlinedFunc == nullptr)
     return false;
if (AnyInline) {
     Cloner.IsFunctionInlined = true;
     if (CalleeEntryCount)
       Cloner.OrigFunc->setEntryCount(CalleeEntryCountV);
   }
 
   return AnyInline;
 }
 
+static bool
+shouldInline(InlineCost IC) {
+  // ORE would shure be useful here ...
+  if(IC.isNever())
+    return false;
+
+  if(IC.isAlways())
+     return true;
+
+  return IC.getCostDelta() >= 0;
+}
+
+bool PartialInlinerImpl::tryNewPartialInline(FunctionCloner &Cloner) {
+  // loop over all the call sites that call the cloned function.
+  assert(Cloner.OrigFunc->user_begin() == Cloner.OrigFunc->user_end() &&
+         "F's users should all be replaced!");
+  assert(Cloner.OutlinedFunctions.size() &&
+         "Expected at least 1 outlined reigon!");
+  std::vector<User *> Users(Cloner.ClonedFunc->user_begin(),
+                            Cloner.ClonedFunc->user_end());
+  OptimizationRemarkEmitter ORE(Cloner.OrigFunc);
+
+  DenseMap<User *, uint64_t> CallSiteToProfCountMap;
+  assert(Cloner.OrigFunc->getEntryCount());
+  computeCallsiteToProfCountMap(Cloner.ClonedFunc, CallSiteToProfCountMap);
+  uint64_t CalleeEntryCount = *Cloner.OrigFunc->getEntryCount();
+
+  bool AnyInlined = false;
+  for (User *User : Users) {
+    CallSite CS = getCallSite(User);
+    Function *Callee = CS.getCalledFunction();
+    auto &CalleeTTI = (*GetTTI)(*Callee);
+    InlineCost IC = getInlineCost(CS, getInlineParams(), CalleeTTI,
+                                  *GetAssumptionCache, GetBFI, PSI, &ORE);
+    if (!shouldInline(IC))
+      continue;
+
+    if (TracePartialInlining || TracePartialInliningSparse) {
+      dbgs() << "Partial Inlined callee " << Callee->getName();
+      dbgs() << " in caller " << User->getName() << "in file ";
+      dbgs() << Callee->getParent()->getName() << "\n";
+    }
+    ++NewNumPartialInlined;
+
+    AnyInlined = true;
+    InlineFunctionInfo IFI(nullptr, GetAssumptionCache, PSI);
+    InlineFunction(CS, IFI);
+
+    if (CallSiteToProfCountMap.count(User))
+      CalleeEntryCount -=
+          std::min(CalleeEntryCount, CallSiteToProfCountMap[User]);
+  }
+
+  if (AnyInlined) {
+    Cloner.IsFunctionInlined = true;
+    Cloner.OrigFunc->setEntryCount(CalleeEntryCount);
+  }
+
+  return AnyInlined;
+}
+
 bool PartialInlinerImpl::run(Module &M) {
   if (DisablePartialInlining)
     return false;
 
   std::vector<Function *> Worklist;
   Worklist.reserve(M.size());
   for (Function &F : M)
     if (!F.use_empty() && !F.isDeclaration())
for (User *U : CurrFunc->users())
       if (Instruction *I = dyn_cast<Instruction>(U))
         if (I->getParent()->getParent() == CurrFunc) {
           Recursive = true;
           break;
         }
     if (Recursive)
       continue;
 
-    if (Function *NewFunc = unswitchFunction(CurrFunc)) {
-      Worklist.push_back(NewFunc);
+    std::pair<bool, Function * > Result = unswitchFunction(CurrFunc);
+    if (Result.second)
+      Worklist.push_back(Result.second);
+    if (Result.first)
       Changed = true;
-    }
-  }
+  if (TracePartialInlining && Changed)
+    dbgs() << "Partially inlined at least one function.\n";
 
   return Changed;
 }
 
 char PartialInlinerLegacyPass::ID = 0;
 INITIALIZE_PASS_BEGIN(PartialInlinerLegacyPass, "partial-inliner",
                       "Partial Inliner", false, false)
 INITIALIZE_PASS_DEPENDENCY(AssumptionCacheTracker)