CodeExtractor : Add ability to preserve profile data.

[rriddle]

Summary

Added ability to estimate the entry count of the extracted function and the branch probabilities of the exit branches.

[rriddle]

Hard exit if an exit frequency is zero.

[rriddle]

Formatting

[davidxl]

include/llvm/Analysis/BlockFrequencyInfo.h

67

nit: perhaps just freqToProfileCount or getProfileCountFromFreq?

include/llvm/Analysis/BlockFrequencyInfoImpl.h

485

Same here

lib/Transforms/Utils/CodeExtractor.cpp

744

clang-format? add space

747

EntryCount --> EntryFreq

806

nit: OptEntryCount --> EntryCount

808

When EntryCount does not exist, skip it instead of setting it to 0.

816

Move this into a helper function

[rriddle]

Updated after Davids comments

[silvas]

Sorry for the delay on getting to this.

lib/Transforms/IPO/PartialInlining.cpp

149

Can you add a comment before the false showing the name of the argument? E.g. CodeExtractor(ToExtract, &DT, /*TheArgName*/false, BFI).extractCodeRegion();

lib/Transforms/Utils/CodeExtractor.cpp

690

Can you explain more how we would get into a situation with a weight of zero and why we need to check it? maybe parts of this should be proper methods on BFI? E.g. below you are directly accessing the impl class so it sounds like something here should maybe be a method on BFI.

also, a nit: can't

742

Should we change the getBPI method to return non-const? I assume the author made them const for a reason. E.g. BFI might be relying on BPI not changing under it.

808

nit: no space after the open paren

test/Transforms/CodeExtractor/2016-07-20-MultipleExitBranchProb.ll

30

Can you do a check with a pattern match and capture a variable, like you do for the other test?

You should be able to do something like test/Transforms/PGOProfile/branch1.ll to check the profile metadata on a particular instruction:

; USE: br i1 %cmp, label %if.then, label %if.end
; USE-SAME: !prof ![[BW_ENTRY:[0-9]+]]
; USE-DAG: ![[BW_ENTRY]] = !{!"branch_weights", i32 2, i32 1}

[davidxl]

lib/Transforms/Utils/CodeExtractor.cpp

693

This check does not look correct. It is possible for the frequency to be zero

805

The Freq > 0 check can be removed. I don't see how it can be zero (it is not actual count).

808

The check should be here: if optEntryCount has value (hasValue() returns true, not >0), set it to the newFunction entry. hasValue is false when PGO is not enabled.

[rriddle]

lib/Transforms/Utils/CodeExtractor.cpp

693

Yeah, that was my understanding. Thanks for the sanity check! I need to add the weights without using the current methods given that they do not allow for weights of 0.

[rriddle]

lib/Transforms/Utils/CodeExtractor.cpp

808

Thanks for the clarification, I was unsure if it should just be for the entry count.

[rriddle]

lib/Transforms/Utils/CodeExtractor.cpp

693

Although, can you help clarify on what I should be doing in the case of a 0 frequency? I am unsure on whether the branch weight can be set to a hard 0 or if it needs to be modified to always have a non zero weight?

[davidxl]

lib/Transforms/Utils/CodeExtractor.cpp

693

zero weight is fine -- BFI will always adds 1 to it if it sees it due to some limitation (it is less than ideal -- a better way to is to fix BPI (which will be fixed in the near future).

One the other hand if the total weight is zero, you may want to treat it as invalid.

[rriddle]

lib/Transforms/Utils/CodeExtractor.cpp

693

Awesome. Thanks for the clarification!

[rriddle]

Updated after Sean and David's comments

[davidxl]

lgtm with the small issues (in new comments) fixed.

lib/Transforms/Utils/CodeExtractor.cpp

742

BPI check is not needed.

HasProfileData --> HasBFI because HasProfileData usually means real profile data exists (BFI can be statically computed).

808

if (OptEntryCount->hasValue())

[rriddle]

Updated after David's comments

[silvas]

Some nits.

lib/Transforms/Utils/CodeExtractor.cpp

697

Missing '.' at the end of this comment.

741

should this be "entry frequency" or "entry count" instead of just "entry"?

782

Add a comment here about what this loop is doing.

805

The function is newly created, so this should probably say "Calculate and set" instead of "Update".

[rriddle]

Updated to reflect Sean's comments

[silvas]

Committed in r277313

[silvas]

Thanks!

[silvas]

I've had to revert this in r277317 due to an LSan failure. Can you take a look?

http://lab.llvm.org:8011/builders/sanitizer-x86_64-linux-fast/builds/15140/steps/check-llvm%20asan/logs/stdio

[rriddle]

Manual computation of BFI and BPI in partial inliner.

include/llvm/Analysis/BlockFrequencyInfo.h

public:
   /// floating points.
   BlockFrequency getBlockFreq(const BasicBlock *BB) const;
 
   /// \brief Returns the estimated profile count of \p BB.
   /// This computes the relative block frequency of \p BB and multiplies it by
   /// the enclosing function's count (if available) and returns the value.
   Optional<uint64_t> getBlockProfileCount(const BasicBlock *BB) const;
 
+  /// \brief Returns the estimated profile count of \p Freq.
+  /// This uses the frequency \p Freq and multiplies it by
+  /// the enclosing function's count (if available) and returns the value.
+  Optional<uint64_t> getProfileCountFromFreq(uint64_t Freq) const;
+
   // Set the frequency of the given basic block.
   void setBlockFreq(const BasicBlock *BB, uint64_t Freq);
 
   /// calculate - compute block frequency info for the given function.
   void calculate(const Function &F, const BranchProbabilityInfo &BPI,
                  const LoopInfo &LI);
 
   // Print the block frequency Freq to OS using the current functions entry

include/llvm/Analysis/BlockFrequencyInfoImpl.h

public:
   virtual raw_ostream &print(raw_ostream &OS) const { return OS; }
   void dump() const { print(dbgs()); }
 
   Scaled64 getFloatingBlockFreq(const BlockNode &Node) const;
 
   BlockFrequency getBlockFreq(const BlockNode &Node) const;
   Optional<uint64_t> getBlockProfileCount(const Function &F,
                                           const BlockNode &Node) const;
+  Optional<uint64_t> getProfileCountFromFreq(const Function &F,
+                                         uint64_t Freq) const;
 
   void setBlockFreq(const BlockNode &Node, uint64_t Freq);
 
   raw_ostream &printBlockFreq(raw_ostream &OS, const BlockNode &Node) const;
   raw_ostream &printBlockFreq(raw_ostream &OS,
                               const BlockFrequency &Freq) const;
 
   uint64_t getEntryFreq() const {
public:
   using BlockFrequencyInfoImplBase::getEntryFreq;
   BlockFrequency getBlockFreq(const BlockT *BB) const {
     return BlockFrequencyInfoImplBase::getBlockFreq(getNode(BB));
   }
   Optional<uint64_t> getBlockProfileCount(const Function &F,
                                           const BlockT *BB) const {
     return BlockFrequencyInfoImplBase::getBlockProfileCount(F, getNode(BB));
   }
+  Optional<uint64_t> getProfileCountFromFreq(const Function &F,
+                                         uint64_t Freq) const {
+    return BlockFrequencyInfoImplBase::getProfileCountFromFreq(F, Freq);
+  }
   void setBlockFreq(const BlockT *BB, uint64_t Freq);
   Scaled64 getFloatingBlockFreq(const BlockT *BB) const {
     return BlockFrequencyInfoImplBase::getFloatingBlockFreq(getNode(BB));
   }
 
   const BranchProbabilityInfoT &getBPI() const { return *BPI; }
 
   /// \brief Print the frequencies for the current function.

include/llvm/CodeGen/MachineBlockFrequencyInfo.h

public:
   /// getblockFreq - Return block frequency. Return 0 if we don't have the
   /// information. Please note that initial frequency is equal to 1024. It means
   /// that we should not rely on the value itself, but only on the comparison to
   /// the other block frequencies. We do this to avoid using of floating points.
   ///
   BlockFrequency getBlockFreq(const MachineBasicBlock *MBB) const;
 
   Optional<uint64_t> getBlockProfileCount(const MachineBasicBlock *MBB) const;
+  Optional<uint64_t> getProfileCountFromFreq(uint64_t Freq) const;
+
 
   const MachineFunction *getFunction() const;
   const MachineBranchProbabilityInfo *getMBPI() const;
   void view() const;
 
   // Print the block frequency Freq to OS using the current functions entry
   // frequency to convert freq into a relative decimal form.
   raw_ostream &printBlockFreq(raw_ostream &OS, const BlockFrequency Freq) const;

include/llvm/Transforms/Utils/CodeExtractor.h

 #ifndef LLVM_TRANSFORMS_UTILS_CODEEXTRACTOR_H
 #define LLVM_TRANSFORMS_UTILS_CODEEXTRACTOR_H
 
 #include "llvm/ADT/SetVector.h"
 
 namespace llvm {
 template <typename T> class ArrayRef;
   class BasicBlock;
+  class BlockFrequency;
+  class BlockFrequencyInfo;
+  class BranchProbabilityInfo;
   class DominatorTree;
   class Function;
   class Loop;
   class Module;
   class RegionNode;
+  class TargetTransformInfo;
   class Type;
   class Value;
 
   /// \brief Utility class for extracting code into a new function.
   ///
   /// This utility provides a simple interface for extracting some sequence of
   /// code into its own function, replacing it with a call to that function. It
   /// also provides various methods to query about the nature and result of
   /// such a transformation.
   ///
   /// The rough algorithm used is:
   /// 1) Find both the inputs and outputs for the extracted region.
   /// 2) Pass the inputs as arguments, remapping them within the extracted
   ///    function to arguments.
   /// 3) Add allocas for any scalar outputs, adding all of the outputs' allocas
   ///    as arguments, and inserting stores to the arguments for any scalars.
   class CodeExtractor {
     typedef SetVector<Value *> ValueSet;
 
     // Various bits of state computed on construction.
     DominatorTree *const DT;
     const bool AggregateArgs;
+    BlockFrequencyInfo *BFI;
 
     // Bits of intermediate state computed at various phases of extraction.
     SetVector<BasicBlock *> Blocks;
     unsigned NumExitBlocks;
     Type *RetTy;
 
   public:
     /// \brief Create a code extractor for a single basic block.
     ///
     /// In this formation, we don't require a dominator tree. The given basic
     /// block is set up for extraction.
-    CodeExtractor(BasicBlock *BB, bool AggregateArgs = false);
+    CodeExtractor(BasicBlock *BB, bool AggregateArgs = false,
+                  BlockFrequencyInfo *BFI = nullptr);
 
     /// \brief Create a code extractor for a sequence of blocks.
     ///
     /// Given a sequence of basic blocks where the first block in the sequence
     /// dominates the rest, prepare a code extractor object for pulling this
     /// sequence out into its new function. When a DominatorTree is also given,
     /// extra checking and transformations are enabled.
     CodeExtractor(ArrayRef<BasicBlock *> BBs, DominatorTree *DT = nullptr,
-                  bool AggregateArgs = false);
+                  bool AggregateArgs = false,
+                  BlockFrequencyInfo *BFI = nullptr);
 
     /// \brief Create a code extractor for a loop body.
     ///
     /// Behaves just like the generic code sequence constructor, but uses the
     /// block sequence of the loop.
-    CodeExtractor(DominatorTree &DT, Loop &L, bool AggregateArgs = false);
+    CodeExtractor(DominatorTree &DT, Loop &L, bool AggregateArgs = false,
+                  BlockFrequencyInfo *BFI = nullptr);
 
     /// \brief Create a code extractor for a region node.
     ///
     /// Behaves just like the generic code sequence constructor, but uses the
     /// block sequence of the region node passed in.
     CodeExtractor(DominatorTree &DT, const RegionNode &RN,
-                  bool AggregateArgs = false);
+                  bool AggregateArgs = false,
+                  BlockFrequencyInfo *BFI = nullptr);
 
     /// \brief Perform the extraction, returning the new function.
     ///
     /// Returns zero when called on a CodeExtractor instance where isEligible
     /// returns false.
     Function *extractCodeRegion();
 
     /// \brief Test whether this code extractor is eligible.
private:
 
     Function *constructFunction(const ValueSet &inputs,
                                 const ValueSet &outputs,
                                 BasicBlock *header,
                                 BasicBlock *newRootNode, BasicBlock *newHeader,
                                 Function *oldFunction, Module *M);
 
     void moveCodeToFunction(Function *newFunction);
+    void updateNewCallTerminatorWeights(BasicBlock *CodeReplacer,
+                                        DenseMap<BasicBlock *,
+                                            BlockFrequency> &ExitWeights,
+                                        BranchProbabilityInfo *BPI);
 
     void emitCallAndSwitchStatement(Function *newFunction,
                                     BasicBlock *newHeader,
                                     ValueSet &inputs,
                                     ValueSet &outputs);
   };
 }
 
 #endif

lib/Analysis/BlockFrequencyInfo.cpp

 Optional<uint64_t>
 BlockFrequencyInfo::getBlockProfileCount(const BasicBlock *BB) const {
   if (!BFI)
     return None;
 
   return BFI->getBlockProfileCount(*getFunction(), BB);
 }
 
+Optional<uint64_t>
+BlockFrequencyInfo::getProfileCountFromFreq(uint64_t Freq) const {
+  if (!BFI)
+    return None;
+
+  return BFI->getProfileCountFromFreq(*getFunction(), Freq);
+}
+
 void BlockFrequencyInfo::setBlockFreq(const BasicBlock *BB, uint64_t Freq) {
   assert(BFI && "Expected analysis to be available");
   BFI->setBlockFreq(BB, Freq);
 }
 
 /// Pop up a ghostview window with the current block frequency propagation
 /// rendered using dot.
 void BlockFrequencyInfo::view() const {

lib/Analysis/BlockFrequencyInfoImpl.cpp

BlockFrequencyInfoImplBase::getBlockFreq(const BlockNode &Node) const {
   if (!Node.isValid())
     return 0;
   return Freqs[Node.Index].Integer;
 }
 
 Optional<uint64_t>
 BlockFrequencyInfoImplBase::getBlockProfileCount(const Function &F,
                                                  const BlockNode &Node) const {
+  return getProfileCountFromFreq(F, getBlockFreq(Node).getFrequency());
+}
+
+Optional<uint64_t>
+BlockFrequencyInfoImplBase::getProfileCountFromFreq(const Function &F,
+                                                uint64_t Freq) const {
   auto EntryCount = F.getEntryCount();
   if (!EntryCount)
     return None;
   // Use 128 bit APInt to do the arithmetic to avoid overflow.
   APInt BlockCount(128, EntryCount.getValue());
-  APInt BlockFreq(128, getBlockFreq(Node).getFrequency());
+  APInt BlockFreq(128, Freq);
   APInt EntryFreq(128, getEntryFreq());
   BlockCount *= BlockFreq;
   BlockCount = BlockCount.udiv(EntryFreq);
   return BlockCount.getLimitedValue();
 }
 
 Scaled64
 BlockFrequencyInfoImplBase::getFloatingBlockFreq(const BlockNode &Node) const {

lib/CodeGen/MachineBlockFrequencyInfo.cpp

 }
 
 Optional<uint64_t> MachineBlockFrequencyInfo::getBlockProfileCount(
     const MachineBasicBlock *MBB) const {
   const Function *F = MBFI->getFunction()->getFunction();
   return MBFI ? MBFI->getBlockProfileCount(*F, MBB) : None;
 }
 
+Optional<uint64_t> MachineBlockFrequencyInfo::getProfileCountFromFreq(
+    uint64_t Freq) const {
+  const Function *F = MBFI->getFunction()->getFunction();
+  return MBFI ? MBFI->getProfileCountFromFreq(*F, Freq) : None;
+}
+
 const MachineFunction *MachineBlockFrequencyInfo::getFunction() const {
   return MBFI ? MBFI->getFunction() : nullptr;
 }
 
 const MachineBranchProbabilityInfo *MachineBlockFrequencyInfo::getMBPI() const {
   return MBFI ? &MBFI->getBPI() : nullptr;
 }
 

lib/Transforms/IPO/PartialInlining.cpp

 //===- PartialInlining.cpp - Inline parts of functions --------------------===//
 //
 //                     The LLVM Compiler Infrastructure
 //
 // This file is distributed under the University of Illinois Open Source
 // License. See LICENSE.TXT for details.
 //
 //===----------------------------------------------------------------------===//
 //
 // This pass performs partial inlining, typically by inlining an if statement
 // that surrounds the body of the function.
 //
 //===----------------------------------------------------------------------===//
 
+#include <llvm/Analysis/BlockFrequencyInfo.h>
 #include "llvm/Transforms/IPO/PartialInlining.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instructions.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 "partialinlining"
 
 STATISTIC(NumPartialInlined, "Number of functions partially inlined");
 
 namespace {
+typedef std::function<BlockFrequencyInfo *(Function &)> GetBlockFreqFn;
 struct PartialInlinerImpl {
-  PartialInlinerImpl(InlineFunctionInfo IFI) : IFI(IFI) {}
+  PartialInlinerImpl(InlineFunctionInfo IFI, GetBlockFreqFn GetBlockFreq)
+      : IFI(IFI), GetBlockFrequency(GetBlockFreq) {}
   bool run(Module &M);
   Function *unswitchFunction(Function *F);
 
 private:
   InlineFunctionInfo IFI;
+  GetBlockFreqFn GetBlockFrequency;
 };
 struct PartialInlinerLegacyPass : public ModulePass {
   static char ID; // Pass identification, replacement for typeid
   PartialInlinerLegacyPass() : ModulePass(ID) {
     initializePartialInlinerLegacyPassPass(*PassRegistry::getPassRegistry());
   }
 
   void getAnalysisUsage(AnalysisUsage &AU) const override {
     AU.addRequired<AssumptionCacheTracker>();
+    AU.addRequired<BlockFrequencyInfoWrapperPass>();
   }
   bool runOnModule(Module &M) override {
     if (skipModule(M))
       return false;
 
     AssumptionCacheTracker *ACT = &getAnalysis<AssumptionCacheTracker>();
     std::function<AssumptionCache &(Function &)> GetAssumptionCache =
         [&ACT](Function &F) -> AssumptionCache & {
       return ACT->getAssumptionCache(F);
     };
+    GetBlockFreqFn GetBlockFreq =
+        [this](Function &F) -> BlockFrequencyInfo *{
+          return &getAnalysis<BlockFrequencyInfoWrapperPass>(F).getBFI();
+        };
     InlineFunctionInfo IFI(nullptr, &GetAssumptionCache);
-    return PartialInlinerImpl(IFI).run(M);
+    return PartialInlinerImpl(IFI, GetBlockFreq).run(M);
   }
 };
 }
 
 Function *PartialInlinerImpl::unswitchFunction(Function *F) {
   // First, verify that this function is an unswitching candidate...
   BasicBlock *EntryBlock = &F->front();
   BranchInst *BR = dyn_cast<BranchInst>(EntryBlock->getTerminator());
for (BasicBlock &BB : *DuplicateFunction)
     if (&BB != NewEntryBlock && &BB != NewReturnBlock &&
         &BB != NewNonReturnBlock)
       ToExtract.push_back(&BB);
 
   // The CodeExtractor needs a dominator tree.
   DominatorTree DT;
   DT.recalculate(*DuplicateFunction);
 
+  BlockFrequencyInfo *BFI = GetBlockFrequency(*DuplicateFunction);
+
   // Extract the body of the if.
   Function *ExtractedFunction =
-      CodeExtractor(ToExtract, &DT).extractCodeRegion();
+      CodeExtractor(ToExtract, &DT, false, BFI).extractCodeRegion();
 
   // Inline the top-level if test into all callers.
   std::vector<User *> Users(DuplicateFunction->user_begin(),
                             DuplicateFunction->user_end());
   for (User *User : Users)
     if (CallInst *CI = dyn_cast<CallInst>(User))
       InlineFunction(CI, IFI);
     else if (InvokeInst *II = dyn_cast<InvokeInst>(User))
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);
+    if (Function *NewFunc = unswitchFunction(CurrFunc)) {
+      Worklist.push_back(NewFunc);
       Changed = true;
     }
   }
 
   return Changed;
 }
 
 char PartialInlinerLegacyPass::ID = 0;

 
 PreservedAnalyses PartialInlinerPass::run(Module &M,
                                           ModuleAnalysisManager &AM) {
   auto &FAM = AM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager();
   std::function<AssumptionCache &(Function &)> GetAssumptionCache =
       [&FAM](Function &F) -> AssumptionCache & {
     return FAM.getResult<AssumptionAnalysis>(F);
   };
+  GetBlockFreqFn GetBlockFreq =
+      [&FAM](Function &F) -> BlockFrequencyInfo * {
+        return &FAM.getResult<BlockFrequencyAnalysis>(F);
+      };
   InlineFunctionInfo IFI(nullptr, &GetAssumptionCache);
-  if (PartialInlinerImpl(IFI).run(M))
+  if (PartialInlinerImpl(IFI, GetBlockFreq).run(M))
     return PreservedAnalyses::none();
   return PreservedAnalyses::all();
 }

lib/Transforms/Utils/CodeExtractor.cpp

 // a call to the new function.
 //
 //===----------------------------------------------------------------------===//
 
 #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/LoopInfo.h"
 #include "llvm/Analysis/RegionInfo.h"
 #include "llvm/Analysis/RegionIterator.h"
+#include "llvm/Analysis/TargetTransformInfo.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"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Verifier.h"
 #include "llvm/Pass.h"
+#include "llvm/Support/BlockFrequency.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include <algorithm>
 #include <set>
 using namespace llvm;
if (!RN.isSubRegion())
     // Just a single BasicBlock.
     return buildExtractionBlockSet(RN.getNodeAs<BasicBlock>());
 
   const Region &R = *RN.getNodeAs<Region>();
 
   return buildExtractionBlockSet(R.block_begin(), R.block_end());
 }
 
-CodeExtractor::CodeExtractor(BasicBlock *BB, bool AggregateArgs)
-  : DT(nullptr), AggregateArgs(AggregateArgs||AggregateArgsOpt),
+CodeExtractor::CodeExtractor(BasicBlock *BB, bool AggregateArgs,
+                             BlockFrequencyInfo *BFI)
+  : DT(nullptr), AggregateArgs(AggregateArgs||AggregateArgsOpt), BFI(BFI),
     Blocks(buildExtractionBlockSet(BB)), NumExitBlocks(~0U) {}
 
 CodeExtractor::CodeExtractor(ArrayRef<BasicBlock *> BBs, DominatorTree *DT,
-                             bool AggregateArgs)
-  : DT(DT), AggregateArgs(AggregateArgs||AggregateArgsOpt),
+                             bool AggregateArgs, BlockFrequencyInfo *BFI)
+  : DT(DT), AggregateArgs(AggregateArgs||AggregateArgsOpt), BFI(BFI),
     Blocks(buildExtractionBlockSet(BBs)), NumExitBlocks(~0U) {}
 
-CodeExtractor::CodeExtractor(DominatorTree &DT, Loop &L, bool AggregateArgs)
-  : DT(&DT), AggregateArgs(AggregateArgs||AggregateArgsOpt),
+CodeExtractor::CodeExtractor(DominatorTree &DT, Loop &L, bool AggregateArgs,
+                             BlockFrequencyInfo *BFI)
+  : DT(&DT), AggregateArgs(AggregateArgs||AggregateArgsOpt), BFI(BFI),
     Blocks(buildExtractionBlockSet(L.getBlocks())), NumExitBlocks(~0U) {}
 
 CodeExtractor::CodeExtractor(DominatorTree &DT, const RegionNode &RN,
-                             bool AggregateArgs)
-  : DT(&DT), AggregateArgs(AggregateArgs||AggregateArgsOpt),
+                             bool AggregateArgs, BlockFrequencyInfo *BFI)
+    : DT(&DT), AggregateArgs(AggregateArgs || AggregateArgsOpt), BFI(BFI),
       Blocks(buildExtractionBlockSet(RN)), NumExitBlocks(~0U) {}
 
 /// definedInRegion - Return true if the specified value is defined in the
 /// extracted region.
 static bool definedInRegion(const SetVector<BasicBlock *> &Blocks, Value *V) {
   if (Instruction *I = dyn_cast<Instruction>(V))
     if (Blocks.count(I->getParent()))
       return true;
   return false;
Function *CodeExtractor::constructFunction(const ValueSet &inputs,
   // Create the new function
   Function *newFunction = Function::Create(funcType,
                                            GlobalValue::InternalLinkage,
                                            oldFunction->getName() + "_" +
                                            header->getName(), M);
   // If the old function is no-throw, so is the new one.
   if (oldFunction->doesNotThrow())
     newFunction->setDoesNotThrow();
-  
   newFunction->getBasicBlockList().push_back(newRootNode);
 
   // Create an iterator to name all of the arguments we inserted.
   Function::arg_iterator AI = newFunction->arg_begin();
 
   // Rewrite all users of the inputs in the extracted region to use the
   // arguments (or appropriate addressing into struct) instead.
   for (unsigned i = 0, e = inputs.size(); i != e; ++i) {

 /// the call instruction, splitting any PHI nodes in the header block as
 /// necessary.
 void CodeExtractor::
 emitCallAndSwitchStatement(Function *newFunction, BasicBlock *codeReplacer,
                            ValueSet &inputs, ValueSet &outputs) {
   // Emit a call to the new function, passing in: *pointer to struct (if
   // aggregating parameters), or plan inputs and allocated memory for outputs
   std::vector<Value*> params, StructValues, ReloadOutputs, Reloads;
-  
   LLVMContext &Context = newFunction->getContext();
 
   // Add inputs as params, or to be filled into the struct
   for (Value *input : inputs)
     if (AggregateArgs)
       StructValues.push_back(input);
     else
       params.push_back(input);
for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i)
               // with an invoke instruction, make sure that we don't emit a
               // store of the invoke value for the unwind block.
               if (!DT && DefBlock != OldTarget)
                 DominatesDef = false;
             }
 
             if (DT) {
               DominatesDef = DT->dominates(DefBlock, OldTarget);
-              
               // If the output value is used by a phi in the target block,
               // then we need to test for dominance of the phi's predecessor
               // instead.  Unfortunately, this a little complicated since we
               // have already rewritten uses of the value to uses of the reload.
-              BasicBlock* pred = FindPhiPredForUseInBlock(Reloads[out], 
+              BasicBlock* pred = FindPhiPredForUseInBlock(Reloads[out],
                                                           OldTarget);
               if (pred && DT && DT->dominates(DefBlock, pred))
                 DominatesDef = true;
             }
 
             if (DominatesDef) {
               if (AggregateArgs) {
                 Value *Idx[2];
case 0:
     if (OldFnRetTy->isVoidTy()) {
       ReturnInst::Create(Context, nullptr, TheSwitch);  // Return void
     } else if (OldFnRetTy == TheSwitch->getCondition()->getType()) {
       // return what we have
       ReturnInst::Create(Context, TheSwitch->getCondition(), TheSwitch);
     } else {
       // Otherwise we must have code extracted an unwind or something, just
       // return whatever we want.
-      ReturnInst::Create(Context, 
+      ReturnInst::Create(Context,
                          Constant::getNullValue(OldFnRetTy), TheSwitch);
     }
 
     TheSwitch->eraseFromParent();
     break;
   case 1:
     // Only a single destination, change the switch into an unconditional
     // branch.
for (BasicBlock *Block : Blocks) {
     // Delete the basic block from the old function, and the list of blocks
     oldBlocks.remove(Block);
 
     // Insert this basic block into the new function
     newBlocks.push_back(Block);
   }
 }
 
+void CodeExtractor::updateNewCallTerminatorWeights(
+    BasicBlock *CodeReplacer,
+    DenseMap<BasicBlock *, BlockFrequency> &ExitWeights,
+    BranchProbabilityInfo *BPI) {
+  // Update the branch weights for the exit block.
+  TerminatorInst *TI = CodeReplacer->getTerminator();
+
+  // Make sure that the weights are valid, we cant have a weight of zero.
+  bool HasValidWeights = true;
+  for (auto &Weight : ExitWeights) {
+    if (Weight.second.getFrequency() == 0) {
+      HasValidWeights = false;
+      break;
+    }
+  }
+  if (!HasValidWeights)
+    return;
+
+  // Block Frequency distribution with dummy node.
+  BlockFrequencyInfoImplBase::Distribution BranchDist;
+
+  // Add each of the frequencies of the successors
+  for (unsigned I = 0, E = TI->getNumSuccessors(); I < E; ++I) {
+    BlockFrequencyInfoImplBase::BlockNode DummyNode(I);
+    BranchDist.addExit(DummyNode,
+                       ExitWeights[TI->getSuccessor(I)].getFrequency());
+  }
+
+  // Normalize the distribution so that they can fit in unsigned.
+  BranchDist.normalize();
+
+  // Create normalized branch weights and set the metadata.
+  SmallVector<unsigned, 8> BranchWeights;
+  BranchWeights.reserve(TI->getNumSuccessors());
+  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.push_back(Weight.Amount);
+    BranchProbability BP(Weight.Amount, BranchDist.Total);
+    BPI->setEdgeProbability(CodeReplacer, I, BP);
+  }
+  TI->setMetadata(
+      LLVMContext::MD_prof,
+      MDBuilder(TI->getContext()).createBranchWeights(BranchWeights));
+}
+
 Function *CodeExtractor::extractCodeRegion() {
   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();
 
+  // Get non-const in case we modify the exit probabilities when the exit count
+  // is
+  //  greater than 1.
+  BranchProbabilityInfo *BPI = nullptr;
+  if (BFI)
+    BPI = const_cast<BranchProbabilityInfo *>(BFI->getBPI());
+
+  // Calculate the entry of the new function before we change the root block.
+  BlockFrequency EntryFreq;
+  if (BFI) {
+    for (BasicBlock *Pred : predecessors(header)) {
+      if (Blocks.count(Pred))
+        continue;
+      EntryFreq +=
+          BFI->getBlockFreq(Pred) * BPI->getEdgeProbability(Pred, header);
+    }
+  }
+
   // If we have to split PHI nodes or the entry block, do so now.
   severSplitPHINodes(header);
 
   // If we have any return instructions in the region, split those blocks so
   // that the return is not in the region.
   splitReturnBlocks();
 
   Function *oldFunction = header->getParent();
 
   // This takes place of the original loop
-  BasicBlock *codeReplacer = BasicBlock::Create(header->getContext(), 
+  BasicBlock *codeReplacer = BasicBlock::Create(header->getContext(),
                                                 "codeRepl", oldFunction,
                                                 header);
 
   // 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(), 
+  BasicBlock *newFuncRoot = BasicBlock::Create(header->getContext(),
                                                "newFuncRoot");
   newFuncRoot->getInstList().push_back(BranchInst::Create(header));
 
   // Find inputs to, outputs from the code region.
   findInputsOutputs(inputs, outputs);
 
+  DenseMap<BasicBlock *, BlockFrequency> ExitWeights;
   SmallPtrSet<BasicBlock *, 1> ExitBlocks;
-  for (BasicBlock *Block : Blocks)
+  for (BasicBlock *Block : Blocks) {
     for (succ_iterator SI = succ_begin(Block), SE = succ_end(Block); SI != SE;
-         ++SI)
-      if (!Blocks.count(*SI))
+         ++SI) {
+      if (!Blocks.count(*SI)) {
+        // Update the branch weight for this successor.
+        if (BPI) {
+          BlockFrequency &BF = ExitWeights[*SI];
+          BF += BFI->getBlockFreq(Block) * BPI->getEdgeProbability(Block, *SI);
+        }
         ExitBlocks.insert(*SI);
+      }
+    }
+  }
   NumExitBlocks = ExitBlocks.size();
 
   // Construct new function based on inputs/outputs & add allocas for all defs.
   Function *newFunction = constructFunction(inputs, outputs, header,
                                             newFuncRoot,
                                             codeReplacer, oldFunction,
                                             oldFunction->getParent());
 
+  // Update the entry count of the function.
+  if (BFI && EntryFreq.getFrequency() > 0) {
+    Optional<uint64_t> OptEntryCount =
+        BFI->getProfileCountFromFreq(EntryFreq.getFrequency());
+    newFunction->setEntryCount( OptEntryCount.getValue());
+    BFI->setBlockFreq(codeReplacer, EntryFreq.getFrequency());
+  }
+
   emitCallAndSwitchStatement(newFunction, codeReplacer, inputs, outputs);
 
   moveCodeToFunction(newFunction);
 
+  // Update the branch weights for the exit block.
+  if (BPI && NumExitBlocks > 1)
+    updateNewCallTerminatorWeights(codeReplacer, ExitWeights, BPI);
+
   // Loop over all of the PHI nodes in the header block, and change any
   // references to the old incoming edge to be the new incoming edge.
   for (BasicBlock::iterator I = header->begin(); isa<PHINode>(I); ++I) {
     PHINode *PN = cast<PHINode>(I);
     for (unsigned i = 0, e = PN->getNumIncomingValues(); i != e; ++i)
       if (!Blocks.count(PN->getIncomingBlock(i)))
         PN->setIncomingBlock(i, newFuncRoot);
   }
for (unsigned i = 0, e = Succs.size(); i != e; ++i)
     }
 
   //cerr << "NEW FUNCTION: " << *newFunction;
   //  verifyFunction(*newFunction);
 
   //  cerr << "OLD FUNCTION: " << *oldFunction;
   //  verifyFunction(*oldFunction);
 
-  DEBUG(if (verifyFunction(*newFunction)) 
+  DEBUG(if (verifyFunction(*newFunction))
         report_fatal_error("verifyFunction failed!"));
   return newFunction;
-}
+}
 No newline at end of file

test/Transforms/CodeExtractor/2016-07-20-ExtractedFnEntryCount.ll

+; RUN: opt < %s -partial-inliner -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 {
+entry:
+  br i1 %cond, label %if.then, label %return, !prof !2
+if.then:
+  br i1 %cond, label %if.then, label %return, !prof !3
+return:             ; preds = %entry
+  ret i32 0
+}
+
+
+define internal i32 @dummyCaller(i1 %cond) !prof !1 {
+entry:
+  %val = call i32 @inlinedFunc(i1 %cond)
+  ret i32 %val
+}
+
+; CHECK: @inlinedFunc.1_if.then(i1 %cond) !prof [[COUNT1:![0-9]+]]
+
+
+!llvm.module.flags = !{!0}
+; CHECK: [[COUNT1]] = !{!"function_entry_count", i64 250}
+!0 = !{i32 1, !"MaxFunctionCount", i32 1000}
+!1 = !{!"function_entry_count", i64 1000}
+!2 = !{!"branch_weights", i32 250, i32 750}
+!3 = !{!"branch_weights", i32 125, i32 125}

test/Transforms/CodeExtractor/2016-07-20-MultipleExitBranchProb.ll

+; RUN: opt < %s -partial-inliner -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:
+  br i1 %cond, label %return, label %return.2, !prof !3
+return.2:
+  ret i32 10
+return:             ; preds = %entry
+  ret i32 0
+}
+
+
+define internal i32 @dummyCaller(i1 %cond) !prof !1 {
+entry:
+%val = call i32 @inlinedFunc(i1 %cond)
+ret i32 %val
+}
+
+!llvm.module.flags = !{!0}
+!0 = !{i32 1, !"MaxFunctionCount", i32 10000}
+!1 = !{!"function_entry_count", i64 10000}
+!2 = !{!"branch_weights", i32 5, i32 5}
+!3 = !{!"branch_weights", i32 4, i32 1}
+
+; CHECK: !4 = !{!"branch_weights", i32 8, i32 31}