1 files changed, 69 insertions, 8 deletions

diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
index 96a3e5b..963c363 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
@@ -91,6 +91,10 @@ static cl::list<std::string> InputFiles(cl::Positional, cl::OneOrMore,
                                         cl::desc("input files"),
                                         cl::cat(JITLinkCategory));
 
+static cl::opt<size_t> MaterializationThreads(
+    "num-threads", cl::desc("Number of materialization threads to use"),
+    cl::init(std::numeric_limits<size_t>::max()), cl::cat(JITLinkCategory));
+
 static cl::list<std::string>
     LibrarySearchPaths("L",
                        cl::desc("Add dir to the list of library search paths"),
@@ -400,6 +404,7 @@ bool lazyLinkingRequested() {
 }
 
 static Error applyHarnessPromotions(Session &S, LinkGraph &G) {
+  std::lock_guard<std::mutex> Lock(S.M);
 
   // If this graph is part of the test harness there's nothing to do.
   if (S.HarnessFiles.empty() || S.HarnessFiles.count(G.getName()))
@@ -450,7 +455,11 @@ static Error applyHarnessPromotions(Session &S, LinkGraph &G) {
   return Error::success();
 }
 
-static void dumpSectionContents(raw_ostream &OS, LinkGraph &G) {
+static void dumpSectionContents(raw_ostream &OS, Session &S, LinkGraph &G) {
+  std::lock_guard<std::mutex> Lock(S.M);
+
+  outs() << "Relocated section contents for " << G.getName() << ":\n";
+
   constexpr orc::ExecutorAddrDiff DumpWidth = 16;
   static_assert(isPowerOf2_64(DumpWidth), "DumpWidth must be a power of two");
 
@@ -842,7 +851,7 @@ static Expected<std::unique_ptr<ExecutorProcessControl>> launchExecutor() {
     S.CreateMemoryManager = createSharedMemoryManager;
 
   return SimpleRemoteEPC::Create<FDSimpleRemoteEPCTransport>(
-      std::make_unique<DynamicThreadPoolTaskDispatcher>(std::nullopt),
+      std::make_unique<DynamicThreadPoolTaskDispatcher>(MaterializationThreads),
       std::move(S), FromExecutor[ReadEnd], ToExecutor[WriteEnd]);
 #endif
 }
@@ -984,10 +993,21 @@ Expected<std::unique_ptr<Session>> Session::Create(Triple TT,
     auto PageSize = sys::Process::getPageSize();
     if (!PageSize)
       return PageSize.takeError();
+    std::unique_ptr<TaskDispatcher> Dispatcher;
+    if (MaterializationThreads == 0)
+      Dispatcher = std::make_unique<InPlaceTaskDispatcher>();
+    else {
+#if LLVM_ENABLE_THREADS
+      Dispatcher = std::make_unique<DynamicThreadPoolTaskDispatcher>(
+          MaterializationThreads);
+#else
+      llvm_unreachable("MaterializationThreads should be 0");
+#endif
+    }
+
     EPC = std::make_unique<SelfExecutorProcessControl>(
-        std::make_shared<SymbolStringPool>(),
-        std::make_unique<InPlaceTaskDispatcher>(), std::move(TT), *PageSize,
-        createInProcessMemoryManager());
+        std::make_shared<SymbolStringPool>(), std::move(Dispatcher),
+        std::move(TT), *PageSize, createInProcessMemoryManager());
   }
 
   Error Err = Error::success();
@@ -1221,6 +1241,7 @@ void Session::modifyPassConfig(LinkGraph &G, PassConfiguration &PassConfig) {
 
   if (ShowGraphsRegex)
     PassConfig.PostFixupPasses.push_back([this](LinkGraph &G) -> Error {
+      std::lock_guard<std::mutex> Lock(M);
       // Print graph if ShowLinkGraphs is specified-but-empty, or if
       // it contains the given graph.
       if (ShowGraphsRegex->match(G.getName())) {
@@ -1239,9 +1260,8 @@ void Session::modifyPassConfig(LinkGraph &G, PassConfiguration &PassConfig) {
       [this](LinkGraph &G) { return applyHarnessPromotions(*this, G); });
 
   if (ShowRelocatedSectionContents)
-    PassConfig.PostFixupPasses.push_back([](LinkGraph &G) -> Error {
-      outs() << "Relocated section contents for " << G.getName() << ":\n";
-      dumpSectionContents(outs(), G);
+    PassConfig.PostFixupPasses.push_back([this](LinkGraph &G) -> Error {
+      dumpSectionContents(outs(), *this, G);
       return Error::success();
     });
 
@@ -1613,6 +1633,46 @@ static Error sanitizeArguments(const Triple &TT, const char *ArgV0) {
     }
   }
 
+#if LLVM_ENABLE_THREADS
+  if (MaterializationThreads == std::numeric_limits<size_t>::max()) {
+    if (auto HC = std::thread::hardware_concurrency())
+      MaterializationThreads = HC;
+    else {
+      errs() << "Warning: std::thread::hardware_concurrency() returned 0, "
+                "defaulting to -num-threads=1.\n";
+      MaterializationThreads = 1;
+    }
+  }
+#else
+  if (MaterializationThreads.getNumOccurrences() &&
+      MaterializationThreads != 0) {
+    errs() << "Warning: -num-threads was set, but LLVM was built with threads "
+              "disabled. Resetting to -num-threads=0\n";
+  }
+  MaterializationThreads = 0;
+#endif
+
+  if (!!OutOfProcessExecutor.getNumOccurrences() ||
+      !!OutOfProcessExecutorConnect.getNumOccurrences()) {
+    if (NoExec)
+      return make_error<StringError>("-noexec cannot be used with " +
+                                         OutOfProcessExecutor.ArgStr + " or " +
+                                         OutOfProcessExecutorConnect.ArgStr,
+                                     inconvertibleErrorCode());
+
+    if (MaterializationThreads == 0)
+      return make_error<StringError>("-threads=0 cannot be used with " +
+                                         OutOfProcessExecutor.ArgStr + " or " +
+                                         OutOfProcessExecutorConnect.ArgStr,
+                                     inconvertibleErrorCode());
+  }
+
+#ifndef NDEBUG
+  if (DebugFlag && MaterializationThreads != 0)
+    errs() << "Warning: debugging output is not thread safe. "
+              "Use -num-threads=0 to stabilize output.\n";
+#endif // NDEBUG
+
   // Only one of -oop-executor and -oop-executor-connect can be used.
   if (!!OutOfProcessExecutor.getNumOccurrences() &&
       !!OutOfProcessExecutorConnect.getNumOccurrences())
@@ -2502,6 +2562,7 @@ int main(int argc, char *argv[]) {
     if (Timers)
       Timers->JITLinkTG.printAll(errs());
     reportLLVMJITLinkError(EntryPoint.takeError());
+    ExitOnErr(S->ES.endSession());
     exit(1);
   }