[MBP] Partial tail duplication into hot predecessors

[Carrot]

Summary

Current tail duplication embedded in MBP duplicates a BB into all or none of its predecessors without too much cost analysis. So sometimes it is duplicated into cold predecessors, and in other cases it may miss the duplication into hot predecessors.

This patch improves tail duplication in 3 aspects:

1. A successor can be duplicated into part of its predecessors.
2. A more fine-grained benefit analysis, combined with 1, now a successor is duplicated into hot predecessors only.
3. If a successor can't be duplicated into one predecessor, it doesn't impact the duplication into other predecessors.

It doesn't impact the performance of spec2006int like many other code layout changes, spec doesn't have large instruction work set.
Performance test with two Google's internal search benchmarks shows obvious improvement on large binaries. I got 0.30% and 0.18% on them.

[xbolva00]

New improvement looks cool, any data for clang itself (large binary too)?

[davidxl]

llvm/lib/CodeGen/MachineBlockPlacement.cpp

1158

should the benefit analysis be done here instead of just returning true?

1820

can you explain this new condition?

3107

Add some description (and examples) showing the algorithm/cost model used here.

[Carrot]

In D73387#1839878, @xbolva00 wrote:

New improvement looks cool, any data for clang itself (large binary too)?

Without this patch
real 139m31.300s
user 130m0.326s
sys 6m33.313s

With this patch
real 139m1.414s
user 129m46.296s
sys 6m30.924s

[xbolva00]

Thanks!

[Carrot]

llvm/lib/CodeGen/MachineBlockPlacement.cpp

1158

It's possible to call findDuplicateCandidates here. But it needs to be called again just before TailDup.tailDuplicateAndUpdate to provide the duplicate candidates.

1820

Previously if function repeatedlyTailDuplicateBlock returns true, BestSucc is duplicated into all predecessors and removed. So there is no BestSucc any more, so we goto the loop header to find another best successor for BB.

Now with partial duplication, BestSucc may or may not be duplicated into BB.

• If BestSucc is duplicated into BB, thus the condition is true, BestSucc is not BB's successor, we need to find another successor of BB.
• If BestSucc is not duplicated into BB, it falls through to the code to layout BestSucc after BB.

[xur]

Thanks for working on this. Current Duplication in block-placement does need improvement. We notice that disable it often improved performance.

llvm/include/llvm/CodeGen/TailDuplicator.h

96

Need to update the comments for this function.
Nit: maybe move CandidatePtr before DuplicatePreds?

llvm/lib/CodeGen/MachineBlockPlacement.cpp

3035

It seems to me that findDuplicateCandidates() will not let CandidatePreds to be full Preds (it will put the first one to fallthrough -- which makes this condition always true).

3072

what does "most possible" mean here?

3084

add "excluding BB".

3104

Add some comments to explain the modeling.
The same value is computed in findDuplicationCandidates() i.e. BBDupThreshold-- should we refactor the code to share the same heuristic?

3107

As David said, need to have cost model description here.

I give my first impression here:
(1) should you check if BB can be placed after a Pred? Pred could be in the middle of the chain.
(2) It seems that you assume the Succ BBs will be chosen as fall-through by different Dupped BBs.
Should you check if SuccIt has already be placed? Also, if SuccIt can be duplicated, the assumption does not to hold.
(3) Pred fall-through can be handled better. this algorithm does not actually try to find the optimal fall through of BB. It passively get the fall-through when we cannot duplicate. Maybe we can refactor code better so the post fixup at line 3171 is not needed.

3160

I don't think this will ever be true.

3192

I think using profile count will get you the same thing here. Profile count is derived from frequency.

[xbolva00]

https://bugs.llvm.org/show_bug.cgi?id=44596

offtopic: Can you look at this issue? MBP fais here?

[Carrot]

llvm/lib/CodeGen/MachineBlockPlacement.cpp

3035

Only when Candidates.size() < Preds.size() is true then it changes the first predecessor to fall through.

If Candidates.size() == Preds.size(), the first predecessor will not be changed to fall through.

3072

It should be changed to "best".

3107

If BB can't be duplicated into Pred, isBestSuccessor checks if Pred is in the middle of a chain.
If BB can be duplicated into Pred, Pred has the only successor BB, it is extremely unlikely in the middle of a chain.

The main purpose of the loop is finding out beneficial duplications. If a predecessor has multiple successors we can't duplicate BB into it, even if the edge is hot. But we can make it fall through to BB to get the same benefit as tail duplication.
Change one duplication into fall through is only a plus for code size.

3160

Good catch.

3192

I mean compare the raw profile count with the hot number from profile summary information.

[davidxl]

llvm/lib/CodeGen/MachineBlockPlacement.cpp

3072

Add a comment explaining the scaling.

3137

it assumes *SuccIt is the layout successor (fall through bb). Is it always the case?

3167

"the benefit of duplicating into a predecessor is .."

3171

laid out

[Carrot]

llvm/lib/CodeGen/MachineBlockPlacement.cpp

3137

No, *SuccIt may have other predecessors can fall through to it.
After sorting Succs, Succs.begin() has the largest probability among edges come out from BB. So *SuccIt has the largest chance to be laid after BB.

Consider other cases may improve the following cost analysis. But I think the extra benefit is very small with much more complicated implementation.

[davidxl]

lgtm

llvm/include/llvm/CodeGen/TailDuplicator.h

public:
   bool canTailDuplicate(MachineBasicBlock *TailBB, MachineBasicBlock *PredBB);
 
   /// Tail duplicate a single basic block into its predecessors, and then clean
   /// up.
   /// If \p DuplicatePreds is not null, it will be updated to contain the list
   /// of predecessors that received a copy of \p MBB.
   /// If \p RemovalCallback is non-null. It will be called before MBB is
   /// deleted.
+  /// If \p CandidatePtr is not null, duplicate into these blocks only.
   bool tailDuplicateAndUpdate(
       bool IsSimple, MachineBasicBlock *MBB,
       MachineBasicBlock *ForcedLayoutPred,
       SmallVectorImpl<MachineBasicBlock*> *DuplicatedPreds = nullptr,
-      function_ref<void(MachineBasicBlock *)> *RemovalCallback = nullptr);
+      function_ref<void(MachineBasicBlock *)> *RemovalCallback = nullptr,
+      SmallVectorImpl<MachineBasicBlock *> *CandidatePtr = nullptr);
 
 private:
   using RegSubRegPair = TargetInstrInfo::RegSubRegPair;
 
   void addSSAUpdateEntry(unsigned OrigReg, unsigned NewReg,
                          MachineBasicBlock *BB);
   void processPHI(MachineInstr *MI, MachineBasicBlock *TailBB,
                   MachineBasicBlock *PredBB,
private:
   bool duplicateSimpleBB(MachineBasicBlock *TailBB,
                          SmallVectorImpl<MachineBasicBlock *> &TDBBs,
                          const DenseSet<unsigned> &RegsUsedByPhi,
                          SmallVectorImpl<MachineInstr *> &Copies);
   bool tailDuplicate(bool IsSimple,
                      MachineBasicBlock *TailBB,
                      MachineBasicBlock *ForcedLayoutPred,
                      SmallVectorImpl<MachineBasicBlock *> &TDBBs,
-                     SmallVectorImpl<MachineInstr *> &Copies);
+                     SmallVectorImpl<MachineInstr *> &Copies,
+                     SmallVectorImpl<MachineBasicBlock *> *CandidatePtr);
   void appendCopies(MachineBasicBlock *MBB,
                  SmallVectorImpl<std::pair<unsigned,RegSubRegPair>> &CopyInfos,
                  SmallVectorImpl<MachineInstr *> &Copies);
 
   void removeDeadBlock(
       MachineBasicBlock *MBB,
       function_ref<void(MachineBasicBlock *)> *RemovalCallback = nullptr);
 };
 
 } // end namespace llvm
 
 #endif // LLVM_CODEGEN_TAILDUPLICATOR_H

llvm/lib/CodeGen/MachineBlockPlacement.cpp

class MachineBlockPlacement : public MachineFunctionPass {
 
   /// Duplicator used to duplicate tails during placement.
   ///
   /// Placement decisions can open up new tail duplication opportunities, but
   /// since tail duplication affects placement decisions of later blocks, it
   /// must be done inline.
   TailDuplicator TailDup;
 
+  /// Partial tail duplication threshold.
+  BlockFrequency DupThreshold;
+
   /// Allocator and owner of BlockChain structures.
   ///
   /// We build BlockChains lazily while processing the loop structure of
   /// a function. To reduce malloc traffic, we allocate them using this
   /// slab-like allocator, and destroy them after the pass completes. An
   /// important guarantee is that this allocator produces stable pointers to
   /// the chains.
   SpecificBumpPtrAllocator<BlockChain> ChainAllocator;

 #ifndef NDEBUG
   /// The set of basic blocks that have terminators that cannot be fully
   /// analyzed.  These basic blocks cannot be re-ordered safely by
   /// MachineBlockPlacement, and we must preserve physical layout of these
   /// blocks and their successors through the pass.
   SmallPtrSet<MachineBasicBlock *, 4> BlocksWithUnanalyzableExits;
 #endif
 
+  /// Scale the DupThreshold according to basic block size.
+  BlockFrequency scaleThreshold(MachineBasicBlock *BB);
+  void initDupThreshold();
+
   /// Decrease the UnscheduledPredecessors count for all blocks in chain, and
   /// if the count goes to 0, add them to the appropriate work list.
   void markChainSuccessors(
       const BlockChain &Chain, const MachineBasicBlock *LoopHeaderBB,
       const BlockFilterSet *BlockFilter = nullptr);
 
   /// Decrease the UnscheduledPredecessors count for a single block, and
   /// if the count goes to 0, add them to the appropriate work list.
   void markBlockSuccessors(
       const BlockChain &Chain, const MachineBasicBlock *BB,
       const MachineBasicBlock *LoopHeaderBB,
       const BlockFilterSet *BlockFilter = nullptr);
 
   BranchProbability
   collectViableSuccessors(
       const MachineBasicBlock *BB, const BlockChain &Chain,
       const BlockFilterSet *BlockFilter,
       SmallVector<MachineBasicBlock *, 4> &Successors);
   bool shouldPredBlockBeOutlined(
       const MachineBasicBlock *BB, const MachineBasicBlock *Succ,
       const BlockChain &Chain, const BlockFilterSet *BlockFilter,
       BranchProbability SuccProb, BranchProbability HotProb);
+  bool isBestSuccessor(MachineBasicBlock *BB, MachineBasicBlock *Pred,
+                       BlockFilterSet *BlockFilter);
+  void findDuplicateCandidates(SmallVectorImpl<MachineBasicBlock *> &Candidates,
+                               MachineBasicBlock *BB,
+                               BlockFilterSet *BlockFilter);
   bool repeatedlyTailDuplicateBlock(
       MachineBasicBlock *BB, MachineBasicBlock *&LPred,
       const MachineBasicBlock *LoopHeaderBB,
       BlockChain &Chain, BlockFilterSet *BlockFilter,
       MachineFunction::iterator &PrevUnplacedBlockIt);
   bool maybeTailDuplicateBlock(
       MachineBasicBlock *BB, MachineBasicBlock *LPred,
       BlockChain &Chain, BlockFilterSet *BlockFilter,
for (MachineBasicBlock *Pred : Succ->predecessors()) {
     }
     NumDup++;
   }
 
   // No possible duplication in current filter set.
   if (NumDup == 0)
     return false;
 
+  // If profile information is available, findDuplicateCandidates can do more
+  // precise benefit analysis.
+  if (F->getFunction().hasProfileData())
+    return true;
+
   // This is mainly for function exit BB.
   // The integrated tail duplication is really designed for increasing
   // fallthrough from predecessors from Succ to its successors. We may need
   // other machanism to handle different cases.
   if (Succ->succ_size() == 0)
     return true;
 
   // Plus the already placed predecessor.
bool MachineBlockPlacement::canTailDuplicateUnplacedPreds(
   //       Succ1  Succ2
   //
   // In this example Dup has 2 successors and 3 predecessors, duplication of Dup
   // can increase the fallthrough from Pred1 to Succ1 and from Pred2 to Succ2,
   // but the duplication into Pred3 can't increase fallthrough.
   //
   // A small number of extra duplication may not hurt too much. We need a better
   // heuristic to handle it.
-  //
-  // FIXME: we should selectively tail duplicate a BB into part of its
-  // predecessors.
   if ((NumDup > Succ->succ_size()) || !Duplicate)
     return false;
 
   return true;
 }
 
 /// Find chains of triangles where we believe it would be profitable to
 /// tail-duplicate them all, but a local analysis would not find them.
if (!BestSucc) {
 
       LLVM_DEBUG(dbgs() << "Unnatural loop CFG detected, forcibly merging the "
                            "layout successor until the CFG reduces\n");
     }
 
     // Placement may have changed tail duplication opportunities.
     // Check for that now.
     if (allowTailDupPlacement() && BestSucc && ShouldTailDup) {
-      // If the chosen successor was duplicated into all its predecessors,
-      // don't bother laying it out, just go round the loop again with BB as
-      // the chain end.
-      if (repeatedlyTailDuplicateBlock(BestSucc, BB, LoopHeaderBB, Chain,
-                                       BlockFilter, PrevUnplacedBlockIt))
+      repeatedlyTailDuplicateBlock(BestSucc, BB, LoopHeaderBB, Chain,
+                                       BlockFilter, PrevUnplacedBlockIt);
+      // If the chosen successor was duplicated into BB, don't bother laying
+      // it out, just go round the loop again with BB as the chain end.
+      if (!BB->isSuccessor(BestSucc))
         continue;
     }
 
     // Place this block, updating the datastructures to reflect its placement.
     BlockChain &SuccChain = *BlockToChain[BestSucc];
     // Zero out UnscheduledPredecessors for the successor we're about to merge in case
     // we selected a successor that didn't fit naturally into the CFG.
     SuccChain.UnscheduledPredecessors = 0;
auto RemovalCallback =
         LLVM_DEBUG(dbgs() << "TailDuplicator deleted block: "
                           << getBlockName(RemBB) << "\n");
       };
   auto RemovalCallbackRef =
       function_ref<void(MachineBasicBlock*)>(RemovalCallback);
 
   SmallVector<MachineBasicBlock *, 8> DuplicatedPreds;
   bool IsSimple = TailDup.isSimpleBB(BB);
-  TailDup.tailDuplicateAndUpdate(IsSimple, BB, LPred,
-                                 &DuplicatedPreds, &RemovalCallbackRef);
+  SmallVector<MachineBasicBlock *, 8> CandidatePreds;
+  SmallVectorImpl<MachineBasicBlock *> *CandidatePtr = nullptr;
+  if (F->getFunction().hasProfileData()) {
+    // We can do partial duplication with precise profile information.
+    findDuplicateCandidates(CandidatePreds, BB, BlockFilter);
+    if (CandidatePreds.size() == 0)
+      return false;
+    if (CandidatePreds.size() < BB->pred_size())
+      CandidatePtr = &CandidatePreds;
+  }
+  TailDup.tailDuplicateAndUpdate(IsSimple, BB, LPred, &DuplicatedPreds,
+                                 &RemovalCallbackRef, CandidatePtr);
 
   // Update UnscheduledPredecessors to reflect tail-duplication.
   DuplicatedToLPred = false;
   for (MachineBasicBlock *Pred : DuplicatedPreds) {
     // We're only looking for unscheduled predecessors that match the filter.
     BlockChain* PredChain = BlockToChain[Pred];
     if (Pred == LPred)
       DuplicatedToLPred = true;
     if (Pred == LPred || (BlockFilter && !BlockFilter->count(Pred))
         || PredChain == &Chain)
       continue;
     for (MachineBasicBlock *NewSucc : Pred->successors()) {
       if (BlockFilter && !BlockFilter->count(NewSucc))
         continue;
       BlockChain *NewChain = BlockToChain[NewSucc];
       if (NewChain != &Chain && NewChain != PredChain)
         NewChain->UnscheduledPredecessors++;
     }
   }
   return Removed;
 }
 
+// Count the number of actual machine instructions.
+static uint64_t countMBBInstruction(MachineBasicBlock *MBB) {
+  uint64_t InstrCount = 0;
+  for (MachineInstr &MI : *MBB) {
+    if (!MI.isPHI() && !MI.isMetaInstruction())
+      InstrCount += 1;
+  }
+  return InstrCount;
+}
+
+// The size cost of duplication is the instruction size of the duplicated block.
+// So we should scale the threshold accordingly. But the instruction size is not
+// available on all targets, so we use the number of instructions instead.
+BlockFrequency MachineBlockPlacement::scaleThreshold(MachineBasicBlock *BB) {
+  return DupThreshold.getFrequency() * countMBBInstruction(BB);
+}
+
+// Returns true if BB is Pred's best successor.
+bool MachineBlockPlacement::isBestSuccessor(MachineBasicBlock *BB,
+                                            MachineBasicBlock *Pred,
+                                            BlockFilterSet *BlockFilter) {
+  if (BB == Pred)
+    return false;
+  if (BlockFilter && !BlockFilter->count(Pred))
+    return false;
+  BlockChain *PredChain = BlockToChain[Pred];
+  if (PredChain && (Pred != *std::prev(PredChain->end())))
+    return false;
+
+  // Find the successor with largest probability excluding BB.
+  BranchProbability BestProb = BranchProbability::getZero();
+  for (MachineBasicBlock *Succ : Pred->successors())
+    if (Succ != BB) {
+      if (BlockFilter && !BlockFilter->count(Succ))
+        continue;
+      BlockChain *SuccChain = BlockToChain[Succ];
+      if (SuccChain && (Succ != *SuccChain->begin()))
+        continue;
+      BranchProbability SuccProb = MBPI->getEdgeProbability(Pred, Succ);
+      if (SuccProb > BestProb)
+        BestProb = SuccProb;
+    }
+
+  BranchProbability BBProb = MBPI->getEdgeProbability(Pred, BB);
+  if (BBProb <= BestProb)
+    return false;
+
+  // Compute the number of reduced taken branches if Pred falls through to BB
+  // instead of another successor. Then compare it with threshold.
+  BlockFrequency PredFreq = MBFI->getBlockFreq(Pred);
+  BlockFrequency Gain = PredFreq * (BBProb - BestProb);
+  return Gain > scaleThreshold(BB);
+}
+
+// Find out the predecessors of BB and BB can be beneficially duplicated into
+// them.
+void MachineBlockPlacement::findDuplicateCandidates(
+    SmallVectorImpl<MachineBasicBlock *> &Candidates,
+    MachineBasicBlock *BB,
+    BlockFilterSet *BlockFilter) {
+  MachineBasicBlock *Fallthrough = nullptr;
+  BranchProbability DefaultBranchProb = BranchProbability::getZero();
+  BlockFrequency BBDupThreshold(scaleThreshold(BB));
+  SmallVector<MachineBasicBlock *, 8> Preds(BB->pred_begin(), BB->pred_end());
+  SmallVector<MachineBasicBlock *, 8> Succs(BB->succ_begin(), BB->succ_end());
+
+  // Sort for highest frequency.
+  auto CmpSucc = [&](MachineBasicBlock *A, MachineBasicBlock *B) {
+    return MBPI->getEdgeProbability(BB, A) > MBPI->getEdgeProbability(BB, B);
+  };
+  auto CmpPred = [&](MachineBasicBlock *A, MachineBasicBlock *B) {
+    return MBFI->getBlockFreq(A) > MBFI->getBlockFreq(B);
+  };
+  llvm::stable_sort(Succs, CmpSucc);
+  llvm::stable_sort(Preds, CmpPred);
+
+  auto SuccIt = Succs.begin();
+  if (SuccIt != Succs.end()) {
+    DefaultBranchProb = MBPI->getEdgeProbability(BB, *SuccIt).getCompl();
+  }
+
+  // For each predecessors of BB, compute the benefit of duplicating BB,
+  // if it is larger than the threshold, add it into Candidates.
+  //
+  // If we have following control flow.
+  //
+  //     PB1 PB2 PB3 PB4
+  //      \   |  /    /\
+  //       \  | /    /  \
+  //        \ |/    /    \
+  //         BB----/     OB
+  //         /\
+  //        /  \
+  //      SB1 SB2
+  //
+  // And it can be partially duplicated as
+  //
+  //   PB2+BB
+  //      |  PB1 PB3 PB4
+  //      |   |  /    /\
+  //      |   | /    /  \
+  //      |   |/    /    \
+  //      |  BB----/     OB
+  //      |\ /|
+  //      | X |
+  //      |/ \|
+  //     SB2 SB1
+  //
+  // The benefit of duplicating into a predecessor is defined as
+  //         Orig_taken_branch - Duplicated_taken_branch
+  //
+  // The Orig_taken_branch is computed with the assumption that predecessor
+  // jumps to BB and the most possible successor is laid out after BB.
+  //
+  // The Duplicated_taken_branch is computed with the assumption that BB is
+  // duplicated into PB, and one successor is layout after it (SB1 for PB1 and
+  // SB2 for PB2 in our case). If there is no available successor, the combined
+  // block jumps to all BB's successor, like PB3 in this example.
+  //
+  // If a predecessor has multiple successors, so BB can't be duplicated into
+  // it. But it can beneficially fall through to BB, and duplicate BB into other
+  // predecessors.
+  for (MachineBasicBlock *Pred : Preds) {
+    BlockFrequency PredFreq = MBFI->getBlockFreq(Pred);
+
+    if (!TailDup.canTailDuplicate(BB, Pred)) {
+      // BB can't be duplicated into Pred, but it is possible to be layout
+      // below Pred.
+      if (!Fallthrough && isBestSuccessor(BB, Pred, BlockFilter)) {
+        Fallthrough = Pred;
+        if (SuccIt != Succs.end())
+          SuccIt++;
+      }
+      continue;
+    }
+
+    BlockFrequency OrigCost = PredFreq + PredFreq * DefaultBranchProb;
+    BlockFrequency DupCost;
+    if (SuccIt == Succs.end()) {
+      // Jump to all successors;
+      if (Succs.size() > 0)
+        DupCost += PredFreq;
+    } else {
+      // Fallthrough to *SuccIt, jump to all other successors;
+      DupCost += PredFreq;
+      DupCost -= PredFreq * MBPI->getEdgeProbability(BB, *SuccIt);
+    }
+
+    assert(OrigCost >= DupCost);
+    OrigCost -= DupCost;
+    if (OrigCost > BBDupThreshold) {
+      Candidates.push_back(Pred);
+      if (SuccIt != Succs.end())
+        SuccIt++;
+    }
+  }
+
+  // No predecessors can optimally fallthrough to BB.
+  // So we can change one duplication into fallthrough.
+  if (!Fallthrough) {
+    if ((Candidates.size() < Preds.size()) && (Candidates.size() > 0)) {
+      Candidates[0] = Candidates.back();
+      Candidates.pop_back();
+    }
+  }
+}
+
+void MachineBlockPlacement::initDupThreshold() {
+  DupThreshold = 0;
+  if (!F->getFunction().hasProfileData())
+    return;
+
+  BlockFrequency MaxFreq = 0;
+  for (MachineBasicBlock &MBB : *F) {
+    BlockFrequency Freq = MBFI->getBlockFreq(&MBB);
+    if (Freq > MaxFreq)
+      MaxFreq = Freq;
+  }
+
+  // FIXME: we may use profile count instead of frequency,
+  // and need more fine tuning.
+  BranchProbability ThresholdProb(TailDupPlacementPenalty, 100);
+  DupThreshold = MaxFreq * ThresholdProb;
+}
+
 bool MachineBlockPlacement::runOnMachineFunction(MachineFunction &MF) {
   if (skipFunction(MF.getFunction()))
     return false;
 
   // Check for single-block functions and skip them.
   if (std::next(MF.begin()) == MF.end())
     return false;
 
   F = &MF;
   MBPI = &getAnalysis<MachineBranchProbabilityInfo>();
   MBFI = std::make_unique<MBFIWrapper>(
       getAnalysis<MachineBlockFrequencyInfo>());
   MLI = &getAnalysis<MachineLoopInfo>();
   TII = MF.getSubtarget().getInstrInfo();
   TLI = MF.getSubtarget().getTargetLowering();
   MPDT = nullptr;
   PSI = &getAnalysis<ProfileSummaryInfoWrapperPass>().getPSI();
 
+  initDupThreshold();
+
   // Initialize PreferredLoopExit to nullptr here since it may never be set if
   // there are no MachineLoops.
   PreferredLoopExit = nullptr;
 
   assert(BlockToChain.empty() &&
          "BlockToChain map should be empty before starting placement.");
   assert(ComputedEdges.empty() &&
          "Computed Edge map should be empty before starting placement.");

llvm/lib/CodeGen/TailDuplicator.cpp

 ///     predecessor, instead of using the ordering in MF
 /// \p DuplicatedPreds - if non-null, \p DuplicatedPreds will contain a list of
 ///     all Preds that received a copy of \p MBB.
 /// \p RemovalCallback - if non-null, called just before MBB is deleted.
 bool TailDuplicator::tailDuplicateAndUpdate(
     bool IsSimple, MachineBasicBlock *MBB,
     MachineBasicBlock *ForcedLayoutPred,
     SmallVectorImpl<MachineBasicBlock*> *DuplicatedPreds,
-    function_ref<void(MachineBasicBlock *)> *RemovalCallback) {
+    function_ref<void(MachineBasicBlock *)> *RemovalCallback,
+    SmallVectorImpl<MachineBasicBlock *> *CandidatePtr) {
   // Save the successors list.
   SmallSetVector<MachineBasicBlock *, 8> Succs(MBB->succ_begin(),
                                                MBB->succ_end());
 
   SmallVector<MachineBasicBlock *, 8> TDBBs;
   SmallVector<MachineInstr *, 16> Copies;
-  if (!tailDuplicate(IsSimple, MBB, ForcedLayoutPred, TDBBs, Copies))
+  if (!tailDuplicate(IsSimple, MBB, ForcedLayoutPred,
+                     TDBBs, Copies, CandidatePtr))
     return false;
 
   ++NumTails;
 
   SmallVector<MachineInstr *, 8> NewPHIs;
   MachineSSAUpdater SSAUpdate(*MF, &NewPHIs);
 
   // TailBB's immediate successors are now successors of those predecessors

 /// \p TailBB   Block to be duplicated.
 /// \p ForcedLayoutPred  When non-null, use this block as the layout predecessor
 ///                      instead of the previous block in MF's order.
 /// \p TDBBs             A vector to keep track of all blocks tail-duplicated
 ///                      into.
 /// \p Copies            A vector of copy instructions inserted. Used later to
 ///                      walk all the inserted copies and remove redundant ones.
 bool TailDuplicator::tailDuplicate(bool IsSimple, MachineBasicBlock *TailBB,
-                                   MachineBasicBlock *ForcedLayoutPred,
-                                   SmallVectorImpl<MachineBasicBlock *> &TDBBs,
-                                   SmallVectorImpl<MachineInstr *> &Copies) {
+                          SmallVectorImpl<MachineInstr *> &Copies,
+                          SmallVectorImpl<MachineBasicBlock *> *CandidatePtr) {
   LLVM_DEBUG(dbgs() << "\n*** Tail-duplicating " << printMBBReference(*TailBB)
                     << '\n');
 
   DenseSet<unsigned> UsedByPhi;
   getRegsUsedByPHIs(*TailBB, &UsedByPhi);
 
   if (IsSimple)
     return duplicateSimpleBB(TailBB, TDBBs, UsedByPhi, Copies);
 
   // Iterate through all the unique predecessors and tail-duplicate this
   // block into them, if possible. Copying the list ahead of time also
   // avoids trouble with the predecessor list reallocating.
   bool Changed = false;
-  SmallSetVector<MachineBasicBlock *, 8> Preds(TailBB->pred_begin(),
-                                               TailBB->pred_end());
+  SmallSetVector<MachineBasicBlock *, 8> Preds;
+  if (CandidatePtr)
+    Preds.insert(CandidatePtr->begin(), CandidatePtr->end());
+  else
+    Preds.insert(TailBB->pred_begin(), TailBB->pred_end());
+
   for (MachineBasicBlock *PredBB : Preds) {
     assert(TailBB != PredBB &&
            "Single-block loop should have been rejected earlier!");
 
     if (!canTailDuplicate(TailBB, PredBB))
       continue;
 
     // Don't duplicate into a fall-through predecessor (at least for now).
+    // If profile is available, findDuplicateCandidates can choose better
+    // fall-through predecessor.
+    if (!(MF->getFunction().hasProfileData() && LayoutMode)) {
       bool IsLayoutSuccessor = false;
       if (ForcedLayoutPred)
         IsLayoutSuccessor = (ForcedLayoutPred == PredBB);
       else if (PredBB->isLayoutSuccessor(TailBB) && PredBB->canFallThrough())
         IsLayoutSuccessor = true;
       if (IsLayoutSuccessor)
         continue;
+    }
 
     LLVM_DEBUG(dbgs() << "\nTail-duplicating into PredBB: " << *PredBB
                       << "From Succ: " << *TailBB);
 
     TDBBs.push_back(PredBB);
 
     // Remove PredBB's unconditional branch.
     TII->removeBranch(*PredBB);

llvm/test/CodeGen/X86/partial-tail-dup.ll

+; RUN: llc < %s -mtriple=x86_64-unknown-linux-gnu | FileCheck %s
+
+
+@gvar = external global i32
+
+; dupbb has two predecessors, p1 and p2. p1 is hot, p2 is cold. So dupbb
+; should be placed after p1, and not duplicated into p2.
+;
+; CHECK-LABEL: test1
+; CHECK:       %p1
+; CHECK:       .LBB0_4: # %dupbb
+; CHECK:       %p2
+; CHECK:       jmp .LBB0_4
+
+define void @test1(i32* %p) !prof !1 {
+entry:
+  br label %header
+
+header:
+  %call = call zeroext i1 @a()
+  br i1 %call, label %p1, label %p2, !prof !2
+
+p1:
+  call void @b()
+  br label %dupbb
+
+p2:
+  call void @c()
+  br label %dupbb
+
+dupbb:
+  %cond = icmp eq i32* @gvar, %p
+  br i1 %cond, label %header, label %latch, !prof !3
+
+latch:
+  %call3 = call zeroext i1 @a()
+  br i1 %call3, label %header, label %end, !prof !2
+
+end:
+  ret void
+}
+
+
+; dupbb has four predecessors p1, p2, p3 and p4. p1 and p2 are hot, p3 and  p4
+; are cold. So dupbb should be placed after p1, duplicated into p2. p3 and p4
+; should jump to dupbb.
+;
+; CHECK-LABEL: test2
+; CHECK:       %p1
+; CHECK:       .LBB1_8: # %dupbb
+;
+; CHECK:       %p2
+; CHECK:       callq c
+; CHECK-NEXT:  cmpq
+; CHECK-NEXT:  je
+; CHECK-NEXT:  jmp
+;
+; CHECK:       %p3
+; CHECK:       jmp .LBB1_8
+; CHECK:       %p4
+; CHECK:       jmp .LBB1_8
+
+define void @test2(i32* %p) !prof !1 {
+entry:
+  br label %header
+
+header:
+  %call = call zeroext i1 @a()
+  br i1 %call, label %bb1, label %bb2, !prof !2
+
+bb1:
+  %call1 = call zeroext i1 @a()
+  br i1 %call1, label %p1, label %p2, !prof !4
+
+bb2:
+  %call2 = call zeroext i1 @a()
+  br i1 %call2, label %p3, label %p4, !prof !4
+
+p1:
+  call void @b()
+  br label %dupbb
+
+p2:
+  call void @c()
+  br label %dupbb
+
+p3:
+  call void @d()
+  br label %dupbb
+
+p4:
+  call void @e()
+  br label %dupbb
+
+dupbb:
+  %cond = icmp eq i32* @gvar, %p
+  br i1 %cond, label %bb3, label %bb4, !prof !4
+
+bb3:
+  call void @b()
+  br label %bb4
+
+bb4:
+  %call4 = call zeroext i1 @a()
+  br i1 %call4, label %header, label %latch, !prof !3
+
+latch:
+  %call3 = call zeroext i1 @a()
+  br i1 %call3, label %header, label %end, !prof !2
+
+end:
+  ret void
+}
+
+
+; dupbb has three predecessors p1, p2 and p3. p3 has two successors, so dupbb
+; can't be duplicated into p3, but it should not block it to be duplicated into
+; other predecessors.
+;
+; CHECK-LABEL: test3
+; CHECK:       %p1
+; CHECK:       .LBB2_6: # %dupbb
+;
+; CHECK:       %p2
+; CHECK:       callq c
+; CHECK:       cmpq
+; CHECK-NEXT:  je
+; CHECK-NEXT:  jmp
+;
+; CHECK:       %p3
+; CHECK:       jne .LBB2_6
+
+define void @test3(i32* %p) !prof !1 {
+entry:
+  br label %header
+
+header:
+  %call = call zeroext i1 @a()
+  br i1 %call, label %bb1, label %p3, !prof !2
+
+bb1:
+  %call1 = call zeroext i1 @a()
+  br i1 %call1, label %p1, label %p2, !prof !4
+
+p1:
+  call void @b()
+  br label %dupbb
+
+p2:
+  call void @c()
+  br label %dupbb
+
+p3:
+  %call2 = call zeroext i1 @a()
+  br i1 %call2, label %dupbb, label %bb4, !prof !4
+
+dupbb:
+  %cond = icmp eq i32* @gvar, %p
+  br i1 %cond, label %bb3, label %bb4, !prof !4
+
+bb3:
+  call void @b()
+  br label %bb4
+
+bb4:
+  %call4 = call zeroext i1 @a()
+  br i1 %call4, label %header, label %latch, !prof !3
+
+latch:
+  %call3 = call zeroext i1 @a()
+  br i1 %call3, label %header, label %end, !prof !2
+
+end:
+  ret void
+}
+
+declare zeroext i1 @a()
+declare void @b()
+declare void @c()
+declare void @d()
+declare void @e()
+declare void @f()
+
+!1 = !{!"function_entry_count", i64 1000}
+!2 = !{!"branch_weights", i32 100, i32 1}
+!3 = !{!"branch_weights", i32 1, i32 100}
+!4 = !{!"branch_weights", i32 60, i32 40}