[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

7 files changed, 156 insertions, 114 deletions

diff --git a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt b/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
index a7500ef20f3..59c9ee7364e 100644
--- a/llvm/lib/ExecutionEngine/Orc/CMakeLists.txt
+++ b/llvm/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/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp b/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp
index b731814293f..a68848f2f30 100644
--- a/llvm/lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp
+++ b/llvm/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;
-
-  {
-    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;
-    });
-
-    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);
-    }
+static ThreadSafeModule extractAndClone(ThreadSafeModule &TSM, StringRef Suffix,
+                                        GVPredicate ShouldCloneDefinition) {
 
-    BitcodeWriter BCWriter(ClonedModuleBuffer);
+  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");
 
-    BCWriter.writeModule(*Tmp);
-    BCWriter.writeSymtab();
-    BCWriter.writeStrtab();
-  }
+    GV.setLinkage(GlobalValue::ExternalLinkage);
+  };
 
-  MemoryBufferRef ClonedModuleBufferRef(
-      StringRef(ClonedModuleBuffer.data(), ClonedModuleBuffer.size()),
-      "cloned module buffer");
+  auto NewTSMod = cloneToNewContext(TSM, ShouldCloneDefinition,
+                                    DeleteClonedDefsAndPromoteDeclLinkages);
+  auto &M = *NewTSMod.getModule();
+  M.setModuleIdentifier((M.getModuleIdentifier() + Suffix).str());
 
-  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/llvm/lib/ExecutionEngine/Orc/IRCompileLayer.cpp b/llvm/lib/ExecutionEngine/Orc/IRCompileLayer.cpp
index 0c17f9b7ad4..5dee1c80e0b 100644
--- a/llvm/lib/ExecutionEngine/Orc/IRCompileLayer.cpp
+++ b/llvm/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/llvm/lib/ExecutionEngine/Orc/IRTransformLayer.cpp b/llvm/lib/ExecutionEngine/Orc/IRTransformLayer.cpp
index 4dd3cfdfe38..ddd5c4a10c1 100644
--- a/llvm/lib/ExecutionEngine/Orc/IRTransformLayer.cpp
+++ b/llvm/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/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp b/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
index 464b1e2413a..c79c47a0e33 100644
--- a/llvm/lib/ExecutionEngine/Orc/LLJIT.cpp
+++ b/llvm/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/llvm/lib/ExecutionEngine/Orc/Layer.cpp b/llvm/lib/ExecutionEngine/Orc/Layer.cpp
index 7e2f830e4cb..f47634f853a 100644
--- a/llvm/lib/ExecutionEngine/Orc/Layer.cpp
+++ b/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) {}
diff --git a/llvm/lib/ExecutionEngine/Orc/ThreadSafeModule.cpp b/llvm/lib/ExecutionEngine/Orc/ThreadSafeModule.cpp
new file mode 100644
index 00000000000..c5a38aec102
--- /dev/null
+++ b/llvm/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