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

696

clang-format? add space

699

EntryCount --> EntryFreq

756

nit: OptEntryCount --> EntryCount

758

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

766

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 ↗

(On Diff #65625)

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

689

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

694

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.

758

nit: no space after the open paren

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

31

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

692

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

755

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

758

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

692

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

758

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

[rriddle]

lib/Transforms/Utils/CodeExtractor.cpp

692

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

692

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

692

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

694

BPI check is not needed.

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

758

if (OptEntryCount->hasValue())

[rriddle]

Updated after David's comments

[silvas]

Some nits.

lib/Transforms/Utils/CodeExtractor.cpp

693

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

696

Missing '.' at the end of this comment.

732

Add a comment here about what this loop is doing.

755

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> getFreqToProfileCount(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> getFreqToProfileCount(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> getFreqToProfileCount(const Function &F,
+                                         uint64_t Freq) const {
+    return BlockFrequencyInfoImplBase::getFreqToProfileCount(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> getFreqToProfileCount(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 BlockFrequencyInfo;
   class DominatorTree;
   class Function;
   class Loop;
   class Module;
   class RegionNode;
   class Type;
   class Value;
 
namespace llvm {
   /// 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:
 
     void emitCallAndSwitchStatement(Function *newFunction,
                                     BasicBlock *newHeader,
                                     ValueSet &inputs,
                                     ValueSet &outputs);
   };
 }
 
-#endif
+#endif
 No newline at end of file

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::getFreqToProfileCount(uint64_t Freq) const {
+  if (!BFI)
+    return None;
+
+  return BFI->getFreqToProfileCount(*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 getFreqToProfileCount(F, getBlockFreq(Node).getFrequency());
+}
+
+Optional<uint64_t>
+BlockFrequencyInfoImplBase::getFreqToProfileCount(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::getFreqToProfileCount(
+    uint64_t Freq) const {
+  const Function *F = MBFI->getFunction()->getFunction();
+  return MBFI ? MBFI->getFreqToProfileCount(*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/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/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;
 
 #define DEBUG_TYPE "code-extractor"
 
 // Provide a command-line option to aggregate function arguments into a struct
 // for functions produced by the code extractor. This is useful when converting
 // extracted functions to pthread-based code, as only one argument (void*) can
 // be passed in to pthread_create().
 static cl::opt<bool>
     AggregateArgsOpt("aggregate-extracted-args", cl::Hidden,
                      cl::desc("Aggregate arguments to code-extracted functions"));
 
 /// \brief Test whether a block is valid for extraction.
 static bool isBlockValidForExtraction(const BasicBlock &BB) {
   // Landing pads must be in the function where they were inserted for cleanup.
   if (BB.isEHPad())
     return false;
 
   // Don't hoist code containing allocas, invokes, or vastarts.
   for (BasicBlock::const_iterator I = BB.begin(), E = BB.end(); I != E; ++I) {
     if (isa<AllocaInst>(I) || isa<InvokeInst>(I))
       return false;
     if (const CallInst *CI = dyn_cast<CallInst>(I))
-      if (const Function *F = CI->getCalledFunction())
-        if (F->getIntrinsicID() == Intrinsic::vastart)
-          return false;
   }
 
   return true;
 }
 
 /// \brief Build a set of blocks to extract if the input blocks are viable.
 template <typename IteratorT>
 static SetVector<BasicBlock *> buildExtractionBlockSet(IteratorT BBBegin,
do {
     if (!isBlockValidForExtraction(**BBBegin)) {
       Result.clear();
       return Result;
     }
   } while (++BBBegin != BBEnd);
 
 #ifndef NDEBUG
   for (SetVector<BasicBlock *>::iterator I = std::next(Result.begin()),
-                                         E = Result.end();
        I != E; ++I)
     for (pred_iterator PI = pred_begin(*I), PE = pred_end(*I);
          PI != PE; ++PI)
       assert(Result.count(*PI) &&
              "No blocks in this region may have entries from outside the region"
                  " except for the first block!");
 #endif
 
   return Result;
 }
 
 /// \brief Helper to call buildExtractionBlockSet with an ArrayRef.
 static SetVector<BasicBlock *>
 buildExtractionBlockSet(ArrayRef<BasicBlock *> BBs) {
   return buildExtractionBlockSet(BBs.begin(), BBs.end());
 }
 
 /// \brief Helper to call buildExtractionBlockSet with a RegionNode.
 static SetVector<BasicBlock *>
 buildExtractionBlockSet(const RegionNode &RN) {
   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;
 }
 
 /// definedInCaller - Return true if the specified value is defined in the
 /// function being code extracted, but not in the region being extracted.
 /// These values must be passed in as live-ins to the function.
 static bool definedInCaller(const SetVector<BasicBlock *> &Blocks, Value *V) {
   if (isa<Argument>(V)) return true;
   if (Instruction *I = dyn_cast<Instruction>(V))
-    if (!Blocks.count(I->getParent()))
-      return true;
   return false;
 }
 
 void CodeExtractor::findInputsOutputs(ValueSet &Inputs,
                                       ValueSet &Outputs) const {
   for (BasicBlock *BB : Blocks) {
     // If a used value is defined outside the region, it's an input.  If an
     // instruction is used outside the region, it's an output.
Function *CodeExtractor::constructFunction(const ValueSet &inputs,
                                            BasicBlock *newHeader,
                                            Function *oldFunction,
                                            Module *M) {
   DEBUG(dbgs() << "inputs: " << inputs.size() << "\n");
   DEBUG(dbgs() << "outputs: " << outputs.size() << "\n");
 
   // This function returns unsigned, outputs will go back by reference.
   switch (NumExitBlocks) {
     case 0:
     case 1: RetTy = Type::getVoidTy(header->getContext()); break;
     case 2: RetTy = Type::getInt1Ty(header->getContext()); break;
     default: RetTy = Type::getInt16Ty(header->getContext()); break;
   }
 
   std::vector<Type*> paramTy;
 
   // Add the types of the input values to the function's argument list
   for (Value *value : inputs) {
     DEBUG(dbgs() << "value used in func: " << *value << "\n");
     paramTy.push_back(value->getType());
   }
 
   // Add the types of the output values to the function's argument list.
   for (Value *output : outputs) {
     DEBUG(dbgs() << "instr used in func: " << *output << "\n");
     if (AggregateArgs)
       paramTy.push_back(output->getType());
     else
       paramTy.push_back(PointerType::getUnqual(output->getType()));
   }
 
   DEBUG({
           dbgs() << "Function type: " << *RetTy << " f(";
           for (Type *i : paramTy)
             dbgs() << *i << ", ";
           dbgs() << ")\n";
         });
 
   StructType *StructTy;
   if (AggregateArgs && (inputs.size() + outputs.size() > 0)) {
     StructTy = StructType::get(M->getContext(), paramTy);
     paramTy.clear();
     paramTy.push_back(PointerType::getUnqual(StructTy));
   }
   FunctionType *funcType =
-                  FunctionType::get(RetTy, paramTy, false);
 
   // 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) {
     Value *RewriteVal;
     if (AggregateArgs) {
       Value *Idx[2];
       Idx[0] = Constant::getNullValue(Type::getInt32Ty(header->getContext()));
       Idx[1] = ConstantInt::get(Type::getInt32Ty(header->getContext()), i);
       TerminatorInst *TI = newFunction->begin()->getTerminator();
       GetElementPtrInst *GEP = GetElementPtrInst::Create(
           StructTy, &*AI, Idx, "gep_" + inputs[i]->getName(), TI);
       RewriteVal = new LoadInst(GEP, "loadgep_" + inputs[i]->getName(), TI);
     } else
       RewriteVal = &*AI++;
 
     std::vector<User*> Users(inputs[i]->user_begin(), inputs[i]->user_end());
     for (User *use : Users)
       if (Instruction *inst = dyn_cast<Instruction>(use))
-        if (Blocks.count(inst->getParent()))
-          inst->replaceUsesOfWith(inputs[i], RewriteVal);
   }
 
   // Set names for input and output arguments.
   if (!AggregateArgs) {
     AI = newFunction->arg_begin();
     for (unsigned i = 0, e = inputs.size(); i != e; ++i, ++AI)
       AI->setName(inputs[i]->getName());
     for (unsigned i = 0, e = outputs.size(); i != e; ++i, ++AI)
       AI->setName(outputs[i]->getName()+".out");
   }
 
   // Rewrite branches to basic blocks outside of the loop to new dummy blocks
   // within the new function. This must be done before we lose track of which
   // blocks were originally in the code region.
   std::vector<User*> Users(header->user_begin(), header->user_end());
   for (unsigned i = 0, e = Users.size(); i != e; ++i)
     // The BasicBlock which contains the branch is not in the region
     // modify the branch target to a new block
     if (TerminatorInst *TI = dyn_cast<TerminatorInst>(Users[i]))
-      if (!Blocks.count(TI->getParent()) &&
-          TI->getParent()->getParent() == oldFunction)
-        TI->replaceUsesOfWith(header, newHeader);
 
   return newFunction;
 }
 
 /// FindPhiPredForUseInBlock - Given a value and a basic block, find a PHI
 /// that uses the value within the basic block, and return the predecessor
 /// block associated with that use, or return 0 if none is found.
 static BasicBlock* FindPhiPredForUseInBlock(Value* Used, BasicBlock* BB) {
   for (Use &U : Used->uses()) {
-     PHINode *P = dyn_cast<PHINode>(U.getUser());
-     if (P && P->getParent() == BB)
-       return P->getIncomingBlock(U);
   }
 
   return nullptr;
 }
 
 /// emitCallAndSwitchStatement - This method sets up the caller side by adding
 /// 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 (Value *output : outputs) {
     }
   }
 
   StructType *StructArgTy = nullptr;
   AllocaInst *Struct = nullptr;
   if (AggregateArgs && (inputs.size() + outputs.size() > 0)) {
     std::vector<Type*> ArgTypes;
     for (ValueSet::iterator v = StructValues.begin(),
              ve = StructValues.end(); v != ve; ++v)
       ArgTypes.push_back((*v)->getType());
 
     // Allocate a struct at the beginning of this function
     StructArgTy = StructType::get(newFunction->getContext(), ArgTypes);
     Struct = new AllocaInst(StructArgTy, nullptr, "structArg",
                             &codeReplacer->getParent()->front().front());
     params.push_back(Struct);
 
for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i)
           // destination, create one now!
           NewTarget = BasicBlock::Create(Context,
                                          OldTarget->getName() + ".exitStub",
                                          newFunction);
           unsigned SuccNum = switchVal++;
 
           Value *brVal = nullptr;
           switch (NumExitBlocks) {
             case 0:
             case 1: break;  // No value needed.
             case 2:         // Conditional branch, return a bool
               brVal = ConstantInt::get(Type::getInt1Ty(Context), !SuccNum);
               break;
             default:
               brVal = ConstantInt::get(Type::getInt16Ty(Context), SuccNum);
               break;
           }
 
           ReturnInst *NTRet = ReturnInst::Create(Context, brVal, NewTarget);
 
           // Update the switch instruction.
           TheSwitch->addCase(ConstantInt::get(Type::getInt16Ty(Context),
                                               SuccNum),
                              OldTarget);
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];
for (unsigned i = 0, e = TI->getNumSuccessors(); i != e; ++i)
         // rewrite the original branch instruction with this new target
         TI->setSuccessor(i, NewTarget);
       }
   }
 
   // Now that we've done the deed, simplify the switch instruction.
   Type *OldFnRetTy = TheSwitch->getParent()->getParent()->getReturnType();
   switch (NumExitBlocks) {
     case 0:
       // There are no successors (the block containing the switch itself), which
       // means that previously this was the last part of the function, and hence
       // this should be rewritten as a `ret'
 
       // Check if the function should return a value
       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.
       BranchInst::Create(TheSwitch->getSuccessor(1), TheSwitch);
       TheSwitch->eraseFromParent();
       break;
     case 2:
       BranchInst::Create(TheSwitch->getSuccessor(1), TheSwitch->getSuccessor(2),
                          call, TheSwitch);
       TheSwitch->eraseFromParent();
       break;
     default:
       // Otherwise, make the default destination of the switch instruction be one
       // of the other successors.
       TheSwitch->setCondition(call);
       TheSwitch->setDefaultDest(TheSwitch->getSuccessor(NumExitBlocks));
       // Remove redundant case
       TheSwitch->removeCase(SwitchInst::CaseIt(TheSwitch, NumExitBlocks-1));
       break;
   }
 }
 
 void CodeExtractor::moveCodeToFunction(Function *newFunction) {
   Function *oldFunc = (*Blocks.begin())->getParent();
   Function::BasicBlockListType &oldBlocks = oldFunc->getBasicBlockList();
   Function::BasicBlockListType &newBlocks = newFunction->getBasicBlockList();
 
   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);
   }
 }
 
 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 EntryCount;
+  if (BFI) {
+    for (BasicBlock *Pred : predecessors(header)) {
+      if (Blocks.count(Pred))
+        continue;
+      EntryCount += 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) {
+    Optional<uint64_t> OptEntryCount =
+        BFI->getFreqToProfileCount(EntryCount.getFrequency());
+    newFunction->setEntryCount(OptEntryCount ? OptEntryCount.getValue() : 0);
+    BFI->setBlockFreq(codeReplacer, EntryCount.getFrequency());
+  }
+
   emitCallAndSwitchStatement(newFunction, codeReplacer, inputs, outputs);
 
   moveCodeToFunction(newFunction);
 
+  // Update the branch weights for the exit block.
+  TerminatorInst *TI = codeReplacer->getTerminator();
+  if (BPI && TI->getNumSuccessors() > 1) {
+    // 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) {
+      // Block Frequency distribution for creating normalized weights.
+      BlockFrequencyInfoImplBase::Distribution BranchDist;
+
+      // Add each of the frequencies of the successors
+      for (unsigned I = 0, E = TI->getNumSuccessors(); I < E; ++I) {
+        // Dummy node to represent this successor.
+        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, Weight.TargetNode.Index, BP);
+      }
+      TI->setMetadata(LLVMContext::MD_prof,
+                      MDBuilder(TI->getContext()).
+                          createBranchWeights(BranchWeights));
+    }
+  }
+
   // 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}