From 7d9758f33aa4ceabd63dc424907ba8628cbc0ae5 Mon Sep 17 00:00:00 2001
From: Lang Hames <lhames@gmail.com>
Date: Wed, 26 Sep 2018 01:24:12 +0000
Subject: [PATCH] [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
---
 .../Orc/CompileOnDemandLayer.h                |  12 +-
 .../llvm/ExecutionEngine/Orc/IRCompileLayer.h |   4 +-
 .../ExecutionEngine/Orc/IRTransformLayer.h    |   9 +-
 include/llvm/ExecutionEngine/Orc/LLJIT.h      |  15 ++-
 include/llvm/ExecutionEngine/Orc/Layer.h      |  33 +++--
 .../ExecutionEngine/Orc/ThreadSafeModule.h    | 124 ++++++++++++++++++
 lib/ExecutionEngine/Orc/CMakeLists.txt        |   1 +
 .../Orc/CompileOnDemandLayer.cpp              | 121 +++++++----------
 lib/ExecutionEngine/Orc/IRCompileLayer.cpp    |  10 +-
 lib/ExecutionEngine/Orc/IRTransformLayer.cpp  |  10 +-
 lib/ExecutionEngine/Orc/LLJIT.cpp             |  32 +++--
 lib/ExecutionEngine/Orc/Layer.cpp             |  31 +++--
 lib/ExecutionEngine/Orc/ThreadSafeModule.cpp  |  65 +++++++++
 tools/lli/lli.cpp                             |  53 ++++----
 .../Orc/RTDyldObjectLinkingLayer2Test.cpp     |  25 ++--
 15 files changed, 364 insertions(+), 181 deletions(-)
 create mode 100644 include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h
 create mode 100644 lib/ExecutionEngine/Orc/ThreadSafeModule.cpp

diff --git a/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h b/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h
index 20382f5086d..592f14843e9 100644
--- a/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h
+++ b/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.
diff --git a/include/llvm/ExecutionEngine/Orc/IRCompileLayer.h b/include/llvm/ExecutionEngine/Orc/IRCompileLayer.h
index ad6481548d5..cb8df26bfdc 100644
--- a/include/llvm/ExecutionEngine/Orc/IRCompileLayer.h
+++ b/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;
diff --git a/include/llvm/ExecutionEngine/Orc/IRTransformLayer.h b/include/llvm/ExecutionEngine/Orc/IRTransformLayer.h
index 266a0f45b3e..1d4d33e5df4 100644
--- a/include/llvm/ExecutionEngine/Orc/IRTransformLayer.h
+++ b/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:
diff --git a/include/llvm/ExecutionEngine/Orc/LLJIT.h b/include/llvm/ExecutionEngine/Orc/LLJIT.h
index 85aec7601b1..57b991f8c75 100644
--- a/include/llvm/ExecutionEngine/Orc/LLJIT.h
+++ b/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);
 
diff --git a/include/llvm/ExecutionEngine/Orc/Layer.h b/include/llvm/ExecutionEngine/Orc/Layer.h
index 013350f6233..ce3ed216d91 100644
--- a/include/llvm/ExecutionEngine/Orc/Layer.h
+++ b/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;
diff --git a/include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h b/include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h
new file mode 100644
index 00000000000..c82fb983bbb
--- /dev/null
+++ b/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
diff --git a/lib/ExecutionEngine/Orc/CMakeLists.txt b/lib/ExecutionEngine/Orc/CMakeLists.txt
index a7500ef20f3..59c9ee7364e 100644
--- a/lib/ExecutionEngine/Orc/CMakeLists.txt
+++ b/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
diff --git a/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp b/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp
index b731814293f..a68848f2f30 100644
--- a/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp
+++ b/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
diff --git a/lib/ExecutionEngine/Orc/IRCompileLayer.cpp b/lib/ExecutionEngine/Orc/IRCompileLayer.cpp
index 0c17f9b7ad4..5dee1c80e0b 100644
--- a/lib/ExecutionEngine/Orc/IRCompileLayer.cpp
+++ b/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 {
diff --git a/lib/ExecutionEngine/Orc/IRTransformLayer.cpp b/lib/ExecutionEngine/Orc/IRTransformLayer.cpp
index 4dd3cfdfe38..ddd5c4a10c1 100644
--- a/lib/ExecutionEngine/Orc/IRTransformLayer.cpp
+++ b/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());
   }
 }
 
diff --git a/lib/ExecutionEngine/Orc/LLJIT.cpp b/lib/ExecutionEngine/Orc/LLJIT.cpp
index 464b1e2413a..c79c47a0e33 100644
--- a/lib/ExecutionEngine/Orc/LLJIT.cpp
+++ b/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.
diff --git a/lib/ExecutionEngine/Orc/Layer.cpp b/lib/ExecutionEngine/Orc/Layer.cpp
index 7e2f830e4cb..f47634f853a 100644
--- a/lib/ExecutionEngine/Orc/Layer.cpp
+++ b/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) {}
diff --git a/lib/ExecutionEngine/Orc/ThreadSafeModule.cpp b/lib/ExecutionEngine/Orc/ThreadSafeModule.cpp
new file mode 100644
index 00000000000..c5a38aec102
--- /dev/null
+++ b/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
diff --git a/tools/lli/lli.cpp b/tools/lli/lli.cpp
index b88036e2c4f..2312d770348 100644
--- a/tools/lli/lli.cpp
+++ b/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.
diff --git a/unittests/ExecutionEngine/Orc/RTDyldObjectLinkingLayer2Test.cpp b/unittests/ExecutionEngine/Orc/RTDyldObjectLinkingLayer2Test.cpp
index b99fc494880..d5d15b5be9b 100644
--- a/unittests/ExecutionEngine/Orc/RTDyldObjectLinkingLayer2Test.cpp
+++ b/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.