[llvm] [CodeGen] Reland "Use SmallMapVector for SpillPlacement::Node::Links" (PR #196270)

Mingjie Xu via llvm-commits llvm-commits at lists.llvm.org
Thu May 7 02:58:00 PDT 2026


https://github.com/Enna1 created https://github.com/llvm/llvm-project/pull/196270

None

>From a0ee51ce686455aa8bb8e24ac057e88b0bb9d9b9 Mon Sep 17 00:00:00 2001
From: Enna1 <xumingjie.enna1 at bytedance.com>
Date: Thu, 7 May 2026 11:45:53 +0800
Subject: [PATCH 1/2] [ADT] apply small-size optimization to SmallMapVector

---
 llvm/include/llvm/ADT/MapVector.h | 99 +++++++++++++++++++++++++++++--
 1 file changed, 93 insertions(+), 6 deletions(-)

diff --git a/llvm/include/llvm/ADT/MapVector.h b/llvm/include/llvm/ADT/MapVector.h
index 2b2f098dd3abf..fa52f806ca579 100644
--- a/llvm/include/llvm/ADT/MapVector.h
+++ b/llvm/include/llvm/ADT/MapVector.h
@@ -18,6 +18,7 @@
 #define LLVM_ADT_MAPVECTOR_H
 
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallVector.h"
 #include <cassert>
 #include <cstddef>
@@ -32,7 +33,8 @@ namespace llvm {
 /// mapping is done with DenseMap from Keys to indexes in that vector.
 template <typename KeyT, typename ValueT,
           typename MapType = DenseMap<KeyT, unsigned>,
-          typename VectorType = SmallVector<std::pair<KeyT, ValueT>, 0>>
+          typename VectorType = SmallVector<std::pair<KeyT, ValueT>, 0>,
+          unsigned N = 0>
 class MapVector {
 public:
   using key_type = KeyT;
@@ -108,6 +110,15 @@ class MapVector {
   [[nodiscard]] ValueT lookup(const KeyT &Key) const {
     static_assert(std::is_copy_constructible_v<ValueT>,
                   "Cannot call lookup() if ValueT is not copyable.");
+    if constexpr (canBeSmall())
+      if (isSmall()) {
+        auto I =
+            find_if(Vector, [&Key](const auto &P) { return P.first == Key; });
+        if (I != Vector.end())
+          return I->second;
+        return ValueT();
+      }
+
     typename MapType::const_iterator Pos = Map.find(Key);
     return Pos == Map.end()? ValueT() : Vector[Pos->second].second;
   }
@@ -144,6 +155,10 @@ class MapVector {
   }
 
   [[nodiscard]] bool contains(const KeyT &Key) const {
+    if constexpr (canBeSmall())
+      if (isSmall())
+        return any_of(Vector, [&Key](const auto &P) { return P.first == Key; });
+
     return Map.find(Key) != Map.end();
   }
 
@@ -152,11 +167,21 @@ class MapVector {
   }
 
   [[nodiscard]] iterator find(const KeyT &Key) {
+    if constexpr (canBeSmall())
+      if (isSmall())
+        return find_if(Vector,
+                       [&Key](const auto &P) { return P.first == Key; });
+
     typename MapType::const_iterator Pos = Map.find(Key);
     return Pos == Map.end() ? Vector.end() : (Vector.begin() + Pos->second);
   }
 
   [[nodiscard]] const_iterator find(const KeyT &Key) const {
+    if constexpr (canBeSmall())
+      if (isSmall())
+        return find_if(Vector,
+                       [&Key](const auto &P) { return P.first == Key; });
+
     typename MapType::const_iterator Pos = Map.find(Key);
     return Pos == Map.end() ? Vector.end() : (Vector.begin() + Pos->second);
   }
@@ -164,6 +189,15 @@ class MapVector {
   /// at - Return the entry for the specified key, or abort if no such
   /// entry exists.
   [[nodiscard]] ValueT &at(const KeyT &Key) {
+    if constexpr (canBeSmall())
+      if (isSmall()) {
+        auto I =
+            find_if(Vector, [&Key](const auto &P) { return P.first == Key; });
+        assert(I != Vector.end() &&
+               "MapVector::at failed due to a missing key");
+        return I->second;
+      }
+
     auto Pos = Map.find(Key);
     assert(Pos != Map.end() && "MapVector::at failed due to a missing key");
     return Vector[Pos->second].second;
@@ -172,6 +206,15 @@ class MapVector {
   /// at - Return the entry for the specified key, or abort if no such
   /// entry exists.
   [[nodiscard]] const ValueT &at(const KeyT &Key) const {
+    if constexpr (canBeSmall())
+      if (isSmall()) {
+        auto I =
+            find_if(Vector, [&Key](const auto &P) { return P.first == Key; });
+        assert(I != Vector.end() &&
+               "MapVector::at failed due to a missing key");
+        return I->second;
+      }
+
     auto Pos = Map.find(Key);
     assert(Pos != Map.end() && "MapVector::at failed due to a missing key");
     return Vector[Pos->second].second;
@@ -179,6 +222,12 @@ class MapVector {
 
   /// Remove the last element from the vector.
   void pop_back() {
+    if constexpr (canBeSmall())
+      if (isSmall()) {
+        Vector.pop_back();
+        return;
+      }
+
     typename MapType::iterator Pos = Map.find(Vector.back().first);
     Map.erase(Pos);
     Vector.pop_back();
@@ -192,6 +241,10 @@ class MapVector {
   /// \note This is a deceivingly expensive operation (linear time).  It's
   /// usually better to use \a remove_if() if possible.
   typename VectorType::iterator erase(typename VectorType::iterator Iterator) {
+    if constexpr (canBeSmall())
+      if (isSmall())
+        return Vector.erase(Iterator);
+
     Map.erase(Iterator->first);
     auto Next = Vector.erase(Iterator);
     if (Next == Vector.end())
@@ -225,15 +278,43 @@ class MapVector {
   template <class Predicate> void remove_if(Predicate Pred);
 
 private:
+  [[nodiscard]] static constexpr bool canBeSmall() { return N != 0; }
+
+  [[nodiscard]] bool isSmall() const { return Map.empty(); }
+
+  void makeBig() {
+    if constexpr (canBeSmall()) {
+      unsigned Index = 0;
+      for (const auto &entry : Vector)
+        Map[entry.first] = Index++;
+    }
+  }
+
   MapType Map;
   VectorType Vector;
 
+  static_assert(N <= 32, "Small size should be less than or equal to 32!");
+
   static_assert(
       std::is_integral_v<typename MapType::mapped_type>,
       "The mapped_type of the specified Map must be an integral type");
 
   template <typename KeyArgT, typename... Ts>
   std::pair<iterator, bool> try_emplace_impl(KeyArgT &&Key, Ts &&...Args) {
+    if constexpr (canBeSmall())
+      if (isSmall()) {
+        auto I =
+            find_if(Vector, [&Key](const auto &P) { return P.first == Key; });
+        if (I != Vector.end())
+          return {I, false};
+        Vector.emplace_back(std::piecewise_construct,
+                            std::forward_as_tuple(std::forward<KeyArgT>(Key)),
+                            std::forward_as_tuple(std::forward<Ts>(Args)...));
+        if (Vector.size() > N)
+          makeBig();
+        return {std::prev(end()), true};
+      }
+
     auto [It, Inserted] = Map.try_emplace(Key);
     if (Inserted) {
       It->second = Vector.size();
@@ -246,9 +327,16 @@ class MapVector {
   }
 };
 
-template <typename KeyT, typename ValueT, typename MapType, typename VectorType>
+template <typename KeyT, typename ValueT, typename MapType, typename VectorType,
+          unsigned N>
 template <class Function>
-void MapVector<KeyT, ValueT, MapType, VectorType>::remove_if(Function Pred) {
+void MapVector<KeyT, ValueT, MapType, VectorType, N>::remove_if(Function Pred) {
+  if constexpr (canBeSmall())
+    if (isSmall()) {
+      Vector.erase(llvm::remove_if(Vector, Pred), Vector.end());
+      return;
+    }
+
   auto O = Vector.begin();
   for (auto I = O, E = Vector.end(); I != E; ++I) {
     if (Pred(*I)) {
@@ -272,9 +360,8 @@ void MapVector<KeyT, ValueT, MapType, VectorType>::remove_if(Function Pred) {
 /// size.
 template <typename KeyT, typename ValueT, unsigned N>
 struct SmallMapVector
-    : MapVector<KeyT, ValueT, SmallDenseMap<KeyT, unsigned, N>,
-                SmallVector<std::pair<KeyT, ValueT>, N>> {
-};
+    : MapVector<KeyT, ValueT, DenseMap<KeyT, unsigned>,
+                SmallVector<std::pair<KeyT, ValueT>, N>, N> {};
 
 } // end namespace llvm
 

>From 5cd23b57566fe9c713c39db30f3842535f2e6a4e Mon Sep 17 00:00:00 2001
From: Enna1 <xumingjie.enna1 at bytedance.com>
Date: Tue, 28 Apr 2026 22:30:54 +0800
Subject: [PATCH 2/2] [CodeGen] Use SmallMapVector for
 SpillPlacement::Node::Links

Previously, `SpillPlacement::Node::Links` was implemented as a `SmallVector`
of `(Weight, BundleNo)` pairs.

This patch replaces the `SmallVector` with a `SmallMapVector<unsigned, BlockFrequency, 4>`,
which stores `(BundleNo, Weight)` pairs. This allows for more efficient
lookups and weight accumulations when multiple links to the same bundle are
added.
---
 llvm/lib/CodeGen/SpillPlacement.cpp | 34 ++++++++++++-----------------
 1 file changed, 14 insertions(+), 20 deletions(-)

diff --git a/llvm/lib/CodeGen/SpillPlacement.cpp b/llvm/lib/CodeGen/SpillPlacement.cpp
index 55a96a22a00ec..1b3cd9ccbf08b 100644
--- a/llvm/lib/CodeGen/SpillPlacement.cpp
+++ b/llvm/lib/CodeGen/SpillPlacement.cpp
@@ -28,6 +28,7 @@
 
 #include "llvm/CodeGen/SpillPlacement.h"
 #include "llvm/ADT/BitVector.h"
+#include "llvm/ADT/MapVector.h"
 #include "llvm/CodeGen/EdgeBundles.h"
 #include "llvm/CodeGen/MachineBasicBlock.h"
 #include "llvm/CodeGen/MachineBlockFrequencyInfo.h"
@@ -81,11 +82,9 @@ struct SpillPlacement::Node {
   /// variable should go in a register through this bundle.
   int Value;
 
-  using LinkVector = SmallVector<std::pair<BlockFrequency, unsigned>, 4>;
-
-  /// Links - (Weight, BundleNo) for all transparent blocks connecting to other
+  /// Links - (BundleNo, Weight) for all transparent blocks connecting to other
   /// bundles. The weights are all positive block frequencies.
-  LinkVector Links;
+  SmallMapVector<unsigned, BlockFrequency, 4> Links;
 
   /// SumLinkWeights - Cached sum of the weights of all links + ThresHold.
   BlockFrequency SumLinkWeights;
@@ -120,13 +119,9 @@ struct SpillPlacement::Node {
     SumLinkWeights += w;
 
     // There can be multiple links to the same bundle, add them up.
-    for (std::pair<BlockFrequency, unsigned> &L : Links)
-      if (L.second == b) {
-        L.first += w;
-        return;
-      }
-    // This must be the first link to b.
-    Links.push_back(std::make_pair(w, b));
+    auto [It, Inserted] = Links.try_emplace(b, w);
+    if (!Inserted)
+      It->second += w;
   }
 
   /// addBias - Bias this node.
@@ -152,11 +147,11 @@ struct SpillPlacement::Node {
     // Compute the weighted sum of inputs.
     BlockFrequency SumN = BiasN;
     BlockFrequency SumP = BiasP;
-    for (std::pair<BlockFrequency, unsigned> &L : Links) {
-      if (nodes[L.second].Value == -1)
-        SumN += L.first;
-      else if (nodes[L.second].Value == 1)
-        SumP += L.first;
+    for (auto [BundleNo, Weight] : Links) {
+      if (nodes[BundleNo].Value == -1)
+        SumN += Weight;
+      else if (nodes[BundleNo].Value == 1)
+        SumP += Weight;
     }
 
     // Each weighted sum is going to be less than the total frequency of the
@@ -179,12 +174,11 @@ struct SpillPlacement::Node {
 
   void getDissentingNeighbors(SparseSet<unsigned> &List,
                               const Node nodes[]) const {
-    for (const auto &Elt : Links) {
-      unsigned n = Elt.second;
+    for (auto [BundleNo, _] : Links) {
       // Neighbors that already have the same value are not going to
       // change because of this node changing.
-      if (Value != nodes[n].Value)
-        List.insert(n);
+      if (Value != nodes[BundleNo].Value)
+        List.insert(BundleNo);
     }
   }
 };



More information about the llvm-commits mailing list