Update the branch weight metadata in JumpThreading pass.

[congh]

Summary

Currently in JumpThreading pass, the branch weight metadata is not updated after CFG modification. Consider the jump threading on PredBB, BB, and SuccBB. After jump threading, the weight on BB->SuccBB should be adjusted as some of it is contributed by the edge PredBB->BB, which doesn't exist anymore. This patch tries to update the edge weight in metadata on BB->SuccBB by scaling it by 1 - Freq(PredBB->BB) / Freq(BB->SuccBB). Two more analyses (BlockFrequencyInfo and BranchProbabilityInfo) are needed then.

[dnovillo]

lib/Transforms/Scalar/JumpThreading.cpp

1522

Shouldn't you be using weights here instead of frequencies? Frequencies can be computed later out of the branch weights.

[davidxl]

Weights are relative values that can only be compared/mixed across
branches sharing the same source BB. For Jump-thread updating, the
split-out edge from NewBB to SuccBB is a new edge without profile
info. We need to use the profile info from edge Pred -> NewBB which
does not share the same source BB with BB->SuccBB. Frequency needs to
be used for scaling purpose here.

David

[davidxl]

lib/Transforms/Scalar/JumpThreading.cpp

172

Ideally the new passes in the pipeline is only needed when profile-instr-use or autofdo is turned on.

[congh]

lib/Transforms/Scalar/JumpThreading.cpp

1522

Weights are only meaningful in terms of the same source BB. Here I need to know how many weights on BB to SuccBB are contributed by the edge PredBB to BB, and I could not compare the weights on those two edges directly as they only make sense for PredBB and BB respectively.

[congh]

This update merges the patch that adds setBlockFreq() interfaces to BlockFrequencyInfo (http://reviews.llvm.org/D11104). As both BFI and BPI can be used without including passes with them, in jump-threading pass, we can update the edge weights only in PGO mode (i.e., only when we have the entry count of a function) without introducing new passes. To make sure both BFI/BPI have up-to-date information, for every CFG change, we need to update BFI and BPI accordingly. Two utility functions (JumpThreading::UpdateEdgeWeight() and JumpThreading::SplitBlockPreds()) are created to facilitate the update.

[davidxl]

include/llvm/Analysis/BlockFrequencyInfoImpl.h

972

It might be slightly better to split the interfaces into two: one is set and the other is update.

lib/Analysis/BranchProbabilityInfo.cpp

622

It is probably better to call this one 'updateEdgeWeight'

lib/Transforms/Scalar/JumpThreading.cpp

150

This function certainly does more than just edge weight update -- a different name may be better.

1562

assert here?

1570

fix typo.

1571

A side question: do we have a bug tracking that data insanity issue?

1585

Should probably assert here that BPI->getEdgeWeight(..) == W

1585

This is not really BBFreq, it should be BB2Succ Edge frequency

1590

Another insanity to track?

1593

I don't think this one is correct. You are using the New BB->Succ Frequency in the update which is wrong.

The invariant is :

PredBBFreq + NewBB2SuccBBEdgeFreq = OldBB2SuccBBEdgeFreq

and the update is:

NewWeight = (NewBB2SuccBBEdgeFreq/OldBB2EdgeFreq)* OldWeight

combine them we have

NewWeight = (1 - PredBBFreq/OldBB2EdgeFreq)*OldWeight

[congh]

Update the patch according to David's comments.

[davidxl]

lib/Transforms/Scalar/JumpThreading.cpp

1545

Can be skipped when Preds' s size == 1

1553

Can be skipped if Preds's size is 1

[congh]

Fixed a bug and let two newly added internal functions be private.

[congh]

Fix precision loss problem when calculating new edge weights.

[davidxl]

I suggest also adding another test case (documented in the PR -- the one that has problem with jump-threading + loop-rotation + jump-threading).

This reminds me of a way to do profile sanity check/validation when debug check is turned on:

1. compute the BB freq before the pass
2. cfg transformation + BB freq update
3. branch probability weight update

4. recompute BFI using updated weight in 3). The result of 4) should match that of 2).

lib/Transforms/Scalar/JumpThreading.cpp

1572

change the name to OrigBBFreq

1578

Is this still an issue that needs the workaround here?

1596

Why this condition?

1599

It is unclear how can this happen.

1614

The scaling code looks unnecessarily complicated. How about this:

a) Pre-normalization

1. first sum up all the weights before the update
2. if the sum is too small, scale it up to a fixed val such as ScaledNumber<64>(1,20)
3. compute the scale factor and rescale the weights

b) update BB->SuccBB with (1- Freq(Pred->BB)/Freq(BB->SuccBB))

c) post-normalization: scale all weights such that the weight sum is 2^20

Some helper routine can be defined.

[congh]

In D10979#208776, @davidxl wrote:

I suggest also adding another test case (documented in the PR -- the one that has problem with jump-threading + loop-rotation + jump-threading).

This reminds me of a way to do profile sanity check/validation when debug check is turned on:

1. compute the BB freq before the pass
2. cfg transformation + BB freq update
3. branch probability weight update

4. recompute BFI using updated weight in 3). The result of 4) should match that of 2).

You mean comparing result from 4) and 1)? What if a CFG node is split? If CFG is modified, we may not be able to find one-to-one BB match. To make a reasonable comparison, we may need to understand what the pass does...

[congh]

lib/Transforms/Scalar/JumpThreading.cpp

1578

Yes. The issue emerges after jump-threading + loop-rotation + jump-threading will lead to this.

1596

I forget to add a comment when I did this. But I think this condition can be removed now.

1599

Consider the following CFG and assume the edge frequency is 100 for each edge:

A   B
 \ /
  C
 / \
D   E

In some pass C may be split into two blocks:

A   B
|   |
C1 C2
| X |
D   E

Then edges from C1/C2 have frequency 50. Now suppose in jump-threading pass we thread B->E across C2 (C2->D is impossible at runtime), then B has frequency 100 but C2->E has frequency 50. This is when this condition is not met.

1614

This method looks good. I have two questions:

1. How to tell if the sum is too small? Can we just normalize all weights such that the sum of them is 2^20?
2. Is c) necessary?

[congh]

Fixed two bugs in the patch.

[davidxl]

lib/Transforms/Scalar/JumpThreading.cpp

1599

What I meant is that at runtime Pred->BB can only be followed by BB->Succ, then Freq(BB->SuccBB) can be proved to be no less than Freq(PredBB->BB) -- so that check is unnecessary.

1614

1. I think it is fine to do normalization uncondionally

2. c is not strictly necessary. However I do expect in the future we can assert that the weight sum at any given time is always normalized.

[congh]

Update the patch by normalizing a list of weights of uint64_t type using MachineBranchProbabilityInfo::normalizeEdgeWeights().

Please take a look.

[davidxl]

As discussed, there is another more complicated jump-threading related
profile updating bug (also involving loop rotation). Do you plan to
address that in a different patch?

David

[davidxl]

After reading the code again, I think the change needs to be restructured. There are a couple of main points we should stick to:

1. When CFG is changed, the code range in which the in-memory profile data is in inconsistent state should be minimized -- ideally immediately after the CFG is changed -- so that the profile state before and after the CFG change is kosher.
2. BFI and BPI update need to be bundled together
3. weight meta data should be a direct translation of updated BPI.

I suggest reorganizing the change like this

1. BFI and BPI update for threading should be in one place (one helper function). a) BBs involved: PredBB, NewBB, BB, and SuccBB b) Edges involved: Pred->BB : removed Pred->NewBB: new NewBB->SuccBB: new BB->SuccBB: need edge freq and bp update.

2. after threading profile update, the meta data update should be simple.

There should be no worry for overflow as subtraction is involved. Due to this, the scaling code can be simple.

lib/Transforms/Scalar/JumpThreading.cpp

1442

Should the NewBB' freq = BFI->getBlockFreq(PredBB)*BPI->getEdgeProbability(PredBB, BB)?

Where is the BPI update for PredBB->NewBB?

1560

Freq(PredBB->NewBB) / Orig_Freq(BB->SuccBB)

which is basically Freq (NewBB)/Orig_Freq(BB->SuccBB)

1574

When NewBB is created, there should be no edge from PredBB to BB anymore. The code here assumes the BPI has not been updated which can be broken.

I think the right way is to use NewBB's frequency which is already computed -- the update interface either needs to pass in NewBB's freq or NewBB (clone of BB).

1580

To avoid ambiguity of the name, it is better to change the name of the variable to BBNewFreq, otherwise it may be thought of as NewBB's frequency where NewBB is the cloned bb of 'BB'.

Similarly change OrigBBFreq to BBOrigFreq.

1584

Here is the case where BB's frequency is updated, but the BPI for BB's successors are not updated. I think these updates needs to be wrapped into its own helper function independent of meta data weights update (mixing them makes the code hard to read and reason).

1590

This whole loop seems unnecessary.

1610

There does not seem to be a need to collect the old weights at all. As long as we have newly updated branch probabilty info (using the updated edge frequency), the new weights can be resynthezied.

[congh]

In D10979#234372, @davidxl wrote:

After reading the code again, I think the change needs to be restructured. There are a couple of main points we should stick to:

1. When CFG is changed, the code range in which the in-memory profile data is in inconsistent state should be minimized -- ideally immediately after the CFG is changed -- so that the profile state before and after the CFG change is kosher.
2. BFI and BPI update need to be bundled together

BFI and BPI are updated in different ways and it is difficult to write a general function doing this. However, I have greatly reduced the size of the function that updates BFI and BPI in this case.

3. weight meta data should be a direct translation of updated BPI.

This is a good idea. I have used the edge frequencies to update edge weights. The only possible issue is that when the source BB's frequency is too low, we may suffer some precision loss. But I think it is OK because BB's frequency is low.

I suggest reorganizing the change like this

1. BFI and BPI update for threading should be in one place (one helper function). a) BBs involved: PredBB, NewBB, BB, and SuccBB b) Edges involved: Pred->BB : removed Pred->NewBB: new NewBB->SuccBB: new BB->SuccBB: need edge freq and bp update.

This is now done in one function.

2. after threading profile update, the meta data update should be simple.

Right.

There should be no worry for overflow as subtraction is involved. Due to this, the scaling code can be simple.

Right. As we now use edge frequency to calculate edge weights, the only think we should take care is that we need to normalize those frequencies to let them fit in 32-bit.

lib/Transforms/Scalar/JumpThreading.cpp

1442

You are right. I have corrected this part.

The BPI update for PredBB->NewBB is unnecessary here as PredBB's old successor BB is replaced by NewBB and PredBB's weights list in BPI can remain unchanged (in BPI each edge weight is mapped with a BB-to-success_index pair).

1574

You are right. Here I should use NewBB to calculate the frequency of the edge from PredBB to BB.

1580

OK.

1590

This is removed now as we are using edge frequencies to calculate edge weights.

1596

It was used to prevent the "Denominator cannot be 0!" error when BB2SuccBBFreq is used to build branch probability. I have also handled this case when scaling weights below.

1610

This is done now.

[congh]

Update the patch by using edge frequencies to update edge weights.

[davidxl]

lib/Support/BlockFrequency.cpp

70 ↗

(On Diff #33467)

No underflow check here?

lib/Transforms/Scalar/JumpThreading.cpp

1573

It is better to call PredBB2BBFreq NewBBFreq (with comment: frequency of the newly split BB)

1583

operator - does not have underflow check (see comment above).

1589

Is this interface available in trunk?

1598

BPI for PredBB also needs to be updated for completeness -- as its successor BB is now NewBB

[congh]

lib/Support/BlockFrequency.cpp

70 ↗

(On Diff #33467)

-= below is using the new operator defined above, so underflow is checked.

lib/Transforms/Scalar/JumpThreading.cpp

1573

OK.

1583

But -= has underflow check which is used in operator -.

1589

Not yet :( Still waiting for the review of that patch.

1598

The old PredBB's successor BB is replaced by NewBB in the same position, so BPI doesn't have to be updated. (In BPI, the successor weights are stores in a map with the index of the successor as the key. The NewBB is set as a successor using the following code:

if (PredTerm->getSuccessor(i) == BB) {
  BB->removePredecessor(PredBB, true);
  PredTerm->setSuccessor(i, NewBB);
}

So the position of this new successor is not changed.)

[davidxl]

lib/Support/BlockFrequency.cpp

70 ↗

(On Diff #33467)

ok

lib/Transforms/Scalar/JumpThreading.cpp

1589

Since the interfaces are under large rework, I suggest simply extract extract the minimal necessary normalization code from that patch here (and abandon that one). Is it doable?

1598

ok then.

[congh]

Update the patch by adding the definition of MachineBranchProbabilityInfo::normalizeEdgeWeights().

include/llvm/Analysis/BlockFrequencyInfo.h

public:
 
   /// getblockFreq - Return block frequency. Return 0 if we don't have the
   /// information. Please note that initial frequency is equal to ENTRY_FREQ. 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 BasicBlock *BB) 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
   // frequency to convert freq into a relative decimal form.
   raw_ostream &printBlockFreq(raw_ostream &OS, const BlockFrequency Freq) const;
 

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;
 
+  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 {
     assert(!Freqs.empty());
     return Freqs[0].Integer;
   }
public:
   void calculate(const FunctionT &F, const BranchProbabilityInfoT &BPI,
                  const LoopInfoT &LI);
   BlockFrequencyInfoImpl() : BPI(nullptr), LI(nullptr), F(nullptr) {}
 
   using BlockFrequencyInfoImplBase::getEntryFreq;
   BlockFrequency getBlockFreq(const BlockT *BB) const {
     return BlockFrequencyInfoImplBase::getBlockFreq(getNode(BB));
   }
+  void setBlockFreq(const BlockT *BB, uint64_t Freq);
   Scaled64 getFloatingBlockFreq(const BlockT *BB) const {
     return BlockFrequencyInfoImplBase::getFloatingBlockFreq(getNode(BB));
   }
 
   /// \brief Print the frequencies for the current function.
   ///
   /// Prints the frequencies for the blocks in the current function.
   ///
void BlockFrequencyInfoImpl<BT>::calculate(const FunctionT &F,
   // Visit loops in post-order to find the local mass distribution, and then do
   // the full function.
   computeMassInLoops();
   computeMassInFunction();
   unwrapLoops();
   finalizeMetrics();
 }
 
+template <class BT>
+void BlockFrequencyInfoImpl<BT>::setBlockFreq(const BlockT *BB, uint64_t Freq) {
+  if (Nodes.count(BB))
+    BlockFrequencyInfoImplBase::setBlockFreq(getNode(BB), Freq);
+  else {
+    // If BB is a newly added block after BFI is done, we need to create a new
+    // BlockNode for it assigned with a new index. The index can be determined
+    // by the size of Freqs.
+    BlockNode NewNode(Freqs.size());
+    Nodes[BB] = NewNode;
+    Freqs.emplace_back();
+    BlockFrequencyInfoImplBase::setBlockFreq(NewNode, Freq);
+  }
+}
+
 template <class BT> void BlockFrequencyInfoImpl<BT>::initializeRPOT() {
   const BlockT *Entry = F->begin();
   RPOT.reserve(F->size());
   std::copy(po_begin(Entry), po_end(Entry), std::back_inserter(RPOT));
   std::reverse(RPOT.begin(), RPOT.end());
 
   assert(RPOT.size() - 1 <= BlockNode::getMaxIndex() &&
          "More nodes in function than Block Frequency Info supports");

include/llvm/Analysis/BranchProbabilityInfo.h

public:
   ///
   /// This allows a pass to explicitly set the edge weight for an edge. It can
   /// be used when updating the CFG to update and preserve the branch
   /// probability information. Read the implementation of how these edge
   /// weights are calculated carefully before using!
   void setEdgeWeight(const BasicBlock *Src, unsigned IndexInSuccessors,
                      uint32_t Weight);
 
+  /// \brief Update the raw edge weight for a given edge.
+  ///
+  /// In this function the index of Dst in Src's successors will be calculated
+  /// on the fly.
+  void updateEdgeWeight(const BasicBlock *Src, const BasicBlock *Dst,
+                        uint32_t Weight);
+
   static uint32_t getBranchWeightStackProtector(bool IsLikely) {
     return IsLikely ? (1u << 20) - 1 : 1;
   }
 
   void calculate(Function &F, const LoopInfo& LI);
 
 private:
   // Since we allow duplicate edges from one basic block to another, we use

lib/Analysis/BlockFrequencyInfo.cpp

if (ViewBlockFreqPropagationDAG != GVDT_None)
     view();
 #endif
 }
 
 BlockFrequency BlockFrequencyInfo::getBlockFreq(const BasicBlock *BB) const {
   return BFI ? BFI->getBlockFreq(BB) : 0;
 }
 
+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 {
 // This code is only for debugging.
 #ifndef NDEBUG
   ViewGraph(const_cast<BlockFrequencyInfo *>(this), "BlockFrequencyDAGs");
 #else
   errs() << "BlockFrequencyInfo::view is only available in debug builds on "

lib/Analysis/BlockFrequencyInfoImpl.cpp

 }
 Scaled64
 BlockFrequencyInfoImplBase::getFloatingBlockFreq(const BlockNode &Node) const {
   if (!Node.isValid())
     return Scaled64::getZero();
   return Freqs[Node.Index].Scaled;
 }
 
+void BlockFrequencyInfoImplBase::setBlockFreq(const BlockNode &Node,
+                                              uint64_t Freq) {
+  assert(Node.isValid() && "Expected valid node");
+  assert(Node.Index < Freqs.size() && "Expected legal index");
+  Freqs[Node.Index].Integer = Freq;
+}
+
 std::string
 BlockFrequencyInfoImplBase::getBlockName(const BlockNode &Node) const {
   return std::string();
 }
 std::string
 BlockFrequencyInfoImplBase::getLoopName(const LoopData &Loop) const {
   return getBlockName(Loop.getHeader()) + (Loop.isIrreducible() ? "**" : "*");
 }

lib/Analysis/BranchProbabilityInfo.cpp

getEdgeWeight(const BasicBlock *Src, const BasicBlock *Dst) const {
   for (succ_const_iterator I = succ_begin(Src), E = succ_end(Src); I != E; ++I)
     if (*I == Dst) {
       MapI = Weights.find(std::make_pair(Src, I.getSuccessorIndex()));
       if (MapI != Weights.end()) {
         FoundWeight = true;
         Weight += MapI->second;
       }
     }
-  return (!FoundWeight) ? DEFAULT_WEIGHT : Weight;
+  return FoundWeight ? Weight : DEFAULT_WEIGHT;
 }
 
 /// Set the edge weight for a given edge specified by PredBlock and an index
 /// to the successors.
 void BranchProbabilityInfo::
 setEdgeWeight(const BasicBlock *Src, unsigned IndexInSuccessors,
               uint32_t Weight) {
   Weights[std::make_pair(Src, IndexInSuccessors)] = Weight;
   DEBUG(dbgs() << "set edge " << Src->getName() << " -> "
                << IndexInSuccessors << " successor weight to "
                << Weight << "\n");
 }
 
+/// Set the edge weight for a given edge specified by Src and Dst.
+void BranchProbabilityInfo::updateEdgeWeight(const BasicBlock *Src,
+                                             const BasicBlock *Dst,
+                                             uint32_t Weight) {
+  unsigned Index;
+  bool FoundEdge = false;
+  for (auto I = succ_begin(Src), E = succ_end(Src); I != E; ++I) {
+    if (*I == Dst) {
+      Index = I.getSuccessorIndex();
+      FoundEdge = true;
+      break;
+    }
+  }
+  assert(FoundEdge && "Could not find Dst in Src's successors");
+  setEdgeWeight(Src, Index, Weight);
+}
+
 /// Get an edge's probability, relative to other out-edges from Src.
 BranchProbability BranchProbabilityInfo::
 getEdgeProbability(const BasicBlock *Src, unsigned IndexInSuccessors) const {
   uint32_t N = getEdgeWeight(Src, IndexInSuccessors);
   uint32_t D = getSumForBlock(Src);
 
   return BranchProbability(N, D);
 }

lib/Transforms/Scalar/JumpThreading.cpp

 #include "llvm/Transforms/Scalar.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallSet.h"
 #include "llvm/ADT/Statistic.h"
 #include "llvm/Analysis/CFG.h"
+#include "llvm/Analysis/BlockFrequencyInfo.h"
+#include "llvm/Analysis/BlockFrequencyInfoImpl.h"
+#include "llvm/Analysis/BranchProbabilityInfo.h"
 #include "llvm/Analysis/ConstantFolding.h"
 #include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/LazyValueInfo.h"
 #include "llvm/Analysis/Loads.h"
+#include "llvm/Analysis/LoopInfo.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/ValueHandle.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
 #include "llvm/Transforms/Utils/Local.h"
namespace {
   ///   if (X < 3) {
   ///
   /// In this case, the unconditional branch at the end of the first if can be
   /// revectored to the false side of the second if.
   ///
   class JumpThreading : public FunctionPass {
     TargetLibraryInfo *TLI;
     LazyValueInfo *LVI;
+    std::unique_ptr<BlockFrequencyInfo> BFI;
+    std::unique_ptr<BranchProbabilityInfo> BPI;
+    bool HasProfileData;
 #ifdef NDEBUG
     SmallPtrSet<BasicBlock*, 16> LoopHeaders;
 #else
     SmallSet<AssertingVH<BasicBlock>, 16> LoopHeaders;
 #endif
     DenseSet<std::pair<Value*, BasicBlock*> > RecursionSet;
 
     unsigned BBDupThreshold;
bool ProcessThreadableEdges(Value *Cond, BasicBlock *BB,
                                 ConstantPreference Preference,
                                 Instruction *CxtI = nullptr);
 
     bool ProcessBranchOnPHI(PHINode *PN);
     bool ProcessBranchOnXOR(BinaryOperator *BO);
 
     bool SimplifyPartiallyRedundantLoad(LoadInst *LI);
     bool TryToUnfoldSelect(CmpInst *CondCmp, BasicBlock *BB);
+
+  private:
+    BasicBlock *SplitBlockPreds(BasicBlock *BB, ArrayRef<BasicBlock *> Preds,
+                                const char *Suffix);
+    void UpdateBlockFreqAndEdgeWeight(BasicBlock *PredBB, BasicBlock *BB,
+                                      BasicBlock *SuccBB);
   };
 }
 
 char JumpThreading::ID = 0;
 INITIALIZE_PASS_BEGIN(JumpThreading, "jump-threading",
                 "Jump Threading", false, false)
 INITIALIZE_PASS_DEPENDENCY(LazyValueInfo)
 INITIALIZE_PASS_DEPENDENCY(TargetLibraryInfoWrapperPass)
 INITIALIZE_PASS_END(JumpThreading, "jump-threading",
                 "Jump Threading", false, false)
 
 // Public interface to the Jump Threading pass
 FunctionPass *llvm::createJumpThreadingPass(int Threshold) { return new JumpThreading(Threshold); }
 
 /// runOnFunction - Top level algorithm.
 ///
 bool JumpThreading::runOnFunction(Function &F) {
   if (skipOptnoneFunction(F))
     return false;
 
   DEBUG(dbgs() << "Jump threading on function '" << F.getName() << "'\n");
   TLI = &getAnalysis<TargetLibraryInfoWrapperPass>().getTLI();
   LVI = &getAnalysis<LazyValueInfo>();
+  BFI.reset();
+  BPI.reset();
+  HasProfileData = F.getEntryCount().hasValue();
+  if (HasProfileData) {
+    LoopInfo LI{DominatorTree(F)};
+    BPI.reset(new BranchProbabilityInfo(F, LI));
+    BFI.reset(new BlockFrequencyInfo(F, *BPI, LI));
+  }
 
   // Remove unreachable blocks from function as they may result in infinite
   // loop. We do threading if we found something profitable. Jump threading a
   // branch can create other opportunities. If these opportunities form a cycle
   // i.e. if any jump treading is undoing previous threading in the path, then
   // we will loop forever. We take care of this issue by not jump threading for
   // back edges. This works for normal cases but not for unreachable blocks as
   // they may have cycle with no back edge.
for (pred_iterator PI = pred_begin(LoadBB), PE = pred_end(LoadBB);
       if (isa<IndirectBrInst>(P->getTerminator()))
         return false;
 
       if (!AvailablePredSet.count(P))
         PredsToSplit.push_back(P);
     }
 
     // Split them out to their own block.
-    UnavailablePred =
-      SplitBlockPredecessors(LoadBB, PredsToSplit, "thread-pre-split");
+    UnavailablePred = SplitBlockPreds(LoadBB, PredsToSplit, "thread-pre-split");
   }
 
   // If the value isn't available in all predecessors, then there will be
   // exactly one where it isn't available.  Insert a load on that edge and add
   // it to the AvailablePreds list.
   if (UnavailablePred) {
     assert(UnavailablePred->getTerminator()->getNumSuccessors() == 1 &&
            "Can't handle critical edge here!");
bool JumpThreading::ThreadEdge(BasicBlock *BB,
 
   // And finally, do it!  Start by factoring the predecessors is needed.
   BasicBlock *PredBB;
   if (PredBBs.size() == 1)
     PredBB = PredBBs[0];
   else {
     DEBUG(dbgs() << "  Factoring out " << PredBBs.size()
           << " common predecessors.\n");
-    PredBB = SplitBlockPredecessors(BB, PredBBs, ".thr_comm");
+    PredBB = SplitBlockPreds(BB, PredBBs, ".thr_comm");
   }
 
   // And finally, do it!
   DEBUG(dbgs() << "  Threading edge from '" << PredBB->getName() << "' to '"
         << SuccBB->getName() << "' with cost: " << JumpThreadCost
         << ", across block:\n    "
         << *BB << "\n");
 
   LVI->threadEdge(PredBB, BB, SuccBB);
 
   // We are going to have to map operands from the original BB block to the new
   // copy of the block 'NewBB'.  If there are PHI nodes in BB, evaluate them to
   // account for entry from PredBB.
   DenseMap<Instruction*, Value*> ValueMapping;
 
   BasicBlock *NewBB = BasicBlock::Create(BB->getContext(),
                                          BB->getName()+".thread",
                                          BB->getParent(), BB);
   NewBB->moveAfter(PredBB);
 
+  // Set the block frequency of NewBB, which is identical to the frequency of
+  // PredBB.
+  if (HasProfileData)
+    BFI->setBlockFreq(NewBB, BFI->getBlockFreq(PredBB).getFrequency());
+
   BasicBlock::iterator BI = BB->begin();
   for (; PHINode *PN = dyn_cast<PHINode>(BI); ++BI)
     ValueMapping[PN] = PN->getIncomingValueForBlock(PredBB);
 
   // Clone the non-phi instructions of BB into NewBB, keeping track of the
   // mapping and using it to remap operands in the cloned instructions.
   for (; !isa<TerminatorInst>(BI); ++BI) {
     Instruction *New = BI->clone();
     New->setName(BI->getName());
     NewBB->getInstList().push_back(New);
     ValueMapping[BI] = New;
 
     // Remap operands to patch up intra-block references.
     for (unsigned i = 0, e = New->getNumOperands(); i != e; ++i)
       if (Instruction *Inst = dyn_cast<Instruction>(New->getOperand(i))) {
         DenseMap<Instruction*, Value*>::iterator I = ValueMapping.find(Inst);
         if (I != ValueMapping.end())
           New->setOperand(i, I->second);
       }
   }
 
   // We didn't copy the terminator from BB over to NewBB, because there is now
   // an unconditional jump to SuccBB.  Insert the unconditional jump.
-  BranchInst *NewBI =BranchInst::Create(SuccBB, NewBB);
+  BranchInst *NewBI = BranchInst::Create(SuccBB, NewBB);
   NewBI->setDebugLoc(BB->getTerminator()->getDebugLoc());
 
   // Check to see if SuccBB has PHI nodes. If so, we need to add entries to the
   // PHI nodes for NewBB now.
   AddPHINodeEntriesForMappedBlock(SuccBB, BB, NewBB, ValueMapping);
 
   // If there were values defined in BB that are used outside the block, then we
   // now have to update all uses of the value to use either the original value,
bool JumpThreading::ThreadEdge(BasicBlock *BB,
   // NewBB instead of BB.  This eliminates predecessors from BB, which requires
   // us to simplify any PHI nodes in BB.
   TerminatorInst *PredTerm = PredBB->getTerminator();
   for (unsigned i = 0, e = PredTerm->getNumSuccessors(); i != e; ++i)
     if (PredTerm->getSuccessor(i) == BB) {
       BB->removePredecessor(PredBB, true);
       PredTerm->setSuccessor(i, NewBB);
     }
 
   // At this point, the IR is fully up to date and consistent.  Do a quick scan
   // over the new instructions and zap any that are constants or dead.  This
   // frequently happens because of phi translation.
   SimplifyInstructionsInBlock(NewBB, TLI);
 
+  // Update the edge weight from BB to SuccBB, which should be less than before.
+  UpdateBlockFreqAndEdgeWeight(PredBB, BB, SuccBB);
+
   // Threaded an edge!
   ++NumThreads;
   return true;
 }
 
+/// Create a new basic block that will be the predecessor of BB and successor of
+/// all blocks in Preds. When profile data is availble, update the frequency of
+/// this new block.
+BasicBlock *JumpThreading::SplitBlockPreds(BasicBlock *BB,
+                                           ArrayRef<BasicBlock *> Preds,
+                                           const char *Suffix) {
+  // Collect the frequencies of all predecessors of BB, which will be used to
+  // update the edge weight on BB->SuccBB.
+  BlockFrequency PredBBFreq(0);
+  if (HasProfileData)
+    for (auto Pred : Preds)
+      PredBBFreq += BFI->getBlockFreq(Pred) * BPI->getEdgeProbability(Pred, BB);
+
+  BasicBlock *PredBB = SplitBlockPredecessors(BB, Preds, Suffix);
+
+  // Set the block frequency of the newly created PredBB, which is the sum of
+  // frequencies of Preds.
+  if (HasProfileData)
+    BFI->setBlockFreq(PredBB, PredBBFreq.getFrequency());
+  return PredBB;
+}
+
+/// Update the block frequency of BB and branch weight and the metadata on the
+/// edge BB->SuccBB. This is done by scaling the weight of BB->SuccBB by 1 -
+/// Freq(PredBB->BB) / Freq(BB->SuccBB).
+void JumpThreading::UpdateBlockFreqAndEdgeWeight(BasicBlock *PredBB,
+                                                 BasicBlock *BB,
+                                                 BasicBlock *SuccBB) {
+  if (!HasProfileData)
+    return;
+
+  assert(BFI && "BFI should have been created here");
+  assert(BPI && "BPI should have been created here");
+
+  // As the edge from PredBB to BB is deleted, we have to update the block
+  // frequency of BB.
+  auto BBFreq = BFI->getBlockFreq(BB);
+  auto PredBBFreq = BFI->getBlockFreq(PredBB);
+  auto BB2SuccBBFreq = BBFreq * BPI->getEdgeProbability(BB, SuccBB);
+  // Theoretically PredBBFreq should not be greater than BBFreq. However, in
+  // some cases this may happen when the frequencies are incorrectly
+  // distributed to CFG nodes. In this case we just set the new frequency of
+  // BB to zero.
+  uint64_t NewBBFreq = (PredBBFreq <= BBFreq)
+                           ? BBFreq.getFrequency() - PredBBFreq.getFrequency()
+                           : 0;
+  BFI->setBlockFreq(BB, NewBBFreq);
+
+  SmallVector<uint32_t, 8> Weights;
+  auto TI = BB->getTerminator();
+  if (MDNode *MD = TI->getMetadata(LLVMContext::MD_prof)) {
+    int SuccIndexToUpdate = -1;
+    for (unsigned MD_i = 1, MD_e = MD->getNumOperands(); MD_i < MD_e; ++MD_i) {
+      ConstantInt *CI = mdconst::extract<ConstantInt>(MD->getOperand(MD_i));
+      auto W = CI->getValue().getZExtValue();
+      if (TI->getSuccessor(MD_i - 1) == SuccBB) {
+        // Note that we are using the old BB's frequency to calculate the old
+        // BB->SuccBB's frequency.
+        assert(BPI->getEdgeWeight(BB, SuccBB) == W &&
+               "The weight on BB->SuccBB is not coherent with it in BPI");
+        if (BB2SuccBBFreq.getFrequency() > 0) {
+          DEBUG(dbgs() << "  Update the weight on the edge from "
+                       << BB->getName() << " to " << SuccBB->getName() << "\n");
+          if (PredBBFreq <= BB2SuccBBFreq) {
+            // Here we record the index of SuccBB in successors of BB and
+            // updates weights of all BB's successors later.
+            SuccIndexToUpdate = MD_i - 1;
+          } else {
+            W = 0;
+            // Update BPI.
+            BPI->updateEdgeWeight(BB, SuccBB, W);
+          }
+        }
+      }
+      Weights.push_back(W);
+    }
+
+    if (SuccIndexToUpdate >= 0) {
+      // We need to update the weights of BB's successors.
+      // Instead of scaling the weight of BB->SuccBB by 1 - Freq(PredBB->BB) /
+      // Freq(BB->SuccBB), which may suffer from precision loss, we scale it by
+      // Freq(BB->SuccBB) - Freq(PredBB->BB), and the weights of other BB's
+      // successors by Freq(BB->SuccBB). We need to avoid potential overflow
+      // here.
+
+      uint64_t ScaleOtherSucc = BB2SuccBBFreq.getFrequency();
+      uint64_t ScaleBB2SuccBB = ScaleOtherSucc - PredBBFreq.getFrequency();
+      SmallVector<uint64_t, 8> Weights64;
+      // A lambda that calculates the number of effective bits of the given
+      // integer.
+      auto GetBits = [](uint64_t Val) {
+        int Bits = 0;
+        while (Val > 0) {
+          Bits++;
+          Val >>= 1;
+        }
+        return Bits;
+      };
+      int MaxBitsNum = 0;
+
+      for (int i = 0, s = Weights.size(); i < s; ++i) {
+        uint64_t Scale =
+            i == SuccIndexToUpdate ? ScaleBB2SuccBB : ScaleOtherSucc;
+        // Use BranchProbability::scaleByInverse to calculate Weights[i] * Scale
+        // to avoid overflow.
+        uint64_t Weight64 =
+            BranchProbability(1, Weights[i]).scaleByInverse(Scale);
+        MaxBitsNum = std::max(MaxBitsNum, GetBits(Weight64));
+        Weights64.push_back(Weight64);
+      }
+
+      // If there is at least one weight that is greater than UINT32_MAX, we
+      // scale all weights by 1 / 2 ^ (MaxBitsNum - 32)/
+      if (MaxBitsNum > 32)
+        for (uint64_t &Weight64 : Weights64)
+          Weight64 >>= MaxBitsNum - 32;
+
+      // Finally, safely convert each 64-bit weight into 32-bit one, and update
+      // the edge weights in BPI.
+      for (int i = 0, s = Weights.size(); i < s; ++i) {
+        Weights[i] = static_cast<uint32_t>(Weights64[i]);
+        BPI->updateEdgeWeight(BB, TI->getSuccessor(i), Weights[i]);
+      }
+    }
+
+    if (Weights.size() >= 2) {
+      TI->setMetadata(LLVMContext::MD_prof,
+                      MDBuilder(TI->getParent()->getContext())
+                          .createBranchWeights(Weights));
+    }
+  }
+}
+
 /// DuplicateCondBranchOnPHIIntoPred - PredBB contains an unconditional branch
 /// to BB which contains an i1 PHI node and a conditional branch on that PHI.
 /// If we can duplicate the contents of BB up into PredBB do so now, this
 /// improves the odds that the branch will be on an analyzable instruction like
 /// a compare.
 bool JumpThreading::DuplicateCondBranchOnPHIIntoPred(BasicBlock *BB,
                                  const SmallVectorImpl<BasicBlock *> &PredBBs) {
   assert(!PredBBs.empty() && "Can't handle an empty set");
bool JumpThreading::DuplicateCondBranchOnPHIIntoPred(BasicBlock *BB,
 
   // And finally, do it!  Start by factoring the predecessors is needed.
   BasicBlock *PredBB;
   if (PredBBs.size() == 1)
     PredBB = PredBBs[0];
   else {
     DEBUG(dbgs() << "  Factoring out " << PredBBs.size()
           << " common predecessors.\n");
-    PredBB = SplitBlockPredecessors(BB, PredBBs, ".thr_comm");
+    PredBB = SplitBlockPreds(BB, PredBBs, ".thr_comm");
   }
 
   // Okay, we decided to do this!  Clone all the instructions in BB onto the end
   // of PredBB.
   DEBUG(dbgs() << "  Duplicating block '" << BB->getName() << "' into end of '"
         << PredBB->getName() << "' to eliminate branch on phi.  Cost: "
         << DuplicationCost << " block is:" << *BB << "\n");
 

test/Transforms/JumpThreading/update-edge-weight.ll

+; RUN: opt -S -jump-threading < %s | FileCheck %s
+
+; Test if edge weights are properly updated after jump threading.
+
+; CHECK: !2 = !{!"branch_weights", i32 210, i32 79}
+define void @foo(i32 %n) !prof !0 {
+entry:
+  %cmp = icmp sgt i32 %n, 10
+  br i1 %cmp, label %if.then.1, label %if.else.1, !prof !1
+
+if.then.1:
+  tail call void @a()
+  br label %if.cond
+
+if.else.1:
+  tail call void @b()
+  br label %if.cond
+
+if.cond:
+  %cmp1 = icmp sgt i32 %n, 5
+  br i1 %cmp1, label %if.then.2, label %if.else.2, !prof !2
+
+if.then.2:
+  tail call void @c()
+  br label %if.end
+
+if.else.2:
+  tail call void @d()
+  br label %if.end
+
+if.end:
+  ret void
+}
+
+declare void @a()
+declare void @b()
+declare void @c()
+declare void @d()
+
+!0 = !{!"function_entry_count", i64 1}
+!1 = !{!"branch_weights", i32 10, i32 5}
+!2 = !{!"branch_weights", i32 10, i32 1}