[llvm] e0b3f45 - [ORC] Automatically suspend and resume lookups that depend on in-use generators.
Lang Hames via llvm-commits
llvm-commits at lists.llvm.org
Mon Jul 31 12:17:24 PDT 2023
Author: Lang Hames
Date: 2023-07-31T12:17:17-07:00
New Revision: e0b3f45d87a6677efb59d8a4f0e6deac9c346c76
URL: https://github.com/llvm/llvm-project/commit/e0b3f45d87a6677efb59d8a4f0e6deac9c346c76
DIFF: https://github.com/llvm/llvm-project/commit/e0b3f45d87a6677efb59d8a4f0e6deac9c346c76.diff
LOG: [ORC] Automatically suspend and resume lookups that depend on in-use generators.
Access to individual DefinitionGenerators is serialized in order to make
generators easier to implement: serializing access means that tryToGenerate
methods don't have to handle concurrent, potentially overlapping, requests.
Prior to this patch serialization was achieved by having each lookup acquire a
lock on each generator, however this causes the lookup thread to block if the
generator is in use. In the common case where many objects reference some
common library symbol that is provided by a generator this may cause many
threads to block concurrently preventing progress on other work.
This patch changes the model so that lookups are automatically suspended if
they need to use a generator that is already in use, and then automatically
resumed once the generator is free. This is achieved by reusing the lookup
suspension machinery that was introduced in 069919c9ba3 for optionally
asynchronous generators.
Added:
Modified:
llvm/include/llvm/ExecutionEngine/Orc/Core.h
llvm/lib/ExecutionEngine/Orc/Core.cpp
llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
Removed:
################################################################################
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/Core.h b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
index 2c755968c45b5f..8851b655b4b286 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/Core.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
@@ -28,6 +28,7 @@
#include "llvm/Support/ExtensibleRTTI.h"
#include <atomic>
+#include <deque>
#include <future>
#include <memory>
#include <vector>
@@ -912,6 +913,8 @@ class LookupState {
/// Definition generators can be attached to JITDylibs to generate new
/// definitions for otherwise unresolved symbols during lookup.
class DefinitionGenerator {
+ friend class ExecutionSession;
+
public:
virtual ~DefinitionGenerator();
@@ -924,6 +927,11 @@ class DefinitionGenerator {
virtual Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
JITDylibLookupFlags JDLookupFlags,
const SymbolLookupSet &LookupSet) = 0;
+
+private:
+ std::mutex M;
+ bool InUse = false;
+ std::deque<LookupState> PendingLookups;
};
/// Represents a JIT'd dynamic library.
@@ -1362,6 +1370,21 @@ class MaterializationTask : public RTTIExtends<MaterializationTask, Task> {
std::unique_ptr<MaterializationResponsibility> MR;
};
+/// Lookups are usually run on the current thread, but in some cases they may
+/// be run as tasks, e.g. if the lookup has been continued from a suspended
+/// state.
+class LookupTask : public RTTIExtends<LookupTask, Task> {
+public:
+ static char ID;
+
+ LookupTask(LookupState LS) : LS(std::move(LS)) {}
+ void printDescription(raw_ostream &OS) override;
+ void run() override;
+
+private:
+ LookupState LS;
+};
+
/// An ExecutionSession represents a running JIT program.
class ExecutionSession {
friend class InProgressLookupFlagsState;
@@ -1718,6 +1741,9 @@ class ExecutionSession {
SymbolLookupSet &Candidates,
SymbolLookupSet *NonCandidates);
+ /// Handle resumption of a lookup after entering a generator.
+ void OL_resumeLookupAfterGeneration(InProgressLookupState &IPLS);
+
/// OL_applyQueryPhase1 is an optionally re-startable loop for triggering
/// definition generation. It is called when a lookup is performed, and again
/// each time that LookupState::continueLookup is called.
diff --git a/llvm/lib/ExecutionEngine/Orc/Core.cpp b/llvm/lib/ExecutionEngine/Orc/Core.cpp
index 95ceba78a2d4d8..f7ebffdc898a2b 100644
--- a/llvm/lib/ExecutionEngine/Orc/Core.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/Core.cpp
@@ -31,6 +31,7 @@ char SymbolsCouldNotBeRemoved::ID = 0;
char MissingSymbolDefinitions::ID = 0;
char UnexpectedSymbolDefinitions::ID = 0;
char MaterializationTask::ID = 0;
+char LookupTask::ID = 0;
RegisterDependenciesFunction NoDependenciesToRegister =
RegisterDependenciesFunction();
@@ -529,11 +530,16 @@ class InProgressLookupState {
SymbolLookupSet LookupSet;
SymbolState RequiredState;
- std::unique_lock<std::mutex> GeneratorLock;
size_t CurSearchOrderIndex = 0;
bool NewJITDylib = true;
SymbolLookupSet DefGeneratorCandidates;
SymbolLookupSet DefGeneratorNonCandidates;
+
+ enum {
+ NotInGenerator, // Not currently using a generator.
+ ResumedForGenerator, // Resumed after being auto-suspended before generator.
+ InGenerator // Currently using generator.
+ } GenState = NotInGenerator;
std::vector<std::weak_ptr<DefinitionGenerator>> CurDefGeneratorStack;
};
@@ -547,15 +553,11 @@ class InProgressLookupFlagsState : public InProgressLookupState {
OnComplete(std::move(OnComplete)) {}
void complete(std::unique_ptr<InProgressLookupState> IPLS) override {
- GeneratorLock = {}; // Unlock and release.
auto &ES = SearchOrder.front().first->getExecutionSession();
ES.OL_completeLookupFlags(std::move(IPLS), std::move(OnComplete));
}
- void fail(Error Err) override {
- GeneratorLock = {}; // Unlock and release.
- OnComplete(std::move(Err));
- }
+ void fail(Error Err) override { OnComplete(std::move(Err)); }
private:
unique_function<void(Expected<SymbolFlagsMap>)> OnComplete;
@@ -574,14 +576,12 @@ class InProgressFullLookupState : public InProgressLookupState {
}
void complete(std::unique_ptr<InProgressLookupState> IPLS) override {
- GeneratorLock = {}; // Unlock and release.
auto &ES = SearchOrder.front().first->getExecutionSession();
ES.OL_completeLookup(std::move(IPLS), std::move(Q),
std::move(RegisterDependencies));
}
void fail(Error Err) override {
- GeneratorLock = {};
Q->detach();
Q->handleFailed(std::move(Err));
}
@@ -638,7 +638,19 @@ void LookupState::continueLookup(Error Err) {
ES.OL_applyQueryPhase1(std::move(IPLS), std::move(Err));
}
-DefinitionGenerator::~DefinitionGenerator() = default;
+DefinitionGenerator::~DefinitionGenerator() {
+ std::deque<LookupState> LookupsToFail;
+ {
+ std::lock_guard<std::mutex> Lock(M);
+ std::swap(PendingLookups, LookupsToFail);
+ InUse = false;
+ }
+
+ for (auto &LS : LookupsToFail)
+ LS.continueLookup(make_error<StringError>(
+ "Query waiting on DefinitionGenerator that was destroyed",
+ inconvertibleErrorCode()));
+}
JITDylib::~JITDylib() {
LLVM_DEBUG(dbgs() << "Destroying JITDylib " << getName() << "\n");
@@ -677,6 +689,10 @@ ResourceTrackerSP JITDylib::createResourceTracker() {
}
void JITDylib::removeGenerator(DefinitionGenerator &G) {
+ // DefGenerator moved into TmpDG to ensure that it's destroyed outside the
+ // session lock (since it may have to send errors to pending queries).
+ std::shared_ptr<DefinitionGenerator> TmpDG;
+
ES.runSessionLocked([&] {
assert(State == Open && "JD is defunct");
auto I = llvm::find_if(DefGenerators,
@@ -684,6 +700,7 @@ void JITDylib::removeGenerator(DefinitionGenerator &G) {
return H.get() == &G;
});
assert(I != DefGenerators.end() && "Generator not found");
+ TmpDG = std::move(*I);
DefGenerators.erase(I);
});
}
@@ -1903,6 +1920,10 @@ void MaterializationTask::printDescription(raw_ostream &OS) {
void MaterializationTask::run() { MU->materialize(std::move(MR)); }
+void LookupTask::printDescription(raw_ostream &OS) { OS << "Lookup task"; }
+
+void LookupTask::run() { LS.continueLookup(Error::success()); }
+
ExecutionSession::ExecutionSession(std::unique_ptr<ExecutorProcessControl> EPC)
: EPC(std::move(EPC)) {
// Associated EPC and this.
@@ -2406,6 +2427,37 @@ Error ExecutionSession::IL_updateCandidatesFor(
});
}
+void ExecutionSession::OL_resumeLookupAfterGeneration(
+ InProgressLookupState &IPLS) {
+
+ assert(IPLS.GenState != InProgressLookupState::NotInGenerator &&
+ "Should not be called for not-in-generator lookups");
+ IPLS.GenState = InProgressLookupState::NotInGenerator;
+
+ LookupState LS;
+
+ if (auto DG = IPLS.CurDefGeneratorStack.back().lock()) {
+ IPLS.CurDefGeneratorStack.pop_back();
+ std::lock_guard<std::mutex> Lock(DG->M);
+
+ // If there are no pending lookups then mark the generator as free and
+ // return.
+ if (DG->PendingLookups.empty()) {
+ DG->InUse = false;
+ return;
+ }
+
+ // Otherwise resume the next lookup.
+ LS = std::move(DG->PendingLookups.front());
+ DG->PendingLookups.pop_front();
+ }
+
+ if (LS.IPLS) {
+ LS.IPLS->GenState = InProgressLookupState::ResumedForGenerator;
+ dispatchTask(std::make_unique<LookupTask>(std::move(LS)));
+ }
+}
+
void ExecutionSession::OL_applyQueryPhase1(
std::unique_ptr<InProgressLookupState> IPLS, Error Err) {
@@ -2422,6 +2474,12 @@ void ExecutionSession::OL_applyQueryPhase1(
<< IPLS->DefGeneratorNonCandidates << "\n";
});
+ if (IPLS->GenState == InProgressLookupState::InGenerator)
+ OL_resumeLookupAfterGeneration(*IPLS);
+
+ assert(IPLS->GenState != InProgressLookupState::InGenerator &&
+ "Lookup should not be in InGenerator state here");
+
// FIXME: We should attach the query as we go: This provides a result in a
// single pass in the common case where all symbols have already reached the
// required state. The query could be detached again in the 'fail' method on
@@ -2447,10 +2505,6 @@ void ExecutionSession::OL_applyQueryPhase1(
// If we've just reached a new JITDylib then perform some setup.
if (IPLS->NewJITDylib) {
-
- // Acquire the generator lock for this JITDylib.
- IPLS->GeneratorLock = std::unique_lock<std::mutex>(JD.GeneratorsMutex);
-
// Add any non-candidates from the last JITDylib (if any) back on to the
// list of definition candidates for this JITDylib, reset definition
// non-candidates to the empty set.
@@ -2488,6 +2542,13 @@ void ExecutionSession::OL_applyQueryPhase1(
dbgs() << " Remaining candidates = " << IPLS->DefGeneratorCandidates
<< "\n";
});
+
+ // If this lookup was resumed after auto-suspension but all candidates
+ // have already been generated (by some previous call to the generator)
+ // treat the lookup as if it had completed generation.
+ if (IPLS->GenState == InProgressLookupState::ResumedForGenerator &&
+ IPLS->DefGeneratorCandidates.empty())
+ OL_resumeLookupAfterGeneration(*IPLS);
});
// If we encountered an error while filtering generation candidates then
@@ -2509,13 +2570,32 @@ void ExecutionSession::OL_applyQueryPhase1(
while (!IPLS->CurDefGeneratorStack.empty() &&
!IPLS->DefGeneratorCandidates.empty()) {
auto DG = IPLS->CurDefGeneratorStack.back().lock();
- IPLS->CurDefGeneratorStack.pop_back();
if (!DG)
return IPLS->fail(make_error<StringError>(
"DefinitionGenerator removed while lookup in progress",
inconvertibleErrorCode()));
+ // At this point the lookup is in either the NotInGenerator state, or in
+ // the ResumedForGenerator state.
+ // If this lookup is in the NotInGenerator state then check whether the
+ // generator is in use. If the generator is not in use then move the
+ // lookup to the InGenerator state and continue. If the generator is
+ // already in use then just add this lookup to the pending lookups list
+ // and bail out.
+ // If this lookup is in the ResumedForGenerator state then just move it
+ // to InGenerator and continue.
+ if (IPLS->GenState == InProgressLookupState::NotInGenerator) {
+ std::lock_guard<std::mutex> Lock(DG->M);
+ if (DG->InUse) {
+ DG->PendingLookups.push_back(std::move(IPLS));
+ return;
+ }
+ DG->InUse = true;
+ }
+
+ IPLS->GenState = InProgressLookupState::InGenerator;
+
auto K = IPLS->K;
auto &LookupSet = IPLS->DefGeneratorCandidates;
@@ -2528,6 +2608,11 @@ void ExecutionSession::OL_applyQueryPhase1(
IPLS = std::move(LS.IPLS);
}
+ // If the lookup returned then pop the generator stack and unblock the
+ // next lookup on this generator (if any).
+ if (IPLS)
+ OL_resumeLookupAfterGeneration(*IPLS);
+
// If there was an error then fail the query.
if (Err) {
LLVM_DEBUG({
diff --git a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
index f0adcae4344801..a136e941fa9488 100644
--- a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
@@ -13,6 +13,7 @@
#include "llvm/ExecutionEngine/Orc/Shared/OrcError.h"
#include "llvm/Testing/Support/Error.h"
+#include <deque>
#include <set>
#include <thread>
@@ -1146,23 +1147,43 @@ TEST_F(CoreAPIsStandardTest, GeneratorTest) {
<< "Expected fallback def for Bar to be equal to BarSym";
}
-TEST_F(CoreAPIsStandardTest, AsynchronousGeneratorTest) {
- class TestGenerator : public DefinitionGenerator {
- public:
- TestGenerator(LookupState &TLS) : TLS(TLS) {}
- Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
- JITDylibLookupFlags JDLookupFlags,
- const SymbolLookupSet &Name) override {
- TLS = std::move(LS);
- return Error::success();
- }
-
- private:
- LookupState &TLS;
+/// By default appends LookupStates to a queue.
+/// Behavior can be overridden by setting TryToGenerateOverride.
+class SimpleAsyncGenerator : public DefinitionGenerator {
+public:
+ struct SuspendedLookupInfo {
+ LookupState LS;
+ LookupKind K;
+ JITDylibSP JD;
+ JITDylibLookupFlags JDLookupFlags;
+ SymbolLookupSet Names;
};
- LookupState LS;
- JD.addGenerator(std::make_unique<TestGenerator>(LS));
+ unique_function<Error(LookupState &, LookupKind, JITDylib &,
+ JITDylibLookupFlags, const SymbolLookupSet &)>
+ TryToGenerateOverride;
+
+ Error tryToGenerate(LookupState &LS, LookupKind K, JITDylib &JD,
+ JITDylibLookupFlags JDLookupFlags,
+ const SymbolLookupSet &Names) override {
+ if (TryToGenerateOverride)
+ return TryToGenerateOverride(LS, K, JD, JDLookupFlags, Names);
+ Lookup = SuspendedLookupInfo{std::move(LS), K, &JD, JDLookupFlags, Names};
+ return Error::success();
+ }
+
+ SuspendedLookupInfo takeLookup() {
+ std::optional<SuspendedLookupInfo> Tmp;
+ std::swap(Tmp, Lookup);
+ return std::move(*Tmp);
+ }
+
+ std::optional<SuspendedLookupInfo> Lookup;
+};
+
+TEST_F(CoreAPIsStandardTest, SimpleAsynchronousGeneratorTest) {
+
+ auto &G = JD.addGenerator(std::make_unique<SimpleAsyncGenerator>());
bool LookupCompleted = false;
@@ -1171,28 +1192,137 @@ TEST_F(CoreAPIsStandardTest, AsynchronousGeneratorTest) {
SymbolState::Ready,
[&](Expected<SymbolMap> Result) {
LookupCompleted = true;
- if (!Result) {
- ADD_FAILURE() << "Lookup failed unexpected";
- logAllUnhandledErrors(Result.takeError(), errs(), "");
- return;
- }
-
- EXPECT_EQ(Result->size(), 1U) << "Unexpected number of results";
- EXPECT_EQ(Result->count(Foo), 1U) << "Expected result for Foo";
- EXPECT_EQ((*Result)[Foo].getAddress(), FooSym.getAddress())
- << "Bad result for Foo";
+ EXPECT_THAT_EXPECTED(Result, Succeeded());
+ if (Result)
+ EXPECT_EQ(*Result, SymbolMap({{Foo, FooSym}}));
},
NoDependenciesToRegister);
EXPECT_FALSE(LookupCompleted);
cantFail(JD.define(absoluteSymbols({{Foo, FooSym}})));
-
- LS.continueLookup(Error::success());
+ G.takeLookup().LS.continueLookup(Error::success());
EXPECT_TRUE(LookupCompleted);
}
+TEST_F(CoreAPIsStandardTest, BlockedGeneratorAutoSuspensionTest) {
+ // Test that repeated lookups while a generator is in use cause automatic
+ // lookup suspension / resumption.
+
+ auto &G = JD.addGenerator(std::make_unique<SimpleAsyncGenerator>());
+
+ bool Lookup1Completed = false;
+ bool Lookup2Completed = false;
+ bool Lookup3Completed = false;
+ bool Lookup4Completed = false;
+
+ // Add lookup 1.
+ //
+ // Tests that tryToGenerate-suspended lookups resume auto-suspended lookups
+ // when the tryToGenerate-suspended lookup continues (i.e. the call to
+ // OL_resumeLookupAfterGeneration at the top of OL_applyQueryPhase1).
+ ES.lookup(
+ LookupKind::Static, makeJITDylibSearchOrder(&JD), SymbolLookupSet(Foo),
+ SymbolState::Ready,
+ [&](Expected<SymbolMap> Result) {
+ Lookup1Completed = true;
+ EXPECT_THAT_EXPECTED(Result, Succeeded());
+ if (Result)
+ EXPECT_EQ(*Result, SymbolMap({{Foo, FooSym}}));
+ },
+ NoDependenciesToRegister);
+
+ // The generator should immediately see the first lookup.
+ EXPECT_NE(G.Lookup, std::nullopt);
+
+ // Add lookup 2.
+ //
+ // Tests that lookups that pass through tryToGenerate without being captured
+ // resume auto-suspended lookups. We set a one-shot TryToGenerateOverride to
+ // prevent capture of lookup 2 by tryToGenerate. This tests the call to
+ // OL_resumeLookupAfterGeneration inside the generator loop.
+ G.TryToGenerateOverride = [&](LookupState &LS, LookupKind K, JITDylib &JD,
+ JITDylibLookupFlags JDLookupFlags,
+ const SymbolLookupSet &Names) -> Error {
+ cantFail(JD.define(absoluteSymbols({{Bar, BarSym}})));
+ G.TryToGenerateOverride = nullptr;
+ return Error::success();
+ };
+
+ ES.lookup(
+ LookupKind::Static, makeJITDylibSearchOrder(&JD), SymbolLookupSet(Bar),
+ SymbolState::Ready,
+ [&](Expected<SymbolMap> Result) {
+ Lookup2Completed = true;
+ EXPECT_THAT_EXPECTED(Result, Succeeded());
+ if (Result)
+ EXPECT_EQ(*Result, SymbolMap({{Bar, BarSym}}));
+ },
+ NoDependenciesToRegister);
+
+ // Add lookup 3.
+ //
+ // Test that if a lookup's symbols have already been generated (and it
+ // consequently skips the generator loop entirely) it still resumes the next
+ // suspended lookup. This tests the call to OL_resumeLookupAfterGeneration
+ // just above the generator loop.
+ ES.lookup(
+ LookupKind::Static, makeJITDylibSearchOrder(&JD), SymbolLookupSet(Bar),
+ SymbolState::Ready,
+ [&](Expected<SymbolMap> Result) {
+ Lookup3Completed = true;
+ EXPECT_THAT_EXPECTED(Result, Succeeded());
+ if (Result)
+ EXPECT_EQ(*Result, SymbolMap({{Bar, BarSym}}));
+ },
+ NoDependenciesToRegister);
+
+ // Add lookup 4.
+ //
+ // This is just used to verify that lookup 3 triggered resumption of the next
+ // lookup as expected.
+ ES.lookup(
+ LookupKind::Static, makeJITDylibSearchOrder(&JD), SymbolLookupSet(Baz),
+ SymbolState::Ready,
+ [&](Expected<SymbolMap> Result) {
+ Lookup4Completed = true;
+ EXPECT_THAT_EXPECTED(Result, Succeeded());
+ if (Result)
+ EXPECT_EQ(*Result, SymbolMap({{Baz, BazSym}}));
+ },
+ NoDependenciesToRegister);
+
+ // All lookups have been started, but none should have been completed yet.
+ EXPECT_FALSE(Lookup1Completed);
+ EXPECT_FALSE(Lookup2Completed);
+ EXPECT_FALSE(Lookup3Completed);
+ EXPECT_FALSE(Lookup4Completed);
+
+ // Start continuing lookups.
+
+ // First Define foo and continue lookup 1. We expect this to complete lookups
+ // 1, 2 and 3: the TryToGenerateOverride set above will define bar, which will
+ // allow both 2 and 3 to complete.
+ cantFail(JD.define(absoluteSymbols({{Foo, FooSym}})));
+ G.takeLookup().LS.continueLookup(Error::success());
+
+ EXPECT_TRUE(Lookup1Completed);
+ EXPECT_TRUE(Lookup2Completed);
+ EXPECT_TRUE(Lookup3Completed);
+ EXPECT_FALSE(Lookup4Completed);
+ EXPECT_NE(G.Lookup, std::nullopt);
+
+ // Check that the most recently captured lookup is lookup 4 (for baz).
+ if (G.Lookup)
+ EXPECT_EQ(G.Lookup->Names.begin()->first, Baz);
+
+ cantFail(JD.define(absoluteSymbols({{Baz, BazSym}})));
+ G.takeLookup().LS.continueLookup(Error::success());
+
+ EXPECT_TRUE(Lookup4Completed);
+}
+
TEST_F(CoreAPIsStandardTest, FailResolution) {
auto MU = std::make_unique<SimpleMaterializationUnit>(
SymbolFlagsMap({{Foo, JITSymbolFlags::Exported | JITSymbolFlags::Weak},
More information about the llvm-commits
mailing list