[llvm] 5344c88 - [ORC] Generalize materialization dispatch to task dispatch.
Lang Hames via llvm-commits
llvm-commits at lists.llvm.org
Sun May 9 19:38:23 PDT 2021
Author: Lang Hames
Date: 2021-05-09T19:19:39-07:00
New Revision: 5344c88dcb2845f6a12cd0992deab1448b4d1419
URL: https://github.com/llvm/llvm-project/commit/5344c88dcb2845f6a12cd0992deab1448b4d1419
DIFF: https://github.com/llvm/llvm-project/commit/5344c88dcb2845f6a12cd0992deab1448b4d1419.diff
LOG: [ORC] Generalize materialization dispatch to task dispatch.
Generalizing this API allows work to be distributed more evenly. In particular,
query callbacks can now be dispatched (rather than running immediately on the
thread that satisfied the query). This avoids the pathalogical case where an
operation on one thread satisfies many queries simultaneously, causing large
amounts of work to be run on that thread while other threads potentially sit
idle.
Added:
Modified:
llvm/include/llvm/ExecutionEngine/Orc/Core.h
llvm/lib/ExecutionEngine/Orc/Core.cpp
llvm/lib/ExecutionEngine/Orc/LLJIT.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 4a4b58ed32e30..c37361fbe57e9 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/Core.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
@@ -22,6 +22,7 @@
#include "llvm/ExecutionEngine/Orc/SymbolStringPool.h"
#include "llvm/ExecutionEngine/OrcV1Deprecation.h"
#include "llvm/Support/Debug.h"
+#include "llvm/Support/ExtensibleRTTI.h"
#include <atomic>
#include <memory>
@@ -637,6 +638,8 @@ class MaterializationUnit {
friend class JITDylib;
public:
+ static char ID;
+
MaterializationUnit(SymbolFlagsMap InitalSymbolFlags,
SymbolStringPtr InitSymbol)
: SymbolFlags(std::move(InitalSymbolFlags)),
@@ -1217,6 +1220,37 @@ class Platform {
const DenseMap<JITDylib *, SymbolLookupSet> &InitSyms);
};
+/// Represents an abstract task for ORC to run.
+class Task : public RTTIExtends<Task, RTTIRoot> {
+public:
+ static char ID;
+
+ /// Description of the task to be performed. Used for logging.
+ virtual void printDescription(raw_ostream &OS) = 0;
+
+ /// Run the task.
+ virtual void run() = 0;
+
+private:
+ void anchor() override;
+};
+
+/// A materialization task.
+class MaterializationTask : public RTTIExtends<MaterializationTask, Task> {
+public:
+ static char ID;
+
+ MaterializationTask(std::unique_ptr<MaterializationUnit> MU,
+ std::unique_ptr<MaterializationResponsibility> MR)
+ : MU(std::move(MU)), MR(std::move(MR)) {}
+ void printDescription(raw_ostream &OS) override;
+ void run() override;
+
+private:
+ std::unique_ptr<MaterializationUnit> MU;
+ std::unique_ptr<MaterializationResponsibility> MR;
+};
+
/// An ExecutionSession represents a running JIT program.
class ExecutionSession {
friend class InProgressLookupFlagsState;
@@ -1230,10 +1264,8 @@ class ExecutionSession {
/// For reporting errors.
using ErrorReporter = std::function<void(Error)>;
- /// For dispatching MaterializationUnit::materialize calls.
- using DispatchMaterializationFunction =
- std::function<void(std::unique_ptr<MaterializationUnit> MU,
- std::unique_ptr<MaterializationResponsibility> MR)>;
+ /// For dispatching ORC tasks (typically materialization tasks).
+ using DispatchTaskFunction = std::function<void(std::unique_ptr<Task> T)>;
/// Construct an ExecutionSession.
///
@@ -1306,10 +1338,9 @@ class ExecutionSession {
/// Unhandled errors can be sent here to log them.
void reportError(Error Err) { ReportError(std::move(Err)); }
- /// Set the materialization dispatch function.
- ExecutionSession &setDispatchMaterialization(
- DispatchMaterializationFunction DispatchMaterialization) {
- this->DispatchMaterialization = std::move(DispatchMaterialization);
+ /// Set the task dispatch function.
+ ExecutionSession &setDispatchTask(DispatchTaskFunction DispatchTask) {
+ this->DispatchTask = std::move(DispatchTask);
return *this;
}
@@ -1384,12 +1415,10 @@ class ExecutionSession {
SymbolState RequiredState = SymbolState::Ready);
/// Materialize the given unit.
- void
- dispatchMaterialization(std::unique_ptr<MaterializationUnit> MU,
- std::unique_ptr<MaterializationResponsibility> MR) {
- assert(MU && "MU must be non-null");
- DEBUG_WITH_TYPE("orc", dumpDispatchInfo(MR->getTargetJITDylib(), *MU));
- DispatchMaterialization(std::move(MU), std::move(MR));
+ void dispatchTask(std::unique_ptr<Task> T) {
+ assert(T && "T must be non-null");
+ DEBUG_WITH_TYPE("orc", dumpDispatchInfo(*T));
+ DispatchTask(std::move(T));
}
/// Dump the state of all the JITDylibs in this session.
@@ -1400,11 +1429,7 @@ class ExecutionSession {
logAllUnhandledErrors(std::move(Err), errs(), "JIT session error: ");
}
- static void materializeOnCurrentThread(
- std::unique_ptr<MaterializationUnit> MU,
- std::unique_ptr<MaterializationResponsibility> MR) {
- MU->materialize(std::move(MR));
- }
+ static void runOnCurrentThread(std::unique_ptr<Task> T) { T->run(); }
void dispatchOutstandingMUs();
@@ -1474,7 +1499,7 @@ class ExecutionSession {
const SymbolDependenceMap &Dependencies);
#ifndef NDEBUG
- void dumpDispatchInfo(JITDylib &JD, MaterializationUnit &MU);
+ void dumpDispatchInfo(Task &T);
#endif // NDEBUG
mutable std::recursive_mutex SessionMutex;
@@ -1482,8 +1507,7 @@ class ExecutionSession {
std::shared_ptr<SymbolStringPool> SSP;
std::unique_ptr<Platform> P;
ErrorReporter ReportError = logErrorsToStdErr;
- DispatchMaterializationFunction DispatchMaterialization =
- materializeOnCurrentThread;
+ DispatchTaskFunction DispatchTask = runOnCurrentThread;
std::vector<ResourceManager *> ResourceManagers;
diff --git a/llvm/lib/ExecutionEngine/Orc/Core.cpp b/llvm/lib/ExecutionEngine/Orc/Core.cpp
index 1e3384d4787dd..4300a0bbd1bc0 100644
--- a/llvm/lib/ExecutionEngine/Orc/Core.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/Core.cpp
@@ -29,6 +29,8 @@ char SymbolsNotFound::ID = 0;
char SymbolsCouldNotBeRemoved::ID = 0;
char MissingSymbolDefinitions::ID = 0;
char UnexpectedSymbolDefinitions::ID = 0;
+char Task::ID = 0;
+char MaterializationTask::ID = 0;
RegisterDependenciesFunction NoDependenciesToRegister =
RegisterDependenciesFunction();
@@ -750,7 +752,8 @@ Error JITDylib::replace(MaterializationResponsibility &FromMR,
if (MustRunMU) {
assert(MustRunMR && "MustRunMU set implies MustRunMR set");
- ES.dispatchMaterialization(std::move(MustRunMU), std::move(MustRunMR));
+ ES.dispatchTask(std::make_unique<MaterializationTask>(
+ std::move(MustRunMU), std::move(MustRunMR)));
} else {
assert(!MustRunMR && "MustRunMU unset implies MustRunMR unset");
}
@@ -1730,6 +1733,15 @@ Expected<DenseMap<JITDylib *, SymbolMap>> Platform::lookupInitSymbols(
return std::move(CompoundResult);
}
+void Task::anchor() {}
+
+void MaterializationTask::printDescription(raw_ostream &OS) {
+ OS << "Materialization task: " << MU->getName() << " in "
+ << MR->getTargetJITDylib().getName() << "\n";
+}
+
+void MaterializationTask::run() { MU->materialize(std::move(MR)); }
+
ExecutionSession::ExecutionSession(std::shared_ptr<SymbolStringPool> SSP)
: SSP(SSP ? std::move(SSP) : std::make_shared<SymbolStringPool>()) {}
@@ -2003,7 +2015,8 @@ void ExecutionSession::dispatchOutstandingMUs() {
assert(JMU->first && "No MU?");
LLVM_DEBUG(dbgs() << " Dispatching \"" << JMU->first->getName() << "\"\n");
- dispatchMaterialization(std::move(JMU->first), std::move(JMU->second));
+ dispatchTask(std::make_unique<MaterializationTask>(std::move(JMU->first),
+ std::move(JMU->second)));
}
LLVM_DEBUG(dbgs() << "Done dispatching MaterializationUnits.\n");
}
@@ -2776,9 +2789,10 @@ void ExecutionSession::OL_addDependenciesForAll(
}
#ifndef NDEBUG
-void ExecutionSession::dumpDispatchInfo(JITDylib &JD, MaterializationUnit &MU) {
+void ExecutionSession::dumpDispatchInfo(Task &T) {
runSessionLocked([&]() {
- dbgs() << "Dispatching " << MU << " for " << JD.getName() << "\n";
+ dbgs() << "Dispatching: ";
+ T.printDescription(dbgs());
});
}
#endif // NDEBUG
diff --git a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
index 9e2177e1eaea0..a79882a04e708 100644
--- a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
@@ -1126,20 +1126,16 @@ LLJIT::LLJIT(LLJITBuilderState &S, Error &Err)
InitHelperTransformLayer->setCloneToNewContextOnEmit(true);
CompileThreads =
std::make_unique<ThreadPool>(hardware_concurrency(S.NumCompileThreads));
- ES->setDispatchMaterialization(
- [this](std::unique_ptr<MaterializationUnit> MU,
- std::unique_ptr<MaterializationResponsibility> MR) {
- // FIXME: We should be able to use move-capture here, but ThreadPool's
- // AsyncTaskTys are std::functions rather than unique_functions
- // (because MSVC's std::packaged_tasks don't support move-only types).
- // Fix this when all the above gets sorted out.
- CompileThreads->async(
- [UnownedMU = MU.release(), UnownedMR = MR.release()]() mutable {
- std::unique_ptr<MaterializationUnit> MU(UnownedMU);
- std::unique_ptr<MaterializationResponsibility> MR(UnownedMR);
- MU->materialize(std::move(MR));
- });
- });
+ ES->setDispatchTask([this](std::unique_ptr<Task> T) {
+ // FIXME: We should be able to use move-capture here, but ThreadPool's
+ // AsyncTaskTys are std::functions rather than unique_functions
+ // (because MSVC's std::packaged_tasks don't support move-only types).
+ // Fix this when all the above gets sorted out.
+ CompileThreads->async([UnownedT = T.release()]() mutable {
+ std::unique_ptr<Task> T(UnownedT);
+ T->run();
+ });
+ });
}
if (S.SetUpPlatform)
diff --git a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
index 5b8d5550e79cd..5128cc9f558ae 100644
--- a/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
+++ b/llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp
@@ -1019,12 +1019,11 @@ TEST_F(CoreAPIsStandardTest, TestBasicWeakSymbolMaterialization) {
TEST_F(CoreAPIsStandardTest, DefineMaterializingSymbol) {
bool ExpectNoMoreMaterialization = false;
- ES.setDispatchMaterialization(
- [&](std::unique_ptr<MaterializationUnit> MU,
- std::unique_ptr<MaterializationResponsibility> MR) {
+ ES.setDispatchTask(
+ [&](std::unique_ptr<Task> T) {
if (ExpectNoMoreMaterialization)
ADD_FAILURE() << "Unexpected materialization";
- MU->materialize(std::move(MR));
+ T->run();
});
auto MU = std::make_unique<SimpleMaterializationUnit>(
@@ -1252,12 +1251,11 @@ TEST_F(CoreAPIsStandardTest, TestLookupWithThreadedMaterialization) {
#if LLVM_ENABLE_THREADS
std::thread MaterializationThread;
- ES.setDispatchMaterialization(
- [&](std::unique_ptr<MaterializationUnit> MU,
- std::unique_ptr<MaterializationResponsibility> MR) {
+ ES.setDispatchTask(
+ [&](std::unique_ptr<Task> T) {
MaterializationThread =
- std::thread([MU = std::move(MU), MR = std::move(MR)]() mutable {
- MU->materialize(std::move(MR));
+ std::thread([T = std::move(T)]() mutable {
+ T->run();
});
});
More information about the llvm-commits
mailing list