[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