[ORC] Introduce LazyReexportsManager, JITLinkTrampolines, ORC-RT base… (

llvm#118923)

…d reentry.

These utilities provide new, more generic and easier to use support for
lazy compilation in ORC.

LazyReexportsManager is an alternative to LazyCallThroughManager. It
takes requests for lazy re-entry points in the form of an alias map:
lazy-reexports = {
  ( <entry point symbol #1>, <implementation symbol #1> ),
  ( <entry point symbol #2>, <implementation symbol #2> ),
  ...
  ( <entry point symbol #n>, <implementation symbol #n> )
}

LazyReexportsManager then:
1. binds the entry points to the implementation names in an internal
table.
2. creates a JIT re-entry trampoline for each entry point.
3. creates a redirectable symbol for each of the entry point name and
binds redirectable symbol to the corresponding reentry trampoline.

When an entry point symbol is first called at runtime (which may be on
any thread of the JIT'd program) it will re-enter the JIT via the
trampoline and trigger a lookup for the implementation symbol stored in
LazyReexportsManager's internal table. When the lookup completes the
entry point symbol will be updated (via the RedirectableSymbolManager)
to point at the implementation symbol, and execution will proceed to the
implementation symbol.

Actual construction of the re-entry trampolines and redirectable symbols
is delegated to an EmitTrampolines functor and the
RedirectableSymbolsManager respectively.

JITLinkReentryTrampolines.h provides a JITLink-based implementation of
the EmitTrampolines functor. (AArch64 only in this patch, but other
architectures will be added in the near future).

Register state save and reentry functionality is added to the ORC
runtime in the __orc_rt_sysv_resolve and __orc_rt_resolve_implementation
functions (the latter is generic, the former will need custom
implementations for each ABI and architecture to be supported, however
this should be much less effort than the existing OrcABISupport
approach, since the ORC runtime allows this code to be written as native
assembly).

The resulting system:
1. Works equally well for in-process and out-of-process JIT'd code.
2. Requires less boilerplate to set up.

Given an ObjectLinkingLayer and PlatformJD (JITDylib containing the ORC
runtime), setup is just:

```c++
auto RSMgr = JITLinkRedirectableSymbolManager::Create(OLL);
if (!RSMgr)
  return RSMgr.takeError();

auto LRMgr = createJITLinkLazyReexportsManager(OLL, **RSMgr, PlatformJD);
if (!LRMgr)
  return LRMgr.takeError();
```

after which lazy reexports can be introduced with:

```c++
JD.define(lazyReexports(LRMgr, <alias map>));
```

LazyObectLinkingLayer is updated to use this new method, but the LLVM-IR
level CompileOnDemandLayer will continue to use LazyCallThroughManager
and OrcABISupport until the new system supports a wider range of
architectures and ABIs.

The llvm-jitlink utility's -lazy option now uses the new scheme. Since
it depends on the ORC runtime, the lazy-link.ll testcase and associated
helpers are moved to the ORC runtime.

compiler-rt/lib/orc/CMakeLists.txt

@@ -51,6 +51,7 @@ if (APPLE)
   set(ORC_ASM_SOURCES
     macho_tlv.x86-64.S
     macho_tlv.arm64.S
+    sysv_reentry.arm64.S
     )
 
   set(ORC_IMPL_HEADERS
@@ -61,6 +62,7 @@ if (APPLE)
   set(ORC_SOURCES
     ${ORC_COMMON_SOURCES}
     macho_platform.cpp
+    sysv_resolve.cpp
     )
 
   add_compiler_rt_object_libraries(RTOrc

compiler-rt/lib/orc/sysv_reentry.arm64.S

@@ -0,0 +1,102 @@
+//===-- sysv_reentry.arm64.s ------------------------------------*- ASM -*-===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file is a part of the ORC runtime support library.
+//
+//===----------------------------------------------------------------------===//
+
+// The content of this file is arm64-only
+#if defined(__arm64__) || defined(__aarch64__)
+
+        .text
+
+        // Saves GPRs, calls __orc_rt_sysv_resolve
+        .globl __orc_rt_sysv_reentry
+__orc_rt_sysv_reentry:
+        // Save register state, set up new stack frome.
+        stp  x27, x28, [sp, #-16]!
+        stp  x25, x26, [sp, #-16]!
+        stp  x23, x24, [sp, #-16]!
+        stp  x21, x22, [sp, #-16]!
+        stp  x19, x20, [sp, #-16]!
+        stp  x14, x15, [sp, #-16]!
+        stp  x12, x13, [sp, #-16]!
+        stp  x10, x11, [sp, #-16]!
+        stp   x8,  x9, [sp, #-16]!
+        stp   x6,  x7, [sp, #-16]!
+        stp   x4,  x5, [sp, #-16]!
+        stp   x2,  x3, [sp, #-16]!
+        stp   x0,  x1, [sp, #-16]!
+        stp  q30, q31, [sp, #-32]!
+        stp  q28, q29, [sp, #-32]!
+        stp  q26, q27, [sp, #-32]!
+        stp  q24, q25, [sp, #-32]!
+        stp  q22, q23, [sp, #-32]!
+        stp  q20, q21, [sp, #-32]!
+        stp  q18, q19, [sp, #-32]!
+        stp  q16, q17, [sp, #-32]!
+        stp  q14, q15, [sp, #-32]!
+        stp  q12, q13, [sp, #-32]!
+        stp  q10, q11, [sp, #-32]!
+        stp   q8,  q9, [sp, #-32]!
+        stp   q6,  q7, [sp, #-32]!
+        stp   q4,  q5, [sp, #-32]!
+        stp   q2,  q3, [sp, #-32]!
+        stp   q0,  q1, [sp, #-32]!
+
+        // Look up the return address and subtract 8 from it (on the
+        // assumption that it's a standard arm64 reentry trampoline) to get
+        // back the trampoline's address.
+        sub   x0, x30, #8
+
+        // Call __orc_rt_sysv_resolve to look up the implementation
+        // corresponding to the calling stub, then store this in x17 (which
+        // we'll return to below.
+#if !defined(__APPLE__)
+        bl    __orc_rt_sysv_resolve
+#else
+        bl    ___orc_rt_sysv_resolve
+#endif
+        mov   x17, x0
+
+        // Restore the register state.
+        ldp   q0,  q1, [sp], #32
+        ldp   q2,  q3, [sp], #32
+        ldp   q4,  q5, [sp], #32
+        ldp   q6,  q7, [sp], #32
+        ldp   q8,  q9, [sp], #32
+        ldp  q10, q11, [sp], #32
+        ldp  q12, q13, [sp], #32
+        ldp  q14, q15, [sp], #32
+        ldp  q16, q17, [sp], #32
+        ldp  q18, q19, [sp], #32
+        ldp  q20, q21, [sp], #32
+        ldp  q22, q23, [sp], #32
+        ldp  q24, q25, [sp], #32
+        ldp  q26, q27, [sp], #32
+        ldp  q28, q29, [sp], #32
+        ldp  q30, q31, [sp], #32
+        ldp   x0,  x1, [sp], #16
+        ldp   x2,  x3, [sp], #16
+        ldp   x4,  x5, [sp], #16
+        ldp   x6,  x7, [sp], #16
+        ldp   x8,  x9, [sp], #16
+        ldp  x10, x11, [sp], #16
+        ldp  x12, x13, [sp], #16
+        ldp  x14, x15, [sp], #16
+        ldp  x19, x20, [sp], #16
+        ldp  x21, x22, [sp], #16
+        ldp  x23, x24, [sp], #16
+        ldp  x25, x26, [sp], #16
+        ldp  x27, x28, [sp], #16
+        ldp  x29, x30, [sp], #16
+
+        // Return to the function implementation (rather than the stub).
+        ret  x17
+
+#endif // defined(__arm64__) || defined(__aarch64__)

compiler-rt/lib/orc/sysv_resolve.cpp

@@ -0,0 +1,49 @@
+//===- sysv_resolve.cpp ---------------------------------------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains a generic "resolver" function compatible with the SysV
+// ABI.
+//
+//===----------------------------------------------------------------------===//
+
+#include "executor_symbol_def.h"
+#include "jit_dispatch.h"
+#include "wrapper_function_utils.h"
+
+#include <stdio.h>
+
+#define DEBUG_TYPE "sysv_resolve"
+
+using namespace orc_rt;
+
+// Declare function tags for functions in the JIT process.
+ORC_RT_JIT_DISPATCH_TAG(__orc_rt_resolve_tag)
+
+// FIXME: Make this configurable via an alias.
+static void __orc_rt_sysv_fail(void *Caller, const char *ErrMsg) {
+  fprintf(stderr, "error resolving implementation for stub %p: %s\n", Caller,
+          ErrMsg);
+  abort();
+}
+
+extern "C" ORC_RT_HIDDEN void *__orc_rt_sysv_resolve(void *Caller) {
+  Expected<ExecutorSymbolDef> Result((ExecutorSymbolDef()));
+  if (auto Err = WrapperFunction<SPSExpected<SPSExecutorSymbolDef>(
+          SPSExecutorAddr)>::call(JITDispatch(&__orc_rt_resolve_tag), Result,
+                                  ExecutorAddr::fromPtr(Caller))) {
+    __orc_rt_sysv_fail(Caller, toString(std::move(Err)).c_str());
+    return nullptr; // Unreachable.
+  }
+
+  if (!Result) {
+    __orc_rt_sysv_fail(Caller, toString(Result.takeError()).c_str());
+    return nullptr; // Unreachable.
+  }
+
+  return Result->getAddress().toPtr<void *>();
+}

llvm/test/ExecutionEngine/JITLink/Generic/lazy-link.ll → compiler-rt/test/orc/TestCases/Generic/lazy-link.ll

@@ -4,13 +4,14 @@
 ; 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: %clang -c -o %t/foo.o %S/Inputs/foo-ret-42.ll
+; RUN: %clang -c -o %t/x.o %S/Inputs/var-x-42.ll
+; RUN: %clang -c -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, target={{arm[^6][^4].*}}, target=powerpc64{{.*}}
+; UNSUPPORTED: system-windows
+; REQUIRES: target={{(arm|aarch)64.*}}
 ;
 ; CHECK: Linking {{.*}}main.o
 ; CHECK-DAG: Linking <indirect stubs graph #1>

compiler-rt/test/orc/lit.cfg.py

@@ -16,6 +16,11 @@
     host_arch_compatible = True
 if host_arch_compatible:
     config.available_features.add("host-arch-compatible")
+
+# If the target OS hasn't been set then assume host.
+if not config.target_os:
+    config.target_os = config.host_os
+
 config.test_target_is_host_executable = (
     config.target_os == config.host_os and host_arch_compatible
 )

llvm/include/llvm/ExecutionEngine/JITLink/aarch64.h

@@ -755,6 +755,32 @@ inline Symbol &createAnonymousPointerJumpStub(LinkGraph &G,
       sizeof(PointerJumpStubContent), true, false);
 }
 
+/// AArch64 reentry trampoline.
+///
+/// Contains the instruction sequence for a trampoline that stores its return
+/// address on the stack and passes its own address in x0:
+///   STP  x29, x30, [sp, #-16]!
+///   BL   <reentry-symbol>
+extern const char ReentryTrampolineContent[8];
+
+/// Create a block of N reentry trampolines.
+inline Block &createReentryTrampolineBlock(LinkGraph &G,
+                                           Section &TrampolineSection,
+                                           Symbol &ReentrySymbol) {
+  auto &B = G.createContentBlock(TrampolineSection, ReentryTrampolineContent,
+                                 orc::ExecutorAddr(~uint64_t(7)), 4, 0);
+  B.addEdge(Branch26PCRel, 4, ReentrySymbol, 0);
+  return B;
+}
+
+inline Symbol &createAnonymousReentryTrampoline(LinkGraph &G,
+                                                Section &TrampolineSection,
+                                                Symbol &ReentrySymbol) {
+  return G.addAnonymousSymbol(
+      createReentryTrampolineBlock(G, TrampolineSection, ReentrySymbol), 0,
+      sizeof(ReentryTrampolineContent), true, false);
+}
+
 /// Global Offset Table Builder.
 class GOTTableManager : public TableManager<GOTTableManager> {
 public:

llvm/include/llvm/ExecutionEngine/Orc/Core.h

@@ -51,6 +51,26 @@ enum class SymbolState : uint8_t;
 using ResourceTrackerSP = IntrusiveRefCntPtr<ResourceTracker>;
 using JITDylibSP = IntrusiveRefCntPtr<JITDylib>;
 
+/// A definition of a Symbol within a JITDylib.
+class SymbolInstance {
+public:
+  using LookupAsyncOnCompleteFn =
+      unique_function<void(Expected<ExecutorSymbolDef>)>;
+
+  SymbolInstance(JITDylibSP JD, SymbolStringPtr Name)
+      : JD(std::move(JD)), Name(std::move(Name)) {}
+
+  const JITDylib &getJITDylib() const { return *JD; }
+  const SymbolStringPtr &getName() const { return Name; }
+
+  Expected<ExecutorSymbolDef> lookup() const;
+  void lookupAsync(LookupAsyncOnCompleteFn OnComplete) const;
+
+private:
+  JITDylibSP JD;
+  SymbolStringPtr Name;
+};
+
 using ResourceKey = uintptr_t;
 
 /// API to remove / transfer ownership of JIT resources.
@@ -550,6 +570,9 @@ class MaterializationResponsibility {
   ///        emitted or notified of an error.
   ~MaterializationResponsibility();
 
+  /// Return the ResourceTracker associated with this instance.
+  const ResourceTrackerSP &getResourceTracker() const { return RT; }
+
   /// Runs the given callback under the session lock, passing in the associated
   /// ResourceKey. This is the safe way to associate resources with trackers.
   template <typename Func> Error withResourceKeyDo(Func &&F) const {
@@ -1748,6 +1771,10 @@ class ExecutionSession {
       JITDispatchHandlers;
 };
 
+inline Expected<ExecutorSymbolDef> SymbolInstance::lookup() const {
+  return JD->getExecutionSession().lookup({JD.get()}, Name);
+}
+
 template <typename Func> Error ResourceTracker::withResourceKeyDo(Func &&F) {
   return getJITDylib().getExecutionSession().runSessionLocked([&]() -> Error {
     if (isDefunct())

llvm/include/llvm/ExecutionEngine/Orc/JITLinkLazyCallThroughManager.h

@@ -0,0 +1,26 @@
+//===- JITLinkLazyCallThroughManager.h - JITLink based laziness -*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Redirectable Symbol Manager implementation using JITLink
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_JITLINKLAZYCALLTHROUGHMANAGER_H
+#define LLVM_EXECUTIONENGINE_ORC_JITLINKLAZYCALLTHROUGHMANAGER_H
+
+#include "llvm/ExecutionEngine/Orc/ObjectLinkingLayer.h"
+#include "llvm/ExecutionEngine/Orc/RedirectionManager.h"
+#include "llvm/Support/StringSaver.h"
+
+#include <atomic>
+
+namespace llvm {
+namespace orc {} // namespace orc
+} // namespace llvm
+
+#endif // LLVM_EXECUTIONENGINE_ORC_JITLINKLAZYCALLTHROUGHMANAGER_H

llvm/include/llvm/ExecutionEngine/Orc/JITLinkRedirectableSymbolManager.h

@@ -39,12 +39,6 @@ class JITLinkRedirectableSymbolManager : public RedirectableSymbolManager {
             ObjLinkingLayer, AnonymousPtrCreator, PtrJumpStubCreator));
   }
 
-  void emitRedirectableSymbols(std::unique_ptr<MaterializationResponsibility> R,
-                               SymbolMap InitialDests) override;
-
-  Error redirect(JITDylib &JD, const SymbolMap &NewDests) override;
-
-private:
   JITLinkRedirectableSymbolManager(
       ObjectLinkingLayer &ObjLinkingLayer,
       jitlink::AnonymousPointerCreator &AnonymousPtrCreator,
@@ -53,6 +47,14 @@ class JITLinkRedirectableSymbolManager : public RedirectableSymbolManager {
         AnonymousPtrCreator(std::move(AnonymousPtrCreator)),
         PtrJumpStubCreator(std::move(PtrJumpStubCreator)) {}
 
+  ObjectLinkingLayer &getObjectLinkingLayer() const { return ObjLinkingLayer; }
+
+  void emitRedirectableSymbols(std::unique_ptr<MaterializationResponsibility> R,
+                               SymbolMap InitialDests) override;
+
+  Error redirect(JITDylib &JD, const SymbolMap &NewDests) override;
+
+private:
   ObjectLinkingLayer &ObjLinkingLayer;
   jitlink::AnonymousPointerCreator AnonymousPtrCreator;
   jitlink::PointerJumpStubCreator PtrJumpStubCreator;