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 ↗

(On Diff #29140)

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 ↗

(On Diff #29140)

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 ↗

(On Diff #29140)

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 ↗

(On Diff #29963)

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

lib/Analysis/BranchProbabilityInfo.cpp

622 ↗

(On Diff #29963)

It is probably better to call this one 'updateEdgeWeight'

lib/Transforms/Scalar/JumpThreading.cpp

150 ↗

(On Diff #29963)

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

1560 ↗

(On Diff #29963)

assert here?

1568 ↗

(On Diff #29963)

fix typo.

1569 ↗

(On Diff #29963)

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

1583 ↗

(On Diff #29963)

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

1583 ↗

(On Diff #29963)

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

1588 ↗

(On Diff #29963)

Another insanity to track?

1591 ↗

(On Diff #29963)

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

1543 ↗

(On Diff #29963)

Can be skipped when Preds' s size == 1

1551 ↗

(On Diff #29963)

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 ↗

(On Diff #30214)

change the name to OrigBBFreq

1578 ↗

(On Diff #30214)

Is this still an issue that needs the workaround here?

1596 ↗

(On Diff #30214)

Why this condition?

1599 ↗

(On Diff #30214)

It is unclear how can this happen.

1614 ↗

(On Diff #30214)

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 ↗

(On Diff #30214)

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

1596 ↗

(On Diff #30214)

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

1599 ↗

(On Diff #30214)

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 ↗

(On Diff #30214)

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 ↗

(On Diff #30214)

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 ↗

(On Diff #30214)

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

1453 ↗

(On Diff #31408)

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

Where is the BPI update for PredBB->NewBB?

1571 ↗

(On Diff #31408)

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

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

1585 ↗

(On Diff #31408)

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).

1591 ↗

(On Diff #31408)

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.

1595 ↗

(On Diff #31408)

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).

1601 ↗

(On Diff #31408)

This whole loop seems unnecessary.

1621 ↗

(On Diff #31408)

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

1453 ↗

(On Diff #31408)

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).

1585 ↗

(On Diff #31408)

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

1591 ↗

(On Diff #31408)

OK.

1601 ↗

(On Diff #31408)

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

1621 ↗

(On Diff #31408)

This is done now.

1596 ↗

(On Diff #30214)

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.

[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

1591 ↗

(On Diff #33467)

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

1601 ↗

(On Diff #33467)

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

1607 ↗

(On Diff #33467)

Is this interface available in trunk?

1616 ↗

(On Diff #33467)

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

1591 ↗

(On Diff #33467)

OK.

1601 ↗

(On Diff #33467)

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

1607 ↗

(On Diff #33467)

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

1616 ↗

(On Diff #33467)

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

1607 ↗

(On Diff #33467)

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?

1616 ↗

(On Diff #33467)

ok then.

[congh]

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

llvm/trunk/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;
 

llvm/trunk/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->front();
   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");

llvm/trunk/include/llvm/CodeGen/MachineBranchProbabilityInfo.h

BranchProbability getEdgeProbability(const MachineBasicBlock *Src,
                                        const MachineBasicBlock *Dst) const;
 
   // Print value between 0 (0% probability) and 1 (100% probability),
   // however the value is never equal to 0, and can be 1 only iff SRC block
   // has only one successor.
   raw_ostream &printEdgeProbability(raw_ostream &OS,
                                     const MachineBasicBlock *Src,
                                     const MachineBasicBlock *Dst) const;
+
+  // Normalize a list of weights by scaling them down so that the sum of them
+  // doesn't exceed UINT32_MAX. Return the scale.
+  template <class WeightListIter>
+  static uint32_t normalizeEdgeWeights(WeightListIter Begin,
+                                       WeightListIter End);
 };
 
+template <class WeightListIter>
+uint32_t
+MachineBranchProbabilityInfo::normalizeEdgeWeights(WeightListIter Begin,
+                                                   WeightListIter End) {
+  // First we compute the sum with 64-bits of precision.
+  uint64_t Sum = std::accumulate(Begin, End, uint64_t(0));
+
+  // If the computed sum fits in 32-bits, we're done.
+  if (Sum <= UINT32_MAX)
+    return 1;
+
+  // Otherwise, compute the scale necessary to cause the weights to fit, and
+  // re-sum with that scale applied.
+  assert((Sum / UINT32_MAX) < UINT32_MAX &&
+         "The sum of weights exceeds UINT32_MAX^2!");
+  uint32_t Scale = (Sum / UINT32_MAX) + 1;
+  for (auto I = Begin; I != End; ++I)
+    *I /= Scale;
+  return Scale;
+}
+
 }
 
 
 #endif

llvm/trunk/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 "

llvm/trunk/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() ? "**" : "*");
 }

llvm/trunk/lib/Transforms/Scalar/JumpThreading.cpp

 #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/GlobalsModRef.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/CodeGen/MachineBranchProbabilityInfo.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"
 #include "llvm/Transforms/Utils/SSAUpdater.h"
+#include <algorithm>
+#include <memory>
 using namespace llvm;
 
 #define DEBUG_TYPE "jump-threading"
 
 STATISTIC(NumThreads, "Number of jumps threaded");
 STATISTIC(NumFolds,   "Number of terminators folded");
 STATISTIC(NumDupes,   "Number of branch blocks duplicated to eliminate phi");
 
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;
public:
 
     void getAnalysisUsage(AnalysisUsage &AU) const override {
       AU.addRequired<LazyValueInfo>();
       AU.addPreserved<LazyValueInfo>();
       AU.addPreserved<GlobalsAAWrapperPass>();
       AU.addRequired<TargetLibraryInfoWrapperPass>();
     }
 
+    void releaseMemory() override {
+      BFI.reset();
+      BPI.reset();
+    }
+
     void FindLoopHeaders(Function &F);
     bool ProcessBlock(BasicBlock *BB);
     bool ThreadEdge(BasicBlock *BB, const SmallVectorImpl<BasicBlock*> &PredBBs,
                     BasicBlock *SuccBB);
     bool DuplicateCondBranchOnPHIIntoPred(BasicBlock *BB,
                                   const SmallVectorImpl<BasicBlock *> &PredBBs);
 
     bool ComputeValueKnownInPredecessors(Value *V, BasicBlock *BB,
                                          PredValueInfo &Result,
                                          ConstantPreference Preference,
                                          Instruction *CxtI = nullptr);
     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 *NewBB, 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();
+  // When profile data is available, we need to update edge weights after
+  // successful jump threading, which requires both BPI and BFI being available.
+  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 threading 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 if 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.
+  if (HasProfileData) {
+    auto NewBBFreq =
+        BFI->getBlockFreq(PredBB) * BPI->getEdgeProbability(PredBB, BB);
+    BFI->setBlockFreq(NewBB, NewBBFreq.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,
if (PredTerm->getSuccessor(i) == BB) {
       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, NewBB, 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 *NewBB,
+                                                 BasicBlock *SuccBB) {
+  if (!HasProfileData)
+    return;
+
+  assert(BFI && BPI && "BFI & 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 BBOrigFreq = BFI->getBlockFreq(BB);
+  auto NewBBFreq = BFI->getBlockFreq(NewBB);
+  auto BB2SuccBBFreq = BBOrigFreq * BPI->getEdgeProbability(BB, SuccBB);
+  auto BBNewFreq = BBOrigFreq - NewBBFreq;
+  BFI->setBlockFreq(BB, BBNewFreq.getFrequency());
+
+  // Collect updated outgoing edges' frequencies from BB and use them to update
+  // edge weights.
+  SmallVector<uint64_t, 4> BBSuccFreq;
+  for (auto I = succ_begin(BB), E = succ_end(BB); I != E; ++I) {
+    auto SuccFreq = (*I == SuccBB)
+                        ? BB2SuccBBFreq - NewBBFreq
+                        : BBOrigFreq * BPI->getEdgeProbability(BB, *I);
+    BBSuccFreq.push_back(SuccFreq.getFrequency());
+  }
+
+  // Normalize edge weights in Weights64 so that the sum of them can fit in
+  MachineBranchProbabilityInfo::normalizeEdgeWeights(BBSuccFreq.begin(),
+                                                     BBSuccFreq.end());
+
+  SmallVector<uint32_t, 4> Weights;
+  for (auto Freq : BBSuccFreq)
+    Weights.push_back(static_cast<uint32_t>(Freq));
+
+  // Update edge weights in BPI.
+  for (int I = 0, E = Weights.size(); I < E; I++)
+    BPI->setEdgeWeight(BB, I, Weights[I]);
+
+  if (Weights.size() >= 2) {
+    auto TI = BB->getTerminator();
+    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 if 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");
 

llvm/trunk/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 22, i32 7}
+
+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}