extending code layout alg

[spupyrev]

Summary

The diff modifies ext-tsp code layout algorithm in the following ways:
(i) fixes merging of cold block chains (this is a port of D129397);
(ii) adjusts the cost model utilized for optimization;
(iii) adjusts some APIs so that the implementation can be used in BOLT; this is
a prerequisite for D129895.

The only non-trivial change is (ii). Here we introduce different weights for
conditional and unconditional branches in the cost model. Based on the new model
it is slightly more important to increase the number of "fall-through
unconditional" jumps, which makes sense, as placing two blocks with an
unconditional jump next to each other reduces the number of jump instructions in
the generated code. Experimentally, this makes a mild impact on the performance;
I've seen up to 0.2%-0.3% perf win on some benchmarks.

[hoy]

llvm/lib/CodeGen/MachineBlockPlacement.cpp

3491

nit: use EdgeT here instead of std::pair<uint64_t, uint64_t>?

Maybe also define a type, e.g, EdgeCountT for std::pair<std::pair<uint64_t, uint64_t>, uint64_t> ?

llvm/lib/Transforms/Utils/CodeLayout.cpp

133

nit: return jumpExtTSPScore(0, 1, Count, IsConditional ? FallthroughWeightCond : FallthroughWeightUncond);

407–410

Why do you need 'const` here?

[spupyrev]

review

[hoy]

LGTM, thanks.

[wenlei]

lgtm with a few nits.

llvm/lib/Transforms/Utils/CodeLayout.cpp

514

nit: SuccNodes.resize(NumNodes) should just work, so you don't invoke copy/move constructor here. Same for PredNodes.

516

nit: std::vector<uint64_t> OutDegree(NumNodes, 0); this is avoids copy/move constructor from the temp object.

Same for the other auto OutDegree and other instances of auto xx = std::vector<..>.

[spupyrev]

more suggested changes + rebase

llvm/include/llvm/Transforms/Utils/CodeLayout.h

 #define LLVM_TRANSFORMS_UTILS_CODELAYOUT_H
 
 #include "llvm/ADT/DenseMap.h"
 
 #include <vector>
 
 namespace llvm {
 
+using EdgeT = std::pair<uint64_t, uint64_t>;
+using EdgeCountT = std::pair<EdgeT, uint64_t>;
+
 /// Find a layout of nodes (basic blocks) of a given CFG optimizing jump
 /// locality and thus processor I-cache utilization. This is achieved via
 /// increasing the number of fall-through jumps and co-locating frequently
 /// executed nodes together.
 /// The nodes are assumed to be indexed by integers from [0, |V|) so that the
 /// current order is the identity permutation.
 /// \p NodeSizes: The sizes of the nodes (in bytes).
 /// \p NodeCounts: The execution counts of the nodes in the profile.
 /// \p EdgeCounts: The execution counts of every edge (jump) in the profile. The
 ///    map also defines the edges in CFG and should include 0-count edges.
 /// \returns The best block order found.
-std::vector<uint64_t> applyExtTspLayout(
-    const std::vector<uint64_t> &NodeSizes,
+std::vector<uint64_t>
+applyExtTspLayout(const std::vector<uint64_t> &NodeSizes,
                   const std::vector<uint64_t> &NodeCounts,
-    const DenseMap<std::pair<uint64_t, uint64_t>, uint64_t> &EdgeCounts);
+                  const std::vector<EdgeCountT> &EdgeCounts);
 
 /// Estimate the "quality" of a given node order in CFG. The higher the score,
 /// the better the order is. The score is designed to reflect the locality of
 /// the given order, which is anti-correlated with the number of I-cache misses
 /// in a typical execution of the function.
-double calcExtTspScore(
-    const std::vector<uint64_t> &Order, const std::vector<uint64_t> &NodeSizes,
+double calcExtTspScore(const std::vector<uint64_t> &Order,
+                       const std::vector<uint64_t> &NodeSizes,
                        const std::vector<uint64_t> &NodeCounts,
-    const DenseMap<std::pair<uint64_t, uint64_t>, uint64_t> &EdgeCounts);
+                       const std::vector<EdgeCountT> &EdgeCounts);
 
 /// Estimate the "quality" of the current node order in CFG.
-double calcExtTspScore(
-    const std::vector<uint64_t> &NodeSizes,
+double calcExtTspScore(const std::vector<uint64_t> &NodeSizes,
                        const std::vector<uint64_t> &NodeCounts,
-    const DenseMap<std::pair<uint64_t, uint64_t>, uint64_t> &EdgeCounts);
+                       const std::vector<EdgeCountT> &EdgeCounts);
 
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_UTILS_CODELAYOUT_H

llvm/lib/CodeGen/MachineBlockPlacement.cpp

void MachineBlockPlacement::applyExtTsp() {
   size_t NumBlocks = 0;
   for (const MachineBasicBlock &MBB : *F) {
     BlockIndex[&MBB] = NumBlocks++;
     CurrentBlockOrder.push_back(&MBB);
   }
 
   auto BlockSizes = std::vector<uint64_t>(F->size());
   auto BlockCounts = std::vector<uint64_t>(F->size());
-  DenseMap<std::pair<uint64_t, uint64_t>, uint64_t> JumpCounts;
+  std::vector<EdgeCountT> JumpCounts;
   for (MachineBasicBlock &MBB : *F) {
     // Getting the block frequency.
     BlockFrequency BlockFreq = MBFI->getBlockFreq(&MBB);
     BlockCounts[BlockIndex[&MBB]] = BlockFreq.getFrequency();
     // Getting the block size:
     // - approximate the size of an instruction by 4 bytes, and
     // - ignore debug instructions.
     // Note: getting the exact size of each block is target-dependent and can be
     // done by extending the interface of MCCodeEmitter. Experimentally we do
     // not see a perf improvement with the exact block sizes.
     auto NonDbgInsts =
         instructionsWithoutDebug(MBB.instr_begin(), MBB.instr_end());
     int NumInsts = std::distance(NonDbgInsts.begin(), NonDbgInsts.end());
     BlockSizes[BlockIndex[&MBB]] = 4 * NumInsts;
     // Getting jump frequencies.
     for (MachineBasicBlock *Succ : MBB.successors()) {
       auto EP = MBPI->getEdgeProbability(&MBB, Succ);
-      BlockFrequency EdgeFreq = BlockFreq * EP;
-      auto Edge = std::make_pair(BlockIndex[&MBB], BlockIndex[Succ]);
-      JumpCounts[Edge] = EdgeFreq.getFrequency();
+      BlockFrequency JumpFreq = BlockFreq * EP;
+      auto Jump = std::make_pair(BlockIndex[&MBB], BlockIndex[Succ]);
+      JumpCounts.push_back(std::make_pair(Jump, JumpFreq.getFrequency()));
     }
   }
 
   LLVM_DEBUG(dbgs() << "Applying ext-tsp layout for |V| = " << F->size()
                     << " with profile = " << F->getFunction().hasProfileData()
                     << " (" << F->getName().str() << ")"
                     << "\n");
   LLVM_DEBUG(

llvm/lib/Transforms/Utils/CodeLayout.cpp

 // consider all possible ways of gluing the three chains (e.g., X1YX2, X1X2Y,
 // X2X1Y, X2YX1, YX1X2, YX2X1) and choose the one producing the largest score.
 // This improves the quality of the final result (the search space is larger)
 // while keeping the implementation sufficiently fast.
 //
 // Reference:
 //   * A. Newell and S. Pupyrev, Improved Basic Block Reordering,
 //     IEEE Transactions on Computers, 2020
+//     https://arxiv.org/abs/1809.04676
 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Transforms/Utils/CodeLayout.h"
 #include "llvm/Support/CommandLine.h"
 
 using namespace llvm;
 #define DEBUG_TYPE "code-layout"
 
 cl::opt<bool> EnableExtTspBlockPlacement(
     "enable-ext-tsp-block-placement", cl::Hidden, cl::init(false),
     cl::desc("Enable machine block placement based on the ext-tsp model, "
              "optimizing I-cache utilization."));
 
 cl::opt<bool> ApplyExtTspWithoutProfile(
     "ext-tsp-apply-without-profile",
     cl::desc("Whether to apply ext-tsp placement for instances w/o profile"),
     cl::init(true), cl::Hidden);
 
-// Algorithm-specific constants. The values are tuned for the best performance
+// Algorithm-specific params. The values are tuned for the best performance
 // of large-scale front-end bound binaries.
-static cl::opt<double>
-    ForwardWeight("ext-tsp-forward-weight", cl::Hidden, cl::init(0.1),
-                  cl::desc("The weight of forward jumps for ExtTSP value"));
-
-static cl::opt<double>
-    BackwardWeight("ext-tsp-backward-weight", cl::Hidden, cl::init(0.1),
-                   cl::desc("The weight of backward jumps for ExtTSP value"));
+static cl::opt<double> ForwardWeightCond(
+    "ext-tsp-forward-weight-cond", cl::ReallyHidden, cl::init(0.1),
+    cl::desc("The weight of conditional forward jumps for ExtTSP value"));
+
+static cl::opt<double> ForwardWeightUncond(
+    "ext-tsp-forward-weight-uncond", cl::ReallyHidden, cl::init(0.1),
+    cl::desc("The weight of unconditional forward jumps for ExtTSP value"));
+
+static cl::opt<double> BackwardWeightCond(
+    "ext-tsp-backward-weight-cond", cl::ReallyHidden, cl::init(0.1),
+    cl::desc("The weight of conditonal backward jumps for ExtTSP value"));
+
+static cl::opt<double> BackwardWeightUncond(
+    "ext-tsp-backward-weight-uncond", cl::ReallyHidden, cl::init(0.1),
+    cl::desc("The weight of unconditonal backward jumps for ExtTSP value"));
+
+static cl::opt<double> FallthroughWeightCond(
+    "ext-tsp-fallthrough-weight-cond", cl::ReallyHidden, cl::init(1.0),
+    cl::desc("The weight of conditional fallthrough jumps for ExtTSP value"));
+
+static cl::opt<double> FallthroughWeightUncond(
+    "ext-tsp-fallthrough-weight-uncond", cl::ReallyHidden, cl::init(1.05),
+    cl::desc("The weight of unconditional fallthrough jumps for ExtTSP value"));
 
 static cl::opt<unsigned> ForwardDistance(
-    "ext-tsp-forward-distance", cl::Hidden, cl::init(1024),
+    "ext-tsp-forward-distance", cl::ReallyHidden, cl::init(1024),
     cl::desc("The maximum distance (in bytes) of a forward jump for ExtTSP"));
 
 static cl::opt<unsigned> BackwardDistance(
-    "ext-tsp-backward-distance", cl::Hidden, cl::init(640),
+    "ext-tsp-backward-distance", cl::ReallyHidden, cl::init(640),
     cl::desc("The maximum distance (in bytes) of a backward jump for ExtTSP"));
 
 // The maximum size of a chain created by the algorithm. The size is bounded
 // so that the algorithm can efficiently process extremely large instance.
 static cl::opt<unsigned>
-    MaxChainSize("ext-tsp-max-chain-size", cl::Hidden, cl::init(4096),
+    MaxChainSize("ext-tsp-max-chain-size", cl::ReallyHidden, cl::init(4096),
                  cl::desc("The maximum size of a chain to create."));
 
 // The maximum size of a chain for splitting. Larger values of the threshold
 // may yield better quality at the cost of worsen run-time.
 static cl::opt<unsigned> ChainSplitThreshold(
-    "ext-tsp-chain-split-threshold", cl::Hidden, cl::init(128),
+    "ext-tsp-chain-split-threshold", cl::ReallyHidden, cl::init(128),
     cl::desc("The maximum size of a chain to apply splitting"));
 
 // The option enables splitting (large) chains along in-coming and out-going
 // jumps. This typically results in a better quality.
 static cl::opt<bool> EnableChainSplitAlongJumps(
-    "ext-tsp-enable-chain-split-along-jumps", cl::Hidden, cl::init(true),
+    "ext-tsp-enable-chain-split-along-jumps", cl::ReallyHidden, cl::init(true),
     cl::desc("The maximum size of a chain to apply splitting"));
 
 namespace {
 
 // Epsilon for comparison of doubles.
 constexpr double EPS = 1e-8;
 
+// Compute the Ext-TSP score for a given jump.
+double jumpExtTSPScore(uint64_t JumpDist, uint64_t JumpMaxDist, uint64_t Count,
+                       double Weight) {
+  if (JumpDist > JumpMaxDist)
+    return 0;
+  double Prob = 1.0 - static_cast<double>(JumpDist) / JumpMaxDist;
+  return Weight * Prob * Count;
+}
+
 // Compute the Ext-TSP score for a jump between a given pair of blocks,
 // using their sizes, (estimated) addresses and the jump execution count.
 double extTSPScore(uint64_t SrcAddr, uint64_t SrcSize, uint64_t DstAddr,
-                   uint64_t Count) {
+                   uint64_t Count, bool IsConditional) {
   // Fallthrough
   if (SrcAddr + SrcSize == DstAddr) {
-    // Assume that FallthroughWeight = 1.0 after normalization
-    return static_cast<double>(Count);
+    return jumpExtTSPScore(0, 1, Count,
+                           IsConditional ? FallthroughWeightCond
+                                         : FallthroughWeightUncond);
   }
   // Forward
   if (SrcAddr + SrcSize < DstAddr) {
-    const auto Dist = DstAddr - (SrcAddr + SrcSize);
-    if (Dist <= ForwardDistance) {
-      double Prob = 1.0 - static_cast<double>(Dist) / ForwardDistance;
-      return ForwardWeight * Prob * Count;
-    }
-    return 0;
+    const uint64_t Dist = DstAddr - (SrcAddr + SrcSize);
+    return jumpExtTSPScore(Dist, ForwardDistance, Count,
+                           IsConditional ? ForwardWeightCond
+                                         : ForwardWeightUncond);
   }
   // Backward
-  const auto Dist = SrcAddr + SrcSize - DstAddr;
-  if (Dist <= BackwardDistance) {
-    double Prob = 1.0 - static_cast<double>(Dist) / BackwardDistance;
-    return BackwardWeight * Prob * Count;
-  }
-  return 0;
+  const uint64_t Dist = SrcAddr + SrcSize - DstAddr;
+  return jumpExtTSPScore(Dist, BackwardDistance, Count,
+                         IsConditional ? BackwardWeightCond
+                                       : BackwardWeightUncond);
 }
 
 /// A type of merging two chains, X and Y. The former chain is split into
 /// X1 and X2 and then concatenated with Y in the order specified by the type.
 enum class MergeTypeTy : int { X_Y, X1_Y_X2, Y_X2_X1, X2_X1_Y };
 
 /// The gain of merging two chains, that is, the Ext-TSP score of the merge
 /// together with the corresponfiding merge 'type' and 'offset'.
public:
   // Forced predecessor of the block in CFG.
   Block *ForcedPred{nullptr};
   // Outgoing jumps from the block.
   std::vector<Jump *> OutJumps;
   // Incoming jumps to the block.
   std::vector<Jump *> InJumps;
 
 public:
-  explicit Block(size_t Index, uint64_t Size_, uint64_t EC)
-      : Index(Index), Size(Size_), ExecutionCount(EC) {}
+  explicit Block(size_t Index, uint64_t Size, uint64_t EC)
+      : Index(Index), Size(Size), ExecutionCount(EC) {}
   bool isEntry() const { return Index == 0; }
 };
 
 /// An arc in the graph, typically corresponding to a jump between two blocks.
 class Jump {
 public:
   Jump(const Jump &) = delete;
   Jump(Jump &&) = default;
   Jump &operator=(const Jump &) = delete;
   Jump &operator=(Jump &&) = default;
 
   // Source block of the jump.
   Block *Source;
   // Target block of the jump.
   Block *Target;
   // Execution count of the arc in the profile data.
   uint64_t ExecutionCount{0};
+  // Whether the jump corresponds to a conditional branch.
+  bool IsConditional{false};
 
 public:
   explicit Jump(Block *Source, Block *Target, uint64_t ExecutionCount)
       : Source(Source), Target(Target), ExecutionCount(ExecutionCount) {}
 };
 
 /// A chain (ordered sequence) of blocks.
 class Chain {
 public:
   Chain(const Chain &) = delete;
   Chain(Chain &&) = default;
   Chain &operator=(const Chain &) = delete;
   Chain &operator=(Chain &&) = default;
 
   explicit Chain(uint64_t Id, Block *Block)
       : Id(Id), Score(0), Blocks(1, Block) {}
 
   uint64_t id() const { return Id; }
 
   bool isEntry() const { return Blocks[0]->Index == 0; }
 
+  bool isCold() const {
+    for (auto Block : Blocks) {
+      if (Block->ExecutionCount > 0)
+        return false;
+    }
+    return true;
+  }
+
   double score() const { return Score; }
 
   void setScore(double NewScore) { Score = NewScore; }
 
   const std::vector<Block *> &blocks() const { return Blocks; }
 
   size_t numBlocks() const { return Blocks.size(); }
 
private:
   bool CacheValidBackward{false};
 };
 
 void Chain::mergeEdges(Chain *Other) {
   assert(this != Other && "cannot merge a chain with itself");
 
   // Update edges adjacent to chain Other
   for (auto EdgeIt : Other->Edges) {
-    const auto DstChain = EdgeIt.first;
-    const auto DstEdge = EdgeIt.second;
-    const auto TargetChain = DstChain == Other ? this : DstChain;
-    auto CurEdge = getEdge(TargetChain);
+    Chain *DstChain = EdgeIt.first;
+    ChainEdge *DstEdge = EdgeIt.second;
+    Chain *TargetChain = DstChain == Other ? this : DstChain;
+    ChainEdge *CurEdge = getEdge(TargetChain);
     if (CurEdge == nullptr) {
       DstEdge->changeEndpoint(Other, this);
       this->addEdge(TargetChain, DstEdge);
       if (DstChain != this && DstChain != Other) {
         DstChain->addEdge(this, DstEdge);
       }
     } else {
       CurEdge->moveJumps(DstEdge);
private:
   BlockIter End2;
   BlockIter Begin3;
   BlockIter End3;
 };
 
 /// The implementation of the ExtTSP algorithm.
 class ExtTSPImpl {
   using EdgeT = std::pair<uint64_t, uint64_t>;
-  using EdgeCountMap = DenseMap<EdgeT, uint64_t>;
+  using EdgeCountMap = std::vector<std::pair<EdgeT, uint64_t>>;
 
 public:
   ExtTSPImpl(size_t NumNodes, const std::vector<uint64_t> &NodeSizes,
              const std::vector<uint64_t> &NodeCounts,
              const EdgeCountMap &EdgeCounts)
       : NumNodes(NumNodes) {
     initialize(NodeSizes, NodeCounts, EdgeCounts);
   }
for (uint64_t Node = 0; Node < NumNodes; Node++) {
       uint64_t ExecutionCount = NodeCounts[Node];
       // The execution count of the entry block is set to at least 1
       if (Node == 0 && ExecutionCount == 0)
         ExecutionCount = 1;
       AllBlocks.emplace_back(Node, Size, ExecutionCount);
     }
 
     // Initialize jumps between blocks
     SuccNodes = std::vector<std::vector<uint64_t>>(NumNodes);
     PredNodes = std::vector<std::vector<uint64_t>>(NumNodes);
+    auto OutDegree = std::vector<uint64_t>(NumNodes, 0);
     AllJumps.reserve(EdgeCounts.size());
     for (auto It : EdgeCounts) {
       auto Pred = It.first.first;
       auto Succ = It.first.second;
+      OutDegree[Pred]++;
       // Ignore self-edges
       if (Pred == Succ)
         continue;
 
       SuccNodes[Pred].push_back(Succ);
       PredNodes[Succ].push_back(Pred);
       auto ExecutionCount = It.second;
       if (ExecutionCount > 0) {
         auto &Block = AllBlocks[Pred];
         auto &SuccBlock = AllBlocks[Succ];
         AllJumps.emplace_back(&Block, &SuccBlock, ExecutionCount);
         SuccBlock.InJumps.push_back(&AllJumps.back());
         Block.OutJumps.push_back(&AllJumps.back());
       }
     }
+    for (auto &Jump : AllJumps) {
+      assert(OutDegree[Jump.Source->Index] > 0);
+      Jump.IsConditional = OutDegree[Jump.Source->Index] > 1;
+    }
 
     // Initialize chains
     AllChains.reserve(NumNodes);
     HotChains.reserve(NumNodes);
-    for (auto &Block : AllBlocks) {
+    for (Block &Block : AllBlocks) {
       AllChains.emplace_back(Block.Index, &Block);
       Block.CurChain = &AllChains.back();
       if (Block.ExecutionCount > 0) {
         HotChains.push_back(&AllChains.back());
       }
     }
 
     // Initialize chain edges
     AllEdges.reserve(AllJumps.size());
-    for (auto &Block : AllBlocks) {
+    for (Block &Block : AllBlocks) {
       for (auto &Jump : Block.OutJumps) {
         auto SuccBlock = Jump->Target;
-        auto CurEdge = Block.CurChain->getEdge(SuccBlock->CurChain);
+        ChainEdge *CurEdge = Block.CurChain->getEdge(SuccBlock->CurChain);
         // this edge is already present in the graph
         if (CurEdge != nullptr) {
           assert(SuccBlock->CurChain->getEdge(Block.CurChain) != nullptr);
           CurEdge->appendJump(Jump);
           continue;
         }
         // this is a new edge
         AllEdges.emplace_back(Jump);
auto compareChainPairs = [](const Chain *A1, const Chain *B1,
       return B1->id() < B2->id();
     };
 
     while (HotChains.size() > 1) {
       Chain *BestChainPred = nullptr;
       Chain *BestChainSucc = nullptr;
       auto BestGain = MergeGainTy();
       // Iterate over all pairs of chains
-      for (auto ChainPred : HotChains) {
+      for (Chain *ChainPred : HotChains) {
         // Get candidates for merging with the current chain
         for (auto EdgeIter : ChainPred->edges()) {
-          auto ChainSucc = EdgeIter.first;
-          auto ChainEdge = EdgeIter.second;
+          Chain *ChainSucc = EdgeIter.first;
+          class ChainEdge *ChainEdge = EdgeIter.second;
           // Ignore loop edges
           if (ChainPred == ChainSucc)
             continue;
 
           // Stop early if the combined chain violates the maximum allowed size
           if (ChainPred->numBlocks() + ChainSucc->numBlocks() >= MaxChainSize)
             continue;
 
           // Compute the gain of merging the two chains
-          auto CurGain = getBestMergeGain(ChainPred, ChainSucc, ChainEdge);
+          MergeGainTy CurGain =
+              getBestMergeGain(ChainPred, ChainSucc, ChainEdge);
           if (CurGain.score() <= EPS)
             continue;
 
           if (BestGain < CurGain ||
               (std::abs(CurGain.score() - BestGain.score()) < EPS &&
                compareChainPairs(ChainPred, ChainSucc, BestChainPred,
                                  BestChainSucc))) {
             BestGain = CurGain;
             BestChainPred = ChainPred;
             BestChainSucc = ChainSucc;
           }
         }
       }
 
       // Stop merging when there is no improvement
       if (BestGain.score() <= EPS)
         break;
 
       // Merge the best pair of chains
       mergeChains(BestChainPred, BestChainSucc, BestGain.mergeOffset(),
                   BestGain.mergeType());
     }
   }
 
-  /// Merge cold blocks to reduce code size.
+  /// Merge remaining blocks into chains w/o taking jump counts into
+  /// consideration. This allows to maintain the original block order in the
+  /// absense of profile data
   void mergeColdChains() {
     for (size_t SrcBB = 0; SrcBB < NumNodes; SrcBB++) {
-      // Iterating over neighbors in the reverse order to make sure original
-      // fallthrough jumps are merged first
+      // Iterating in reverse order to make sure original fallthrough jumps are
+      // merged first; this might be beneficial for code size.
       size_t NumSuccs = SuccNodes[SrcBB].size();
       for (size_t Idx = 0; Idx < NumSuccs; Idx++) {
         auto DstBB = SuccNodes[SrcBB][NumSuccs - Idx - 1];
         auto SrcChain = AllBlocks[SrcBB].CurChain;
         auto DstChain = AllBlocks[DstBB].CurChain;
         if (SrcChain != DstChain && !DstChain->isEntry() &&
             SrcChain->blocks().back()->Index == SrcBB &&
-            DstChain->blocks().front()->Index == DstBB) {
+            DstChain->blocks().front()->Index == DstBB &&
+            SrcChain->isCold() == DstChain->isCold()) {
           mergeChains(SrcChain, DstChain, 0, MergeTypeTy::X_Y);
         }
       }
     }
   }
 
   /// Compute the Ext-TSP score for a given block order and a list of jumps.
   double extTSPScore(const MergedChain &MergedBlocks,
                      const std::vector<Jump *> &Jumps) const {
     if (Jumps.empty())
       return 0.0;
     uint64_t CurAddr = 0;
     MergedBlocks.forEach([&](const Block *BB) {
       BB->EstimatedAddr = CurAddr;
       CurAddr += BB->Size;
     });
 
     double Score = 0;
-    for (const auto &Jump : Jumps) {
-      const auto SrcBlock = Jump->Source;
-      const auto DstBlock = Jump->Target;
+    for (auto &Jump : Jumps) {
+      const Block *SrcBlock = Jump->Source;
+      const Block *DstBlock = Jump->Target;
       Score += ::extTSPScore(SrcBlock->EstimatedAddr, SrcBlock->Size,
-                             DstBlock->EstimatedAddr, Jump->ExecutionCount);
+                             DstBlock->EstimatedAddr, Jump->ExecutionCount,
+                             Jump->IsConditional);
     }
     return Score;
   }
 
   /// Compute the gain of merging two chains.
   ///
   /// The function considers all possible ways of merging two chains and
   /// computes the one having the largest increase in ExtTSP objective. The
   /// result is a pair with the first element being the gain and the second
   /// element being the corresponding merging type.
   MergeGainTy getBestMergeGain(Chain *ChainPred, Chain *ChainSucc,
                                ChainEdge *Edge) const {
     if (Edge->hasCachedMergeGain(ChainPred, ChainSucc)) {
       return Edge->getCachedMergeGain(ChainPred, ChainSucc);
     }
 
     // Precompute jumps between ChainPred and ChainSucc
     auto Jumps = Edge->jumps();
-    auto EdgePP = ChainPred->getEdge(ChainPred);
+    ChainEdge *EdgePP = ChainPred->getEdge(ChainPred);
     if (EdgePP != nullptr) {
       Jumps.insert(Jumps.end(), EdgePP->jumps().begin(), EdgePP->jumps().end());
     }
     assert(!Jumps.empty() && "trying to merge chains w/o jumps");
 
     // The object holds the best currently chosen gain of merging the two chains
     MergeGainTy Gain = MergeGainTy();
 
private:
 
   /// Merge chain From into chain Into, update the list of active chains,
   /// adjacency information, and the corresponding cached values.
   void mergeChains(Chain *Into, Chain *From, size_t MergeOffset,
                    MergeTypeTy MergeType) {
     assert(Into != From && "a chain cannot be merged with itself");
 
     // Merge the blocks
-    auto MergedBlocks =
+    MergedChain MergedBlocks =
         mergeBlocks(Into->blocks(), From->blocks(), MergeOffset, MergeType);
     Into->merge(From, MergedBlocks.getBlocks());
     Into->mergeEdges(From);
     From->clear();
 
     // Update cached ext-tsp score for the new chain
-    auto SelfEdge = Into->getEdge(Into);
+    ChainEdge *SelfEdge = Into->getEdge(Into);
     if (SelfEdge != nullptr) {
       MergedBlocks = MergedChain(Into->blocks().begin(), Into->blocks().end());
       Into->setScore(extTSPScore(MergedBlocks, SelfEdge->jumps()));
     }
 
     // Remove chain From from the list of active chains
     llvm::erase_value(HotChains, From);
 
for (auto &Chain : AllChains) {
         assert(Size > 0 && "a chain of zero size");
         ChainDensity[&Chain] = ExecutionCount / Size;
       }
     }
 
     // Sorting chains by density in the decreasing order
     std::stable_sort(SortedChains.begin(), SortedChains.end(),
                      [&](const Chain *C1, const Chain *C2) {
-                       // Makre sure the original entry block is at the
+                       // Make sure the original entry block is at the
                        // beginning of the order
                        if (C1->isEntry() != C2->isEntry()) {
                          return C1->isEntry();
                        }
 
                        const double D1 = ChainDensity[C1];
                        const double D2 = ChainDensity[C2];
                        // Compare by density and break ties by chain identifiers
                        return (D1 != D2) ? (D1 > D2) : (C1->id() < C2->id());
                      });
 
     // Collect the blocks in the order specified by their chains
     Order.reserve(NumNodes);
-    for (auto Chain : SortedChains) {
-      for (auto Block : Chain->blocks()) {
+    for (Chain *Chain : SortedChains) {
+      for (Block *Block : Chain->blocks()) {
         Order.push_back(Block->Index);
       }
     }
   }
 
 private:
   /// The number of nodes in the graph.
   const size_t NumNodes;
private:
   std::vector<Chain *> HotChains;
 };
 
 } // end of anonymous namespace
 
 std::vector<uint64_t> llvm::applyExtTspLayout(
     const std::vector<uint64_t> &NodeSizes,
     const std::vector<uint64_t> &NodeCounts,
-    const DenseMap<std::pair<uint64_t, uint64_t>, uint64_t> &EdgeCounts) {
+    const std::vector<std::pair<EdgeT, uint64_t>> &EdgeCounts) {
   size_t NumNodes = NodeSizes.size();
 
   // Verify correctness of the input data.
   assert(NodeCounts.size() == NodeSizes.size() && "Incorrect input");
   assert(NumNodes > 2 && "Incorrect input");
 
   // Apply the reordering algorithm.
   auto Alg = ExtTSPImpl(NumNodes, NodeSizes, NodeCounts, EdgeCounts);
   std::vector<uint64_t> Result;
   Alg.run(Result);
 
   // Verify correctness of the output.
   assert(Result.front() == 0 && "Original entry point is not preserved");
   assert(Result.size() == NumNodes && "Incorrect size of reordered layout");
   return Result;
 }
 
 double llvm::calcExtTspScore(
     const std::vector<uint64_t> &Order, const std::vector<uint64_t> &NodeSizes,
     const std::vector<uint64_t> &NodeCounts,
-    const DenseMap<std::pair<uint64_t, uint64_t>, uint64_t> &EdgeCounts) {
+    const std::vector<std::pair<EdgeT, uint64_t>> &EdgeCounts) {
   // Estimate addresses of the blocks in memory
   auto Addr = std::vector<uint64_t>(NodeSizes.size(), 0);
   for (size_t Idx = 1; Idx < Order.size(); Idx++) {
     Addr[Order[Idx]] = Addr[Order[Idx - 1]] + NodeSizes[Order[Idx - 1]];
   }
+  auto OutDegree = std::vector<uint64_t>(NodeSizes.size(), 0);
+  for (auto It : EdgeCounts) {
+    auto Pred = It.first.first;
+    OutDegree[Pred]++;
+  }
 
   // Increase the score for each jump
   double Score = 0;
   for (auto It : EdgeCounts) {
     auto Pred = It.first.first;
     auto Succ = It.first.second;
     uint64_t Count = It.second;
-    Score += extTSPScore(Addr[Pred], NodeSizes[Pred], Addr[Succ], Count);
+    bool IsConditional = OutDegree[Pred] > 1;
+    Score += ::extTSPScore(Addr[Pred], NodeSizes[Pred], Addr[Succ], Count,
+                           IsConditional);
   }
   return Score;
 }
 
 double llvm::calcExtTspScore(
     const std::vector<uint64_t> &NodeSizes,
     const std::vector<uint64_t> &NodeCounts,
-    const DenseMap<std::pair<uint64_t, uint64_t>, uint64_t> &EdgeCounts) {
+    const std::vector<std::pair<EdgeT, uint64_t>> &EdgeCounts) {
   auto Order = std::vector<uint64_t>(NodeSizes.size());
   for (size_t Idx = 0; Idx < NodeSizes.size(); Idx++) {
     Order[Idx] = Idx;
   }
   return calcExtTspScore(Order, NodeSizes, NodeCounts, EdgeCounts);
 }

llvm/test/CodeGen/X86/code_placement_ext_tsp_large.ll

 ;   +-------------------------------> |                                          |
 ;                                     +------------------------------------------+
 ;
 ; An expected output with a large chain-split-threshold -- the layout score is
 ; increased by ~17%
 ;
 ; CHECK-LABEL: Applying ext-tsp layout
 ; CHECK:   original  layout score: 9171074274.27
-; CHECK:   optimized layout score: 10756755324.57
+; CHECK:   optimized layout score: 10844307310.87
 ; CHECK: b0
 ; CHECK: b2
 ; CHECK: b3
 ; CHECK: b4
 ; CHECK: b5
 ; CHECK: b8
 ; CHECK: b1
 ; CHECK: b6