[ORC] Add an initial implementation of a replacement CompileOnDemandLayer.

CompileOnDemandLayer2 is a replacement for CompileOnDemandLayer built on the ORC
Core APIs. Functions in added modules are extracted and compiled lazily.
CompileOnDemandLayer2 supports multithreaded JIT'd code, and compilation on
multiple threads.

llvm-svn: 334967

include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h

@@ -20,9 +20,9 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/ExecutionEngine/JITSymbol.h"
-#include "llvm/ExecutionEngine/Orc/Core.h"
 #include "llvm/ExecutionEngine/Orc/IndirectionUtils.h"
 #include "llvm/ExecutionEngine/Orc/LambdaResolver.h"
+#include "llvm/ExecutionEngine/Orc/Layer.h"
 #include "llvm/ExecutionEngine/Orc/OrcError.h"
 #include "llvm/ExecutionEngine/RuntimeDyld.h"
 #include "llvm/IR/Attributes.h"
@@ -57,6 +57,66 @@ class Value;
 
 namespace orc {
 
+class ExtractingIRMaterializationUnit;
+
+class CompileOnDemandLayer2 : public IRLayer {
+  friend class ExtractingIRMaterializationUnit;
+
+public:
+  /// Builder for IndirectStubsManagers.
+  using IndirectStubsManagerBuilder =
+      std::function<std::unique_ptr<IndirectStubsManager>()>;
+
+  /// Retrieve symbol resolver for the given VModuleKey.
+  using GetSymbolResolverFunction =
+      std::function<std::shared_ptr<SymbolResolver>(VModuleKey K)>;
+
+  /// Set the symbol resolver for the given VModuleKey.
+  using SetSymbolResolverFunction =
+      std::function<void(VModuleKey K, std::shared_ptr<SymbolResolver> R)>;
+
+  using GetAvailableContextFunction = std::function<LLVMContext &()>;
+
+  CompileOnDemandLayer2(ExecutionSession &ES, IRLayer &BaseLayer,
+                        JITCompileCallbackManager &CCMgr,
+                        IndirectStubsManagerBuilder BuildIndirectStubsManager,
+                        GetSymbolResolverFunction GetSymbolResolver,
+                        SetSymbolResolverFunction SetSymbolResolver,
+                        GetAvailableContextFunction GetAvailableContext);
+
+  Error add(VSO &V, VModuleKey K, std::unique_ptr<Module> M) override;
+
+  void emit(MaterializationResponsibility R, VModuleKey K,
+            std::unique_ptr<Module> M) override;
+
+private:
+  using StubManagersMap =
+      std::map<const VSO *, std::unique_ptr<IndirectStubsManager>>;
+
+  using SymbolNameToDefinitionMap =
+      IRMaterializationUnit::SymbolNameToDefinitionMap;
+
+  IndirectStubsManager &getStubsManager(const VSO &V);
+
+  std::unique_ptr<Module>
+  extractFunctions(Module &M, const SymbolNameSet &SymbolNames,
+                   const SymbolNameToDefinitionMap &SymbolToDefiniton);
+
+  void emitExtractedFunctionsModule(MaterializationResponsibility R,
+                                    std::unique_ptr<Module> M,
+                                    std::shared_ptr<SymbolResolver> Resolver);
+
+  mutable std::mutex CODLayerMutex;
+
+  IRLayer &BaseLayer;
+  JITCompileCallbackManager &CCMgr;
+  IndirectStubsManagerBuilder BuildIndirectStubsManager;
+  StubManagersMap StubsMgrs;
+  GetSymbolResolverFunction GetSymbolResolver;
+  SetSymbolResolverFunction SetSymbolResolver;
+  GetAvailableContextFunction GetAvailableContext;
+};
+
 /// Compile-on-demand layer.
 ///
 ///   When a module is added to this layer a stub is created for each of its

include/llvm/ExecutionEngine/Orc/Core.h

@@ -115,7 +115,7 @@ public:
 
   /// Returns the target VSO that these symbols are being materialized
   ///        into.
-  const VSO &getTargetVSO() const { return V; }
+  VSO &getTargetVSO() const { return V; }
 
   /// Returns the names of any symbols covered by this
   /// MaterializationResponsibility object that have queries pending. This

lib/ExecutionEngine/Orc/CMakeLists.txt

@@ -1,4 +1,5 @@
 add_llvm_library(LLVMOrcJIT
+  CompileOnDemandLayer.cpp
   Core.cpp
   ExecutionUtils.cpp
   IndirectionUtils.cpp

lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp

@@ -0,0 +1,328 @@
+//===----- CompileOnDemandLayer.cpp - Lazily emit IR on first call --------===//
+//
+//                     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/CompileOnDemandLayer.h"
+#include "llvm/IR/Mangler.h"
+#include "llvm/IR/Module.h"
+#include "llvm/Support/raw_ostream.h"
+
+using namespace llvm;
+using namespace llvm::orc;
+
+namespace {
+
+template <typename MaterializerFtor>
+class LambdaValueMaterializer final : public ValueMaterializer {
+public:
+  LambdaValueMaterializer(MaterializerFtor M) : M(std::move(M)) {}
+
+  Value *materialize(Value *V) final { return M(V); }
+
+private:
+  MaterializerFtor M;
+};
+
+template <typename MaterializerFtor>
+LambdaValueMaterializer<MaterializerFtor>
+createLambdaValueMaterializer(MaterializerFtor M) {
+  return LambdaValueMaterializer<MaterializerFtor>(std::move(M));
+}
+} // namespace
+
+static std::unique_ptr<Module> extractGlobals(Module &M) {
+  // FIXME: Add alias support.
+
+  if (M.global_empty() && M.alias_empty() && !M.getModuleFlagsMetadata())
+    return nullptr;
+
+  auto GlobalsModule = llvm::make_unique<Module>(
+      (M.getName() + ".globals").str(), M.getContext());
+  GlobalsModule->setDataLayout(M.getDataLayout());
+
+  ValueToValueMapTy VMap;
+
+  for (auto &GV : M.globals())
+    if (!GV.isDeclaration() && !VMap.count(&GV))
+      cloneGlobalVariableDecl(*GlobalsModule, GV, &VMap);
+
+  // Clone the module flags.
+  cloneModuleFlagsMetadata(*GlobalsModule, M, VMap);
+
+  auto Materializer = createLambdaValueMaterializer([&](Value *V) -> Value * {
+    if (auto *F = dyn_cast<Function>(V))
+      return cloneFunctionDecl(*GlobalsModule, *F);
+    return nullptr;
+  });
+
+  // Move the global variable initializers.
+  for (auto &GV : M.globals()) {
+    if (!GV.isDeclaration())
+      moveGlobalVariableInitializer(GV, VMap, &Materializer);
+    GV.setInitializer(nullptr);
+  }
+
+  return GlobalsModule;
+}
+
+namespace llvm {
+namespace orc {
+
+class ExtractingIRMaterializationUnit : public IRMaterializationUnit {
+public:
+  ExtractingIRMaterializationUnit(
+      ExecutionSession &ES, CompileOnDemandLayer2 &Parent,
+      std::unique_ptr<Module> M,
+      std::shared_ptr<SymbolResolver> BackingResolver)
+      : IRMaterializationUnit(ES, std::move(M)), Parent(Parent),
+        BackingResolver(std::move(BackingResolver)) {}
+
+  ExtractingIRMaterializationUnit(
+      std::unique_ptr<Module> M, SymbolFlagsMap SymbolFlags,
+      SymbolNameToDefinitionMap SymbolToDefinition,
+      CompileOnDemandLayer2 &Parent,
+      std::shared_ptr<SymbolResolver> BackingResolver)
+      : IRMaterializationUnit(std::move(M), std::move(SymbolFlags),
+                              std::move(SymbolToDefinition)),
+        Parent(Parent), BackingResolver(std::move(BackingResolver)) {}
+
+private:
+  void materialize(MaterializationResponsibility R) override {
+    // FIXME: Need a 'notify lazy-extracting/emitting' callback to tie the
+    //        extracted module key, extracted module, and source module key
+    //        together. This could be used, for example, to provide a specific
+    //        memory manager instance to the linking layer.
+
+    // FIXME: The derived constructor should *only* look for the names of
+    //        original function definitions in the target VSO. All other
+    //        symbols should be looked up in the backing resolver.
+
+    // Find the functions that have been requested.
+    auto RequestedSymbols = R.getRequestedSymbols();
+
+    // Extract them into a new module.
+    auto ExtractedFunctionsModule =
+        Parent.extractFunctions(*M, RequestedSymbols, SymbolToDefinition);
+
+    // Build a new ExtractingIRMaterializationUnit to delegate the unrequested
+    // symbols to.
+    SymbolFlagsMap DelegatedSymbolFlags;
+    IRMaterializationUnit::SymbolNameToDefinitionMap
+        DelegatedSymbolToDefinition;
+    for (auto &KV : SymbolToDefinition) {
+      if (RequestedSymbols.count(KV.first))
+        continue;
+      DelegatedSymbolFlags[KV.first] =
+          JITSymbolFlags::fromGlobalValue(*KV.second);
+      DelegatedSymbolToDefinition[KV.first] = KV.second;
+    }
+
+    if (!DelegatedSymbolFlags.empty()) {
+      assert(DelegatedSymbolFlags.size() ==
+                 DelegatedSymbolToDefinition.size() &&
+             "SymbolFlags and SymbolToDefinition should have the same number "
+             "of entries");
+      R.delegate(llvm::make_unique<ExtractingIRMaterializationUnit>(
+          std::move(M), std::move(DelegatedSymbolFlags),
+          std::move(DelegatedSymbolToDefinition), Parent, BackingResolver));
+    }
+
+    Parent.emitExtractedFunctionsModule(
+        std::move(R), std::move(ExtractedFunctionsModule), BackingResolver);
+  }
+
+  void discard(const VSO &V, SymbolStringPtr Name) override {
+    // All original symbols were materialized by the CODLayer and should be
+    // final. The function bodies provided by M should never be overridden.
+    llvm_unreachable("Discard should never be called on an "
+                     "ExtractingIRMaterializationUnit");
+  }
+
+  CompileOnDemandLayer2 &Parent;
+  std::shared_ptr<SymbolResolver> BackingResolver;
+};
+
+CompileOnDemandLayer2::CompileOnDemandLayer2(
+    ExecutionSession &ES, IRLayer &BaseLayer, JITCompileCallbackManager &CCMgr,
+    IndirectStubsManagerBuilder BuildIndirectStubsManager,
+    GetSymbolResolverFunction GetSymbolResolver,
+    SetSymbolResolverFunction SetSymbolResolver,
+    GetAvailableContextFunction GetAvailableContext)
+    : IRLayer(ES), BaseLayer(BaseLayer), CCMgr(CCMgr),
+      BuildIndirectStubsManager(std::move(BuildIndirectStubsManager)),
+      GetSymbolResolver(std::move(GetSymbolResolver)),
+      SetSymbolResolver(std::move(SetSymbolResolver)),
+      GetAvailableContext(std::move(GetAvailableContext)) {}
+
+Error CompileOnDemandLayer2::add(VSO &V, VModuleKey K,
+                                 std::unique_ptr<Module> M) {
+  makeAllSymbolsExternallyAccessible(*M);
+  return IRLayer::add(V, K, std::move(M));
+}
+
+void CompileOnDemandLayer2::emit(MaterializationResponsibility R, VModuleKey K,
+                                 std::unique_ptr<Module> M) {
+  auto &ES = getExecutionSession();
+  assert(M && "M should not be null");
+
+  for (auto &GV : M->global_values())
+    if (GV.hasWeakLinkage())
+      GV.setLinkage(GlobalValue::ExternalLinkage);
+
+  auto GlobalsModule = extractGlobals(*M);
+
+  MangleAndInterner Mangle(ES, M->getDataLayout());
+
+  // Delete the bodies of any available externally functions, rename the
+  // rest, and build the compile callbacks.
+  std::map<SymbolStringPtr, std::pair<JITTargetAddress, JITSymbolFlags>>
+      StubCallbacksAndLinkages;
+  auto &TargetVSO = R.getTargetVSO();
+
+  for (auto &F : M->functions()) {
+    if (F.isDeclaration())
+      continue;
+
+    if (F.hasAvailableExternallyLinkage()) {
+      F.deleteBody();
+      continue;
+    }
+
+    assert(F.hasName() && "Function should have a name");
+    auto StubName = Mangle(F.getName());
+    F.setName(F.getName() + "$body");
+    auto BodyName = Mangle(F.getName());
+    if (auto CallbackAddr = CCMgr.getCompileCallback(
+            [BodyName, &TargetVSO, &ES]() -> JITTargetAddress {
+              if (auto Sym = lookup({&TargetVSO}, BodyName))
+                return Sym->getAddress();
+              else {
+                ES.reportError(Sym.takeError());
+                return 0;
+              }
+            })) {
+      auto Flags = JITSymbolFlags::fromGlobalValue(F);
+      Flags &= ~JITSymbolFlags::Weak;
+      StubCallbacksAndLinkages[std::move(StubName)] =
+          std::make_pair(*CallbackAddr, Flags);
+    } else {
+      ES.reportError(CallbackAddr.takeError());
+      R.failMaterialization();
+      return;
+    }
+  }
+
+  // Build the stub inits map.
+  IndirectStubsManager::StubInitsMap StubInits;
+  for (auto &KV : StubCallbacksAndLinkages)
+    StubInits[*KV.first] = KV.second;
+
+  // Build the function-body-extracting materialization unit.
+  if (auto Err = R.getTargetVSO().define(
+          llvm::make_unique<ExtractingIRMaterializationUnit>(
+              ES, *this, std::move(M), GetSymbolResolver(K)))) {
+    ES.reportError(std::move(Err));
+    R.failMaterialization();
+    return;
+  }
+
+  // Build the stubs.
+  // FIXME: Remove function bodies materialization unit if stub creation fails.
+  auto &StubsMgr = getStubsManager(TargetVSO);
+  if (auto Err = StubsMgr.createStubs(StubInits)) {
+    ES.reportError(std::move(Err));
+    R.failMaterialization();
+    return;
+  }
+
+  // Resolve and finalize stubs.
+  SymbolMap ResolvedStubs;
+  for (auto &KV : StubCallbacksAndLinkages) {
+    if (auto Sym = StubsMgr.findStub(*KV.first, false))
+      ResolvedStubs[KV.first] = Sym;
+    else
+      llvm_unreachable("Stub went missing");
+  }
+
+  R.resolve(ResolvedStubs);
+
+  BaseLayer.emit(std::move(R), std::move(K), std::move(GlobalsModule));
+}
+
+IndirectStubsManager &CompileOnDemandLayer2::getStubsManager(const VSO &V) {
+  std::lock_guard<std::mutex> Lock(CODLayerMutex);
+  StubManagersMap::iterator I = StubsMgrs.find(&V);
+  if (I == StubsMgrs.end())
+    I = StubsMgrs.insert(std::make_pair(&V, BuildIndirectStubsManager())).first;
+  return *I->second;
+}
+
+std::unique_ptr<Module> CompileOnDemandLayer2::extractFunctions(
+    Module &M, const SymbolNameSet &SymbolNames,
+    const SymbolNameToDefinitionMap &SymbolToDefinition) {
+  assert(!SymbolNames.empty() && "Can not extract an empty function set");
+
+  std::string ExtractedModName;
+  {
+    raw_string_ostream ExtractedModNameStream(ExtractedModName);
+    ExtractedModNameStream << M.getName();
+    for (auto &Name : SymbolNames)
+      ExtractedModNameStream << "." << *Name;
+  }
+
+  auto ExtractedFunctionsModule =
+      llvm::make_unique<Module>(ExtractedModName, GetAvailableContext());
+  ExtractedFunctionsModule->setDataLayout(M.getDataLayout());
+
+  ValueToValueMapTy VMap;
+
+  auto Materializer = createLambdaValueMaterializer([&](Value *V) -> Value * {
+    if (auto *F = dyn_cast<Function>(V))
+      return cloneFunctionDecl(*ExtractedFunctionsModule, *F);
+    else if (auto *GV = dyn_cast<GlobalVariable>(V))
+      return cloneGlobalVariableDecl(*ExtractedFunctionsModule, *GV);
+    return nullptr;
+  });
+
+  std::vector<std::pair<Function *, Function *>> OrigToNew;
+  for (auto &FunctionName : SymbolNames) {
+    assert(SymbolToDefinition.count(FunctionName) &&
+           "No definition for symbol");
+    auto *OrigF = cast<Function>(SymbolToDefinition.find(FunctionName)->second);
+    auto *NewF = cloneFunctionDecl(*ExtractedFunctionsModule, *OrigF, &VMap);
+    OrigToNew.push_back(std::make_pair(OrigF, NewF));
+  }
+
+  for (auto &KV : OrigToNew)
+    moveFunctionBody(*KV.first, VMap, &Materializer, KV.second);
+
+  return ExtractedFunctionsModule;
+}
+
+void CompileOnDemandLayer2::emitExtractedFunctionsModule(
+    MaterializationResponsibility R, std::unique_ptr<Module> M,
+    std::shared_ptr<SymbolResolver> Resolver) {
+  auto &TargetVSO = R.getTargetVSO();
+  auto K = getExecutionSession().allocateVModule();
+
+  auto ExtractedFunctionsResolver = createSymbolResolver(
+      [=](SymbolFlagsMap &Flags, const SymbolNameSet &Symbols) {
+        return Resolver->lookupFlags(Flags, Symbols);
+      },
+      [=, &TargetVSO](std::shared_ptr<AsynchronousSymbolQuery> Query,
+                      SymbolNameSet Symbols) {
+        auto RemainingSymbols = TargetVSO.lookup(Query, std::move(Symbols));
+        return Resolver->lookup(std::move(Query), std::move(RemainingSymbols));
+      });
+
+  SetSymbolResolver(K, std::move(ExtractedFunctionsResolver));
+  BaseLayer.emit(std::move(R), std::move(K), std::move(M));
+}
+
+} // end namespace orc
+} // end namespace llvm

lib/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.cpp

@@ -54,9 +54,24 @@ void RTDyldObjectLinkingLayer2::emit(MaterializationResponsibility R,
   auto Info = RTDyld->loadObject(**ObjFile);
 
   {
+    std::set<StringRef> InternalSymbols;
+    for (auto &Sym : (*ObjFile)->symbols()) {
+      if (!(Sym.getFlags() & object::BasicSymbolRef::SF_Global)) {
+        if (auto SymName = Sym.getName())
+          InternalSymbols.insert(*SymName);
+        else {
+          ES.reportError(SymName.takeError());
+          R.failMaterialization();
+          return;
+        }
+      }
+    }
+
     SymbolMap Symbols;
     for (auto &KV : RTDyld->getSymbolTable())
-      Symbols[ES.getSymbolStringPool().intern(KV.first)] = KV.second;
+      if (!InternalSymbols.count(KV.first))
+        Symbols[ES.getSymbolStringPool().intern(KV.first)] = KV.second;
+
     R.resolve(Symbols);
   }
 
@@ -74,6 +89,7 @@ void RTDyldObjectLinkingLayer2::emit(MaterializationResponsibility R,
     ES.reportError(make_error<StringError>(RTDyld->getErrorString(),
                                            inconvertibleErrorCode()));
     R.failMaterialization();
+    return;
   }
 
   R.finalize();