[PartialInlining] Add frequency based cost analysis

[davidxl]

Summary

In this patch, the runtime and size cost of making calls to outlined function is computed and compared with potential savings of partial inlining.

[davidxl]

Update test case -- added a cold region test.

[eraman]

I am not fully sure if this cost model - first check if the function is a suitable candidate for outlining and then use inliner's cost analysis to inline the original function to its callsites - is the right one. Specifically, I'm thinking if both should be done on a per-callsite basis.

lib/Transforms/IPO/PartialInlining.cpp

468

This should ideally be in InlineCost. Not all instructions have the same cost (eg. switch). At the least it should be documented here.

485

nit: This is also in multiples of InstrCost like SizeCost above. Using a consistent naming will make this more readable.

515

I wonder if it is better to first clone the function and then call this isPartialInliningBeneficial. We might end up cloning it unnecessarily, but could reuse all these analyses which are computed on the duplicate function later.

525

std::max?

530

NonWeightedSavings above only includes the savings due to elimination of call and arg setup. Since inlining also can eliminate instructions in a given context, we are likely undercounting the savings here.

lib/Transforms/Utils/CodeExtractor.cpp

701

This should be multiplied by InstrCost.

705

This is different from how it is used in InlineCost. In inline cost analysis, there is no separate static and runtime costs and even CallPenalty is meant to model static cost (saving and restoring registers etc). It is clear that they need to be modeled separately, but for now I prefer replicating what is in InlineCost.cpp and may be assign both the same value.

[davidxl]

Addressed review feedbacks by eraman

[davidxl]

lib/Transforms/IPO/PartialInlining.cpp

468

Added comment.

530

possibly. Since the inlined region is just a set of guarding blocks, the addition benefit is probably not that high.

lib/Transforms/Utils/CodeExtractor.cpp

701

removed the function. Use the actual extracted function to compute instead.

[eraman]

lib/Transforms/IPO/PartialInlining.cpp

132–174

Add an empty line between functions (her and in other places)

448

It wasn't obvious to me why callee entry freq >= outlining call's frequency. I think I understand it now - we ensure the non-return block has a single predecessor and hence can't be in a loop, but a comment above would help.

467

This is not in multiples of InstrCost (and also inliner's switch cost computation is more sophisticated now, but you've the TODO above)

499

may be add an assert that OutlinedFunctionCost >= OutlinedRegionCost to potentially catch any bugs.

659

nullptr instead of 0

lib/Transforms/Utils/CodeExtractor.cpp

23

Why do you need this?

754

Why not grab the (only)user to this function in unswitchFunction and get the basic block from there. It adds some overhead, but IMO is cleaner than this.

[davidxl]

1. Addressed review feedbacks
2. More refactorings and naming cleanups (for consistency)
3. Added handling of function entry count update and a new test case for it.

[davidxl]

lib/Transforms/Utils/CodeExtractor.cpp

23

for getCallsiteCost and getInlineCost interface.

754

can you clarify on this?

[eraman]

lib/Transforms/Utils/CodeExtractor.cpp

23

I don't see any use of those interfaces in CodeExtractor.cpp. Am I missing something?

754

Something like cast<CallInst>(*F->user_begin())->getParent() in unswitchFunction should work right?

[davidxl]

Address feedback -- revert change in CodeExtractor.h and .cc file

Also let the cost analysis only trust PGO and user annotation data. Without PGO, add an option to control the aggressiveness.

[davidxl]

Further tuned cost/benefit computation with SPEC 2k/06 tuning. Get speedups in a few benchmarks and fixed the regression in perlbmk (many cases with early return branch which is never taken).

[dberris]

Drive-by-comment.

lib/Transforms/IPO/PartialInlining.cpp

755–758

nit: No need for {}'s in single-line bodies.

[davidxl]

Remove { }

[eraman]

In terms of correctness and heuristics this looks good. I'v a few comments that are mostly stylistic and documentation issues.

lib/Transforms/IPO/PartialInlining.cpp

49

make this cl::ReallyHidden since it is only used in testing

66

The description string and comment does not make it clear whether this is the upper limit (I think it is) or not.

120

typo: successfully -> successful

132–174

an empty line below

152

Empty line?

166

Not sure if this has been clang-formatted because it looks like the next two lines can be merged into one (I may be wrong)

169

Inline -> InlineCost. usedto->used to

401

I prefer auto here and below since the get methods clearly convey the type.

410

This certainly need some more comments. Essentially you don't want to over-estimate the savings and hence you assume that outlinined region is at least OutlineRegionFreqPercent likely even if BFI tells otherwise.

411

Flip the branch condition and so an early return?

412

std::min

421

auto here and in the next line.

527

move the BFI computation to a lambda maybe?

721

Instead of using two variables with similar names, why not initialize the uint64_t to 0 above the previous if and inside the if do F->getEntryCount()->getValue() ?

[davidxl]

lib/Transforms/IPO/PartialInlining.cpp

66

Updated the comment.

166

The ; will be @column 80 which clang-format does not like.

721

Both vars are used later. I changed the second var name to make it look friendiler.

[davidxl]

Address review feedbacks.

[eraman]

LGTM

lib/Transforms/IPO/PartialInlining.cpp

421

I think you can use auto for NormWeightedRCost also.

include/llvm/Transforms/Utils/CodeExtractor.h

public:
     /// block sequence of the loop.
     CodeExtractor(DominatorTree &DT, Loop &L, bool AggregateArgs = false,
                   BlockFrequencyInfo *BFI = nullptr,
                   BranchProbabilityInfo *BPI = nullptr);
 
     /// \brief Perform the extraction, returning the new function.
     ///
     /// Returns zero when called on a CodeExtractor instance where isEligible
-    /// returns false.
-    Function *extractCodeRegion();
+    /// returns false. If codeReplacer is not null, the pointer to the basic
+    /// block in which the call sequence to the extracted function will be
+    /// stored in *codeReplacer.
+    Function *extractCodeRegion(BasicBlock **codeReplacer = nullptr);
 
     /// \brief Test whether this code extractor is eligible.
     ///
     /// Based on the blocks used when constructing the code extractor,
     /// determine whether it is eligible for extraction.
     bool isEligible() const { return !Blocks.empty(); }
 
     /// \brief Compute the set of input values and output values for the code.

lib/Transforms/IPO/PartialInlining.cpp

 // that surrounds the body of the function.
 //
 //===----------------------------------------------------------------------===//
 
 #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/InlineCost.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/OptimizationDiagnosticInfo.h"
 #include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/DiagnosticInfo.h"

 
 STATISTIC(NumPartialInlined,
           "Number of callsites functions partially inlined into.");
 
 // 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"));
+// This is an option used by testing:
+static cl::opt<bool> SkipCostAnalysis("skip-partial-inlining-cost-analysis",
+                                      cl::init(false), cl::ZeroOrMore,
+                                      cl::Hidden,
+                                      cl::desc("Skip Cost Analysis"));
 
 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"));
 
 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;
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::unique_ptr<FunctionOutliningInfo> computeOutliningInfo(Function *F);
-
 private:
   int NumPartialInlining = 0;
   std::function<AssumptionCache &(Function &)> *GetAssumptionCache;
   std::function<TargetTransformInfo &(Function &)> *GetTTI;
   Optional<function_ref<BlockFrequencyInfo &(Function &)>> GetBFI;
   ProfileSummaryInfo *PSI;
 
-  bool shouldPartialInline(CallSite CS, OptimizationRemarkEmitter &ORE);
+  // Return the frequency of the OutlininingBB relative to F's entry point.
+  // The result is no larger than 1 and is represented using BP.
+  // (Note that the outlined region's 'head' block can only have incoming
+  // edges from the guarding entry blocks).
+  BranchProbability getOutliningCallBBRelativeFreq(Function *DuplicateFunction,
+                                                   BlockFrequencyInfo *BFI,
+                                                   BasicBlock *OutliningCallBB);
+
+  // Return true if the callee of CS should be partially inlined with
+  // profit.
+  bool shouldPartialInline(CallSite CS, BlockFrequencyInfo *CalleeBFI,
+                           BasicBlock *OutliningCallBB,
+                           int OutliningCallOverhead,
+                           OptimizationRemarkEmitter &ORE);
+
+  // Try to inline DuplicateFunction (cloned from F with call to
+  // the OutlinedFunction into its callers. Return true
+  // if there is any successfully inlining.
+  bool tryPartialInline(Function *DuplicateFunction,
+                        Function *F, /*orignal function */
+                        FunctionOutliningInfo *OI, Function *OutlinedFunction,
+                        BasicBlock *OutliningCallBB,
+                        BlockFrequencyInfo *CalleeBFI);
+
+  // 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);
   }
+  CallSite getCallSite(User *U) {
+    CallSite CS;
+    if (CallInst *CI = dyn_cast<CallInst>(U))
+      CS = CallSite(CI);
+    else if (InvokeInst *II = dyn_cast<InvokeInst>(U))
+      CS = CallSite(II);
+    else
+      llvm_unreachable("All uses must be calls");
+    return CS;
+  }
+
+  CallSite getOneCallSiteTo(Function *F) {
+    User *User = *F->user_begin();
+    return getCallSite(User);
+  }
+
+  std::tuple<DebugLoc, BasicBlock *> getOneDebugLoc(Function *F) {
+    CallSite CS = getOneCallSiteTo(F);
+
+    DebugLoc DLoc = CS.getInstruction()->getDebugLoc();
+    BasicBlock *Block = CS.getParent();
+    return std::make_tuple(DLoc, Block);
+  }
+
+  // Returns the costs associated with function outlining:
+  // - The first value is the non-weighted runtime cost for making the call
+  //   to the outlined function 'OutlinedFunction', including the addtional
+  //   setup cost in the outlined function itself;
+  // - The second value is the estimated size of the new call sequence in
+  //   basic block 'OutliningCallBB';
+  // - The third value is the estimated size of the original code from
+  //   function 'F' that is extracted into the outlined function.
+  std::tuple<int, int, int>
+  computeOutliningCosts(Function *F, const FunctionOutliningInfo *OutliningInfo,
+                        Function *OutlinedFunction,
+                        BasicBlock *OutliningCallBB);
+  // Compute the 'Inline' of block BB. InlineCost is a proxy usedto approximate
+  // both the size and runtime cost (Note that in the current inline cost
+  // analysis, there is no clear distinction there either).
+  int computeBBInlineCost(BasicBlock *BB);
+
+  std::unique_ptr<FunctionOutliningInfo> computeOutliningInfo(Function *F);
 };
 
 struct PartialInlinerLegacyPass : public ModulePass {
   static char ID; // Pass identification, replacement for typeid
   PartialInlinerLegacyPass() : ModulePass(ID) {
     initializePartialInlinerLegacyPassPass(*PassRegistry::getPassRegistry());
   }
 
PartialInlinerImpl::computeOutliningInfo(Function *F) {
   } while (true);
 
   if (!CandidateFound)
     return std::unique_ptr<FunctionOutliningInfo>();
 
   // Do sanity check of the entries: threre should not
   // be any successors (not in the entry set) other than
   // {ReturnBlock, NonReturnBlock}
-  assert(OutliningInfo->Entries[0] == &F->front());
+  assert(OutliningInfo->Entries[0] == &F->front() &&
+         "Function Entry must be the first in Entries vector");
   DenseSet<BasicBlock *> Entries;
   for (BasicBlock *E : OutliningInfo->Entries)
     Entries.insert(E);
 
   // Returns true of BB has Predecessor which is not
   // in Entries set.
   auto HasNonEntryPred = [Entries](BasicBlock *BB) {
     for (auto Pred : predecessors(BB)) {
while (OutliningInfo->GetNumInlinedBlocks() < MaxNumInlineBlocks) {
     OutliningInfo->NonReturnBlock = NonReturnBlock;
     OutliningInfo->ReturnBlockPreds.push_back(Cand);
     Entries.insert(Cand);
   }
 
   return OutliningInfo;
 }
 
+BranchProbability PartialInlinerImpl::getOutliningCallBBRelativeFreq(
+    Function *DuplicateFunction, BlockFrequencyInfo *BFI,
+    BasicBlock *OutliningCallBB) {
+  BlockFrequency EntryFreq =
+      BFI->getBlockFreq(&DuplicateFunction->getEntryBlock());
+  BlockFrequency OutliningCallFreq = BFI->getBlockFreq(OutliningCallBB);
+  return BranchProbability::getBranchProbability(
+      OutliningCallFreq.getFrequency(), EntryFreq.getFrequency());
+}
+
 bool PartialInlinerImpl::shouldPartialInline(CallSite CS,
+                                             BlockFrequencyInfo *CalleeBFI,
+                                             BasicBlock *OutliningCallBB,
+                                             int NonWeightedOutliningRcost,
                                              OptimizationRemarkEmitter &ORE) {
-  // TODO : more sharing with shouldInline in Inliner.cpp
   using namespace ore;
+  if (SkipCostAnalysis)
+    return true;
+
   Instruction *Call = CS.getInstruction();
   Function *Callee = CS.getCalledFunction();
   Function *Caller = CS.getCaller();
   auto &CalleeTTI = (*GetTTI)(*Callee);
   InlineCost IC = getInlineCost(CS, getInlineParams(), CalleeTTI,
                                 *GetAssumptionCache, GetBFI, PSI);
 
   if (IC.isAlways()) {
     ORE.emit(OptimizationRemarkAnalysis(DEBUG_TYPE, "AlwaysInline", Call)
              << NV("Callee", Callee)
              << " should always be fully inlined, not partially");
     return false;
   }
 
   if (IC.isNever()) {
     ORE.emit(OptimizationRemarkMissed(DEBUG_TYPE, "NeverInline", Call)
              << NV("Callee", Callee) << " not partially inlined into "
              << NV("Caller", Caller)
              << " because it should never be inlined (cost=never)");
     return false;
   }
 
   if (!IC) {
-    ORE.emit(OptimizationRemarkMissed(DEBUG_TYPE, "TooCostly", Call)
+    ORE.emit(OptimizationRemarkAnalysis(DEBUG_TYPE, "TooCostly", Call)
              << NV("Callee", Callee) << " not partially inlined into "
              << NV("Caller", Caller) << " because too costly to inline (cost="
              << NV("Cost", IC.getCost()) << ", threshold="
              << NV("Threshold", IC.getCostDelta() + IC.getCost()) << ")");
     return false;
   }
+  const DataLayout &DL = Caller->getParent()->getDataLayout();
+  // The savings of eliminating the call:
+  int NonWeightedSavings = getCallsiteCost(CS, DL);
+  BlockFrequency NormWeightedSavings(NonWeightedSavings);
+
+  BranchProbability RelativeFreq =
+      getOutliningCallBBRelativeFreq(Callee, CalleeBFI, OutliningCallBB);
+  BlockFrequency NormWeightedRcost =
+      BlockFrequency(NonWeightedOutliningRcost) * RelativeFreq;
+
+  // Weighted saving is smaller than weighted cost, return false
+  if (NormWeightedSavings < NormWeightedRcost) {
+    ORE.emit(
+        OptimizationRemarkAnalysis(DEBUG_TYPE, "OutliningCallcostTooHigh", Call)
+        << NV("Callee", Callee) << " not partially inlined into "
+        << NV("Caller", Caller) << " runtime overhead (overhead="
+        << NV("Overhead", (unsigned)NormWeightedRcost.getFrequency())
+        << ", savings="
+        << NV("Savings", (unsigned)NormWeightedSavings.getFrequency()) << ")"
+        << " of making the outlined call is too high");
+
+    return false;
+  }
 
   ORE.emit(OptimizationRemarkAnalysis(DEBUG_TYPE, "CanBePartiallyInlined", Call)
            << NV("Callee", Callee) << " can be partially inlined into "
            << NV("Caller", Caller) << " with cost=" << NV("Cost", IC.getCost())
            << " (threshold="
            << NV("Threshold", IC.getCostDelta() + IC.getCost()) << ")");
   return true;
 }
 
+// TODO: Ideally  we should share Inliner's InlineCost Analysis code.
+// For now use a simplified version. The returned 'InlineCost' will be used
+// to esimate the size cost as well as runtime cost of the BB.
+int PartialInlinerImpl::computeBBInlineCost(BasicBlock *BB) {
+  int InlineCost = 0;
+  for (BasicBlock::iterator I = BB->begin(), E = BB->end(); I != E; ++I) {
+    if (isa<DbgInfoIntrinsic>(I))
+      continue;
+
+    if (dyn_cast<CallInst>(I)) {
+      InlineCost += InlineConstants::CallPenalty;
+      continue;
+    }
+    if (dyn_cast<InvokeInst>(I)) {
+      InlineCost += InlineConstants::CallPenalty;
+      continue;
+    }
+
+    if (SwitchInst *SI = dyn_cast<SwitchInst>(I)) {
+      InlineCost += (SI->getNumCases() + 1) * InlineConstants::InstrCost;
+      continue;
+    }
+    InlineCost += InlineConstants::InstrCost;
+  }
+  return InlineCost;
+}
+
+std::tuple<int, int, int> PartialInlinerImpl::computeOutliningCosts(
+    Function *F, const FunctionOutliningInfo *OI, Function *OutlinedFunction,
+    BasicBlock *OutliningCallBB) {
+  // First compute the cost of the outlined region 'OI' in the original
+  // function 'F':
+  int OutlinedRegionCost = 0;
+  for (BasicBlock &BB : *F) {
+    if (&BB != OI->ReturnBlock &&
+        // Assuming Entry set is small -- do a linear search here:
+        std::find(OI->Entries.begin(), OI->Entries.end(), &BB) ==
+            OI->Entries.end()) {
+      OutlinedRegionCost += computeBBInlineCost(&BB);
+    }
+  }
+
+  // Now compute the cost of the call sequence to the outlined function
+  // 'OutlinedFunction' in BB 'OutliningCallBB':
+  int OutliningFuncCallCost = computeBBInlineCost(OutliningCallBB);
+
+  // Now compute the cost of the extracted/outlined function itself:
+  int OutlinedFunctionCost = 0;
+  for (BasicBlock &BB : *OutlinedFunction) {
+    OutlinedFunctionCost += computeBBInlineCost(&BB);
+  }
+
+  assert(OutlinedFunctionCost >= OutlinedRegionCost &&
+         "Outlined function cost should be no less than the outlined region");
+  int OutliningRuntimeOverhead =
+      OutliningFuncCallCost + (OutlinedFunctionCost - OutlinedRegionCost);
+
+  return std::make_tuple(OutliningFuncCallCost, OutliningRuntimeOverhead,
+                         OutlinedRegionCost);
+}
+
+// Create the callsite to profile count map which is
+// used to update the original function's entry count,
+// after the function is partially inlined into the callsite.
+void PartialInlinerImpl::computeCallsiteToProfCountMap(
+    Function *DuplicateFunction,
+    DenseMap<User *, uint64_t> &CallSiteToProfCountMap) {
+  std::vector<User *> Users(DuplicateFunction->user_begin(),
+                            DuplicateFunction->user_end());
+  Function *CurrentCaller = nullptr;
+  BlockFrequencyInfo *CurrentCallerBFI = nullptr;
+  for (User *User : Users) {
+    CallSite CS = getCallSite(User);
+    Function *Caller = CS.getCaller();
+    if (CurrentCaller != Caller) {
+      CurrentCaller = Caller;
+      // For the old pass manager:
+      if (!GetBFI) {
+        if (CurrentCallerBFI)
+          delete CurrentCallerBFI;
+        DominatorTree DT(*Caller);
+        LoopInfo LI(DT);
+        BranchProbabilityInfo BPI(*Caller, LI);
+        CurrentCallerBFI = new BlockFrequencyInfo(*Caller, BPI, LI);
+      } else {
+        // New pass manager:
+        CurrentCallerBFI = &(*GetBFI)(*Caller);
+      }
+    } else {
+      assert(CurrentCallerBFI && "CallerBFI is not set");
+    }
+    BasicBlock *CallBB = CS.getInstruction()->getParent();
+    auto Count = CurrentCallerBFI->getBlockProfileCount(CallBB);
+    if (Count)
+      CallSiteToProfCountMap[User] = *Count;
+    else
+      CallSiteToProfCountMap[User] = 0;
+  }
+}
+
 Function *PartialInlinerImpl::unswitchFunction(Function *F) {
 
   if (F->hasAddressTaken())
     return nullptr;
 
   // Let inliner handle it
   if (F->hasFnAttribute(Attribute::AlwaysInline))
     return nullptr;
 
   if (F->hasFnAttribute(Attribute::NoInline))
     return nullptr;
 
   if (PSI->isFunctionEntryCold(F))
     return nullptr;
 
-  std::unique_ptr<FunctionOutliningInfo> OutliningInfo =
-      computeOutliningInfo(F);
+  if (F->user_begin() == F->user_end())
+    return nullptr;
 
-  if (!OutliningInfo)
+  std::unique_ptr<FunctionOutliningInfo> OI = computeOutliningInfo(F);
+
+  if (!OI)
     return nullptr;
 
   // Clone the function, so that we can hack away on it.
   ValueToValueMapTy VMap;
   Function *DuplicateFunction = CloneFunction(F, VMap);
-  BasicBlock *NewReturnBlock =
-      cast<BasicBlock>(VMap[OutliningInfo->ReturnBlock]);
-  BasicBlock *NewNonReturnBlock =
-      cast<BasicBlock>(VMap[OutliningInfo->NonReturnBlock]);
+  BasicBlock *NewReturnBlock = cast<BasicBlock>(VMap[OI->ReturnBlock]);
+  BasicBlock *NewNonReturnBlock = cast<BasicBlock>(VMap[OI->NonReturnBlock]);
   DenseSet<BasicBlock *> NewEntries;
-  for (BasicBlock *BB : OutliningInfo->Entries) {
+  for (BasicBlock *BB : OI->Entries) {
     NewEntries.insert(cast<BasicBlock>(VMap[BB]));
   }
 
   // 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(DuplicateFunction);
 
   auto getFirstPHI = [](BasicBlock *BB) {
Function *PartialInlinerImpl::unswitchFunction(Function *F) {
   };
   // 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 = NewReturnBlock;
   // only split block when necessary:
   PHINode *FirstPhi = getFirstPHI(PreReturn);
-  unsigned NumPredsFromEntries = OutliningInfo->ReturnBlockPreds.size();
+  unsigned NumPredsFromEntries = OI->ReturnBlockPreds.size();
   if (FirstPhi && FirstPhi->getNumIncomingValues() > NumPredsFromEntries + 1) {
 
     NewReturnBlock = NewReturnBlock->splitBasicBlock(
         NewReturnBlock->getFirstNonPHI()->getIterator());
     BasicBlock::iterator I = PreReturn->begin();
     Instruction *Ins = &NewReturnBlock->front();
     while (I != PreReturn->end()) {
       PHINode *OldPhi = dyn_cast<PHINode>(I);
       if (!OldPhi)
         break;
 
       PHINode *RetPhi =
           PHINode::Create(OldPhi->getType(), NumPredsFromEntries + 1, "", Ins);
       OldPhi->replaceAllUsesWith(RetPhi);
       Ins = NewReturnBlock->getFirstNonPHI();
 
       RetPhi->addIncoming(&*I, PreReturn);
-      for (BasicBlock *E : OutliningInfo->ReturnBlockPreds) {
+      for (BasicBlock *E : OI->ReturnBlockPreds) {
         BasicBlock *NewE = cast<BasicBlock>(VMap[E]);
         RetPhi->addIncoming(OldPhi->getIncomingValueForBlock(NewE), NewE);
         OldPhi->removeIncomingValue(NewE);
       }
       ++I;
     }
-    for (auto E : OutliningInfo->ReturnBlockPreds) {
+    for (auto E : OI->ReturnBlockPreds) {
       BasicBlock *NewE = cast<BasicBlock>(VMap[E]);
       NewE->getTerminator()->replaceUsesOfWith(PreReturn, NewReturnBlock);
     }
   }
 
   // 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 = [=](BasicBlock *BB) {
Function *PartialInlinerImpl::unswitchFunction(Function *F) {
   DT.recalculate(*DuplicateFunction);
 
   // Manually calculate a BlockFrequencyInfo and BranchProbabilityInfo.
   LoopInfo LI(DT);
   BranchProbabilityInfo BPI(*DuplicateFunction, LI);
   BlockFrequencyInfo BFI(*DuplicateFunction, BPI, LI);
 
   // Extract the body of the if.
-  Function *ExtractedFunction =
+  BasicBlock *OutliningCallBB = nullptr;
+  Function *OutlinedFunction =
       CodeExtractor(ToExtract, &DT, /*AggregateArgs*/ false, &BFI, &BPI)
-          .extractCodeRegion();
+          .extractCodeRegion(&OutliningCallBB);
+
+  bool AnyInline = tryPartialInline(DuplicateFunction, F, OI.get(),
+                                    OutlinedFunction, OutliningCallBB, &BFI);
+
+  // Ditch the duplicate, since we're done with it, and rewrite all remaining
+  // users (function pointers, etc.) back to the original function.
+  DuplicateFunction->replaceAllUsesWith(F);
+  DuplicateFunction->eraseFromParent();
+  if (!AnyInline && OutlinedFunction)
+    OutlinedFunction->eraseFromParent();
+  return OutlinedFunction;
+}
 
-  // Inline the top-level if test into all callers.
+bool PartialInlinerImpl::tryPartialInline(Function *DuplicateFunction,
+                                          Function *F,
+                                          FunctionOutliningInfo *OI,
+                                          Function *OutlinedFunction,
+                                          BasicBlock *OutliningCallBB,
+                                          BlockFrequencyInfo *CalleeBFI) {
+  if (OutlinedFunction == nullptr)
+    return false;
+
+  int NonWeightedRcost;
+  int SizeCost;
+  int OutlinedRegionSizeCost;
+
+  std::tie(SizeCost, NonWeightedRcost, OutlinedRegionSizeCost) =
+      computeOutliningCosts(F, OI, OutlinedFunction, OutliningCallBB);
+
+  // The call sequence to the outlined function is larger than the original
+  // outlined region size, it does not increase the chances of inlining
+  // 'F' with outlining (The inliner usies the size increase to model the
+  // the cost of inlining a callee).
+  if (!SkipCostAnalysis && OutlinedRegionSizeCost < SizeCost) {
+    OptimizationRemarkEmitter ORE(F);
+    DebugLoc DLoc;
+    BasicBlock *Block;
+    std::tie(DLoc, Block) = getOneDebugLoc(DuplicateFunction);
+    ORE.emit(OptimizationRemarkAnalysis(DEBUG_TYPE, "OutlineRegionTooSmall",
+                                        DLoc, Block)
+             << ore::NV("Function", F)
+             << " not partially inlined into callers (Original Size = "
+             << ore::NV("OutlinedRegionOriginalSize", OutlinedRegionSizeCost)
+             << ", Size of call sequence to outlined function = "
+             << ore::NV("NewSize", SizeCost) << ")");
+    return false;
+  }
+
+  assert(F->user_begin() == OrigninalFunction->user_end() &&
+         "F's users should all be replaced!");
   std::vector<User *> Users(DuplicateFunction->user_begin(),
                             DuplicateFunction->user_end());
 
+  DenseMap<User *, uint64_t> CallSiteToProfCountMap;
+  if (F->getEntryCount())
+    computeCallsiteToProfCountMap(DuplicateFunction, CallSiteToProfCountMap);
+
+  auto CalleeEntryCount = F->getEntryCount();
+  uint64_t CalleeEntryC = (CalleeEntryCount ? *CalleeEntryCount : 0);
+  bool AnyInline = false;
   for (User *User : Users) {
-    CallSite CS;
-    if (CallInst *CI = dyn_cast<CallInst>(User))
-      CS = CallSite(CI);
-    else if (InvokeInst *II = dyn_cast<InvokeInst>(User))
-      CS = CallSite(II);
-    else
-      llvm_unreachable("All uses must be calls");
+    CallSite CS = getCallSite(User);
 
     if (IsLimitReached())
       continue;
 
     OptimizationRemarkEmitter ORE(CS.getCaller());
-    if (!shouldPartialInline(CS, ORE))
+
+    if (!shouldPartialInline(CS, CalleeBFI, OutliningCallBB, NonWeightedRcost,
+                             ORE))
       continue;
 
-    DebugLoc DLoc = CS.getInstruction()->getDebugLoc();
-    BasicBlock *Block = CS.getParent();
-    ORE.emit(OptimizationRemark(DEBUG_TYPE, "PartiallyInlined", DLoc, Block)
-             << ore::NV("Callee", F) << " partially inlined into "
-             << ore::NV("Caller", CS.getCaller()));
+    ORE.emit(
+        OptimizationRemark(DEBUG_TYPE, "PartiallyInlined", CS.getInstruction())
 
     InlineFunctionInfo IFI(nullptr, GetAssumptionCache);
     InlineFunction(CS, IFI);
+
+    // Now update the entry count:
+    if (CalleeEntryC && CallSiteToProfCountMap.count(User)) {
+      uint64_t CallSiteCount = CallSiteToProfCountMap[User];
+      CalleeEntryC -= std::min(CalleeEntryC, CallSiteCount);
+    }
+
+    AnyInline = true;
     NumPartialInlining++;
-    // update stats
+    // Update the stats
     NumPartialInlined++;
   }
 
-  // Ditch the duplicate, since we're done with it, and rewrite all remaining
-  // users (function pointers, etc.) back to the original function.
-  DuplicateFunction->replaceAllUsesWith(F);
-  DuplicateFunction->eraseFromParent();
-
+  if (AnyInline && CalleeEntryCount) {
+    F->setEntryCount(CalleeEntryC);
+  }
 
-  return ExtractedFunction;
+  return AnyInline;
 }
 
 bool PartialInlinerImpl::run(Module &M) {
   if (DisablePartialInlining)
     return false;
 
   std::vector<Function *> Worklist;
   Worklist.reserve(M.size());

lib/Transforms/Utils/CodeExtractor.cpp

 
 #include "llvm/Transforms/Utils/CodeExtractor.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/Analysis/BlockFrequencyInfo.h"
 #include "llvm/Analysis/BlockFrequencyInfoImpl.h"
 #include "llvm/Analysis/BranchProbabilityInfo.h"
+#include "llvm/Analysis/InlineCost.h"
 #include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/RegionInfo.h"
 #include "llvm/Analysis/RegionIterator.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DerivedTypes.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/Intrinsics.h"

     else
       BPI->setEdgeProbability(CodeReplacer, i, BranchProbability::getZero());
   }
 
   // Check for no total weight.
   if (BranchDist.Total == 0)
     return;
 
   // Normalize the distribution so that they can fit in unsigned.
   BranchDist.normalize();
 
   // Create normalized branch weights and set the metadata.
   for (unsigned I = 0, E = BranchDist.Weights.size(); I < E; ++I) {
     const auto &Weight = BranchDist.Weights[I];
 
     // Get the weight and update the current BFI.
     BranchWeights[Weight.TargetNode.Index] = Weight.Amount;
     BranchProbability BP(Weight.Amount, BranchDist.Total);
     BPI->setEdgeProbability(CodeReplacer, Weight.TargetNode.Index, BP);
   }
   TI->setMetadata(
       LLVMContext::MD_prof,
       MDBuilder(TI->getContext()).createBranchWeights(BranchWeights));
 }
 
-Function *CodeExtractor::extractCodeRegion() {
+Function *CodeExtractor::extractCodeRegion(BasicBlock **codeReplacerP) {
   if (!isEligible())
     return nullptr;
 
   ValueSet inputs, outputs;
 
   // Assumption: this is a single-entry code region, and the header is the first
   // block in the region.
   BasicBlock *header = *Blocks.begin();

   splitReturnBlocks();
 
   Function *oldFunction = header->getParent();
 
   // This takes place of the original loop
   BasicBlock *codeReplacer = BasicBlock::Create(header->getContext(), 
                                                 "codeRepl", oldFunction,
                                                 header);
+  if (codeReplacerP)
+    *codeReplacerP = codeReplacer;
 
   // The new function needs a root node because other nodes can branch to the
   // head of the region, but the entry node of a function cannot have preds.
   BasicBlock *newFuncRoot = BasicBlock::Create(header->getContext(), 
                                                "newFuncRoot");
   newFuncRoot->getInstList().push_back(BranchInst::Create(header));
 
   // Find inputs to, outputs from the code region.

test/Transforms/CodeExtractor/ExtractedFnEntryCount.ll

-; RUN: opt < %s -partial-inliner -S | FileCheck %s
+; RUN: opt < %s -partial-inliner -skip-partial-inlining-cost-analysis -S | FileCheck %s
 
 ; This test checks to make sure that the CodeExtractor
 ;  properly sets the entry count for the function that is
 ;  extracted based on the root block being extracted and also
 ;  takes into consideration if the block has edges coming from
 ;  a block that is also being extracted.
 
 define i32 @inlinedFunc(i1 %cond) !prof !1 {

test/Transforms/CodeExtractor/MultipleExitBranchProb.ll

-; RUN: opt < %s -partial-inliner -max-num-inline-blocks=2 -S | FileCheck %s
+; RUN: opt < %s -partial-inliner -max-num-inline-blocks=2 -skip-partial-inlining-cost-analysis -S | FileCheck %s
 
 ; This test checks to make sure that CodeExtractor updates
 ;  the exit branch probabilities for multiple exit blocks.
 
 define i32 @inlinedFunc(i1 %cond) !prof !1 {
 entry:
   br i1 %cond, label %if.then, label %return, !prof !2
 if.then:

test/Transforms/CodeExtractor/PartialInlineEntryUpdate.ll

+; RUN: opt < %s -skip-partial-inlining-cost-analysis -partial-inliner -S  | FileCheck %s
+; RUN: opt < %s -skip-partial-inlining-cost-analysis -passes=partial-inliner -S  | FileCheck %s
+
+define i32 @Func(i1 %cond, i32* align 4 %align.val) !prof !1 {
+; CHECK: @Func({{.*}}) !prof [[REMAINCOUNT:![0-9]+]]
+entry:
+  br i1 %cond, label %if.then, label %return
+if.then:
+  ; Dummy store to have more than 0 uses
+  store i32 10, i32* %align.val, align 4
+  br label %return
+return:             ; preds = %entry
+  ret i32 0
+}
+
+define internal i32 @Caller1(i1 %cond, i32* align 2 %align.val) !prof !3{
+entry:
+; CHECK-LABEL: @Caller1
+; CHECK: br
+; CHECK: call void @Func.1_ 
+; CHECK: br
+; CHECK: call void @Func.1_ 
+  %val = call i32 @Func(i1 %cond, i32* %align.val)
+  %val2 = call i32 @Func(i1 %cond, i32* %align.val)
+  ret i32 %val
+}
+
+define internal i32 @Caller2(i1 %cond, i32* align 2 %align.val) !prof !2{
+entry:
+; CHECK-LABEL: @Caller2
+; CHECK: br
+; CHECK: call void @Func.1_ 
+  %val = call i32 @Func(i1 %cond, i32* %align.val)
+  ret i32 %val
+}
+
+; CHECK: [[REMAINCOUNT]] = !{!"function_entry_count", i64 150}
+!1 = !{!"function_entry_count", i64 200}
+!2 = !{!"function_entry_count", i64 10}
+!3 = !{!"function_entry_count", i64 20}
+

test/Transforms/CodeExtractor/PartialInlineHighCost.ll

+; The outlined region has high frequency  and the outlining
+; call sequence is expensive (input, output, multiple exit etc)
+; RUN: opt < %s -partial-inliner -max-num-inline-blocks=2 -S | FileCheck %s
+; RUN: opt < %s -passes=partial-inliner -max-num-inline-blocks=2 -S | FileCheck %s
+; RUN: opt < %s -partial-inliner -skip-partial-inlining-cost-analysis -max-num-inline-blocks=2 -S | FileCheck --check-prefix=NOCOST %s
+; RUN: opt < %s -passes=partial-inliner -skip-partial-inlining-cost-analysis -max-num-inline-blocks=2 -S | FileCheck  --check-prefix=NOCOST %s
+
+
+; Function Attrs: nounwind
+define i32 @bar_hot_outline_region(i32 %arg) local_unnamed_addr #0 {
+bb:
+  %tmp = icmp slt i32 %arg, 0
+  br i1 %tmp, label %bb1, label %bb16, !prof !1
+
+bb1:                                              ; preds = %bb
+  %tmp2 = tail call i32 (...) @foo() #0
+  %tmp3 = tail call i32 (...) @foo() #0
+  %tmp4 = tail call i32 (...) @foo() #0
+  %tmp5 = tail call i32 (...) @foo() #0
+  %tmp6 = tail call i32 (...) @foo() #0
+  %tmp7 = tail call i32 (...) @foo() #0
+  %tmp8 = add nsw i32 %arg, 1
+  %tmp9 = tail call i32 @goo(i32 %tmp8) #0
+  %tmp10 = tail call i32 (...) @foo() #0
+  %tmp11 = icmp eq i32 %tmp10, 0
+  br i1 %tmp11, label %bb12, label %bb16
+
+bb12:                                             ; preds = %bb1
+  %tmp13 = tail call i32 (...) @foo() #0
+  %tmp14 = icmp eq i32 %tmp13, 0
+  %tmp15 = select i1 %tmp14, i32 0, i32 3
+  br label %bb16
+
+bb16:                                             ; preds = %bb12, %bb1, %bb
+  %tmp17 = phi i32 [ 2, %bb1 ], [ %tmp15, %bb12 ], [ 0, %bb ]
+  ret i32 %tmp17
+}
+
+define i32 @bar_cold_outline_region(i32 %arg) local_unnamed_addr #0 {
+bb:
+  %tmp = icmp slt i32 %arg, 0
+  br i1 %tmp, label %bb1, label %bb16, !prof !2
+
+bb1:                                              ; preds = %bb
+  %tmp2 = tail call i32 (...) @foo() #0
+  %tmp3 = tail call i32 (...) @foo() #0
+  %tmp4 = tail call i32 (...) @foo() #0
+  %tmp5 = tail call i32 (...) @foo() #0
+  %tmp6 = tail call i32 (...) @foo() #0
+  %tmp7 = tail call i32 (...) @foo() #0
+  %tmp8 = add nsw i32 %arg, 1
+  %tmp9 = tail call i32 @goo(i32 %tmp8) #0
+  %tmp10 = tail call i32 (...) @foo() #0
+  %tmp11 = icmp eq i32 %tmp10, 0
+  br i1 %tmp11, label %bb12, label %bb16
+
+bb12:                                             ; preds = %bb1
+  %tmp13 = tail call i32 (...) @foo() #0
+  %tmp14 = icmp eq i32 %tmp13, 0
+  %tmp15 = select i1 %tmp14, i32 0, i32 3
+  br label %bb16
+
+bb16:                                             ; preds = %bb12, %bb1, %bb
+  %tmp17 = phi i32 [ 2, %bb1 ], [ %tmp15, %bb12 ], [ 0, %bb ]
+  ret i32 %tmp17
+}
+
+; Function Attrs: nounwind
+declare i32 @foo(...) local_unnamed_addr #0
+
+; Function Attrs: nounwind
+declare i32 @goo(i32) local_unnamed_addr #0
+
+; Function Attrs: nounwind
+define i32 @dummy_caller(i32 %arg) local_unnamed_addr #0 {
+bb:
+; CHECK-LABEL: @dummy_caller
+; CHECK-NOT: br i1
+; CHECK-NOT: call{{.*}}bar_hot_outline_region. 
+; NOCOST-LABEL: @dummy_caller
+; NOCOST: br i1
+; NOCOST: call{{.*}}bar_hot_outline_region.
+
+  %tmp = tail call i32 @bar_hot_outline_region(i32 %arg)
+  ret i32 %tmp
+}
+
+define i32 @dummy_caller2(i32 %arg) local_unnamed_addr #0 {
+bb:
+; CHECK-LABEL: @dummy_caller2
+; CHECK: br i1
+; CHECK: call{{.*}}bar_cold_outline_region.
+; NOCOST-LABEL: @dummy_caller2
+; NOCOST: br i1
+; NOCOST: call{{.*}}bar_cold_outline_region.
+
+  %tmp = tail call i32 @bar_cold_outline_region(i32 %arg)
+  ret i32 %tmp
+}
+
+attributes #0 = { nounwind }
+
+!llvm.ident = !{!0}
+
+!0 = !{!"clang version 5.0.0 (trunk 301898)"}
+!1 = !{!"branch_weights", i32 2000, i32 1}
+!2 = !{!"branch_weights", i32 1, i32 100}

test/Transforms/CodeExtractor/PartialInlineOrAnd.ll

 ; RUN: opt < %s -partial-inliner -S | FileCheck %s
 ; RUN: opt < %s -passes=partial-inliner -S | FileCheck %s
-; RUN: opt < %s -partial-inliner -max-num-inline-blocks=3 -S | FileCheck --check-prefix=LIMIT3 %s
-; RUN: opt < %s -passes=partial-inliner -max-num-inline-blocks=3 -S | FileCheck  --check-prefix=LIMIT3 %s
+; RUN: opt < %s -partial-inliner -max-num-inline-blocks=3 -skip-partial-inlining-cost-analysis  -S | FileCheck --check-prefix=LIMIT3 %s
+; RUN: opt < %s -passes=partial-inliner -max-num-inline-blocks=3 -skip-partial-inlining-cost-analysis -S | FileCheck  --check-prefix=LIMIT3 %s
 ; RUN: opt < %s -partial-inliner -max-num-inline-blocks=2 -S | FileCheck --check-prefix=LIMIT2 %s
 ; RUN: opt < %s -passes=partial-inliner -max-num-inline-blocks=2 -S | FileCheck  --check-prefix=LIMIT2 %s
 
 
 ; Function Attrs: nounwind uwtable
 define i32 @bar(i32 %arg) local_unnamed_addr #0 {
 bb:
   %tmp = icmp slt i32 %arg, 0

test/Transforms/CodeExtractor/SingleCondition.ll

-; RUN: opt < %s -partial-inliner -S  | FileCheck %s
-; RUN: opt < %s -passes=partial-inliner -S  | FileCheck %s
+; RUN: opt < %s -skip-partial-inlining-cost-analysis -partial-inliner -S  | FileCheck %s
+; RUN: opt < %s -skip-partial-inlining-cost-analysis -passes=partial-inliner -S  | FileCheck %s
 
 define internal i32 @inlinedFunc(i1 %cond, i32* align 4 %align.val) {
 entry:
   br i1 %cond, label %if.then, label %return
 if.then:
   ; Dummy store to have more than 0 uses
   store i32 10, i32* %align.val, align 4
   br label %return

test/Transforms/CodeExtractor/X86/InheritTargetAttributes.ll

-; RUN: opt < %s -partial-inliner | llc -filetype=null
-; RUN: opt < %s -partial-inliner -S | FileCheck %s
+; RUN: opt < %s -partial-inliner -skip-partial-inlining-cost-analysis | llc -filetype=null
+; RUN: opt < %s -partial-inliner -skip-partial-inlining-cost-analysis -S | FileCheck %s
 ; This testcase checks to see if CodeExtractor properly inherits
 ;   target specific attributes for the extracted function. This can
 ;   cause certain instructions that depend on the attributes to not
 ;   be lowered. Like in this test where we try to 'select' the blendvps
 ;   intrinsic on x86 that requires the +sse4.1 target feature.
 
 target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"