[ORC] Add LazyObjectLinkingLayer, lazy-linking support to llvm-jitlink (#116002)

LazyObjectLinkingLayer can be used to add object files that will not be linked
into the executor unless some function that they define is called at runtime.
(References to data members defined by these objects will still trigger
immediate linking)

To implement lazy linking, LazyObjectLinkingLayer uses the lazyReexports
utility to construct stubs for each function in a given object file, and an
ObjectLinkingLayer::Plugin to rename the function bodies at link-time. (Data
symbols are not renamed)

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.

Author: lhames

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

@@ -29,6 +29,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
+}