[CodeLayout] cache-directed sort: limit max chain size

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

@@ -65,6 +65,8 @@ struct CDSortConfig {
   unsigned CacheEntries = 16;
   /// The size of a line in the cache.
   unsigned CacheSize = 2048;
+  /// The maximum size of a chain to create.
+  unsigned MaxChainSize = 128;
   /// The power exponent for the distance-based locality.
   double DistancePower = 0.25;
   /// The scale factor for the frequency-based locality.

llvm/lib/Transforms/Utils/CodeLayout.cpp

@@ -123,6 +123,10 @@ static cl::opt<unsigned> CacheEntries("cds-cache-entries", cl::ReallyHidden,
 static cl::opt<unsigned> CacheSize("cds-cache-size", cl::ReallyHidden,
                                    cl::desc("The size of a line in the cache"));
 
+static cl::opt<unsigned>
+    CDMaxChainSize("cdsort-max-chain-size", cl::ReallyHidden,
+                   cl::desc("The maximum size of a chain to create"));
+
 static cl::opt<double> DistancePower(
     "cds-distance-power", cl::ReallyHidden,
     cl::desc("The power exponent for the distance-based locality"));
@@ -1156,6 +1160,9 @@ class CDSortImpl {
         // Ignore loop edges.
         if (Edge->isSelfEdge())
           continue;
+        if (Edge->srcChain()->numBlocks() + Edge->dstChain()->numBlocks() >
+            Config.MaxChainSize)
+          continue;
 
         // Compute the gain of merging the two chains.
         MergeGainT Gain = getBestMergeGain(Edge);
@@ -1452,6 +1459,8 @@ std::vector<uint64_t> codelayout::computeCacheDirectedLayout(
     Config.CacheEntries = CacheEntries;
   if (CacheSize.getNumOccurrences() > 0)
     Config.CacheSize = CacheSize;
+  if (CDMaxChainSize.getNumOccurrences() > 0)
+    Config.MaxChainSize = CDMaxChainSize;
   if (DistancePower.getNumOccurrences() > 0)
     Config.DistancePower = DistancePower;
   if (FrequencyScale.getNumOccurrences() > 0)

llvm/unittests/Transforms/Utils/CodeLayoutTest.cpp

@@ -40,6 +40,14 @@ TEST(CodeLayout, HotChain) {
     const std::vector<uint64_t> CallOffsets(std::size(Edges), 5);
     auto Order = computeCacheDirectedLayout(Sizes, Counts, Edges, CallOffsets);
     EXPECT_THAT(Order, ElementsAreArray({0, 3, 4, 2, 1}));
+
+    // -cdsort-max-chain-size disables forming a larger chain and therefore may
+    // change the result.
+    CDSortConfig Config;
+    Config.MaxChainSize = 3;
+    Order =
+        computeCacheDirectedLayout(Config, Sizes, Counts, Edges, CallOffsets);
+    EXPECT_THAT(Order, ElementsAreArray({0, 3, 4, 1, 2}));
   }
 }