[ORC] Add ThreadSafeModule and ThreadSafeContext wrappers to support concurrent

compilation of IR in the JIT.

ThreadSafeContext is a pair of an LLVMContext and a mutex that can be used to
lock that context when it needs to be accessed from multiple threads.

ThreadSafeModule is a pair of a unique_ptr<Module> and a
shared_ptr<ThreadSafeContext>. This allows the lifetime of a ThreadSafeContext
to be managed automatically in terms of the ThreadSafeModules that refer to it:
Once all modules using a ThreadSafeContext are destructed, and providing the
client has not held on to a copy of shared context pointer, the context will be
automatically destructed.

This scheme is necessary due to the following constraits: (1) We need multiple
contexts for multithreaded compilation (at least one per compile thread plus
one to store any IR not currently being compiled, though one context per module
is simpler). (2) We need to free contexts that are no longer being used so that
the JIT does not leak memory over time. (3) Module lifetimes are not
predictable (modules are compiled as needed depending on the flow of JIT'd
code) so there is no single point where contexts could be reclaimed.

JIT clients not using concurrency can safely use one ThreadSafeContext for all
ThreadSafeModules.

JIT clients who want to be able to compile concurrently should use a different
ThreadSafeContext for each module, or call setCloneToNewContextOnEmit on their
top-level IRLayer. The former reduces compile latency (since no clone step is
needed) at the cost of additional memory overhead for uncompiled modules (as
every uncompiled module will duplicate the LLVM types, constants and metadata
that have been shared).

llvm-svn: 343055

include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h

@@ -68,17 +68,14 @@ public:
   using IndirectStubsManagerBuilder =
       std::function<std::unique_ptr<IndirectStubsManager>()>;
 
-  using GetAvailableContextFunction = std::function<LLVMContext &()>;
-
   CompileOnDemandLayer2(ExecutionSession &ES, IRLayer &BaseLayer,
                         JITCompileCallbackManager &CCMgr,
-                        IndirectStubsManagerBuilder BuildIndirectStubsManager,
-                        GetAvailableContextFunction GetAvailableContext);
+                        IndirectStubsManagerBuilder BuildIndirectStubsManager);
 
-  Error add(JITDylib &V, VModuleKey K, std::unique_ptr<Module> M) override;
+  Error add(JITDylib &V, VModuleKey K, ThreadSafeModule TSM) override;
 
   void emit(MaterializationResponsibility R, VModuleKey K,
-            std::unique_ptr<Module> M) override;
+            ThreadSafeModule TSM) override;
 
 private:
   using StubManagersMap =
@@ -87,7 +84,7 @@ private:
   IndirectStubsManager &getStubsManager(const JITDylib &JD);
 
   void emitExtractedFunctionsModule(MaterializationResponsibility R,
-                                    std::unique_ptr<Module> M);
+                                    ThreadSafeModule TSM);
 
   mutable std::mutex CODLayerMutex;
 
@@ -95,7 +92,6 @@ private:
   JITCompileCallbackManager &CCMgr;
   IndirectStubsManagerBuilder BuildIndirectStubsManager;
   StubManagersMap StubsMgrs;
-  GetAvailableContextFunction GetAvailableContext;
 };
 
 /// Compile-on-demand layer.

include/llvm/ExecutionEngine/Orc/IRCompileLayer.h

@@ -34,7 +34,7 @@ public:
       std::function<Expected<std::unique_ptr<MemoryBuffer>>(Module &)>;
 
   using NotifyCompiledFunction =
-      std::function<void(VModuleKey K, std::unique_ptr<Module>)>;
+      std::function<void(VModuleKey K, ThreadSafeModule TSM)>;
 
   IRCompileLayer2(ExecutionSession &ES, ObjectLayer &BaseLayer,
                   CompileFunction Compile);
@@ -42,7 +42,7 @@ public:
   void setNotifyCompiled(NotifyCompiledFunction NotifyCompiled);
 
   void emit(MaterializationResponsibility R, VModuleKey K,
-            std::unique_ptr<Module> M) override;
+            ThreadSafeModule TSM) override;
 
 private:
   mutable std::mutex IRLayerMutex;

include/llvm/ExecutionEngine/Orc/IRTransformLayer.h

@@ -25,9 +25,8 @@ namespace orc {
 
 class IRTransformLayer2 : public IRLayer {
 public:
-
   using TransformFunction =
-    std::function<Expected<std::unique_ptr<Module>>(std::unique_ptr<Module>)>;
+      std::function<Expected<ThreadSafeModule>(ThreadSafeModule)>;
 
   IRTransformLayer2(ExecutionSession &ES, IRLayer &BaseLayer,
                     TransformFunction Transform = identityTransform);
@@ -37,10 +36,10 @@ public:
   }
 
   void emit(MaterializationResponsibility R, VModuleKey K,
-            std::unique_ptr<Module> M) override;
+            ThreadSafeModule TSM) override;
 
-  static std::unique_ptr<Module> identityTransform(std::unique_ptr<Module> M) {
-    return M;
+  static ThreadSafeModule identityTransform(ThreadSafeModule TSM) {
+    return TSM;
   }
 
 private:

include/llvm/ExecutionEngine/Orc/LLJIT.h

@@ -21,6 +21,7 @@
 #include "llvm/ExecutionEngine/Orc/IRTransformLayer.h"
 #include "llvm/ExecutionEngine/Orc/ObjectTransformLayer.h"
 #include "llvm/ExecutionEngine/Orc/RTDyldObjectLinkingLayer.h"
+#include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h"
 #include "llvm/Target/TargetMachine.h"
 
 namespace llvm {
@@ -43,11 +44,11 @@ public:
   Error defineAbsolute(StringRef Name, JITEvaluatedSymbol Address);
 
   /// Adds an IR module to the given JITDylib.
-  Error addIRModule(JITDylib &JD, std::unique_ptr<Module> M);
+  Error addIRModule(JITDylib &JD, ThreadSafeModule TSM);
 
   /// Adds an IR module to the Main JITDylib.
-  Error addIRModule(std::unique_ptr<Module> M) {
-    return addIRModule(Main, std::move(M));
+  Error addIRModule(ThreadSafeModule TSM) {
+    return addIRModule(Main, std::move(TSM));
   }
 
   /// Adds an object file to the given JITDylib.
@@ -119,7 +120,7 @@ class LLLazyJIT : public LLJIT {
 public:
   /// Create an LLLazyJIT instance.
   static Expected<std::unique_ptr<LLLazyJIT>>
-  Create(std::unique_ptr<TargetMachine> TM, DataLayout DL, LLVMContext &Ctx);
+  Create(std::unique_ptr<TargetMachine> TM, DataLayout DL);
 
   /// Set an IR transform (e.g. pass manager pipeline) to run on each function
   /// when it is compiled.
@@ -128,16 +129,16 @@ public:
   }
 
   /// Add a module to be lazily compiled to JITDylib JD.
-  Error addLazyIRModule(JITDylib &JD, std::unique_ptr<Module> M);
+  Error addLazyIRModule(JITDylib &JD, ThreadSafeModule M);
 
   /// Add a module to be lazily compiled to the main JITDylib.
-  Error addLazyIRModule(std::unique_ptr<Module> M) {
+  Error addLazyIRModule(ThreadSafeModule M) {
     return addLazyIRModule(Main, std::move(M));
   }
 
 private:
   LLLazyJIT(std::unique_ptr<ExecutionSession> ES,
-            std::unique_ptr<TargetMachine> TM, DataLayout DL, LLVMContext &Ctx,
+            std::unique_ptr<TargetMachine> TM, DataLayout DL,
             std::unique_ptr<JITCompileCallbackManager> CCMgr,
             std::function<std::unique_ptr<IndirectStubsManager>()> ISMBuilder);
 

include/llvm/ExecutionEngine/Orc/Layer.h

@@ -15,6 +15,7 @@
 #define LLVM_EXECUTIONENGINE_ORC_LAYER_H
 
 #include "llvm/ExecutionEngine/Orc/Core.h"
+#include "llvm/ExecutionEngine/Orc/ThreadSafeModule.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Support/MemoryBuffer.h"
 
@@ -30,21 +31,36 @@ public:
   /// Returns the ExecutionSession for this layer.
   ExecutionSession &getExecutionSession() { return ES; }
 
+  /// Sets the CloneToNewContextOnEmit flag (false by default).
+  ///
+  /// When set, IR modules added to this layer will be cloned on to a new
+  /// context before emit is called. This can be used by clients who want
+  /// to load all IR using one LLVMContext (to save memory via type and
+  /// constant uniquing), but want to move Modules to fresh contexts before
+  /// compiling them to enable concurrent compilation.
+  /// Single threaded clients, or clients who load every module on a new
+  /// context, need not set this.
+  void setCloneToNewContextOnEmit(bool CloneToNewContextOnEmit);
+
+  /// Returns the current value of the CloneToNewContextOnEmit flag.
+  bool getCloneToNewContextOnEmit() const { return CloneToNewContextOnEmit; }
+
   /// Adds a MaterializationUnit representing the given IR to the given
   /// JITDylib.
-  virtual Error add(JITDylib &JD, VModuleKey K, std::unique_ptr<Module> M);
+  virtual Error add(JITDylib &JD, VModuleKey K, ThreadSafeModule TSM);
 
   /// Adds a MaterializationUnit representing the given IR to the main
   /// JITDylib.
-  Error add(VModuleKey K, std::unique_ptr<Module> M) {
-    return add(ES.getMainJITDylib(), K, std::move(M));
+  Error add(VModuleKey K, ThreadSafeModule TSM) {
+    return add(ES.getMainJITDylib(), K, std::move(TSM));
   }
 
   /// Emit should materialize the given IR.
   virtual void emit(MaterializationResponsibility R, VModuleKey K,
-                    std::unique_ptr<Module> M) = 0;
+                    ThreadSafeModule TSM) = 0;
 
 private:
+  bool CloneToNewContextOnEmit = false;
   ExecutionSession &ES;
 };
 
@@ -58,18 +74,18 @@ public:
 
   /// Create an IRMaterializationLayer. Scans the module to build the
   /// SymbolFlags and SymbolToDefinition maps.
-  IRMaterializationUnit(ExecutionSession &ES, std::unique_ptr<Module> M);
+  IRMaterializationUnit(ExecutionSession &ES, ThreadSafeModule TSM);
 
   /// Create an IRMaterializationLayer from a module, and pre-existing
   /// SymbolFlags and SymbolToDefinition maps. The maps must provide
   /// entries for each definition in M.
   /// This constructor is useful for delegating work from one
   /// IRMaterializationUnit to another.
-  IRMaterializationUnit(std::unique_ptr<Module> M, SymbolFlagsMap SymbolFlags,
+  IRMaterializationUnit(ThreadSafeModule TSM, SymbolFlagsMap SymbolFlags,
                         SymbolNameToDefinitionMap SymbolToDefinition);
 
 protected:
-  std::unique_ptr<Module> M;
+  ThreadSafeModule TSM;
   SymbolNameToDefinitionMap SymbolToDefinition;
 
 private:
@@ -81,7 +97,8 @@ private:
 class BasicIRLayerMaterializationUnit : public IRMaterializationUnit {
 public:
   BasicIRLayerMaterializationUnit(IRLayer &L, VModuleKey K,
-                                  std::unique_ptr<Module> M);
+                                  ThreadSafeModule TSM);
+
 private:
 
   void materialize(MaterializationResponsibility R) override;

include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h

@@ -0,0 +1,124 @@
+//===----------- ThreadSafeModule.h -- Layer interfaces ---------*- C++ -*-===//
+//
+//                     The LLVM Compiler Infrastructure
+//
+// This file is distributed under the University of Illinois Open Source
+// License. See LICENSE.TXT for details.
+//
+//===----------------------------------------------------------------------===//
+//
+// Thread safe wrappers and utilities for Module and LLVMContext.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_EXECUTIONENGINE_ORC_THREADSAFEMODULEWRAPPER_H
+#define LLVM_EXECUTIONENGINE_ORC_THREADSAFEMODULEWRAPPER_H
+
+#include "llvm/IR/LLVMContext.h"
+#include "llvm/IR/Module.h"
+
+namespace llvm {
+namespace orc {
+
+/// An LLVMContext together with an associated mutex that can be used to lock
+/// the context to prevent concurrent access by other threads.
+class ThreadSafeContext {
+private:
+
+  struct State {
+    State(std::unique_ptr<LLVMContext> Ctx)
+      : Ctx(std::move(Ctx)) {}
+
+    std::unique_ptr<LLVMContext> Ctx;
+    std::recursive_mutex Mutex;
+  };
+
+public:
+
+  // RAII based lock for ThreadSafeContext.
+  class Lock {
+  private:
+    using UnderlyingLock = std::lock_guard<std::recursive_mutex>;
+  public:
+
+    Lock(std::shared_ptr<State> S)
+      : S(std::move(S)),
+        L(llvm::make_unique<UnderlyingLock>(this->S->Mutex)) {}
+  private:
+    std::shared_ptr<State> S;
+    std::unique_ptr<UnderlyingLock> L;
+  };
+
+  /// Construct a null context.
+  ThreadSafeContext() = default;
+
+  /// Construct a ThreadSafeContext from the given LLVMContext.
+  ThreadSafeContext(std::unique_ptr<LLVMContext> NewCtx)
+      : S(std::make_shared<State>(std::move(NewCtx))) {
+    assert(S->Ctx != nullptr &&
+           "Can not construct a ThreadSafeContext from a nullptr");
+  }
+
+  /// Returns a pointer to the LLVMContext that was used to construct this
+  /// instance, or null if the instance was default constructed.
+  LLVMContext* getContext() {
+    return S ? S->Ctx.get() : nullptr;
+  }
+
+  Lock getLock() {
+    assert(S && "Can not lock an empty ThreadSafeContext");
+    return Lock(S);
+  }
+
+private:
+  std::shared_ptr<State> S;
+};
+
+/// An LLVM Module together with a shared ThreadSafeContext.
+class ThreadSafeModule {
+public:
+  /// Default construct a ThreadSafeModule. This results in a null module and
+  /// null context.
+  ThreadSafeModule() = default;
+
+  /// Construct a ThreadSafeModule from a unique_ptr<Module> and a
+  /// unique_ptr<LLVMContext>. This creates a new ThreadSafeContext from the
+  /// given context.
+  ThreadSafeModule(std::unique_ptr<Module> M,
+                   std::unique_ptr<LLVMContext> Ctx)
+    : M(std::move(M)), TSCtx(std::move(Ctx)) {}
+
+  ThreadSafeModule(std::unique_ptr<Module> M,
+                   ThreadSafeContext TSCtx)
+    : M(std::move(M)), TSCtx(std::move(TSCtx)) {}
+
+  Module* getModule() { return M.get(); }
+
+  ThreadSafeContext::Lock getContextLock() { return TSCtx.getLock(); }
+
+  explicit operator bool() {
+    if (M) {
+      assert(TSCtx.getContext() && "Non-null module must have non-null context");
+      return true;
+    }
+    return false;
+  }
+
+private:
+  std::unique_ptr<Module> M;
+  ThreadSafeContext TSCtx;
+};
+
+using GVPredicate = std::function<bool(const GlobalValue&)>;
+using GVModifier = std::function<void(GlobalValue&)>;
+
+/// Clones the given module on to a new context.
+ThreadSafeModule
+cloneToNewContext(ThreadSafeModule &TSMW,
+                  GVPredicate ShouldCloneDef = GVPredicate(),
+                  GVModifier UpdateClonedDefSource = GVModifier());
+
+} // End namespace orc
+} // End namespace llvm
+
+#endif // LLVM_EXECUTIONENGINE_ORC_THREADSAFEMODULEWRAPPER_H

lib/ExecutionEngine/Orc/CMakeLists.txt

@@ -16,6 +16,7 @@ add_llvm_library(LLVMOrcJIT
   OrcMCJITReplacement.cpp
   RPCUtils.cpp
   RTDyldObjectLinkingLayer.cpp
+  ThreadSafeModule.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${LLVM_MAIN_INCLUDE_DIR}/llvm/ExecutionEngine/Orc

lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp

@@ -8,12 +8,8 @@
 //===----------------------------------------------------------------------===//
 
 #include "llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h"
-#include "llvm/Bitcode/BitcodeReader.h"
-#include "llvm/Bitcode/BitcodeWriter.h"
 #include "llvm/IR/Mangler.h"
 #include "llvm/IR/Module.h"
-#include "llvm/Support/raw_ostream.h"
-#include "llvm/Transforms/Utils/Cloning.h"
 
 using namespace llvm;
 using namespace llvm::orc;
@@ -71,56 +67,33 @@ static void extractAliases(MaterializationResponsibility &R, Module &M,
   R.replace(symbolAliases(std::move(Aliases)));
 }
 
-static std::unique_ptr<Module>
-extractAndClone(Module &M, LLVMContext &NewContext, StringRef Suffix,
-                function_ref<bool(const GlobalValue *)> ShouldCloneDefinition) {
-  SmallVector<char, 1> ClonedModuleBuffer;
+static ThreadSafeModule extractAndClone(ThreadSafeModule &TSM, StringRef Suffix,
+                                        GVPredicate ShouldCloneDefinition) {
 
-  {
-    std::set<GlobalValue *> ClonedDefsInSrc;
-    ValueToValueMapTy VMap;
-    auto Tmp = CloneModule(M, VMap, [&](const GlobalValue *GV) {
-      if (ShouldCloneDefinition(GV)) {
-        ClonedDefsInSrc.insert(const_cast<GlobalValue *>(GV));
-        return true;
-      }
-      return false;
-    });
+  auto DeleteClonedDefsAndPromoteDeclLinkages = [](GlobalValue &GV) {
+    // Delete the definition and bump the linkage in the source module.
+    if (isa<Function>(GV)) {
+      auto &F = cast<Function>(GV);
+      F.deleteBody();
+      F.setPersonalityFn(nullptr);
+    } else if (isa<GlobalVariable>(GV)) {
+      cast<GlobalVariable>(GV).setInitializer(nullptr);
+    } else
+      llvm_unreachable("Unsupported global type");
 
-    for (auto *GV : ClonedDefsInSrc) {
-      // Delete the definition and bump the linkage in the source module.
-      if (isa<Function>(GV)) {
-        auto &F = *cast<Function>(GV);
-        F.deleteBody();
-        F.setPersonalityFn(nullptr);
-      } else if (isa<GlobalVariable>(GV)) {
-        cast<GlobalVariable>(GV)->setInitializer(nullptr);
-      } else
-        llvm_unreachable("Unsupported global type");
+    GV.setLinkage(GlobalValue::ExternalLinkage);
+  };
 
-      GV->setLinkage(GlobalValue::ExternalLinkage);
-    }
+  auto NewTSMod = cloneToNewContext(TSM, ShouldCloneDefinition,
+                                    DeleteClonedDefsAndPromoteDeclLinkages);
+  auto &M = *NewTSMod.getModule();
+  M.setModuleIdentifier((M.getModuleIdentifier() + Suffix).str());
 
-    BitcodeWriter BCWriter(ClonedModuleBuffer);
-
-    BCWriter.writeModule(*Tmp);
-    BCWriter.writeSymtab();
-    BCWriter.writeStrtab();
-  }
-
-  MemoryBufferRef ClonedModuleBufferRef(
-      StringRef(ClonedModuleBuffer.data(), ClonedModuleBuffer.size()),
-      "cloned module buffer");
-
-  auto ClonedModule =
-      cantFail(parseBitcodeFile(ClonedModuleBufferRef, NewContext));
-  ClonedModule->setModuleIdentifier((M.getName() + Suffix).str());
-  return ClonedModule;
+  return NewTSMod;
 }
 
-static std::unique_ptr<Module> extractGlobals(Module &M,
-                                              LLVMContext &NewContext) {
-  return extractAndClone(M, NewContext, ".globals", [](const GlobalValue *GV) {
+static ThreadSafeModule extractGlobals(ThreadSafeModule &TSM) {
+  return extractAndClone(TSM, ".globals", [](const GlobalValue &GV) {
     return isa<GlobalVariable>(GV);
   });
 }
@@ -132,14 +105,14 @@ class ExtractingIRMaterializationUnit : public IRMaterializationUnit {
 public:
   ExtractingIRMaterializationUnit(ExecutionSession &ES,
                                   CompileOnDemandLayer2 &Parent,
-                                  std::unique_ptr<Module> M)
-      : IRMaterializationUnit(ES, std::move(M)), Parent(Parent) {}
+                                  ThreadSafeModule TSM)
+      : IRMaterializationUnit(ES, std::move(TSM)), Parent(Parent) {}
 
-  ExtractingIRMaterializationUnit(std::unique_ptr<Module> M,
+  ExtractingIRMaterializationUnit(ThreadSafeModule TSM,
                                   SymbolFlagsMap SymbolFlags,
                                   SymbolNameToDefinitionMap SymbolToDefinition,
                                   CompileOnDemandLayer2 &Parent)
-      : IRMaterializationUnit(std::move(M), std::move(SymbolFlags),
+      : IRMaterializationUnit(std::move(TSM), std::move(SymbolFlags),
                               std::move(SymbolToDefinition)),
         Parent(Parent) {}
 
@@ -153,7 +126,7 @@ private:
     auto RequestedSymbols = R.getRequestedSymbols();
 
     // Extract the requested functions into a new module.
-    std::unique_ptr<Module> ExtractedFunctionsModule;
+    ThreadSafeModule ExtractedFunctionsModule;
     if (!RequestedSymbols.empty()) {
       std::string Suffix;
       std::set<const GlobalValue *> FunctionsToClone;
@@ -168,10 +141,9 @@ private:
 
       std::lock_guard<std::mutex> Lock(SourceModuleMutex);
       ExtractedFunctionsModule =
-          extractAndClone(*M, Parent.GetAvailableContext(), Suffix,
-                          [&](const GlobalValue *GV) -> bool {
-                            return FunctionsToClone.count(GV);
-                          });
+          extractAndClone(TSM, Suffix, [&](const GlobalValue &GV) -> bool {
+            return FunctionsToClone.count(&GV);
+          });
     }
 
     // Build a new ExtractingIRMaterializationUnit to delegate the unrequested
@@ -193,7 +165,7 @@ private:
              "SymbolFlags and SymbolToDefinition should have the same number "
              "of entries");
       R.replace(llvm::make_unique<ExtractingIRMaterializationUnit>(
-          std::move(M), std::move(DelegatedSymbolFlags),
+          std::move(TSM), std::move(DelegatedSymbolFlags),
           std::move(DelegatedSymbolToDefinition), Parent));
     }
 
@@ -215,31 +187,30 @@ private:
 
 CompileOnDemandLayer2::CompileOnDemandLayer2(
     ExecutionSession &ES, IRLayer &BaseLayer, JITCompileCallbackManager &CCMgr,
-    IndirectStubsManagerBuilder BuildIndirectStubsManager,
-    GetAvailableContextFunction GetAvailableContext)
+    IndirectStubsManagerBuilder BuildIndirectStubsManager)
     : IRLayer(ES), BaseLayer(BaseLayer), CCMgr(CCMgr),
-      BuildIndirectStubsManager(std::move(BuildIndirectStubsManager)),
-      GetAvailableContext(std::move(GetAvailableContext)) {}
+      BuildIndirectStubsManager(std::move(BuildIndirectStubsManager)) {}
 
 Error CompileOnDemandLayer2::add(JITDylib &V, VModuleKey K,
-                                 std::unique_ptr<Module> M) {
-  return IRLayer::add(V, K, std::move(M));
+                                 ThreadSafeModule TSM) {
+  return IRLayer::add(V, K, std::move(TSM));
 }
 
 void CompileOnDemandLayer2::emit(MaterializationResponsibility R, VModuleKey K,
-                                 std::unique_ptr<Module> M) {
+                                 ThreadSafeModule TSM) {
   auto &ES = getExecutionSession();
-  assert(M && "M should not be null");
+  assert(TSM && "M should not be null");
+  auto &M = *TSM.getModule();
 
-  for (auto &GV : M->global_values())
+  for (auto &GV : M.global_values())
     if (GV.hasWeakLinkage())
       GV.setLinkage(GlobalValue::ExternalLinkage);
 
-  MangleAndInterner Mangle(ES, M->getDataLayout());
+  MangleAndInterner Mangle(ES, M.getDataLayout());
 
-  extractAliases(R, *M, Mangle);
+  extractAliases(R, *TSM.getModule(), Mangle);
 
-  auto GlobalsModule = extractGlobals(*M, GetAvailableContext());
+  auto GlobalsModule = extractGlobals(TSM);
 
   // Delete the bodies of any available externally functions, rename the
   // rest, and build the compile callbacks.
@@ -247,7 +218,7 @@ void CompileOnDemandLayer2::emit(MaterializationResponsibility R, VModuleKey K,
       StubCallbacksAndLinkages;
   auto &TargetJD = R.getTargetJITDylib();
 
-  for (auto &F : M->functions()) {
+  for (auto &F : M.functions()) {
     if (F.isDeclaration())
       continue;
 
@@ -260,7 +231,7 @@ void CompileOnDemandLayer2::emit(MaterializationResponsibility R, VModuleKey K,
     assert(F.hasName() && "Function should have a name");
     std::string StubUnmangledName = F.getName();
     F.setName(F.getName() + "$body");
-    auto StubDecl = cloneFunctionDecl(*M, F);
+    auto StubDecl = cloneFunctionDecl(*TSM.getModule(), F);
     StubDecl->setName(StubUnmangledName);
     StubDecl->setPersonalityFn(nullptr);
     StubDecl->setLinkage(GlobalValue::ExternalLinkage);
@@ -296,7 +267,7 @@ void CompileOnDemandLayer2::emit(MaterializationResponsibility R, VModuleKey K,
   // Build the function-body-extracting materialization unit.
   if (auto Err = R.getTargetJITDylib().define(
           llvm::make_unique<ExtractingIRMaterializationUnit>(ES, *this,
-                                                             std::move(M)))) {
+                                                             std::move(TSM)))) {
     ES.reportError(std::move(Err));
     R.failMaterialization();
     return;
@@ -335,9 +306,9 @@ CompileOnDemandLayer2::getStubsManager(const JITDylib &V) {
 }
 
 void CompileOnDemandLayer2::emitExtractedFunctionsModule(
-    MaterializationResponsibility R, std::unique_ptr<Module> M) {
+    MaterializationResponsibility R, ThreadSafeModule TSM) {
   auto K = getExecutionSession().allocateVModule();
-  BaseLayer.emit(std::move(R), std::move(K), std::move(M));
+  BaseLayer.emit(std::move(R), std::move(K), std::move(TSM));
 }
 
 } // end namespace orc

lib/ExecutionEngine/Orc/IRCompileLayer.cpp

@@ -22,16 +22,16 @@ void IRCompileLayer2::setNotifyCompiled(NotifyCompiledFunction NotifyCompiled) {
 }
 
 void IRCompileLayer2::emit(MaterializationResponsibility R, VModuleKey K,
-                           std::unique_ptr<Module> M) {
-  assert(M && "Module must not be null");
+                           ThreadSafeModule TSM) {
+  assert(TSM.getModule() && "Module must not be null");
 
-  if (auto Obj = Compile(*M)) {
+  if (auto Obj = Compile(*TSM.getModule())) {
     {
       std::lock_guard<std::mutex> Lock(IRLayerMutex);
       if (NotifyCompiled)
-        NotifyCompiled(K, std::move(M));
+        NotifyCompiled(K, std::move(TSM));
       else
-        M = nullptr;
+        TSM = ThreadSafeModule();
     }
     BaseLayer.emit(std::move(R), std::move(K), std::move(*Obj));
   } else {

lib/ExecutionEngine/Orc/IRTransformLayer.cpp

@@ -19,14 +19,14 @@ IRTransformLayer2::IRTransformLayer2(ExecutionSession &ES,
     : IRLayer(ES), BaseLayer(BaseLayer), Transform(std::move(Transform)) {}
 
 void IRTransformLayer2::emit(MaterializationResponsibility R, VModuleKey K,
-                             std::unique_ptr<Module> M) {
-  assert(M && "Module must not be null");
+                             ThreadSafeModule TSM) {
+  assert(TSM.getModule() && "Module must not be null");
 
-  if (auto TransformedMod = Transform(std::move(M)))
-    BaseLayer.emit(std::move(R), std::move(K), std::move(*TransformedMod));
+  if (auto TransformedTSM = Transform(std::move(TSM)))
+    BaseLayer.emit(std::move(R), std::move(K), std::move(*TransformedTSM));
   else {
     R.failMaterialization();
-    getExecutionSession().reportError(TransformedMod.takeError());
+    getExecutionSession().reportError(TransformedTSM.takeError());
   }
 }
 

lib/ExecutionEngine/Orc/LLJIT.cpp

@@ -27,14 +27,14 @@ Error LLJIT::defineAbsolute(StringRef Name, JITEvaluatedSymbol Sym) {
   return Main.define(absoluteSymbols(std::move(Symbols)));
 }
 
-Error LLJIT::addIRModule(JITDylib &JD, std::unique_ptr<Module> M) {
-  assert(M && "Can not add null module");
+Error LLJIT::addIRModule(JITDylib &JD, ThreadSafeModule TSM) {
+  assert(TSM && "Can not add null module");
 
-  if (auto Err = applyDataLayout(*M))
+  if (auto Err = applyDataLayout(*TSM.getModule()))
     return Err;
 
   auto K = ES->allocateVModule();
-  return CompileLayer.add(JD, K, std::move(M));
+  return CompileLayer.add(JD, K, std::move(TSM));
 }
 
 Error LLJIT::addObjectFile(JITDylib &JD, std::unique_ptr<MemoryBuffer> Obj) {
@@ -90,8 +90,7 @@ void LLJIT::recordCtorDtors(Module &M) {
 }
 
 Expected<std::unique_ptr<LLLazyJIT>>
-LLLazyJIT::Create(std::unique_ptr<TargetMachine> TM, DataLayout DL,
-                  LLVMContext &Ctx) {
+LLLazyJIT::Create(std::unique_ptr<TargetMachine> TM, DataLayout DL) {
   auto ES = llvm::make_unique<ExecutionSession>();
 
   const Triple &TT = TM->getTargetTriple();
@@ -109,33 +108,32 @@ LLLazyJIT::Create(std::unique_ptr<TargetMachine> TM, DataLayout DL,
         inconvertibleErrorCode());
 
   return std::unique_ptr<LLLazyJIT>(
-      new LLLazyJIT(std::move(ES), std::move(TM), std::move(DL), Ctx,
+      new LLLazyJIT(std::move(ES), std::move(TM), std::move(DL),
                     std::move(CCMgr), std::move(ISMBuilder)));
 }
 
-Error LLLazyJIT::addLazyIRModule(JITDylib &JD, std::unique_ptr<Module> M) {
-  assert(M && "Can not add null module");
+Error LLLazyJIT::addLazyIRModule(JITDylib &JD, ThreadSafeModule TSM) {
+  assert(TSM && "Can not add null module");
 
-  if (auto Err = applyDataLayout(*M))
+  if (auto Err = applyDataLayout(*TSM.getModule()))
     return Err;
 
-  makeAllSymbolsExternallyAccessible(*M);
+  makeAllSymbolsExternallyAccessible(*TSM.getModule());
 
-  recordCtorDtors(*M);
+  recordCtorDtors(*TSM.getModule());
 
   auto K = ES->allocateVModule();
-  return CODLayer.add(JD, K, std::move(M));
+  return CODLayer.add(JD, K, std::move(TSM));
 }
 
 LLLazyJIT::LLLazyJIT(
     std::unique_ptr<ExecutionSession> ES, std::unique_ptr<TargetMachine> TM,
-    DataLayout DL, LLVMContext &Ctx,
-    std::unique_ptr<JITCompileCallbackManager> CCMgr,
+    DataLayout DL, std::unique_ptr<JITCompileCallbackManager> CCMgr,
     std::function<std::unique_ptr<IndirectStubsManager>()> ISMBuilder)
     : LLJIT(std::move(ES), std::move(TM), std::move(DL)),
       CCMgr(std::move(CCMgr)), TransformLayer(*this->ES, CompileLayer),
-      CODLayer(*this->ES, TransformLayer, *this->CCMgr, std::move(ISMBuilder),
-               [&]() -> LLVMContext & { return Ctx; }) {}
+      CODLayer(*this->ES, TransformLayer, *this->CCMgr, std::move(ISMBuilder)) {
+}
 
 } // End namespace orc.
 } // End namespace llvm.

lib/ExecutionEngine/Orc/Layer.cpp

@@ -16,17 +16,19 @@ namespace orc {
 IRLayer::IRLayer(ExecutionSession &ES) : ES(ES) {}
 IRLayer::~IRLayer() {}
 
-Error IRLayer::add(JITDylib &JD, VModuleKey K, std::unique_ptr<Module> M) {
+Error IRLayer::add(JITDylib &JD, VModuleKey K, ThreadSafeModule TSM) {
   return JD.define(llvm::make_unique<BasicIRLayerMaterializationUnit>(
-      *this, std::move(K), std::move(M)));
+      *this, std::move(K), std::move(TSM)));
 }
 
 IRMaterializationUnit::IRMaterializationUnit(ExecutionSession &ES,
-                                             std::unique_ptr<Module> M)
-  : MaterializationUnit(SymbolFlagsMap()), M(std::move(M)) {
+                                             ThreadSafeModule TSM)
+    : MaterializationUnit(SymbolFlagsMap()), TSM(std::move(TSM)) {
 
-  MangleAndInterner Mangle(ES, this->M->getDataLayout());
-  for (auto &G : this->M->global_values()) {
+  assert(this->TSM && "Module must not be null");
+
+  MangleAndInterner Mangle(ES, this->TSM.getModule()->getDataLayout());
+  for (auto &G : this->TSM.getModule()->global_values()) {
     if (G.hasName() && !G.isDeclaration() && !G.hasLocalLinkage() &&
         !G.hasAvailableExternallyLinkage() && !G.hasAppendingLinkage()) {
       auto MangledName = Mangle(G.getName());
@@ -37,9 +39,9 @@ IRMaterializationUnit::IRMaterializationUnit(ExecutionSession &ES,
 }
 
 IRMaterializationUnit::IRMaterializationUnit(
-    std::unique_ptr<Module> M, SymbolFlagsMap SymbolFlags,
+    ThreadSafeModule TSM, SymbolFlagsMap SymbolFlags,
     SymbolNameToDefinitionMap SymbolToDefinition)
-    : MaterializationUnit(std::move(SymbolFlags)), M(std::move(M)),
+    : MaterializationUnit(std::move(SymbolFlags)), TSM(std::move(TSM)),
       SymbolToDefinition(std::move(SymbolToDefinition)) {}
 
 void IRMaterializationUnit::discard(const JITDylib &JD, SymbolStringPtr Name) {
@@ -53,13 +55,18 @@ void IRMaterializationUnit::discard(const JITDylib &JD, SymbolStringPtr Name) {
 }
 
 BasicIRLayerMaterializationUnit::BasicIRLayerMaterializationUnit(
-    IRLayer &L, VModuleKey K, std::unique_ptr<Module> M)
-  : IRMaterializationUnit(L.getExecutionSession(), std::move(M)),
-      L(L), K(std::move(K)) {}
+    IRLayer &L, VModuleKey K, ThreadSafeModule TSM)
+    : IRMaterializationUnit(L.getExecutionSession(), std::move(TSM)), L(L),
+      K(std::move(K)) {}
 
 void BasicIRLayerMaterializationUnit::materialize(
     MaterializationResponsibility R) {
-  L.emit(std::move(R), std::move(K), std::move(M));
+
+  if (L.getCloneToNewContextOnEmit())
+    TSM = cloneToNewContext(TSM);
+
+  auto Lock = TSM.getContextLock();
+  L.emit(std::move(R), std::move(K), std::move(TSM));
 }
 
 ObjectLayer::ObjectLayer(ExecutionSession &ES) : ES(ES) {}

lib/ExecutionEngine/Orc/ThreadSafeModule.cpp

@@ -0,0 +1,65 @@
+//===-- ThreadSafeModule.cpp - Thread safe Module, Context, and Utilities h-===//
+//
+//                     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/ThreadSafeModule.h"
+#include "llvm/Bitcode/BitcodeReader.h"
+#include "llvm/Bitcode/BitcodeWriter.h"
+#include "llvm/Transforms/Utils/Cloning.h"
+
+namespace llvm {
+namespace orc {
+
+ThreadSafeModule cloneToNewContext(ThreadSafeModule &TSM,
+                                   GVPredicate ShouldCloneDef,
+                                   GVModifier UpdateClonedDefSource) {
+  assert(TSM && "Can not clone null module");
+
+  if (!ShouldCloneDef)
+    ShouldCloneDef = [](const GlobalValue&) { return true; };
+
+  auto Lock = TSM.getContextLock();
+
+  SmallVector<char, 1> ClonedModuleBuffer;
+
+  {
+    std::vector<GlobalValue *> ClonedDefsInSrc;
+    ValueToValueMapTy VMap;
+    auto Tmp = CloneModule(*TSM.getModule(), VMap,
+                           [&](const GlobalValue *GV) {
+      if (ShouldCloneDef(*GV)) {
+        ClonedDefsInSrc.push_back(const_cast<GlobalValue *>(GV));
+        return true;
+      }
+      return false;
+    });
+
+    if (UpdateClonedDefSource)
+      for (auto *GV : ClonedDefsInSrc)
+        UpdateClonedDefSource(*GV);
+
+    BitcodeWriter BCWriter(ClonedModuleBuffer);
+
+    BCWriter.writeModule(*Tmp);
+    BCWriter.writeSymtab();
+    BCWriter.writeStrtab();
+  }
+
+  MemoryBufferRef ClonedModuleBufferRef(
+      StringRef(ClonedModuleBuffer.data(), ClonedModuleBuffer.size()),
+      "cloned module buffer");
+  ThreadSafeContext NewTSCtx(llvm::make_unique<LLVMContext>());
+
+  auto ClonedModule =
+      cantFail(parseBitcodeFile(ClonedModuleBufferRef, *NewTSCtx.getContext()));
+  ClonedModule->setModuleIdentifier(TSM.getModule()->getName());
+  return ThreadSafeModule(std::move(ClonedModule), std::move(NewTSCtx));
+}
+
+} // end namespace orc
+} // end namespace llvm

tools/lli/lli.cpp

@@ -678,13 +678,13 @@ int main(int argc, char **argv, char * const *envp) {
 static orc::IRTransformLayer2::TransformFunction createDebugDumper() {
   switch (OrcDumpKind) {
   case DumpKind::NoDump:
-    return [](std::unique_ptr<Module> M) { return M; };
+    return [](orc::ThreadSafeModule TSM) { return TSM; };
 
   case DumpKind::DumpFuncsToStdOut:
-    return [](std::unique_ptr<Module> M) {
+    return [](orc::ThreadSafeModule TSM) {
       printf("[ ");
 
-      for (const auto &F : *M) {
+      for (const auto &F : *TSM.getModule()) {
         if (F.isDeclaration())
           continue;
 
@@ -696,28 +696,29 @@ static orc::IRTransformLayer2::TransformFunction createDebugDumper() {
       }
 
       printf("]\n");
-      return M;
+      return TSM;
     };
 
   case DumpKind::DumpModsToStdOut:
-    return [](std::unique_ptr<Module> M) {
+    return [](orc::ThreadSafeModule TSM) {
       outs() << "----- Module Start -----\n"
-             << *M << "----- Module End -----\n";
+             << *TSM.getModule() << "----- Module End -----\n";
 
-      return M;
+      return TSM;
     };
 
   case DumpKind::DumpModsToDisk:
-    return [](std::unique_ptr<Module> M) {
+    return [](orc::ThreadSafeModule TSM) {
       std::error_code EC;
-      raw_fd_ostream Out(M->getModuleIdentifier() + ".ll", EC, sys::fs::F_Text);
+      raw_fd_ostream Out(TSM.getModule()->getModuleIdentifier() + ".ll", EC,
+                         sys::fs::F_Text);
       if (EC) {
-        errs() << "Couldn't open " << M->getModuleIdentifier()
+        errs() << "Couldn't open " << TSM.getModule()->getModuleIdentifier()
                << " for dumping.\nError:" << EC.message() << "\n";
         exit(1);
       }
-      Out << *M;
-      return M;
+      Out << *TSM.getModule();
+      return TSM;
     };
   }
   llvm_unreachable("Unknown DumpKind");
@@ -736,13 +737,14 @@ int runOrcLazyJIT(const char *ProgName) {
   }
 
   // Parse the main module.
-  LLVMContext Ctx;
+  orc::ThreadSafeContext TSCtx(llvm::make_unique<LLVMContext>());
   SMDiagnostic Err;
-  auto MainModule = parseIRFile(InputFile, Err, Ctx);
+  auto MainModule = orc::ThreadSafeModule(
+      parseIRFile(InputFile, Err, *TSCtx.getContext()), TSCtx);
   if (!MainModule)
     reportError(Err, ProgName);
 
-  const auto &TT = MainModule->getTargetTriple();
+  const auto &TT = MainModule.getModule()->getTargetTriple();
   orc::JITTargetMachineBuilder TMD =
       TT.empty() ? ExitOnErr(orc::JITTargetMachineBuilder::detectHost())
                  : orc::JITTargetMachineBuilder(Triple(TT));
@@ -758,18 +760,17 @@ int runOrcLazyJIT(const char *ProgName) {
                         : None);
   auto TM = ExitOnErr(TMD.createTargetMachine());
   auto DL = TM->createDataLayout();
-  auto J = ExitOnErr(orc::LLLazyJIT::Create(std::move(TM), DL, Ctx));
+  auto J = ExitOnErr(orc::LLLazyJIT::Create(std::move(TM), DL));
 
   auto Dump = createDebugDumper();
 
-  J->setLazyCompileTransform(
-    [&](std::unique_ptr<Module> M) {
-      if (verifyModule(*M, &dbgs())) {
-        dbgs() << "Bad module: " << *M << "\n";
-        exit(1);
-      }
-      return Dump(std::move(M));
-    });
+  J->setLazyCompileTransform([&](orc::ThreadSafeModule TSM) {
+    if (verifyModule(*TSM.getModule(), &dbgs())) {
+      dbgs() << "Bad module: " << *TSM.getModule() << "\n";
+      exit(1);
+    }
+    return Dump(std::move(TSM));
+  });
   J->getMainJITDylib().setFallbackDefinitionGenerator(
       orc::DynamicLibraryFallbackGenerator(
           std::move(LibLLI), DL, [](orc::SymbolStringPtr) { return true; }));
@@ -783,12 +784,12 @@ int runOrcLazyJIT(const char *ProgName) {
 
   // Add any extra modules.
   for (auto &ModulePath : ExtraModules) {
-    auto M = parseIRFile(ModulePath, Err, Ctx);
+    auto M = parseIRFile(ModulePath, Err, *TSCtx.getContext());
     if (!M)
       reportError(Err, ProgName);
 
     orc::makeAllSymbolsExternallyAccessible(*M);
-    ExitOnErr(J->addLazyIRModule(std::move(M)));
+    ExitOnErr(J->addLazyIRModule(orc::ThreadSafeModule(std::move(M), TSCtx)));
   }
 
   // Add the objects.

unittests/ExecutionEngine/Orc/RTDyldObjectLinkingLayer2Test.cpp

@@ -126,23 +126,25 @@ TEST(RTDyldObjectLinkingLayer2Test, TestOverrideObjectFlags) {
   };
 
   // Create a module with two void() functions: foo and bar.
-  LLVMContext Context;
-  std::unique_ptr<Module> M;
+  ThreadSafeContext TSCtx(llvm::make_unique<LLVMContext>());
+  ThreadSafeModule M;
   {
-    ModuleBuilder MB(Context, TM->getTargetTriple().str(), "dummy");
+    ModuleBuilder MB(*TSCtx.getContext(), TM->getTargetTriple().str(), "dummy");
     MB.getModule()->setDataLayout(TM->createDataLayout());
 
     Function *FooImpl = MB.createFunctionDecl<void()>("foo");
-    BasicBlock *FooEntry = BasicBlock::Create(Context, "entry", FooImpl);
+    BasicBlock *FooEntry =
+        BasicBlock::Create(*TSCtx.getContext(), "entry", FooImpl);
     IRBuilder<> B1(FooEntry);
     B1.CreateRetVoid();
 
     Function *BarImpl = MB.createFunctionDecl<void()>("bar");
-    BasicBlock *BarEntry = BasicBlock::Create(Context, "entry", BarImpl);
+    BasicBlock *BarEntry =
+        BasicBlock::Create(*TSCtx.getContext(), "entry", BarImpl);
     IRBuilder<> B2(BarEntry);
     B2.CreateRetVoid();
 
-    M = MB.takeModule();
+    M = ThreadSafeModule(MB.takeModule(), std::move(TSCtx));
   }
 
   // Create a simple stack and set the override flags option.
@@ -192,18 +194,19 @@ TEST(RTDyldObjectLinkingLayer2Test, TestAutoClaimResponsibilityForSymbols) {
   };
 
   // Create a module with two void() functions: foo and bar.
-  LLVMContext Context;
-  std::unique_ptr<Module> M;
+  ThreadSafeContext TSCtx(llvm::make_unique<LLVMContext>());
+  ThreadSafeModule M;
   {
-    ModuleBuilder MB(Context, TM->getTargetTriple().str(), "dummy");
+    ModuleBuilder MB(*TSCtx.getContext(), TM->getTargetTriple().str(), "dummy");
     MB.getModule()->setDataLayout(TM->createDataLayout());
 
     Function *FooImpl = MB.createFunctionDecl<void()>("foo");
-    BasicBlock *FooEntry = BasicBlock::Create(Context, "entry", FooImpl);
+    BasicBlock *FooEntry =
+        BasicBlock::Create(*TSCtx.getContext(), "entry", FooImpl);
     IRBuilder<> B(FooEntry);
     B.CreateRetVoid();
 
-    M = MB.takeModule();
+    M = ThreadSafeModule(MB.takeModule(), std::move(TSCtx));
   }
 
   // Create a simple stack and set the override flags option.