[ORC] Add a "lazy call-through" utility based on the same underlying trampoline

implementation as lazy compile callbacks, and a "lazy re-exports" utility that
builds lazy call-throughs.

Lazy call-throughs are similar to lazy compile callbacks (and are based on the
same underlying state saving/restoring trampolines) but resolve their targets
by performing a standard ORC lookup rather than invoking a user supplied
compiler callback. This allows them to inherit the thread-safety of ORC lookups
while blocking only the calling thread (whereas compile callbacks also block one
compile thread).

Lazy re-exports provide a simple way of building lazy call-throughs. Unlike a
regular re-export, a lazy re-export generates a new address (a stub entry point)
that will act like the re-exported symbol when called. The first call via a
lazy re-export will trigger compilation of the re-exported symbol before calling
through to it.

llvm-svn: 343061

include/llvm/ExecutionEngine/Orc/LazyReexports.h

@@ -0,0 +1,191 @@
+//===------ LazyReexports.h -- Utilities for lazy reexports -----*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// Lazy re-exports are similar to normal re-exports, except that for callable
+// symbols the definitions are replaced with trampolines that will look up and
+// call through to the re-exported symbol at runtime. This can be used to
+// enable lazy compilation.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_LAZYREEXPORTS_H
+#define LLVM_EXECUTIONENGINE_ORC_LAZYREEXPORTS_H
+
+#include "llvm/ExecutionEngine/Orc/Core.h"
+#include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
+
+namespace llvm {
+
+class Triple;
+
+namespace orc {
+
+/// Manages a set of 'lazy call-through' trampolines. These are compiler
+/// re-entry trampolines that are pre-bound to look up a given symbol in a given
+/// JITDylib, then jump to that address. Since compilation of symbols is
+/// triggered on first lookup, these call-through trampolines can be used to
+/// implement lazy compilation.
+///
+/// The easiest way to construct these call-throughs is using the lazyReexport
+/// function.
+class LazyCallThroughManager {
+public:
+  /// Clients will want to take some action on first resolution, e.g. updating
+  /// a stub pointer. Instances of this class can be used to implement this.
+  class NotifyResolvedFunction {
+  public:
+    virtual ~NotifyResolvedFunction() {}
+
+    /// Called the first time a lazy call through is executed and the target
+    /// symbol resolved.
+    virtual Error operator()(JITDylib &SourceJD,
+                             const SymbolStringPtr &SymbolName,
+                             JITTargetAddress ResolvedAddr) = 0;
+
+  private:
+    virtual void anchor();
+  };
+
+  template <typename NotifyResolvedImpl>
+  class NotifyResolvedFunctionImpl : public NotifyResolvedFunction {
+  public:
+    NotifyResolvedFunctionImpl(NotifyResolvedImpl NotifyResolved)
+        : NotifyResolved(std::move(NotifyResolved)) {}
+    Error operator()(JITDylib &SourceJD, const SymbolStringPtr &SymbolName,
+                     JITTargetAddress ResolvedAddr) {
+      return NotifyResolved(SourceJD, SymbolName, ResolvedAddr);
+    }
+
+  private:
+    NotifyResolvedImpl NotifyResolved;
+  };
+
+  /// Create a shared NotifyResolvedFunction from a given type that is
+  /// callable with the correct signature.
+  template <typename NotifyResolvedImpl>
+  static std::unique_ptr<NotifyResolvedFunction>
+  createNotifyResolvedFunction(NotifyResolvedImpl NotifyResolved) {
+    return llvm::make_unique<NotifyResolvedFunctionImpl<NotifyResolvedImpl>>(
+        std::move(NotifyResolved));
+  };
+
+  // Return a free call-through trampoline and bind it to look up and call
+  // through to the given symbol.
+  Expected<JITTargetAddress> getCallThroughTrampoline(
+      JITDylib &SourceJD, SymbolStringPtr SymbolName,
+      std::shared_ptr<NotifyResolvedFunction> NotifyResolved);
+
+protected:
+  LazyCallThroughManager(ExecutionSession &ES,
+                         JITTargetAddress ErrorHandlerAddr,
+                         std::unique_ptr<TrampolinePool> TP);
+
+  JITTargetAddress callThroughToSymbol(JITTargetAddress TrampolineAddr);
+
+  void setTrampolinePool(std::unique_ptr<TrampolinePool> TP) {
+    this->TP = std::move(TP);
+  }
+
+private:
+  using ReexportsMap =
+      std::map<JITTargetAddress, std::pair<JITDylib *, SymbolStringPtr>>;
+
+  using NotifiersMap =
+      std::map<JITTargetAddress, std::shared_ptr<NotifyResolvedFunction>>;
+
+  std::mutex LCTMMutex;
+  ExecutionSession &ES;
+  JITTargetAddress ErrorHandlerAddr;
+  std::unique_ptr<TrampolinePool> TP;
+  ReexportsMap Reexports;
+  NotifiersMap Notifiers;
+};
+
+/// A lazy call-through manager that builds trampolines in the current process.
+class LocalLazyCallThroughManager : public LazyCallThroughManager {
+private:
+  LocalLazyCallThroughManager(ExecutionSession &ES,
+                              JITTargetAddress ErrorHandlerAddr)
+      : LazyCallThroughManager(ES, ErrorHandlerAddr, nullptr) {}
+
+  template <typename ORCABI> Error init() {
+    auto TP = LocalTrampolinePool<ORCABI>::Create(
+        [this](JITTargetAddress TrampolineAddr) {
+          return callThroughToSymbol(TrampolineAddr);
+        });
+
+    if (!TP)
+      return TP.takeError();
+
+    setTrampolinePool(std::move(*TP));
+    return Error::success();
+  }
+
+public:
+  /// Create a LocalLazyCallThroughManager using the given ABI. See
+  /// createLocalLazyCallThroughManager.
+  template <typename ORCABI>
+  static Expected<std::unique_ptr<LocalLazyCallThroughManager>>
+  Create(ExecutionSession &ES, JITTargetAddress ErrorHandlerAddr) {
+    auto LLCTM = std::unique_ptr<LocalLazyCallThroughManager>(
+        new LocalLazyCallThroughManager(ES, ErrorHandlerAddr));
+
+    if (auto Err = LLCTM->init<ORCABI>())
+      return std::move(Err);
+
+    return std::move(LLCTM);
+  }
+};
+
+/// Create a LocalLazyCallThroughManager from the given triple and execution
+/// session.
+Expected<std::unique_ptr<LazyCallThroughManager>>
+createLocalLazyCallThroughManager(const Triple &T, ExecutionSession &ES,
+                                  JITTargetAddress ErrorHandlerAddr);
+
+/// A materialization unit that builds lazy re-exports. These are callable
+/// entry points that call through to the given symbols.
+/// Unlike a 'true' re-export, the address of the lazy re-export will not
+/// match the address of the re-exported symbol, but calling it will behave
+/// the same as calling the re-exported symbol.
+class LazyReexportsMaterializationUnit : public MaterializationUnit {
+public:
+  LazyReexportsMaterializationUnit(LazyCallThroughManager &LCTManager,
+                                   IndirectStubsManager &ISManager,
+                                   JITDylib &SourceJD,
+                                   SymbolAliasMap CallableAliases);
+
+private:
+  void materialize(MaterializationResponsibility R) override;
+  void discard(const JITDylib &JD, SymbolStringPtr Name) override;
+  static SymbolFlagsMap extractFlags(const SymbolAliasMap &Aliases);
+
+  LazyCallThroughManager &LCTManager;
+  IndirectStubsManager &ISManager;
+  JITDylib &SourceJD;
+  SymbolAliasMap CallableAliases;
+  std::shared_ptr<LazyCallThroughManager::NotifyResolvedFunction>
+      NotifyResolved;
+};
+
+/// Define lazy-reexports based on the given SymbolAliasMap. Each lazy re-export
+/// is a callable symbol that will look up and dispatch to the given aliasee on
+/// first call. All subsequent calls will go directly to the aliasee.
+inline std::unique_ptr<LazyReexportsMaterializationUnit>
+lazyReexports(LazyCallThroughManager &LCTManager,
+              IndirectStubsManager &ISManager, JITDylib &SourceJD,
+              SymbolAliasMap CallableAliases) {
+  return llvm::make_unique<LazyReexportsMaterializationUnit>(
+      LCTManager, ISManager, SourceJD, std::move(CallableAliases));
+}
+
+} // End namespace orc
+} // End namespace llvm
+
+#endif // LLVM_EXECUTIONENGINE_ORC_LAZYREEXPORTS_H

lib/ExecutionEngine/Orc/CMakeLists.txt

@@ -5,6 +5,7 @@ add_llvm_library(LLVMOrcJIT
   IndirectionUtils.cpp
   IRCompileLayer.cpp
   IRTransformLayer.cpp
+  LazyReexports.cpp
   Legacy.cpp
   Layer.cpp
   LLJIT.cpp

lib/ExecutionEngine/Orc/LazyReexports.cpp

@@ -0,0 +1,204 @@
+//===---------- LazyReexports.cpp - Utilities for lazy reexports ----------===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/ExecutionEngine/Orc/LazyReexports.h"
+
+#include "llvm/ADT/Triple.h"
+#include "llvm/ExecutionEngine/Orc/OrcABISupport.h"
+
+#define DEBUG_TYPE "orc"
+
+namespace llvm {
+namespace orc {
+
+void LazyCallThroughManager::NotifyResolvedFunction::anchor() {}
+
+LazyCallThroughManager::LazyCallThroughManager(
+    ExecutionSession &ES, JITTargetAddress ErrorHandlerAddr,
+    std::unique_ptr<TrampolinePool> TP)
+    : ES(ES), ErrorHandlerAddr(ErrorHandlerAddr), TP(std::move(TP)) {}
+
+Expected<JITTargetAddress> LazyCallThroughManager::getCallThroughTrampoline(
+    JITDylib &SourceJD, SymbolStringPtr SymbolName,
+    std::shared_ptr<NotifyResolvedFunction> NotifyResolved) {
+  std::lock_guard<std::mutex> Lock(LCTMMutex);
+  auto Trampoline = TP->getTrampoline();
+
+  if (!Trampoline)
+    return Trampoline.takeError();
+
+  Reexports[*Trampoline] = std::make_pair(&SourceJD, std::move(SymbolName));
+  Notifiers[*Trampoline] = std::move(NotifyResolved);
+  return *Trampoline;
+}
+
+JITTargetAddress
+LazyCallThroughManager::callThroughToSymbol(JITTargetAddress TrampolineAddr) {
+  JITDylib *SourceJD = nullptr;
+  SymbolStringPtr SymbolName;
+
+  {
+    std::lock_guard<std::mutex> Lock(LCTMMutex);
+    auto I = Reexports.find(TrampolineAddr);
+    if (I == Reexports.end())
+      return ErrorHandlerAddr;
+    SourceJD = I->second.first;
+    SymbolName = I->second.second;
+  }
+
+  auto LookupResult =
+      ES.lookup({SourceJD}, {SymbolName}, NoDependenciesToRegister);
+
+  if (!LookupResult) {
+    ES.reportError(LookupResult.takeError());
+    return ErrorHandlerAddr;
+  }
+
+  assert(LookupResult->size() == 1 && "Unexpected number of results");
+  assert(LookupResult->count(SymbolName) && "Unexpected result");
+
+  auto ResolvedAddr = LookupResult->begin()->second.getAddress();
+
+  std::shared_ptr<NotifyResolvedFunction> NotifyResolved = nullptr;
+  {
+    std::lock_guard<std::mutex> Lock(LCTMMutex);
+    auto I = Notifiers.find(TrampolineAddr);
+    if (I != Notifiers.end()) {
+      NotifyResolved = I->second;
+      Notifiers.erase(I);
+    }
+  }
+
+  if (NotifyResolved) {
+    if (auto Err = (*NotifyResolved)(*SourceJD, SymbolName, ResolvedAddr)) {
+      ES.reportError(std::move(Err));
+      return ErrorHandlerAddr;
+    }
+  }
+
+  return ResolvedAddr;
+}
+
+Expected<std::unique_ptr<LazyCallThroughManager>>
+createLocalLazyCallThroughManager(const Triple &T, ExecutionSession &ES,
+                                  JITTargetAddress ErrorHandlerAddr) {
+  switch (T.getArch()) {
+  default:
+    return make_error<StringError>(
+        std::string("No callback manager available for ") + T.str(),
+        inconvertibleErrorCode());
+
+  case Triple::aarch64:
+    return LocalLazyCallThroughManager::Create<OrcAArch64>(ES,
+                                                           ErrorHandlerAddr);
+
+  case Triple::x86:
+    return LocalLazyCallThroughManager::Create<OrcI386>(ES, ErrorHandlerAddr);
+
+  case Triple::mips:
+    return LocalLazyCallThroughManager::Create<OrcMips32Be>(ES,
+                                                            ErrorHandlerAddr);
+
+  case Triple::mipsel:
+    return LocalLazyCallThroughManager::Create<OrcMips32Le>(ES,
+                                                            ErrorHandlerAddr);
+
+  case Triple::mips64:
+  case Triple::mips64el:
+    return LocalLazyCallThroughManager::Create<OrcMips64>(ES, ErrorHandlerAddr);
+
+  case Triple::x86_64:
+    if (T.getOS() == Triple::OSType::Win32)
+      return LocalLazyCallThroughManager::Create<OrcX86_64_Win32>(
+          ES, ErrorHandlerAddr);
+    else
+      return LocalLazyCallThroughManager::Create<OrcX86_64_SysV>(
+          ES, ErrorHandlerAddr);
+  }
+}
+
+LazyReexportsMaterializationUnit::LazyReexportsMaterializationUnit(
+    LazyCallThroughManager &LCTManager, IndirectStubsManager &ISManager,
+    JITDylib &SourceJD, SymbolAliasMap CallableAliases)
+    : MaterializationUnit(extractFlags(CallableAliases)),
+      LCTManager(LCTManager), ISManager(ISManager), SourceJD(SourceJD),
+      CallableAliases(std::move(CallableAliases)),
+      NotifyResolved(LazyCallThroughManager::createNotifyResolvedFunction(
+          [&ISManager](JITDylib &JD, const SymbolStringPtr &SymbolName,
+                       JITTargetAddress ResolvedAddr) {
+            return ISManager.updatePointer(*SymbolName, ResolvedAddr);
+          })) {}
+
+void LazyReexportsMaterializationUnit::materialize(
+    MaterializationResponsibility R) {
+  auto RequestedSymbols = R.getRequestedSymbols();
+
+  SymbolAliasMap RequestedAliases;
+  for (auto &RequestedSymbol : RequestedSymbols) {
+    auto I = CallableAliases.find(RequestedSymbol);
+    assert(I != CallableAliases.end() && "Symbol not found in alias map?");
+    RequestedAliases[I->first] = std::move(I->second);
+    CallableAliases.erase(I);
+  }
+
+  if (!CallableAliases.empty())
+    R.replace(lazyReexports(LCTManager, ISManager, SourceJD,
+                            std::move(CallableAliases)));
+
+  IndirectStubsManager::StubInitsMap StubInits;
+  for (auto &Alias : RequestedAliases) {
+
+    auto CallThroughTrampoline = LCTManager.getCallThroughTrampoline(
+        SourceJD, Alias.second.Aliasee, NotifyResolved);
+
+    if (!CallThroughTrampoline) {
+      SourceJD.getExecutionSession().reportError(
+          CallThroughTrampoline.takeError());
+      R.failMaterialization();
+      return;
+    }
+
+    StubInits[*Alias.first] =
+        std::make_pair(*CallThroughTrampoline, Alias.second.AliasFlags);
+  }
+
+  if (auto Err = ISManager.createStubs(StubInits)) {
+    SourceJD.getExecutionSession().reportError(std::move(Err));
+    R.failMaterialization();
+    return;
+  }
+
+  SymbolMap Stubs;
+  for (auto &Alias : RequestedAliases)
+    Stubs[Alias.first] = ISManager.findStub(*Alias.first, false);
+
+  R.resolve(Stubs);
+  R.emit();
+}
+
+void LazyReexportsMaterializationUnit::discard(const JITDylib &JD,
+                                               SymbolStringPtr Name) {
+  assert(CallableAliases.count(Name) &&
+         "Symbol not covered by this MaterializationUnit");
+  CallableAliases.erase(Name);
+}
+
+SymbolFlagsMap
+LazyReexportsMaterializationUnit::extractFlags(const SymbolAliasMap &Aliases) {
+  SymbolFlagsMap SymbolFlags;
+  for (auto &KV : Aliases) {
+    assert(KV.second.AliasFlags.isCallable() &&
+           "Lazy re-exports must be callable symbols");
+    SymbolFlags[KV.first] = KV.second.AliasFlags;
+  }
+  return SymbolFlags;
+}
+
+} // End namespace orc.
+} // End namespace llvm.

unittests/ExecutionEngine/Orc/CMakeLists.txt

@@ -14,6 +14,7 @@ add_llvm_unittest(OrcJITTests
   CoreAPIsTest.cpp
   IndirectionUtilsTest.cpp
   GlobalMappingLayerTest.cpp
+  LazyCallThroughAndReexportsTest.cpp
   LazyEmittingLayerTest.cpp
   LegacyAPIInteropTest.cpp
   ObjectTransformLayerTest.cpp

unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp

@@ -22,44 +22,6 @@ class CoreAPIsStandardTest : public CoreAPIsBasedStandardTest {};
 
 namespace {
 
-class SimpleMaterializationUnit : public MaterializationUnit {
-public:
-  using MaterializeFunction =
-      std::function<void(MaterializationResponsibility)>;
-  using DiscardFunction =
-      std::function<void(const JITDylib &, SymbolStringPtr)>;
-  using DestructorFunction = std::function<void()>;
-
-  SimpleMaterializationUnit(
-      SymbolFlagsMap SymbolFlags, MaterializeFunction Materialize,
-      DiscardFunction Discard = DiscardFunction(),
-      DestructorFunction Destructor = DestructorFunction())
-      : MaterializationUnit(std::move(SymbolFlags)),
-        Materialize(std::move(Materialize)), Discard(std::move(Discard)),
-        Destructor(std::move(Destructor)) {}
-
-  ~SimpleMaterializationUnit() override {
-    if (Destructor)
-      Destructor();
-  }
-
-  void materialize(MaterializationResponsibility R) override {
-    Materialize(std::move(R));
-  }
-
-  void discard(const JITDylib &JD, SymbolStringPtr Name) override {
-    if (Discard)
-      Discard(JD, std::move(Name));
-    else
-      llvm_unreachable("Discard not supported");
-  }
-
-private:
-  MaterializeFunction Materialize;
-  DiscardFunction Discard;
-  DestructorFunction Destructor;
-};
-
 TEST_F(CoreAPIsStandardTest, BasicSuccessfulLookup) {
   bool OnResolutionRun = false;
   bool OnReadyRun = false;

unittests/ExecutionEngine/Orc/LazyCallThroughAndReexportsTest.cpp

@@ -0,0 +1,75 @@
+#include "OrcTestCommon.h"
+#include "llvm/ExecutionEngine/Orc/ExecutionUtils.h"
+#include "llvm/ExecutionEngine/Orc/LazyReexports.h"
+#include "gtest/gtest.h"
+
+using namespace llvm;
+using namespace llvm::orc;
+
+class LazyReexportsTest : public CoreAPIsBasedStandardTest {};
+
+static int dummyTarget() { return 42; }
+
+TEST_F(LazyReexportsTest, BasicLocalCallThroughManagerOperation) {
+  // Create a callthrough manager for the host (if possible) and verify that
+  // a call to the lazy call-through:
+  // (1) Materializes the MU. This verifies that the symbol was looked up, and
+  //     that we didn't arrive at the target via some other path
+  // (2) Returns the expected value (which we take as proof that the call
+  //     reached the target).
+
+  auto JTMB = JITTargetMachineBuilder::detectHost();
+
+  // Bail out if we can not detect the host.
+  if (!JTMB) {
+    consumeError(JTMB.takeError());
+    return;
+  }
+
+  // Bail out if we can not build a local call-through manager.
+  auto LCTM = createLocalLazyCallThroughManager(JTMB->getTargetTriple(), ES, 0);
+  if (!LCTM) {
+    consumeError(LCTM.takeError());
+    return;
+  }
+
+  auto DummyTarget = ES.getSymbolStringPool().intern("DummyTarget");
+
+  bool DummyTargetMaterialized = false;
+
+  cantFail(JD.define(llvm::make_unique<SimpleMaterializationUnit>(
+      SymbolFlagsMap({{DummyTarget, JITSymbolFlags::Exported}}),
+      [&](MaterializationResponsibility R) {
+        DummyTargetMaterialized = true;
+        R.resolve(
+            {{DummyTarget,
+              JITEvaluatedSymbol(static_cast<JITTargetAddress>(
+                                     reinterpret_cast<uintptr_t>(&dummyTarget)),
+                                 JITSymbolFlags::Exported)}});
+        R.emit();
+      })));
+
+  unsigned NotifyResolvedCount = 0;
+  auto NotifyResolved = LazyCallThroughManager::createNotifyResolvedFunction(
+      [&](JITDylib &JD, const SymbolStringPtr &SymbolName,
+          JITTargetAddress ResolvedAddr) {
+        ++NotifyResolvedCount;
+        return Error::success();
+      });
+
+  auto CallThroughTrampoline = cantFail((*LCTM)->getCallThroughTrampoline(
+      JD, DummyTarget, std::move(NotifyResolved)));
+
+  auto CTTPtr = reinterpret_cast<int (*)()>(
+      static_cast<uintptr_t>(CallThroughTrampoline));
+
+  // Call twice to verify nothing unexpected happens on redundant calls.
+  auto Result = CTTPtr();
+  (void)CTTPtr();
+
+  EXPECT_TRUE(DummyTargetMaterialized)
+      << "CallThrough did not materialize target";
+  EXPECT_EQ(NotifyResolvedCount, 1U)
+      << "CallThrough should have generated exactly one 'NotifyResolved' call";
+  EXPECT_EQ(Result, 42) << "Failed to call through to target";
+}

unittests/ExecutionEngine/Orc/OrcTestCommon.h

@@ -85,6 +85,44 @@ private:
   static bool NativeTargetInitialized;
 };
 
+class SimpleMaterializationUnit : public orc::MaterializationUnit {
+public:
+  using MaterializeFunction =
+      std::function<void(orc::MaterializationResponsibility)>;
+  using DiscardFunction =
+      std::function<void(const orc::JITDylib &, orc::SymbolStringPtr)>;
+  using DestructorFunction = std::function<void()>;
+
+  SimpleMaterializationUnit(
+      orc::SymbolFlagsMap SymbolFlags, MaterializeFunction Materialize,
+      DiscardFunction Discard = DiscardFunction(),
+      DestructorFunction Destructor = DestructorFunction())
+      : MaterializationUnit(std::move(SymbolFlags)),
+        Materialize(std::move(Materialize)), Discard(std::move(Discard)),
+        Destructor(std::move(Destructor)) {}
+
+  ~SimpleMaterializationUnit() override {
+    if (Destructor)
+      Destructor();
+  }
+
+  void materialize(orc::MaterializationResponsibility R) override {
+    Materialize(std::move(R));
+  }
+
+  void discard(const orc::JITDylib &JD, orc::SymbolStringPtr Name) override {
+    if (Discard)
+      Discard(JD, std::move(Name));
+    else
+      llvm_unreachable("Discard not supported");
+  }
+
+private:
+  MaterializeFunction Materialize;
+  DiscardFunction Discard;
+  DestructorFunction Destructor;
+};
+
 // Base class for Orc tests that will execute code.
 class OrcExecutionTest {
 public: