1 files changed, 24 insertions, 4 deletions
diff --git a/llvm/lib/Transforms/Utils/CodeLayout.cpp b/llvm/lib/Transforms/Utils/CodeLayout.cpp
index dfb9f608eab2..1ff0f148b3a9 100644
--- a/llvm/lib/Transforms/Utils/CodeLayout.cpp
+++ b/llvm/lib/Transforms/Utils/CodeLayout.cpp
@@ -40,11 +40,20 @@
 
 #include "llvm/Transforms/Utils/CodeLayout.h"
 #include "llvm/Support/CommandLine.h"
-#include "llvm/Support/Debug.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
 // of large-scale front-end bound binaries.
 static cl::opt<double>
@@ -63,6 +72,12 @@ static cl::opt<unsigned> BackwardDistance(
     "ext-tsp-backward-distance", cl::Hidden, 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),
+                 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(
@@ -115,7 +130,7 @@ enum class MergeTypeTy : int { X_Y, X1_Y_X2, Y_X2_X1, X2_X1_Y };
 /// together with the corresponfiding merge 'type' and 'offset'.
 class MergeGainTy {
 public:
-  explicit MergeGainTy() {}
+  explicit MergeGainTy() = default;
   explicit MergeGainTy(double Score, size_t MergeOffset, MergeTypeTy MergeType)
       : Score(Score), MergeOffset(MergeOffset), MergeType(MergeType) {}
 
@@ -142,7 +157,6 @@ private:
   MergeTypeTy MergeType{MergeTypeTy::X_Y};
 };
 
-class Block;
 class Jump;
 class Chain;
 class ChainEdge;
@@ -223,6 +237,8 @@ public:
 
   const std::vector<Block *> &blocks() const { return Blocks; }
 
+  size_t numBlocks() const { return Blocks.size(); }
+
   const std::vector<std::pair<Chain *, ChainEdge *>> &edges() const {
     return Edges;
   }
@@ -499,7 +515,7 @@ private:
     AllEdges.reserve(AllJumps.size());
     for (auto &Block : AllBlocks) {
       for (auto &Jump : Block.OutJumps) {
-        const auto SuccBlock = Jump->Target;
+        auto SuccBlock = Jump->Target;
         auto CurEdge = Block.CurChain->getEdge(SuccBlock->CurChain);
         // this edge is already present in the graph
         if (CurEdge != nullptr) {
@@ -589,6 +605,10 @@ private:
           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);
           if (CurGain.score() <= EPS)