[llvm] [ORC] Add LazyObjectLinkingLayer, add lazy-linking support to llvm-ji… (PR #116002)

Lang Hames via llvm-commits llvm-commits at lists.llvm.org
Wed Nov 13 15:06:30 PST 2024


https://github.com/lhames updated https://github.com/llvm/llvm-project/pull/116002

>From 2130c1e34a95c471a5cfdf231f8e0741a0eb7535 Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames at gmail.com>
Date: Wed, 13 Nov 2024 16:37:08 +1100
Subject: [PATCH 1/2] [ORC] Add LazyObjectLinkingLayer, add lazy-linking
 support to llvm-jitlink.

LazyObjectLinkingLayer can be used to add object files that will not be linked
in unless they're actually called at runtime.

LazyObjectLinkingLayer uses lazyReexports to insert stubs for each function in
an object file, and an ObjectLinkingLayer::Plugin to rename the function bodies
at link-time.

The llvm-jitlink utility is extended with a -lazy option that can be passed
before input files or archives to add them using the lazy linking layer rather
than the base ObjectLinkingLayer.
---
 .../Orc/LazyObjectLinkingLayer.h              |  47 ++++++++
 llvm/lib/ExecutionEngine/Orc/CMakeLists.txt   |   1 +
 .../Orc/JITLinkRedirectableSymbolManager.cpp  |   2 +-
 .../Orc/LazyObjectLinkingLayer.cpp            | 114 ++++++++++++++++++
 .../JITLink/Generic/Inputs/foo-ret-42.ll      |   4 +
 .../JITLink/Generic/Inputs/var-x-42.ll        |   1 +
 .../JITLink/Generic/lazy-link.ll              |  27 +++++
 llvm/tools/llvm-jitlink/llvm-jitlink.cpp      | 111 +++++++++++++++--
 llvm/tools/llvm-jitlink/llvm-jitlink.h        |  32 ++++-
 9 files changed, 324 insertions(+), 15 deletions(-)
 create mode 100644 llvm/include/llvm/ExecutionEngine/Orc/LazyObjectLinkingLayer.h
 create mode 100644 llvm/lib/ExecutionEngine/Orc/LazyObjectLinkingLayer.cpp
 create mode 100644 llvm/test/ExecutionEngine/JITLink/Generic/Inputs/foo-ret-42.ll
 create mode 100644 llvm/test/ExecutionEngine/JITLink/Generic/Inputs/var-x-42.ll
 create mode 100644 llvm/test/ExecutionEngine/JITLink/Generic/lazy-link.ll

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/LazyObjectLinkingLayer.h b/llvm/include/llvm/ExecutionEngine/Orc/LazyObjectLinkingLayer.h
new file mode 100644
index 00000000000000..800f6773f16f55
--- /dev/null
+++ b/llvm/include/llvm/ExecutionEngine/Orc/LazyObjectLinkingLayer.h
@@ -0,0 +1,47 @@
+//===- RedirectionManager.h - Redirection manager interface -----*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Redirection manager interface that redirects a call to symbol to another.
+//
+//===----------------------------------------------------------------------===//
+#ifndef LLVM_EXECUTIONENGINE_ORC_LAZYOBJECTLINKINGLAYER_H
+#define LLVM_EXECUTIONENGINE_ORC_LAZYOBJECTLINKINGLAYER_H
+
+#include "llvm/ExecutionEngine/Orc/Core.h"
+#include "llvm/ExecutionEngine/Orc/Layer.h"
+
+namespace llvm::orc {
+
+class ObjectLinkingLayer;
+class LazyCallThroughManager;
+class RedirectableSymbolManager;
+
+class LazyObjectLinkingLayer : public ObjectLayer {
+public:
+  LazyObjectLinkingLayer(ObjectLinkingLayer &BaseLayer,
+                         LazyCallThroughManager &LCTMgr,
+                         RedirectableSymbolManager &RSMgr);
+
+  llvm::Error add(llvm::orc::ResourceTrackerSP RT,
+                  std::unique_ptr<llvm::MemoryBuffer> O,
+                  llvm::orc::MaterializationUnit::Interface I) override;
+
+  void emit(std::unique_ptr<MaterializationResponsibility> R,
+            std::unique_ptr<MemoryBuffer> O) override;
+
+private:
+  class RenamerPlugin;
+
+  ObjectLinkingLayer &BaseLayer;
+  LazyCallThroughManager &LCTMgr;
+  RedirectableSymbolManager &RSMgr;
+};
+
+} // namespace llvm::orc
+
+#endif // LLVM_EXECUTIONENGINE_ORC_LAZYOBJECTLINKINGLAYER_H
diff --git a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
index 008875118fdeff..50f15fb6d34cb9 100644
--- a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
+++ b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
@@ -28,6 +28,7 @@ add_llvm_component_library(LLVMOrcJIT
   IRTransformLayer.cpp
   IRPartitionLayer.cpp
   JITTargetMachineBuilder.cpp
+  LazyObjectLinkingLayer.cpp
   LazyReexports.cpp
   Layer.cpp
   LoadLinkableFile.cpp
diff --git a/llvm/lib/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.cpp b/llvm/lib/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.cpp
index c186a947b3bb54..9fefa76ed72476 100644
--- a/llvm/lib/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.cpp
@@ -29,7 +29,7 @@ void JITLinkRedirectableSymbolManager::emitRedirectableSymbols(
   Triple TT = ES.getTargetTriple();
 
   auto G = std::make_unique<jitlink::LinkGraph>(
-      ("<INDIRECT STUBS #" + Twine(++StubGraphIdx) + ">").str(), TT,
+      ("<indirect stubs graph #" + Twine(++StubGraphIdx) + ">").str(), TT,
       TT.isArch64Bit() ? 8 : 4,
       TT.isLittleEndian() ? endianness::little : endianness::big,
       jitlink::getGenericEdgeKindName);
diff --git a/llvm/lib/ExecutionEngine/Orc/LazyObjectLinkingLayer.cpp b/llvm/lib/ExecutionEngine/Orc/LazyObjectLinkingLayer.cpp
new file mode 100644
index 00000000000000..4dc2f2efd7e47e
--- /dev/null
+++ b/llvm/lib/ExecutionEngine/Orc/LazyObjectLinkingLayer.cpp
@@ -0,0 +1,114 @@
+//===---------- LazyReexports.cpp - Utilities for lazy reexports ----------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ExecutionEngine/Orc/LazyObjectLinkingLayer.h"
+
+#include "llvm/ExecutionEngine/Orc/LazyReexports.h"
+#include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
+#include "llvm/ExecutionEngine/Orc/RedirectionManager.h"
+
+using namespace llvm;
+using namespace llvm::jitlink;
+
+namespace {
+
+constexpr StringRef FnBodySuffix = "$orc_fnbody";
+
+} // anonymous namespace
+
+namespace llvm::orc {
+
+class LazyObjectLinkingLayer::RenamerPlugin
+    : public ObjectLinkingLayer::Plugin {
+public:
+  void modifyPassConfig(MaterializationResponsibility &MR,
+                        jitlink::LinkGraph &LG,
+                        jitlink::PassConfiguration &Config) override {
+    // We need to insert this before the mark-live pass to ensure that we don't
+    // delete the bodies (their names won't match the responsibility set until
+    // after this pass completes.
+    Config.PrePrunePasses.insert(
+        Config.PrePrunePasses.begin(),
+        [&MR](LinkGraph &G) { return renameFunctionBodies(G, MR); });
+  }
+
+  Error notifyFailed(MaterializationResponsibility &MR) override {
+    return Error::success();
+  }
+
+  Error notifyRemovingResources(JITDylib &JD, ResourceKey K) override {
+    return Error::success();
+  }
+
+  void notifyTransferringResources(JITDylib &JD, ResourceKey DstKey,
+                                   ResourceKey SrcKey) override {}
+
+private:
+  static Error renameFunctionBodies(LinkGraph &G,
+                                    MaterializationResponsibility &MR) {
+    DenseMap<StringRef, NonOwningSymbolStringPtr> SymsToRename;
+    for (auto &[Name, Flags] : MR.getSymbols())
+      if ((*Name).ends_with(FnBodySuffix))
+        SymsToRename[(*Name).drop_back(FnBodySuffix.size())] =
+            NonOwningSymbolStringPtr(Name);
+
+    for (auto *Sym : G.defined_symbols()) {
+      if (!Sym->hasName())
+        continue;
+      auto I = SymsToRename.find(Sym->getName());
+      if (I == SymsToRename.end())
+        continue;
+      Sym->setName(G.allocateName(*I->second));
+    }
+
+    return Error::success();
+  }
+};
+
+LazyObjectLinkingLayer::LazyObjectLinkingLayer(ObjectLinkingLayer &BaseLayer,
+                                               LazyCallThroughManager &LCTMgr,
+                                               RedirectableSymbolManager &RSMgr)
+    : ObjectLayer(BaseLayer.getExecutionSession()), BaseLayer(BaseLayer),
+      LCTMgr(LCTMgr), RSMgr(RSMgr) {
+  BaseLayer.addPlugin(std::make_unique<RenamerPlugin>());
+}
+
+Error LazyObjectLinkingLayer::add(ResourceTrackerSP RT,
+                                  std::unique_ptr<MemoryBuffer> O,
+                                  MaterializationUnit::Interface I) {
+
+  // Object files with initializer symbols can't be lazy.
+  if (I.InitSymbol)
+    return BaseLayer.add(std::move(RT), std::move(O), std::move(I));
+
+  auto &ES = getExecutionSession();
+  SymbolAliasMap LazySymbols;
+  for (auto &[Name, Flags] : I.SymbolFlags)
+    if (Flags.isCallable())
+      LazySymbols[Name] = {ES.intern((*Name + FnBodySuffix).str()), Flags};
+
+  for (auto &[Name, AI] : LazySymbols) {
+    I.SymbolFlags.erase(Name);
+    I.SymbolFlags[AI.Aliasee] = AI.AliasFlags;
+  }
+
+  if (auto Err = BaseLayer.add(RT, std::move(O), std::move(I)))
+    return Err;
+
+  auto &JD = RT->getJITDylib();
+  return JD.define(lazyReexports(LCTMgr, RSMgr, JD, std::move(LazySymbols)),
+                   std::move(RT));
+}
+
+void LazyObjectLinkingLayer::emit(
+    std::unique_ptr<MaterializationResponsibility> MR,
+    std::unique_ptr<MemoryBuffer> Obj) {
+  return BaseLayer.emit(std::move(MR), std::move(Obj));
+}
+
+} // namespace llvm::orc
diff --git a/llvm/test/ExecutionEngine/JITLink/Generic/Inputs/foo-ret-42.ll b/llvm/test/ExecutionEngine/JITLink/Generic/Inputs/foo-ret-42.ll
new file mode 100644
index 00000000000000..2d6b3393375528
--- /dev/null
+++ b/llvm/test/ExecutionEngine/JITLink/Generic/Inputs/foo-ret-42.ll
@@ -0,0 +1,4 @@
+define i32 @foo()  {
+entry:
+  ret i32 42
+}
diff --git a/llvm/test/ExecutionEngine/JITLink/Generic/Inputs/var-x-42.ll b/llvm/test/ExecutionEngine/JITLink/Generic/Inputs/var-x-42.ll
new file mode 100644
index 00000000000000..ee479a1ffa4f40
--- /dev/null
+++ b/llvm/test/ExecutionEngine/JITLink/Generic/Inputs/var-x-42.ll
@@ -0,0 +1 @@
+ at x = global i32 42
diff --git a/llvm/test/ExecutionEngine/JITLink/Generic/lazy-link.ll b/llvm/test/ExecutionEngine/JITLink/Generic/lazy-link.ll
new file mode 100644
index 00000000000000..0b5092ab54d61e
--- /dev/null
+++ b/llvm/test/ExecutionEngine/JITLink/Generic/lazy-link.ll
@@ -0,0 +1,27 @@
+; Check that files passed with the -lazy option aren't linked unless they're
+; needed. The foo-ret-42.ll file, which contains only code, should not be
+; needed in this -noexec case, whereas x.o, which contains a global variable
+; referenced by main, should be linked (despite being passed with -lazy).
+;
+; RUN: rm -rf %t && mkdir -p %t
+; RUN: llc -filetype=obj -o %t/foo.o %S/Inputs/foo-ret-42.ll
+; RUN: llc -filetype=obj -o %t/x.o %S/Inputs/var-x-42.ll
+; RUN: llc -filetype=obj -o %t/main.o %s
+; RUN: llvm-jitlink -noexec -show-linked-files %t/main.o -lazy %t/foo.o \
+; RUN:     -lazy %t/x.o | FileCheck %s
+
+; CHECK: Linking {{.*}}main.o
+; CHECK-DAG: Linking <indirect stubs graph #1>
+; CHECK-DAG: Linking {{.*}}x.o
+; CHECK-NOT: Linking {{.*}}foo.o
+
+declare i32 @foo()
+ at x = external global i32
+
+define i32 @main(i32 %argc, ptr %argv) {
+entry:
+  %foo_result = call i32 @foo()
+  %x_val = load i32, ptr @x
+  %result = add nsw i32 %foo_result, %x_val
+  ret i32 %result
+}
diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
index 261daae5e67f94..79e064a0a549cc 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink.cpp
@@ -28,6 +28,7 @@
 #include "llvm/ExecutionEngine/Orc/EPCEHFrameRegistrar.h"
 #include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
 #include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
+#include "llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h"
 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
 #include "llvm/ExecutionEngine/Orc/LoadLinkableFile.h"
 #include "llvm/ExecutionEngine/Orc/MachO.h"
@@ -80,6 +81,10 @@ using namespace llvm::orc;
 
 static cl::OptionCategory JITLinkCategory("JITLink Options");
 
+static cl::list<bool> LazyLink("lazy",
+                               cl::desc("Link the following file lazily"),
+                               cl::cat(JITLinkCategory));
+
 static cl::list<std::string> InputFiles(cl::Positional, cl::OneOrMore,
                                         cl::desc("input files"),
                                         cl::cat(JITLinkCategory));
@@ -182,6 +187,11 @@ static cl::list<std::string> TestHarnesses("harness", cl::Positional,
                                            cl::PositionalEatsArgs,
                                            cl::cat(JITLinkCategory));
 
+static cl::opt<bool>
+    ShowLinkedFiles("show-linked-files",
+                    cl::desc("List each file/graph name if/when it is linked"),
+                    cl::init(false), cl::cat(JITLinkCategory));
+
 static cl::opt<bool> ShowInitialExecutionSessionState(
     "show-init-es",
     cl::desc("Print ExecutionSession state before resolving entry point"),
@@ -380,6 +390,13 @@ operator<<(raw_ostream &OS, const Session::FileInfoMap &FIM) {
   return OS;
 }
 
+bool lazyLinkingRequested() {
+  for (auto LL : LazyLink)
+    if (LL)
+      return true;
+  return false;
+}
+
 static Error applyHarnessPromotions(Session &S, LinkGraph &G) {
 
   // If this graph is part of the test harness there's nothing to do.
@@ -931,6 +948,43 @@ class PhonyExternalsGenerator : public DefinitionGenerator {
   }
 };
 
+static void handleLazyCallFailure() {
+  dbgs() << "ERROR: failure to materialize lazy call-through target.\n";
+  exit(1);
+}
+
+static void *reenter(void *Ctx, void *TrampolineAddr) {
+  std::promise<void *> LandingAddressP;
+  auto LandingAddressF = LandingAddressP.get_future();
+
+  auto *EPCIU = static_cast<EPCIndirectionUtils *>(Ctx);
+  EPCIU->getLazyCallThroughManager().resolveTrampolineLandingAddress(
+      ExecutorAddr::fromPtr(TrampolineAddr), [&](ExecutorAddr LandingAddress) {
+        LandingAddressP.set_value(LandingAddress.toPtr<void *>());
+      });
+  return LandingAddressF.get();
+}
+
+Expected<std::unique_ptr<Session::LazyLinkingSupport>>
+createLazyLinkingSupport(ObjectLinkingLayer &OLL) {
+  auto EPCIU = EPCIndirectionUtils::Create(OLL.getExecutionSession());
+  if (!EPCIU)
+    return EPCIU.takeError();
+  if (auto Err = (*EPCIU)
+                     ->writeResolverBlock(ExecutorAddr::fromPtr(&reenter),
+                                          ExecutorAddr::fromPtr(EPCIU->get()))
+                     .takeError())
+    return Err;
+  (*EPCIU)->createLazyCallThroughManager(
+      OLL.getExecutionSession(), ExecutorAddr::fromPtr(handleLazyCallFailure));
+  auto RSMgr = JITLinkRedirectableSymbolManager::Create(OLL);
+  if (!RSMgr)
+    return RSMgr.takeError();
+
+  return std::make_unique<Session::LazyLinkingSupport>(std::move(*EPCIU),
+                                                       std::move(*RSMgr), OLL);
+}
+
 Expected<std::unique_ptr<Session>> Session::Create(Triple TT,
                                                    SubtargetFeatures Features) {
 
@@ -963,6 +1017,14 @@ Expected<std::unique_ptr<Session>> Session::Create(Triple TT,
   if (Err)
     return std::move(Err);
   S->Features = std::move(Features);
+
+  if (lazyLinkingRequested()) {
+    if (auto LazyLinking = createLazyLinkingSupport(S->ObjLayer))
+      S->LazyLinking = std::move(*LazyLinking);
+    else
+      return LazyLinking.takeError();
+  }
+
   return std::move(S);
 }
 
@@ -982,7 +1044,7 @@ Session::Session(std::unique_ptr<ExecutorProcessControl> EPC, Error &Err)
     JITLinkSessionPlugin(Session &S) : S(S) {}
     void modifyPassConfig(MaterializationResponsibility &MR, LinkGraph &G,
                           PassConfiguration &PassConfig) override {
-      S.modifyPassConfig(G.getTargetTriple(), PassConfig);
+      S.modifyPassConfig(G, PassConfig);
     }
 
     Error notifyFailed(MaterializationResponsibility &MR) override {
@@ -1157,8 +1219,11 @@ void Session::dumpSessionInfo(raw_ostream &OS) {
   OS << "Registered addresses:\n" << SymbolInfos << FileInfos;
 }
 
-void Session::modifyPassConfig(const Triple &TT,
-                               PassConfiguration &PassConfig) {
+void Session::modifyPassConfig(LinkGraph &G, PassConfiguration &PassConfig) {
+
+  if (ShowLinkedFiles)
+    outs() << "Linking " << G.getName() << "\n";
+
   if (!CheckFiles.empty())
     PassConfig.PostFixupPasses.push_back([this](LinkGraph &G) {
       if (ES.getTargetTriple().getObjectFormat() == Triple::ELF)
@@ -1575,6 +1640,11 @@ static Error sanitizeArguments(const Triple &TT, const char *ArgV0) {
     OutOfProcessExecutor = OOPExecutorPath.str().str();
   }
 
+  if (lazyLinkingRequested() && !TestHarnesses.empty())
+    return make_error<StringError>(
+        "Lazy linking cannot be used with -harness mode",
+        inconvertibleErrorCode());
+
   return Error::success();
 }
 
@@ -1796,7 +1866,8 @@ static Error addTestHarnesses(Session &S) {
 }
 
 static Error addObjects(Session &S,
-                        const std::map<unsigned, JITDylib *> &IdxToJD) {
+                        const std::map<unsigned, JITDylib *> &IdxToJD,
+                        const DenseSet<unsigned> &LazyLinkIdxs) {
 
   // Load each object into the corresponding JITDylib..
   LLVM_DEBUG(dbgs() << "Adding objects...\n");
@@ -1809,15 +1880,18 @@ static Error addObjects(Session &S,
         StringRef(InputFile).ends_with(".lib"))
       continue;
     auto &JD = *std::prev(IdxToJD.lower_bound(InputFileArgIdx))->second;
-    LLVM_DEBUG(dbgs() << "  " << InputFileArgIdx << ": \"" << InputFile
-                      << "\" to " << JD.getName() << "\n";);
+    bool AddLazy = LazyLinkIdxs.count(InputFileArgIdx);
+    LLVM_DEBUG(dbgs() << "  " << InputFileArgIdx << ": \"" << InputFile << "\" "
+                      << (AddLazy ? " (lazy-linked)" : "") << " to "
+                      << JD.getName() << "\n";);
     auto ObjBuffer = loadLinkableFile(InputFile, S.ES.getTargetTriple(),
                                       LoadArchives::Never);
     if (!ObjBuffer)
       return ObjBuffer.takeError();
 
     if (S.HarnessFiles.empty()) {
-      if (auto Err = S.ObjLayer.add(JD, std::move(ObjBuffer->first)))
+      if (auto Err =
+              S.getLinkLayer(AddLazy).add(JD, std::move(ObjBuffer->first)))
         return Err;
     } else {
       // We're in -harness mode. Use a custom interface for this
@@ -1826,6 +1900,7 @@ static Error addObjects(Session &S,
           getTestObjectFileInterface(S, ObjBuffer->first->getMemBufferRef());
       if (!ObjInterface)
         return ObjInterface.takeError();
+
       if (auto Err = S.ObjLayer.add(JD, std::move(ObjBuffer->first),
                                     std::move(*ObjInterface)))
         return Err;
@@ -1858,7 +1933,8 @@ static SmallVector<StringRef, 5> getSearchPathsFromEnvVar(Session &S) {
 }
 
 static Error addLibraries(Session &S,
-                          const std::map<unsigned, JITDylib *> &IdxToJD) {
+                          const std::map<unsigned, JITDylib *> &IdxToJD,
+                          const DenseSet<unsigned> &LazyLinkIdxs) {
 
   // 1. Collect search paths for each JITDylib.
   DenseMap<const JITDylib *, SmallVector<StringRef, 2>> JDSearchPaths;
@@ -1982,15 +2058,17 @@ static Error addLibraries(Session &S,
       break;
     }
 
+    auto &LinkLayer = S.getLinkLayer(LazyLinkIdxs.count(LL.Position));
+
     StaticLibraryDefinitionGenerator::VisitMembersFunction VisitMembers;
     if (AllLoad)
       VisitMembers = StaticLibraryDefinitionGenerator::loadAllObjectFileMembers(
-          S.ObjLayer, JD);
+          LinkLayer, JD);
     else if (S.ES.getTargetTriple().isOSBinFormatMachO() && ForceLoadObjC)
-      VisitMembers = ForceLoadMachOArchiveMembers(S.ObjLayer, JD, true);
+      VisitMembers = ForceLoadMachOArchiveMembers(LinkLayer, JD, true);
 
     auto G = StaticLibraryDefinitionGenerator::Load(
-        S.ObjLayer, Path, std::move(VisitMembers),
+        LinkLayer, Path, std::move(VisitMembers),
         std::move(GetObjFileInterface));
     if (!G)
       return G.takeError();
@@ -2131,6 +2209,13 @@ static Error addLibraries(Session &S,
 
 static Error addSessionInputs(Session &S) {
   std::map<unsigned, JITDylib *> IdxToJD;
+  DenseSet<unsigned> LazyLinkIdxs;
+
+  for (auto LLItr = LazyLink.begin(), LLEnd = LazyLink.end(); LLItr != LLEnd;
+       ++LLItr) {
+    if (*LLItr)
+      LazyLinkIdxs.insert(LazyLink.getPosition(LLItr - LazyLink.begin()) + 1);
+  }
 
   if (auto Err = createJITDylibs(S, IdxToJD))
     return Err;
@@ -2148,10 +2233,10 @@ static Error addSessionInputs(Session &S) {
     if (auto Err = addTestHarnesses(S))
       return Err;
 
-  if (auto Err = addObjects(S, IdxToJD))
+  if (auto Err = addObjects(S, IdxToJD, LazyLinkIdxs))
     return Err;
 
-  if (auto Err = addLibraries(S, IdxToJD))
+  if (auto Err = addLibraries(S, IdxToJD, LazyLinkIdxs))
     return Err;
 
   return Error::success();
diff --git a/llvm/tools/llvm-jitlink/llvm-jitlink.h b/llvm/tools/llvm-jitlink/llvm-jitlink.h
index 07b56680719f69..40fac2fe6888fe 100644
--- a/llvm/tools/llvm-jitlink/llvm-jitlink.h
+++ b/llvm/tools/llvm-jitlink/llvm-jitlink.h
@@ -15,8 +15,11 @@
 
 #include "llvm/ADT/StringSet.h"
 #include "llvm/ExecutionEngine/Orc/Core.h"
+#include "llvm/ExecutionEngine/Orc/EPCIndirectionUtils.h"
 #include "llvm/ExecutionEngine/Orc/ExecutorProcessControl.h"
+#include "llvm/ExecutionEngine/Orc/LazyObjectLinkingLayer.h"
 #include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
+#include "llvm/ExecutionEngine/Orc/RedirectionManager.h"
 #include "llvm/ExecutionEngine/Orc/SimpleRemoteEPC.h"
 #include "llvm/ExecutionEngine/RuntimeDyldChecker.h"
 #include "llvm/Support/Error.h"
@@ -29,11 +32,30 @@ namespace llvm {
 
 struct Session {
 
+  struct LazyLinkingSupport {
+    LazyLinkingSupport(std::unique_ptr<orc::EPCIndirectionUtils> EPCIU,
+                       std::unique_ptr<orc::RedirectableSymbolManager> RSMgr,
+                       orc::ObjectLinkingLayer &ObjLinkingLayer)
+        : EPCIU(std::move(EPCIU)), RSMgr(std::move(RSMgr)),
+          LazyObjLinkingLayer(ObjLinkingLayer,
+                              this->EPCIU->getLazyCallThroughManager(),
+                              *this->RSMgr) {}
+    ~LazyLinkingSupport() {
+      if (auto Err = EPCIU->cleanup())
+        LazyObjLinkingLayer.getExecutionSession().reportError(std::move(Err));
+    }
+
+    std::unique_ptr<orc::EPCIndirectionUtils> EPCIU;
+    std::unique_ptr<orc::RedirectableSymbolManager> RSMgr;
+    orc::LazyObjectLinkingLayer LazyObjLinkingLayer;
+  };
+
   orc::ExecutionSession ES;
   orc::JITDylib *MainJD = nullptr;
   orc::JITDylib *ProcessSymsJD = nullptr;
   orc::JITDylib *PlatformJD = nullptr;
   orc::ObjectLinkingLayer ObjLayer;
+  std::unique_ptr<LazyLinkingSupport> LazyLinking;
   orc::JITDylibSearchOrder JDSearchOrder;
   SubtargetFeatures Features;
 
@@ -42,7 +64,7 @@ struct Session {
   static Expected<std::unique_ptr<Session>> Create(Triple TT,
                                                    SubtargetFeatures Features);
   void dumpSessionInfo(raw_ostream &OS);
-  void modifyPassConfig(const Triple &FTT,
+  void modifyPassConfig(jitlink::LinkGraph &G,
                         jitlink::PassConfiguration &PassConfig);
 
   using MemoryRegionInfo = RuntimeDyldChecker::MemoryRegionInfo;
@@ -72,6 +94,14 @@ struct Session {
   Expected<orc::JITDylib *> getOrLoadDynamicLibrary(StringRef LibPath);
   Error loadAndLinkDynamicLibrary(orc::JITDylib &JD, StringRef LibPath);
 
+  orc::ObjectLayer &getLinkLayer(bool Lazy) {
+    assert((!Lazy || LazyLinking) &&
+           "Lazy linking requested but not available");
+    return Lazy ? static_cast<orc::ObjectLayer &>(
+                      LazyLinking->LazyObjLinkingLayer)
+                : static_cast<orc::ObjectLayer &>(ObjLayer);
+  }
+
   Expected<FileInfo &> findFileInfo(StringRef FileName);
   Expected<MemoryRegionInfo &> findSectionInfo(StringRef FileName,
                                                StringRef SectionName);

>From 239744c1c55218eaa1aad0db683c34d4bcf38b46 Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames at gmail.com>
Date: Thu, 14 Nov 2024 09:57:02 +1100
Subject: [PATCH 2/2] Disable the lazy-linking testcase on Windows -- it's
 failing there due to a missing __ImageBase symbol for reasons that aren't
 clear. Landing this patch will allow others to debug the issue in-tree.

---
 llvm/test/ExecutionEngine/JITLink/Generic/lazy-link.ll | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/llvm/test/ExecutionEngine/JITLink/Generic/lazy-link.ll b/llvm/test/ExecutionEngine/JITLink/Generic/lazy-link.ll
index 0b5092ab54d61e..88bf19cc6e7b1a 100644
--- a/llvm/test/ExecutionEngine/JITLink/Generic/lazy-link.ll
+++ b/llvm/test/ExecutionEngine/JITLink/Generic/lazy-link.ll
@@ -9,7 +9,9 @@
 ; RUN: llc -filetype=obj -o %t/main.o %s
 ; RUN: llvm-jitlink -noexec -show-linked-files %t/main.o -lazy %t/foo.o \
 ; RUN:     -lazy %t/x.o | FileCheck %s
-
+;
+; UNSUPPORTED: system-windows
+;
 ; CHECK: Linking {{.*}}main.o
 ; CHECK-DAG: Linking <indirect stubs graph #1>
 ; CHECK-DAG: Linking {{.*}}x.o



More information about the llvm-commits mailing list