[llvm] [ORC] Refactor some WaitingOnGraph loops into named methods. NFCI. (PR #184775)
Lang Hames via llvm-commits
llvm-commits at lists.llvm.org
Thu Mar 5 04:09:32 PST 2026
https://github.com/lhames updated https://github.com/llvm/llvm-project/pull/184775
>From ea898779840f7a2d305c6eb30ad2a4c5fbe69f2d Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames at gmail.com>
Date: Thu, 5 Mar 2026 22:03:26 +1100
Subject: [PATCH 1/2] [ORC] Refactor some WaitingOnGraph loops into named
methods. NFCI.
Adds new methods to SuperNode for updating ElemToSuperNodeMap values,
use it to tidy up some loops.
---
.../llvm/ExecutionEngine/Orc/WaitingOnGraph.h | 82 +++++++++++++------
1 file changed, 55 insertions(+), 27 deletions(-)
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/WaitingOnGraph.h b/llvm/include/llvm/ExecutionEngine/Orc/WaitingOnGraph.h
index 88371b557fe7c..0f3f67f8c4fd5 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/WaitingOnGraph.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/WaitingOnGraph.h
@@ -88,6 +88,13 @@ template <typename ContainerIdT, typename ElementIdT> class WaitingOnGraph {
using ElementSet = DenseSet<ElementId>;
using ContainerElementsMap = DenseMap<ContainerId, ElementSet>;
+ class SuperNode;
+
+private:
+ using ElemToSuperNodeMap =
+ DenseMap<ContainerId, DenseMap<ElementId, SuperNode *>>;
+
+public:
class SuperNode {
friend class WaitingOnGraph;
friend class WaitingOnGraphTest;
@@ -103,12 +110,50 @@ template <typename ContainerIdT, typename ElementIdT> class WaitingOnGraph {
private:
ContainerElementsMap Defs;
ContainerElementsMap Deps;
+
+ ElemToSuperNodeMap *RegisteredElemToSN = nullptr;
+
+ /// Add a mapping from the Defs in this SuperNode to SN (which may or may
+ /// not be the same as this).
+ void mapDefsTo(ElemToSuperNodeMap &ElemToSN, SuperNode *SN) {
+ assert(!Defs.empty() && "Empty defs!?");
+ for (auto &[Container, Elements] : Defs) {
+ assert(!Elements.empty() && "Empty elements for container?");
+ auto &ContainerElemToSN = ElemToSN[Container];
+ for (auto &Elem : Elements)
+ ContainerElemToSN[Elem] = SN;
+ }
+ assert((!SN->RegisteredElemToSN || SN->RegisteredElemToSN == &ElemToSN) &&
+ "SN defs split across maps");
+ SN->RegisteredElemToSN = &ElemToSN;
+ }
+
+ /// Add a mapping from the Defs in this SuperNode to this.
+ /// (Equivalent to `SN.mapDefsTo(ElemToSN, &SN);`)
+ void mapDefsToThis(ElemToSuperNodeMap &ElemToSN) {
+ mapDefsTo(ElemToSN, this);
+ }
+
+ /// Remove a mapping from the Defs in this SuperNode from the registered
+ /// ElemToSuperNodeMap. The mapping must already exist.
+ void unmapDefsFromThis() {
+ assert(RegisteredElemToSN && "No registered ElemToSuperNodeMap");
+ for (auto &[Container, Elements] : Defs) {
+ auto I = RegisteredElemToSN->find(Container);
+ assert(I != RegisteredElemToSN->end() && "Container not in map");
+ auto &ContainerElemToSN = I->second;
+ for (auto &Elem : Elements) {
+ assert(ContainerElemToSN[Elem] == this && "Mapping not present");
+ ContainerElemToSN.erase(Elem);
+ }
+ if (ContainerElemToSN.empty())
+ RegisteredElemToSN->erase(I);
+ }
+ RegisteredElemToSN = nullptr;
+ }
};
private:
- using ElemToSuperNodeMap =
- DenseMap<ContainerId, DenseMap<ElementId, SuperNode *>>;
-
using SuperNodeDepsMap = DenseMap<SuperNode *, DenseSet<SuperNode *>>;
class Coalescer {
@@ -143,12 +188,9 @@ template <typename ContainerIdT, typename ElementIdT> class WaitingOnGraph {
"Elements of SNs should be new to the coalescer");
auto H = getHash(SN->Deps);
if (auto *CanonicalSN = findCanonicalSuperNode(H, SN->Deps)) {
- for (auto &[Container, Elems] : SN->Defs) {
+ SN->mapDefsTo(ElemToSN, CanonicalSN);
+ for (auto &[Container, Elems] : SN->Defs)
CanonicalSN->Defs[Container].insert(Elems.begin(), Elems.end());
- auto &ContainerElemToSN = ElemToSN[Container];
- for (auto &Elem : Elems)
- ContainerElemToSN[Elem] = CanonicalSN;
- }
std::swap(SN, SNs.back());
SNs.pop_back();
} else {
@@ -286,13 +328,8 @@ template <typename ContainerIdT, typename ElementIdT> class WaitingOnGraph {
static SimplifyResult simplify(std::vector<std::unique_ptr<SuperNode>> SNs) {
// Build ElemToSN map.
ElemToSuperNodeMap ElemToSN;
- for (auto &SN : SNs) {
- for (auto &[Container, Elements] : SN->Defs) {
- auto &ContainerElemToSN = ElemToSN[Container];
- for (auto &E : Elements)
- ContainerElemToSN[E] = SN.get();
- }
- }
+ for (auto &SN : SNs)
+ SN->mapDefsToThis(ElemToSN);
SuperNodeDepsMap SuperNodeDeps;
hoistDeps(SuperNodeDeps, SNs, ElemToSN);
@@ -380,11 +417,7 @@ template <typename ContainerIdT, typename ElementIdT> class WaitingOnGraph {
// Update ElemToPendingSN for the remaining elements.
for (auto &SN : NewSNs) {
- for (auto &[Container, Elems] : SN->Defs) {
- auto &Row = ElemToPendingSN[Container];
- for (auto &Elem : Elems)
- Row[Elem] = SN.get();
- }
+ SN->mapDefsToThis(ElemToPendingSN);
PendingSNs.push_back(std::move(SN));
}
@@ -692,13 +725,8 @@ template <typename ContainerIdT, typename ElementIdT> class WaitingOnGraph {
}
// Update ElemToSNs (if passed) to remove elements pointing at SN.
- for (auto *SN : ToRemoveFromElemToSNs) {
- for (auto &[Container, Elems] : SN->defs()) {
- auto &Row = (*ElemToSNs)[Container];
- for (auto &Elem : Elems)
- Row.erase(Elem);
- }
- }
+ for (auto *SN : ToRemoveFromElemToSNs)
+ SN->unmapDefsFromThis();
}
std::vector<std::unique_ptr<SuperNode>> PendingSNs;
>From 15e65007483592f81fbc94ec13bf6f094c839cd4 Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames at gmail.com>
Date: Thu, 5 Mar 2026 23:08:58 +1100
Subject: [PATCH 2/2] Allow assertion to be relaxed when we're abandoning a
stale mapping.
---
.../llvm/ExecutionEngine/Orc/WaitingOnGraph.h | 13 ++++++++-----
1 file changed, 8 insertions(+), 5 deletions(-)
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/WaitingOnGraph.h b/llvm/include/llvm/ExecutionEngine/Orc/WaitingOnGraph.h
index 0f3f67f8c4fd5..f95838b430342 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/WaitingOnGraph.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/WaitingOnGraph.h
@@ -115,7 +115,8 @@ template <typename ContainerIdT, typename ElementIdT> class WaitingOnGraph {
/// Add a mapping from the Defs in this SuperNode to SN (which may or may
/// not be the same as this).
- void mapDefsTo(ElemToSuperNodeMap &ElemToSN, SuperNode *SN) {
+ void mapDefsTo(ElemToSuperNodeMap &ElemToSN, SuperNode *SN,
+ bool AbandonOldMapping = false) {
assert(!Defs.empty() && "Empty defs!?");
for (auto &[Container, Elements] : Defs) {
assert(!Elements.empty() && "Empty elements for container?");
@@ -123,15 +124,17 @@ template <typename ContainerIdT, typename ElementIdT> class WaitingOnGraph {
for (auto &Elem : Elements)
ContainerElemToSN[Elem] = SN;
}
- assert((!SN->RegisteredElemToSN || SN->RegisteredElemToSN == &ElemToSN) &&
+ assert((AbandonOldMapping || !SN->RegisteredElemToSN ||
+ SN->RegisteredElemToSN == &ElemToSN) &&
"SN defs split across maps");
SN->RegisteredElemToSN = &ElemToSN;
}
/// Add a mapping from the Defs in this SuperNode to this.
/// (Equivalent to `SN.mapDefsTo(ElemToSN, &SN);`)
- void mapDefsToThis(ElemToSuperNodeMap &ElemToSN) {
- mapDefsTo(ElemToSN, this);
+ void mapDefsToThis(ElemToSuperNodeMap &ElemToSN,
+ bool AbandonOldMapping = false) {
+ mapDefsTo(ElemToSN, this, AbandonOldMapping);
}
/// Remove a mapping from the Defs in this SuperNode from the registered
@@ -417,7 +420,7 @@ template <typename ContainerIdT, typename ElementIdT> class WaitingOnGraph {
// Update ElemToPendingSN for the remaining elements.
for (auto &SN : NewSNs) {
- SN->mapDefsToThis(ElemToPendingSN);
+ SN->mapDefsToThis(ElemToPendingSN, /* AbandonOldMapping = */ true);
PendingSNs.push_back(std::move(SN));
}
More information about the llvm-commits
mailing list