[llvm] [CodeGen] Reland "Use SmallMapVector for SpillPlacement::Node::Links" (PR #196270)
via llvm-commits
llvm-commits at lists.llvm.org
Fri May 8 00:10:45 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-llvm-adt
Author: Mingjie Xu (Enna1)
<details>
<summary>Changes</summary>
The original change https://github.com/llvm/llvm-project/pull/194653 caused a compile-time regression.
Reland the replacement of SpillPlacement::Node::Links from a SmallVector of
(Weight, BundleNo) pairs to a SmallMapVector keyed by BundleNo, and increase the small size threshold to 8.
Depends on https://github.com/llvm/llvm-project/pull/196473, which adds a vector-only small mode for SmallMapVector.
---
Full diff: https://github.com/llvm/llvm-project/pull/196270.diff
3 Files Affected:
- (modified) llvm/include/llvm/ADT/MapVector.h (+125-6)
- (modified) llvm/lib/CodeGen/SpillPlacement.cpp (+14-20)
- (modified) llvm/unittests/ADT/MapVectorTest.cpp (+20)
``````````diff
diff --git a/llvm/include/llvm/ADT/MapVector.h b/llvm/include/llvm/ADT/MapVector.h
index 2b2f098dd3abf..423e2a0fd6145 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>
@@ -27,12 +28,26 @@
namespace llvm {
+namespace detail {
+template <typename MapT> struct MapVectorKeyInfo {};
+
+template <typename KeyT, typename ValueT, typename KeyInfoT, typename BucketT>
+struct MapVectorKeyInfo<DenseMap<KeyT, ValueT, KeyInfoT, BucketT>> {
+ using type = KeyInfoT;
+};
+
+template <typename MapT>
+using map_vector_key_info_t = typename MapVectorKeyInfo<MapT>::type;
+
+} // namespace detail
+
/// This class implements a map that also provides access to all stored values
/// in a deterministic order. The values are kept in a SmallVector<*, 0> and the
/// 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 +123,14 @@ 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 = findInVector(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 +167,10 @@ class MapVector {
}
[[nodiscard]] bool contains(const KeyT &Key) const {
+ if constexpr (canBeSmall())
+ if (isSmall())
+ return findInVector(Key) != Vector.end();
+
return Map.find(Key) != Map.end();
}
@@ -152,11 +179,19 @@ class MapVector {
}
[[nodiscard]] iterator find(const KeyT &Key) {
+ if constexpr (canBeSmall())
+ if (isSmall())
+ return findInVector(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 findInVector(Key);
+
typename MapType::const_iterator Pos = Map.find(Key);
return Pos == Map.end() ? Vector.end() : (Vector.begin() + Pos->second);
}
@@ -164,6 +199,14 @@ 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 = findInVector(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 +215,14 @@ 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 = findInVector(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 +230,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 +249,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 +286,67 @@ class MapVector {
template <class Predicate> void remove_if(Predicate Pred);
private:
+ template <typename LookupKeyT, typename StoredKeyT>
+ [[nodiscard]] static bool keyEqual(const LookupKeyT &LookupKey,
+ const StoredKeyT &StoredKey) {
+ if constexpr (is_detected<detail::map_vector_key_info_t, MapType>::value) {
+ using KeyInfo = detail::map_vector_key_info_t<MapType>;
+ return KeyInfo::isEqual(LookupKey, StoredKey);
+ } else {
+ static_assert(has_equality_comparison_v<LookupKeyT, StoredKeyT>,
+ "MapVector small mode requires MapType key info or "
+ "operator== for key comparison");
+ return LookupKey == StoredKey;
+ }
+ }
+
+ template <typename T> [[nodiscard]] iterator findInVector(const T &Key) {
+ return find_if(Vector,
+ [&Key](const auto &P) { return keyEqual(Key, P.first); });
+ }
+
+ template <typename T>
+ [[nodiscard]] const_iterator findInVector(const T &Key) const {
+ return find_if(Vector,
+ [&Key](const auto &P) { return keyEqual(Key, P.first); });
+ }
+
+ [[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 = findInVector(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 +359,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)) {
@@ -271,9 +391,8 @@ void MapVector<KeyT, ValueT, MapType, VectorType>::remove_if(Function Pred) {
/// A MapVector that performs no allocations if smaller than a certain
/// size.
template <typename KeyT, typename ValueT, unsigned N>
-struct SmallMapVector
- : MapVector<KeyT, ValueT, SmallDenseMap<KeyT, unsigned, N>,
- SmallVector<std::pair<KeyT, ValueT>, N>> {
+struct SmallMapVector : MapVector<KeyT, ValueT, DenseMap<KeyT, unsigned>,
+ SmallVector<std::pair<KeyT, ValueT>, N>, N> {
};
} // end namespace llvm
diff --git a/llvm/lib/CodeGen/SpillPlacement.cpp b/llvm/lib/CodeGen/SpillPlacement.cpp
index 55a96a22a00ec..80107b8767ca9 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, 8> 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);
}
}
};
diff --git a/llvm/unittests/ADT/MapVectorTest.cpp b/llvm/unittests/ADT/MapVectorTest.cpp
index b11d4603b90b7..b1d7b9bb09e29 100644
--- a/llvm/unittests/ADT/MapVectorTest.cpp
+++ b/llvm/unittests/ADT/MapVectorTest.cpp
@@ -482,6 +482,26 @@ TEST(SmallMapVectorSmallTest, NonCopyable) {
ASSERT_EQ(*MV.find(2)->second, 2);
}
+TEST(SmallMapVectorSmallTest, UsesDenseMapInfoEquality) {
+ SmallMapVector<A, int, 4> MV;
+
+ auto R0 = MV.try_emplace(A(0), 1);
+ EXPECT_TRUE(R0.second);
+
+ auto R1 = MV.try_emplace(A(0), 2);
+ EXPECT_FALSE(R1.second);
+ EXPECT_EQ(R1.first, R0.first);
+ EXPECT_EQ(R1.first->second, 1);
+ EXPECT_EQ(MV.size(), 1u);
+
+ EXPECT_TRUE(MV.contains(A(0)));
+ EXPECT_EQ(MV.find(A(0)), MV.begin());
+ EXPECT_EQ(MV.lookup(A(0)), 1);
+ EXPECT_EQ(MV.at(A(0)), 1);
+ EXPECT_EQ(MV.erase(A(0)), 1u);
+ EXPECT_TRUE(MV.empty());
+}
+
TEST(SmallMapVectorLargeTest, insert_pop) {
SmallMapVector<int, int, 1> MV;
std::pair<SmallMapVector<int, int, 1>::iterator, bool> R;
``````````
</details>
https://github.com/llvm/llvm-project/pull/196270
More information about the llvm-commits
mailing list