[llvm] [CodeLayout] CDSortImpl: remove linear-time erase_value from mergeChains (PR #69276)

[Fangrui Song via llvm-commits]

https://github.com/MaskRay updated https://github.com/llvm/llvm-project/pull/69276

>From 299bbc698e8dbdcf2163ed6ba213291f7a1b9505 Mon Sep 17 00:00:00 2001
From: Fangrui Song <i at maskray.me>
Date: Mon, 16 Oct 2023 18:54:12 -0700
Subject: [PATCH] [CodeLayout] CDSortImpl: remove HotChains and remove
 linear-time erase_value from mergeChains

After mergeChainPairs initializes a priority queue, HotChains is unused
except a HotChains.size() use in LLVM_DEBUG. Optimize it out.
---
 llvm/lib/Transforms/Utils/CodeLayout.cpp | 25 +++++++++++-------------
 1 file changed, 11 insertions(+), 14 deletions(-)

diff --git a/llvm/lib/Transforms/Utils/CodeLayout.cpp b/llvm/lib/Transforms/Utils/CodeLayout.cpp
index d9e302d8b4fa54d..eb763b886e6ab31 100644
--- a/llvm/lib/Transforms/Utils/CodeLayout.cpp
+++ b/llvm/lib/Transforms/Utils/CodeLayout.cpp
@@ -1025,9 +1025,6 @@ class CDSortImpl {
     // Merge pairs of chains while improving the objective.
     mergeChainPairs();
 
-    LLVM_DEBUG(dbgs() << "Cache-directed function sorting reduced the number"
-                      << " of chains from " << NumNodes << " to "
-                      << HotChains.size() << "\n");
 
     // Collect nodes from all the chains.
     return concatChains();
@@ -1074,7 +1071,6 @@ class CDSortImpl {
 
     // Initialize chains.
     AllChains.reserve(NumNodes);
-    HotChains.reserve(NumNodes);
     for (NodeT &Node : AllNodes) {
       // Adjust execution counts.
       Node.ExecutionCount = std::max(Node.ExecutionCount, Node.inCount());
@@ -1082,8 +1078,6 @@ class CDSortImpl {
       // Create chain.
       AllChains.emplace_back(Node.Index, &Node);
       Node.CurChain = &AllChains.back();
-      if (Node.ExecutionCount > 0)
-        HotChains.push_back(&AllChains.back());
     }
 
     // Initialize chain edges.
@@ -1116,8 +1110,12 @@ class CDSortImpl {
     std::set<ChainEdge *, decltype(GainComparator)> Queue(GainComparator);
 
     // Insert the edges into the queue.
-    for (ChainT *ChainPred : HotChains) {
-      for (const auto &[_, Edge] : ChainPred->Edges) {
+    size_t NumActiveChains = 0;
+    for (NodeT &Node : AllNodes) {
+      if (Node.ExecutionCount == 0)
+        continue;
+      ++NumActiveChains;
+      for (const auto &[_, Edge] : Node.CurChain->Edges) {
         // Ignore self-edges.
         if (Edge->isSelfEdge())
           continue;
@@ -1152,6 +1150,7 @@ class CDSortImpl {
       MergeGainT BestGain = BestEdge->getMergeGain();
       mergeChains(BestSrcChain, BestDstChain, BestGain.mergeOffset(),
                   BestGain.mergeType());
+      --NumActiveChains;
 
       // Insert newly created edges into the queue.
       for (const auto &[_, Edge] : BestSrcChain->Edges) {
@@ -1167,6 +1166,10 @@ class CDSortImpl {
           Queue.insert(Edge);
       }
     }
+
+    LLVM_DEBUG(dbgs() << "Cache-directed function sorting reduced the number"
+                      << " of chains from " << NumNodes << " to "
+                      << NumActiveChains << "\n");
   }
 
   /// Compute the gain of merging two chains.
@@ -1301,9 +1304,6 @@ class CDSortImpl {
     // Merge the edges.
     Into->mergeEdges(From);
     From->clear();
-
-    // Remove the chain from the list of active chains.
-    llvm::erase_value(HotChains, From);
   }
 
   /// Concatenate all chains into the final order.
@@ -1370,9 +1370,6 @@ class CDSortImpl {
   /// All edges between the chains.
   std::vector<ChainEdge> AllEdges;
 
-  /// Active chains. The vector gets updated at runtime when chains are merged.
-  std::vector<ChainT *> HotChains;
-
   /// The total number of samples in the graph.
   uint64_t TotalSamples{0};