[ORC] Add LazyObjectLinkingLayer, add lazy-linking support to llvm-ji…

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

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

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);

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

llvm/test/ExecutionEngine/JITLink/Generic/Inputs/foo-ret-42.ll

@@ -0,0 +1,4 @@
+define i32 @foo()  {
+entry:
+  ret i32 42
+}

llvm/test/ExecutionEngine/JITLink/Generic/Inputs/var-x-42.ll

@@ -0,0 +1 @@
+@x = global i32 42

llvm/test/ExecutionEngine/JITLink/Generic/lazy-link.ll

@@ -0,0 +1,29 @@
+; 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
+;
+; UNSUPPORTED: system-windows
+;
+; CHECK: Linking {{.*}}main.o
+; CHECK-DAG: Linking <indirect stubs graph #1>
+; CHECK-DAG: Linking {{.*}}x.o
+; CHECK-NOT: Linking {{.*}}foo.o
+
+declare i32 @foo()
+@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
+}