[llvm-jitlink] Use concurrent linking by default.

Adds a -threads option to llvm-jitlink. By default llvm-jitlink will now use a DynamicThreadPoolTaskDispatcher with the number of materialization threads set to whatever is returned by std::hardware_concurrency(). This brings the default in-place linking behavior in line with the concurrent linking that is used for -oop-executor and -oop-executor-connect mode. In-place linking on the main thread can be forced by passing -threads=0.
intel · Dec 23, 2024 · edca1d9 · edca1d9
1 parent 2b3aff8
commit edca1d9
Show file tree

Hide file tree

Showing 5 changed files with 55 additions and 8 deletions.
diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink-coff.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink-coff.cpp
@@ -66,6 +66,8 @@ static Expected<Symbol &> getCOFFStubTarget(LinkGraph &G, Block &B) {
 
 namespace llvm {
 Error registerCOFFGraphInfo(Session &S, LinkGraph &G) {
+  std::lock_guard<std::mutex> Lock(S.M);
+
   auto FileName = sys::path::filename(G.getName());
   if (S.FileInfos.count(FileName)) {
     return make_error<StringError>("When -check is passed, file names must be "

diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink-elf.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink-elf.cpp
@@ -101,6 +101,8 @@ static Error registerSymbol(LinkGraph &G, Symbol &Sym, Session::FileInfo &FI,
 namespace llvm {
 
 Error registerELFGraphInfo(Session &S, LinkGraph &G) {
+  std::lock_guard<std::mutex> Lock(S.M);
+
   auto FileName = sys::path::filename(G.getName());
   if (S.FileInfos.count(FileName)) {
     return make_error<StringError>("When -check is passed, file names must be "

diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink-macho.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink-macho.cpp
@@ -69,6 +69,8 @@ static Expected<Symbol &> getMachOStubTarget(LinkGraph &G, Block &B) {
 namespace llvm {
 
 Error registerMachOGraphInfo(Session &S, LinkGraph &G) {
+  std::lock_guard<std::mutex> Lock(S.M);
+
   auto FileName = sys::path::filename(G.getName());
   if (S.FileInfos.count(FileName)) {
     return make_error<StringError>("When -check is passed, file names must be "

diff --git 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));
 
+cl::opt<size_t> MaterializationThreads(
+    "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,16 @@ 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
+      Dispatcher = std::make_unique<DynamicThreadPoolTaskDispatcher>(
+          MaterializationThreads);
+
     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 +1236,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())) {
@@ -1234,9 +1250,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();
     });
 
@@ -1601,6 +1616,31 @@ static Error sanitizeArguments(const Triple &TT, const char *ArgV0) {
     }
   }
 
+  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 -threads=1.\n";
+      MaterializationThreads = 1;
+    }
+  }
+
+  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());
+  }
+
   // Only one of -oop-executor and -oop-executor-connect can be used.
   if (!!OutOfProcessExecutor.getNumOccurrences() &&
       !!OutOfProcessExecutorConnect.getNumOccurrences())

diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink.h b/llvm/tools/llvm-jitlink/llvm-jitlink.h
@@ -110,6 +110,7 @@ struct Session {
 
   DynLibJDMap DynLibJDs;
 
+  std::mutex M;
   SymbolInfoMap SymbolInfos;
   FileInfoMap FileInfos;