[llvm] [ORC] Make ElementSet, ContainerElementsMap inner classes. (PR #184955)
Lang Hames via llvm-commits
llvm-commits at lists.llvm.org
Thu Mar 5 22:14:58 PST 2026
https://github.com/lhames created https://github.com/llvm/llvm-project/pull/184955
ElementSet and ContainerElementsMap were type aliases inside WaitingOnGraph.
This commit replaces the aliases with classes deriving from DenseSet and DenseMap, with convenience operations added for WaitingOnGraph (merge, remove, remove_if, and visit). These convenience functions are used to simplify the implementation of various parts of WaitingOnGraph.
Unit tests are added for the convenience operations to improve test coverage.
In addition to improving readability of the main WaitingOnGraph operations, this will make it easier to experiment with other underlying representations for these types (e.g. sorted vectors).
>From 85d7a4336b41afd395730518c07978098798f305 Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames at gmail.com>
Date: Fri, 6 Mar 2026 17:07:31 +1100
Subject: [PATCH] [ORC] Make ElementSet, ContainerElementsMap inner classes.
ElementSet and ContainerElementsMap were type aliases inside WaitingOnGraph.
This commit replaces the aliases with classes deriving from DenseSet and
DenseMap, with convenience operations added for WaitingOnGraph (merge, remove,
remove_if, and visit). These convenience functions are used to simplify the
implementation of various parts of WaitingOnGraph.
Unit tests are added for the convenience operations to improve test coverage.
In addition to improving readability of the main WaitingOnGraph operations,
this will make it easier to experiment with other underlying representations
for these types (e.g. sorted vectors).
---
.../llvm/ExecutionEngine/Orc/WaitingOnGraph.h | 210 ++++++++++++------
.../Orc/WaitingOnGraphTest.cpp | 189 ++++++++++++++++
2 files changed, 326 insertions(+), 73 deletions(-)
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/WaitingOnGraph.h b/llvm/include/llvm/ExecutionEngine/Orc/WaitingOnGraph.h
index 6c506d7eb5b7e..8249f192321af 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/WaitingOnGraph.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/WaitingOnGraph.h
@@ -85,8 +85,121 @@ template <typename ContainerIdT, typename ElementIdT> class WaitingOnGraph {
public:
using ContainerId = ContainerIdT;
using ElementId = ElementIdT;
- using ElementSet = DenseSet<ElementId>;
- using ContainerElementsMap = DenseMap<ContainerId, ElementSet>;
+
+ class ElementSet : public DenseSet<ElementId> {
+ friend class ElementSetTest;
+
+ public:
+ using DenseSet<ElementId>::DenseSet;
+
+ /// Merge the elements of Other into this set. Returns true if any new
+ /// elements are added.
+ bool merge(const ElementSet &Other) {
+ size_t OrigSize = this->size();
+ this->insert(Other.begin(), Other.end());
+ return this->size() != OrigSize;
+ }
+
+ /// Remove all elements in Other from this set. Returns true if any
+ /// elements were removed.
+ bool remove(const ElementSet &Other) {
+ size_t OrigSize = this->size();
+
+ // Early out for empty sets.
+ if (OrigSize == 0 || Other.empty())
+ return false;
+
+ // TODO: Tweak condition to account for SmallVector cost. We may want to
+ // prefer iterating over elements if the size difference is small.
+ if (OrigSize > Other.size()) {
+ for (auto &Elem : Other)
+ this->erase(Elem);
+ } else {
+ SmallVector<ElementId> ToRemove;
+ for (auto &Elem : *this)
+ if (Other.count(Elem))
+ ToRemove.push_back(Elem);
+ for (auto &Elem : ToRemove)
+ this->erase(Elem);
+ }
+ return this->size() < OrigSize;
+ }
+
+ /// Remove all elements for which Pred returns true.
+ /// Returns true if any elements were removed.
+ template <typename Pred> bool remove_if(Pred &&P) {
+ if (this->empty())
+ return false;
+
+ SmallVector<ElementId> ToRemove;
+ for (auto &Elem : *this)
+ if (P(Elem))
+ ToRemove.push_back(Elem);
+
+ for (auto &Elem : ToRemove)
+ this->erase(Elem);
+
+ return !ToRemove.empty();
+ }
+ };
+
+ class ContainerElementsMap : public DenseMap<ContainerId, ElementSet> {
+ friend class ContainerElementsMapTest;
+
+ public:
+ using DenseMap<ContainerId, ElementSet>::DenseMap;
+
+ /// Merge the elements of Other into this map. Returns true if any new
+ /// elements are added.
+ bool merge(const ContainerElementsMap &Other) {
+ bool Changed = false;
+ for (auto &[Container, Elements] : Other)
+ Changed |= (*this)[Container].merge(Elements);
+ return Changed;
+ }
+
+ /// Remove all elements in Other from this map. Returns true if any
+ /// elements were removed.
+ bool remove(const ContainerElementsMap &Other) {
+ bool Changed = false;
+ for (auto &[Container, Elements] : Other) {
+ assert(!Elements.empty() && "Stale row for Container in Other");
+ auto I = this->find(Container);
+ if (I == this->end())
+ continue;
+ Changed |= I->second.remove(Elements);
+ if (I->second.empty())
+ this->erase(Container);
+ }
+ return Changed;
+ }
+
+ /// Call V on each (Container, Elements) pair in this map.
+ ///
+ /// V should return true if it modifies any elements.
+ ///
+ /// Returns true if V returns true for any pair.
+ template <typename Visitor> bool visit(Visitor &&V) {
+ if (this->empty())
+ return false;
+
+ bool Changed = false;
+ SmallVector<ContainerId> ToRemove;
+ for (auto &[Container, Elements] : *this) {
+ assert(!Elements.empty() && "empty row for container");
+ if (V(Container, Elements)) {
+ Changed = true;
+ if (Elements.empty())
+ ToRemove.push_back(Container);
+ }
+ }
+
+ for (auto &Container : ToRemove)
+ this->erase(Container);
+
+ return Changed;
+ }
+ };
class SuperNode;
@@ -164,56 +277,29 @@ template <typename ContainerIdT, typename ElementIdT> class WaitingOnGraph {
/// Returns true if SuperNodeDeps was changed.
bool hoistDeps(SuperNodeDepsMap &SuperNodeDeps,
ElemToSuperNodeMap &ElemToSN) {
- bool Changed = false;
+ bool SuperNodeDepsChanged = false;
- SmallVector<ContainerId> ContainersToRemove;
- for (auto &[DepContainer, DepElems] : Deps) {
- auto I = ElemToSN.find(DepContainer);
+ Deps.visit([&](ContainerId &Container, ElementSet &Elements) {
+ auto I = ElemToSN.find(Container);
if (I == ElemToSN.end())
- continue;
- auto &ContainerElemToSN = I->second;
+ return false;
- // ElemToSN includes SuperNodes that define elements in DepContainer.
- // We need to iterate over ContainerElemToSN or DepElems: we pick the
- // smaller to minimize the cost.
- if (ContainerElemToSN.size() < DepElems.size()) {
- for (auto &[DefElem, DefSN] : ContainerElemToSN) {
- if (DepElems.erase(DefElem) && DefSN != this) {
- Changed = true;
- SuperNodeDeps[DefSN].insert(this);
- }
- }
- } else {
- SmallVector<ElementId> ElemsToRemove;
- for (auto &DepElem : DepElems) {
- auto J = ContainerElemToSN.find(DepElem);
- if (J == ContainerElemToSN.end())
- continue;
- ElemsToRemove.push_back(DepElem);
- SuperNode *DefSN = J->second;
- if (DefSN != this) {
- Changed = true;
- SuperNodeDeps[DefSN].insert(this);
- }
+ auto &ContainerElemToSN = I->second;
+ return Elements.remove_if([&](const ElementId &Elem) {
+ auto J = ContainerElemToSN.find(Elem);
+ if (J == ContainerElemToSN.end())
+ return false;
+
+ auto *DefSN = J->second;
+ if (DefSN != this) {
+ SuperNodeDepsChanged = true;
+ SuperNodeDeps[DefSN].insert(this);
}
+ return true;
+ });
+ });
- for (auto &DepElem : ElemsToRemove)
- DepElems.erase(DepElem);
- }
-
- // If DepElems has become empty then add DepContainer to the list of
- // containers to remove.
- if (DepElems.empty())
- ContainersToRemove.push_back(DepContainer);
- }
-
- for (auto &DepContainer : ContainersToRemove) {
- assert(Deps.count(DepContainer) && "already removed?");
- assert(Deps[DepContainer].empty() && "non empty?");
- Deps.erase(DepContainer);
- }
-
- return Changed;
+ return SuperNodeDepsChanged;
}
};
@@ -252,8 +338,7 @@ template <typename ContainerIdT, typename ElementIdT> class WaitingOnGraph {
auto H = getHash(SN->Deps);
if (auto *CanonicalSN = findCanonicalSuperNode(H, SN->Deps)) {
SN->mapDefsTo(ElemToSN, CanonicalSN, AbandonOldMapping);
- for (auto &[Container, Elems] : SN->Defs)
- CanonicalSN->Defs[Container].insert(Elems.begin(), Elems.end());
+ CanonicalSN->Defs.merge(SN->Defs);
std::swap(SN, SNs.back());
SNs.pop_back();
} else {
@@ -342,21 +427,7 @@ template <typename ContainerIdT, typename ElementIdT> class WaitingOnGraph {
void add(ContainerElementsMap Defs, ContainerElementsMap Deps) {
if (Defs.empty())
return;
- // Remove any self-reference.
- SmallVector<ContainerId> ToRemove;
- for (auto &[Container, Elems] : Defs) {
- assert(!Elems.empty() && "Defs for container must not be empty");
- auto I = Deps.find(Container);
- if (I == Deps.end())
- continue;
- auto &DepsForContainer = I->second;
- for (auto &Elem : Elems)
- DepsForContainer.erase(Elem);
- if (DepsForContainer.empty())
- ToRemove.push_back(Container);
- }
- for (auto &Container : ToRemove)
- Deps.erase(Container);
+ Deps.remove(Defs); // Remove any self-reference.
if (auto SN = C.addOrCreateSuperNode(std::move(Defs), std::move(Deps)))
SNs.push_back(std::move(SN));
}
@@ -626,16 +697,9 @@ template <typename ContainerIdT, typename ElementIdT> class WaitingOnGraph {
if (I == SuperNodeDeps.end())
continue;
- for (auto *DependantSN : I->second) {
- bool Changed = false;
- for (auto &[DepContainer, DepElems] : SN->Deps) {
- auto &DepSNContainerElems = DependantSN->Deps[DepContainer];
- for (auto &DepElem : DepElems)
- Changed |= DepSNContainerElems.insert(DepElem).second;
- }
- if (Changed)
+ for (auto *DependantSN : I->second)
+ if (DependantSN->Deps.merge(SN->Deps))
ToVisitNext.insert(DependantSN);
- }
}
if (ToVisitNext.empty())
diff --git a/llvm/unittests/ExecutionEngine/Orc/WaitingOnGraphTest.cpp b/llvm/unittests/ExecutionEngine/Orc/WaitingOnGraphTest.cpp
index 0fae82745bbfe..84d5a5fd9e3ed 100644
--- a/llvm/unittests/ExecutionEngine/Orc/WaitingOnGraphTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/WaitingOnGraphTest.cpp
@@ -11,6 +11,44 @@
namespace llvm::orc::detail {
+class ElementSetTest : public testing::Test {
+public:
+ using TestElementSet = WaitingOnGraph<uintptr_t, uintptr_t>::ElementSet;
+
+ bool merge(TestElementSet &S, const TestElementSet &Other) {
+ return S.merge(Other);
+ }
+
+ bool remove(TestElementSet &S, const TestElementSet &Other) {
+ return S.remove(Other);
+ }
+
+ template <typename Pred> bool remove_if(TestElementSet &S, Pred &&P) {
+ return S.remove_if(std::forward<Pred>(P));
+ }
+};
+
+class ContainerElementsMapTest : public testing::Test {
+public:
+ using TestContainerElementsMap =
+ WaitingOnGraph<uintptr_t, uintptr_t>::ContainerElementsMap;
+
+ bool merge(TestContainerElementsMap &M,
+ const TestContainerElementsMap &Other) {
+ return M.merge(Other);
+ }
+
+ bool remove(TestContainerElementsMap &M,
+ const TestContainerElementsMap &Other) {
+ return M.remove(Other);
+ }
+
+ template <typename Visitor>
+ bool visit(TestContainerElementsMap &M, Visitor &&V) {
+ return M.visit(std::forward<Visitor>(V));
+ }
+};
+
class WaitingOnGraphTest : public testing::Test {
public:
using TestGraph = WaitingOnGraph<uintptr_t, uintptr_t>;
@@ -119,6 +157,157 @@ using namespace llvm;
using namespace llvm::orc;
using namespace llvm::orc::detail;
+TEST_F(ElementSetTest, Merge) {
+ // Merge into empty set.
+ TestElementSet S;
+ TestElementSet Other({1, 2, 3});
+ EXPECT_TRUE(merge(S, Other));
+ EXPECT_EQ(S, Other);
+
+ // Merge with all-duplicate elements -- no change.
+ EXPECT_FALSE(merge(S, Other));
+ EXPECT_EQ(S, Other);
+
+ // Merge empty into non-empty -- no change.
+ EXPECT_FALSE(merge(S, TestElementSet()));
+ EXPECT_EQ(S, Other);
+
+ // Merge with partial overlap.
+ EXPECT_TRUE(merge(S, TestElementSet({3, 4, 5})));
+ EXPECT_EQ(S, TestElementSet({1, 2, 3, 4, 5}));
+}
+
+TEST_F(ElementSetTest, Remove) {
+ // Remove from empty set.
+ TestElementSet S;
+ EXPECT_FALSE(remove(S, TestElementSet({1, 2})));
+ EXPECT_TRUE(S.empty());
+
+ // Remove empty from non-empty -- no change.
+ S = TestElementSet({1, 2, 3});
+ EXPECT_FALSE(remove(S, TestElementSet()));
+ EXPECT_EQ(S, TestElementSet({1, 2, 3}));
+
+ // Remove with no overlap -- no change.
+ EXPECT_FALSE(remove(S, TestElementSet({4, 5})));
+ EXPECT_EQ(S, TestElementSet({1, 2, 3}));
+
+ // Remove with partial overlap (|this| > |Other| path).
+ S = TestElementSet({1, 2, 3, 4, 5});
+ EXPECT_TRUE(remove(S, TestElementSet({2, 4})));
+ EXPECT_EQ(S, TestElementSet({1, 3, 5}));
+
+ // Remove with partial overlap (|this| <= |Other| path).
+ S = TestElementSet({1, 2});
+ EXPECT_TRUE(remove(S, TestElementSet({2, 3, 4, 5})));
+ EXPECT_EQ(S, TestElementSet({1}));
+
+ // Remove all elements.
+ S = TestElementSet({1, 2});
+ EXPECT_TRUE(remove(S, TestElementSet({1, 2})));
+ EXPECT_TRUE(S.empty());
+}
+
+TEST_F(ElementSetTest, RemoveIf) {
+ // RemoveIf on empty set.
+ TestElementSet S;
+ EXPECT_FALSE(remove_if(S, [](const auto &) { return true; }));
+
+ // RemoveIf with predicate matching nothing.
+ S = TestElementSet({1, 2, 3});
+ EXPECT_FALSE(remove_if(S, [](const auto &) { return false; }));
+ EXPECT_EQ(S, TestElementSet({1, 2, 3}));
+
+ // RemoveIf with predicate matching some elements.
+ EXPECT_TRUE(remove_if(S, [](const auto &E) { return E % 2 == 0; }));
+ EXPECT_EQ(S, TestElementSet({1, 3}));
+
+ // RemoveIf with predicate matching all elements.
+ EXPECT_TRUE(remove_if(S, [](const auto &) { return true; }));
+ EXPECT_TRUE(S.empty());
+}
+
+TEST_F(ContainerElementsMapTest, Merge) {
+ // Merge into empty map.
+ TestContainerElementsMap M;
+ TestContainerElementsMap Other({{0, {1, 2}}, {1, {3}}});
+ EXPECT_TRUE(merge(M, Other));
+ EXPECT_EQ(M, Other);
+
+ // Merge with all-duplicate entries -- no change.
+ EXPECT_FALSE(merge(M, Other));
+ EXPECT_EQ(M, Other);
+
+ // Merge empty -- no change.
+ EXPECT_FALSE(merge(M, TestContainerElementsMap()));
+ EXPECT_EQ(M, Other);
+
+ // Merge with disjoint containers.
+ EXPECT_TRUE(merge(M, TestContainerElementsMap({{2, {4}}})));
+ EXPECT_EQ(M, TestContainerElementsMap({{0, {1, 2}}, {1, {3}}, {2, {4}}}));
+
+ // Merge with overlapping container, new elements.
+ EXPECT_TRUE(merge(M, TestContainerElementsMap({{0, {3}}})));
+ EXPECT_EQ(M, TestContainerElementsMap({{0, {1, 2, 3}}, {1, {3}}, {2, {4}}}));
+}
+
+TEST_F(ContainerElementsMapTest, Remove) {
+ // Remove from empty map.
+ TestContainerElementsMap M;
+ EXPECT_FALSE(remove(M, TestContainerElementsMap({{0, {1}}})));
+ EXPECT_TRUE(M.empty());
+
+ // Remove with no matching container.
+ M = TestContainerElementsMap({{0, {1, 2}}, {1, {3}}});
+ EXPECT_FALSE(remove(M, TestContainerElementsMap({{2, {1}}})));
+ EXPECT_EQ(M, TestContainerElementsMap({{0, {1, 2}}, {1, {3}}}));
+
+ // Remove with no overlap within matching container.
+ EXPECT_FALSE(remove(M, TestContainerElementsMap({{0, {5}}})));
+ EXPECT_EQ(M, TestContainerElementsMap({{0, {1, 2}}, {1, {3}}}));
+
+ // Remove partial elements from a container.
+ EXPECT_TRUE(remove(M, TestContainerElementsMap({{0, {1}}})));
+ EXPECT_EQ(M, TestContainerElementsMap({{0, {2}}, {1, {3}}}));
+
+ // Remove all elements from a container -- container should be cleaned up.
+ EXPECT_TRUE(remove(M, TestContainerElementsMap({{0, {2}}})));
+ EXPECT_EQ(M, TestContainerElementsMap({{1, {3}}}));
+}
+
+TEST_F(ContainerElementsMapTest, Visit) {
+ // Visit empty map -- no-op.
+ TestContainerElementsMap M;
+ EXPECT_FALSE(visit(M, [](auto &, auto &) { return false; }));
+
+ // Visit with no modifications.
+ M = TestContainerElementsMap({{0, {1, 2}}, {1, {3}}});
+ EXPECT_FALSE(visit(M, [](auto &, auto &) { return false; }));
+ EXPECT_EQ(M, TestContainerElementsMap({{0, {1, 2}}, {1, {3}}}));
+
+ // Visit that removes some elements from one container.
+ M = TestContainerElementsMap({{0, {1, 2, 3}}, {1, {4}}});
+ EXPECT_TRUE(visit(M, [](auto &Container, auto &Elements) {
+ if (Container == 0) {
+ Elements.erase(2);
+ return true;
+ }
+ return false;
+ }));
+ EXPECT_EQ(M, TestContainerElementsMap({{0, {1, 3}}, {1, {4}}}));
+
+ // Visit that empties a container -- container should be removed.
+ M = TestContainerElementsMap({{0, {1}}, {1, {2}}});
+ EXPECT_TRUE(visit(M, [](auto &Container, auto &Elements) {
+ if (Container == 0) {
+ Elements.clear();
+ return true;
+ }
+ return false;
+ }));
+ EXPECT_EQ(M, TestContainerElementsMap({{1, {2}}}));
+}
+
TEST_F(WaitingOnGraphTest, ConstructAndDestroyEmpty) {
// Nothing to do here -- we're just testing construction and destruction
// of the WaitingOnGraphTest::G member.
More information about the llvm-commits
mailing list