[ORC] Change the locking scheme for ThreadSafeModule.

ThreadSafeModule/ThreadSafeContext are used to manage lifetimes and locking
for LLVMContexts in ORCv2. Prior to this patch contexts were locked as soon
as an associated Module was emitted (to be compiled and linked), and were not
unlocked until the emit call returned. This could lead to deadlocks if
interdependent modules that shared contexts were compiled on different threads:
when, during emission of the first module, the dependence was discovered the
second module (which would provide the required symbol) could not be emitted as
the thread emitting the first module still held the lock.

This patch eliminates this possibility by moving to a finer-grained locking
scheme. Each client holds the module lock only while they are actively operating
on it. To make this finer grained locking simpler/safer to implement this patch
removes the explicit lock method, 'getContextLock', from ThreadSafeModule and
replaces it with a new method, 'withModuleDo', that implicitly locks the context,
calls a user-supplied function object to operate on the Module, then implicitly
unlocks the context before returning the result.

ThreadSafeModule TSM = getModule(...);
size_t NumFunctions = TSM.withModuleDo(
    [](Module &M) { // <- context locked before entry to lambda.
      return M.size();
    });

Existing ORCv2 layers that operate on ThreadSafeModules are updated to use the
new method.

This method is used to introduce Module locking into each of the existing
layers.

llvm-svn: 367686

docs/ORCv2.rst

@@ -153,7 +153,7 @@ Design Overview
 ORC's JIT'd program model aims to emulate the linking and symbol resolution
 rules used by the static and dynamic linkers. This allows ORC to JIT
 arbitrary LLVM IR, including IR produced by an ordinary static compiler (e.g.
-clang) that uses constructs like symbol linkage and visibility, and weak [4]_
+clang) that uses constructs like symbol linkage and visibility, and weak [3]_
 and common symbol definitions.
 
 To see how this works, imagine a program ``foo`` which links against a pair
@@ -441,7 +441,7 @@ ThreadSafeModule and ThreadSafeContext are wrappers around Modules and
 LLVMContexts respectively. A ThreadSafeModule is a pair of a
 std::unique_ptr<Module> and a (possibly shared) ThreadSafeContext value. A
 ThreadSafeContext is a pair of a std::unique_ptr<LLVMContext> and a lock.
-This design serves two purposes: providing both a locking scheme and lifetime
+This design serves two purposes: providing a locking scheme and lifetime
 management for LLVMContexts. The ThreadSafeContext may be locked to prevent
 accidental concurrent access by two Modules that use the same LLVMContext.
 The underlying LLVMContext is freed once all ThreadSafeContext values pointing
@@ -471,33 +471,49 @@ Before using a ThreadSafeContext, clients should ensure that either the context
 is only accessible on the current thread, or that the context is locked. In the
 example above (where the context is never locked) we rely on the fact that both
 ``TSM1`` and ``TSM2``, and TSCtx are all created on one thread. If a context is
-going to be shared between threads then it must be locked before the context,
-or any Modules attached to it, are accessed. When code is added to in-tree IR
-layers this locking is is done automatically by the
-``BasicIRLayerMaterializationUnit::materialize`` method. In all other
-situations, for example when writing a custom IR materialization unit, or
-constructing a new ThreadSafeModule from higher-level program representations,
-locking must be done explicitly:
+going to be shared between threads then it must be locked before any accessing
+or creating any Modules attached to it. E.g.
 
   .. code-block:: c++
 
-    void HighLevelRepresentationLayer::emit(MaterializationResponsibility R,
-                                            HighLevelProgramRepresentation H) {
-      // Get or create a context value that may be shared between threads.
-      ThreadSafeContext TSCtx = getContext();
 
-      // Lock the context to prevent concurrent access.
-      auto Lock = TSCtx.getLock();
+  ThreadSafeContext TSCtx(llvm::make_unique<LLVMContext>());
 
-      // IRGen a module onto the locked Context.
-      ThreadSafeModule TSM(IRGen(H, *TSCtx.getContext()), TSCtx);
+  ThreadPool TP(NumThreads);
+  JITStack J;
 
-      // Emit the module to the base layer with the context still locked.
-      BaseIRLayer.emit(std::move(R), std::move(TSM));
-    }
+  for (auto &ModulePath : ModulePaths) {
+    TP.async(
+      [&]() {
+        auto Lock = TSCtx.getLock();
+
+        auto M = loadModuleOnContext(ModulePath, TSCtx.getContext());
+
+        J.addModule(ThreadSafeModule(std::move(M), TSCtx));
+      });
+  }
+
+  TP.wait();
+
+To make exclusive access to Modules easier to manage the ThreadSafeModule class
+provides a convenince function, ``withModuleDo``, that implicitly (1) locks the
+associated context, (2) runs a given function object, (3) unlocks the context,
+and (3) returns the result generated by the function object. E.g.
+
+  .. code-block:: c++
+
+    ThreadSafeModule TSM = getModule(...);
+
+    // Dump the module:
+    size_t NumFunctionsInModule =
+      TSM.withModuleDo(
+        [](Module &M) { // <- Context locked before entering lambda.
+          return M.size();
+        } // <- Context unlocked after leaving.
+      );
 
 Clients wishing to maximize possibilities for concurrent compilation will want
-to create every new ThreadSafeModule on a new ThreadSafeContext [3]_. For this
+to create every new ThreadSafeModule on a new ThreadSafeContext. For this
 reason a convenience constructor for ThreadSafeModule is provided that implicitly
 constructs a new ThreadSafeContext value from a std::unique_ptr<LLVMContext>:
 
@@ -620,13 +636,7 @@ TBD: Speculative compilation. Object Caches.
        across processes, however this functionality appears not to have been
        used.
 
-.. [3] Sharing ThreadSafeModules in a concurrent compilation can be dangerous:
-       if interdependent modules are loaded on the same context, but compiled
-       on different threads a deadlock may occur, with each compile waiting for
-       the other to complete, and the other unable to proceed because the
-       context is locked.
-
-.. [4] Weak definitions are currently handled correctly within dylibs, but if
+.. [3] Weak definitions are currently handled correctly within dylibs, but if
        multiple dylibs provide a weak definition of a symbol then each will end
        up with its own definition (similar to how weak definitions are handled
        in Windows DLLs). This will be fixed in the future.

examples/Kaleidoscope/BuildingAJIT/Chapter2/KaleidoscopeJIT.h

@@ -87,20 +87,22 @@ public:
 private:
   static Expected<ThreadSafeModule>
   optimizeModule(ThreadSafeModule TSM, const MaterializationResponsibility &R) {
-    // Create a function pass manager.
-    auto FPM = llvm::make_unique<legacy::FunctionPassManager>(TSM.getModule());
+    TSM.withModuleDo([](Module &M) {
+      // Create a function pass manager.
+      auto FPM = llvm::make_unique<legacy::FunctionPassManager>(&M);
 
-    // Add some optimizations.
-    FPM->add(createInstructionCombiningPass());
-    FPM->add(createReassociatePass());
-    FPM->add(createGVNPass());
-    FPM->add(createCFGSimplificationPass());
-    FPM->doInitialization();
+      // Add some optimizations.
+      FPM->add(createInstructionCombiningPass());
+      FPM->add(createReassociatePass());
+      FPM->add(createGVNPass());
+      FPM->add(createCFGSimplificationPass());
+      FPM->doInitialization();
 
-    // Run the optimizations over all functions in the module being added to
-    // the JIT.
-    for (auto &F : *TSM.getModule())
-      FPM->run(F);
+      // Run the optimizations over all functions in the module being added to
+      // the JIT.
+      for (auto &F : M)
+        FPM->run(F);
+    });
 
     return TSM;
   }

include/llvm/ExecutionEngine/Orc/IRTransformLayer.h

@@ -22,6 +22,9 @@ namespace llvm {
 class Module;
 namespace orc {
 
+/// A layer that applies a transform to emitted modules.
+/// The transform function is responsible for locking the ThreadSafeContext
+/// before operating on the module.
 class IRTransformLayer : public IRLayer {
 public:
   using TransformFunction = std::function<Expected<ThreadSafeModule>(

include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h

@@ -69,7 +69,7 @@ public:
   /// instance, or null if the instance was default constructed.
   const LLVMContext *getContext() const { return S ? S->Ctx.get() : nullptr; }
 
-  Lock getLock() {
+  Lock getLock() const {
     assert(S && "Can not lock an empty ThreadSafeContext");
     return Lock(S);
   }
@@ -95,7 +95,7 @@ public:
     // We also need to lock the context to make sure the module tear-down
     // does not overlap any other work on the context.
     if (M) {
-      auto L = getContextLock();
+      auto L = TSCtx.getLock();
       M = nullptr;
     }
     M = std::move(Other.M);
@@ -117,23 +117,14 @@ public:
   ~ThreadSafeModule() {
     // We need to lock the context while we destruct the module.
     if (M) {
-      auto L = getContextLock();
+      auto L = TSCtx.getLock();
       M = nullptr;
     }
   }
 
-  /// Get the module wrapped by this ThreadSafeModule.
-  Module *getModule() { return M.get(); }
-
-  /// Get the module wrapped by this ThreadSafeModule.
-  const Module *getModule() const { return M.get(); }
-
-  /// Take out a lock on the ThreadSafeContext for this module.
-  ThreadSafeContext::Lock getContextLock() { return TSCtx.getLock(); }
-
   /// Boolean conversion: This ThreadSafeModule will evaluate to true if it
   /// wraps a non-null module.
-  explicit operator bool() {
+  explicit operator bool() const {
     if (M) {
       assert(TSCtx.getContext() &&
              "Non-null module must have non-null context");
@@ -142,6 +133,33 @@ public:
     return false;
   }
 
+  /// Locks the associated ThreadSafeContext and calls the given function
+  /// on the contained Module.
+  template <typename Func>
+  auto withModuleDo(Func &&F) -> decltype(F(std::declval<Module &>())) {
+    assert(M && "Can not call on null module");
+    auto Lock = TSCtx.getLock();
+    return F(*M);
+  }
+
+  /// Locks the associated ThreadSafeContext and calls the given function
+  /// on the contained Module.
+  template <typename Func>
+  auto withModuleDo(Func &&F) const
+      -> decltype(F(std::declval<const Module &>())) {
+    auto Lock = TSCtx.getLock();
+    return F(*M);
+  }
+
+  /// Get a raw pointer to the contained module without locking the context.
+  Module *getModuleUnlocked() { return M.get(); }
+
+  /// Get a raw pointer to the contained module without locking the context.
+  const Module *getModuleUnlocked() const { return M.get(); }
+
+  /// Returns the context for this ThreadSafeModule.
+  ThreadSafeContext getContext() const { return TSCtx; }
+
 private:
   std::unique_ptr<Module> M;
   ThreadSafeContext TSCtx;

lib/ExecutionEngine/Orc/CompileOnDemandLayer.cpp

@@ -54,11 +54,12 @@ static ThreadSafeModule extractSubModule(ThreadSafeModule &TSM,
       llvm_unreachable("Unsupported global type");
   };
 
-  auto NewTSMod = cloneToNewContext(TSM, ShouldExtract, DeleteExtractedDefs);
-  auto &M = *NewTSMod.getModule();
-  M.setModuleIdentifier((M.getModuleIdentifier() + Suffix).str());
+  auto NewTSM = cloneToNewContext(TSM, ShouldExtract, DeleteExtractedDefs);
+  NewTSM.withModuleDo([&](Module &M) {
+    M.setModuleIdentifier((M.getModuleIdentifier() + Suffix).str());
+  });
 
-  return NewTSMod;
+  return NewTSM;
 }
 
 namespace llvm {
@@ -119,32 +120,34 @@ void CompileOnDemandLayer::setPartitionFunction(PartitionFunction Partition) {
 
 void CompileOnDemandLayer::emit(MaterializationResponsibility R,
                                 ThreadSafeModule TSM) {
-  assert(TSM.getModule() && "Null module");
+  assert(TSM && "Null module");
 
   auto &ES = getExecutionSession();
-  auto &M = *TSM.getModule();
 
-  // First, do some cleanup on the module:
-  cleanUpModule(M);
-
-  // Now sort the callables and non-callables, build re-exports and lodge the
+  // Sort the callables and non-callables, build re-exports and lodge the
   // actual module with the implementation dylib.
   auto &PDR = getPerDylibResources(R.getTargetJITDylib());
 
-  MangleAndInterner Mangle(ES, M.getDataLayout());
   SymbolAliasMap NonCallables;
   SymbolAliasMap Callables;
-  for (auto &GV : M.global_values()) {
-    if (GV.isDeclaration() || GV.hasLocalLinkage() || GV.hasAppendingLinkage())
-      continue;
+  TSM.withModuleDo([&](Module &M) {
+    // First, do some cleanup on the module:
+    cleanUpModule(M);
 
-    auto Name = Mangle(GV.getName());
-    auto Flags = JITSymbolFlags::fromGlobalValue(GV);
-    if (Flags.isCallable())
-      Callables[Name] = SymbolAliasMapEntry(Name, Flags);
-    else
-      NonCallables[Name] = SymbolAliasMapEntry(Name, Flags);
-  }
+    MangleAndInterner Mangle(ES, M.getDataLayout());
+    for (auto &GV : M.global_values()) {
+      if (GV.isDeclaration() || GV.hasLocalLinkage() ||
+          GV.hasAppendingLinkage())
+        continue;
+
+      auto Name = Mangle(GV.getName());
+      auto Flags = JITSymbolFlags::fromGlobalValue(GV);
+      if (Flags.isCallable())
+        Callables[Name] = SymbolAliasMapEntry(Name, Flags);
+      else
+        NonCallables[Name] = SymbolAliasMapEntry(Name, Flags);
+    }
+  });
 
   // Create a partitioning materialization unit and lodge it with the
   // implementation dylib.
@@ -239,14 +242,16 @@ void CompileOnDemandLayer::emitPartition(
   //        memory manager instance to the linking layer.
 
   auto &ES = getExecutionSession();
-
   GlobalValueSet RequestedGVs;
   for (auto &Name : R.getRequestedSymbols()) {
     assert(Defs.count(Name) && "No definition for symbol");
     RequestedGVs.insert(Defs[Name]);
   }
 
-  auto GVsToExtract = Partition(RequestedGVs);
+  /// Perform partitioning with the context lock held, since the partition
+  /// function is allowed to access the globals to compute the partition.
+  auto GVsToExtract =
+      TSM.withModuleDo([&](Module &M) { return Partition(RequestedGVs); });
 
   // Take a 'None' partition to mean the whole module (as opposed to an empty
   // partition, which means "materialize nothing"). Emit the whole module
@@ -265,37 +270,46 @@ void CompileOnDemandLayer::emitPartition(
   }
 
   // Ok -- we actually need to partition the symbols. Promote the symbol
-  // linkages/names.
-  // FIXME: We apply this once per partitioning. It's safe, but overkill.
-  {
-    auto PromotedGlobals = PromoteSymbols(*TSM.getModule());
-    if (!PromotedGlobals.empty()) {
-      MangleAndInterner Mangle(ES, TSM.getModule()->getDataLayout());
-      SymbolFlagsMap SymbolFlags;
-      for (auto &GV : PromotedGlobals)
-        SymbolFlags[Mangle(GV->getName())] =
-            JITSymbolFlags::fromGlobalValue(*GV);
-      if (auto Err = R.defineMaterializing(SymbolFlags)) {
-        ES.reportError(std::move(Err));
-        R.failMaterialization();
-        return;
-      }
-    }
+  // linkages/names, expand the partition to include any required symbols
+  // (i.e. symbols that can't be separated from our partition), and
+  // then extract the partition.
+  //
+  // FIXME: We apply this promotion once per partitioning. It's safe, but
+  // overkill.
+
+  auto ExtractedTSM =
+      TSM.withModuleDo([&](Module &M) -> Expected<ThreadSafeModule> {
+        auto PromotedGlobals = PromoteSymbols(M);
+        if (!PromotedGlobals.empty()) {
+          MangleAndInterner Mangle(ES, M.getDataLayout());
+          SymbolFlagsMap SymbolFlags;
+          for (auto &GV : PromotedGlobals)
+            SymbolFlags[Mangle(GV->getName())] =
+                JITSymbolFlags::fromGlobalValue(*GV);
+          if (auto Err = R.defineMaterializing(SymbolFlags))
+            return std::move(Err);
+        }
+
+        expandPartition(*GVsToExtract);
+
+        // Extract the requested partiton (plus any necessary aliases) and
+        // put the rest back into the impl dylib.
+        auto ShouldExtract = [&](const GlobalValue &GV) -> bool {
+          return GVsToExtract->count(&GV);
+        };
+
+        return extractSubModule(TSM, ".submodule", ShouldExtract);
+      });
+
+  if (!ExtractedTSM) {
+    ES.reportError(ExtractedTSM.takeError());
+    R.failMaterialization();
+    return;
   }
 
-  expandPartition(*GVsToExtract);
-
-  // Extract the requested partiton (plus any necessary aliases) and
-  // put the rest back into the impl dylib.
-  auto ShouldExtract = [&](const GlobalValue &GV) -> bool {
-    return GVsToExtract->count(&GV);
-  };
-
-  auto ExtractedTSM = extractSubModule(TSM, ".submodule", ShouldExtract);
   R.replace(llvm::make_unique<PartitioningIRMaterializationUnit>(
       ES, std::move(TSM), R.getVModuleKey(), *this));
-
-  BaseLayer.emit(std::move(R), std::move(ExtractedTSM));
+  BaseLayer.emit(std::move(R), std::move(*ExtractedTSM));
 }
 
 } // end namespace orc

lib/ExecutionEngine/Orc/IRCompileLayer.cpp

@@ -22,9 +22,9 @@ void IRCompileLayer::setNotifyCompiled(NotifyCompiledFunction NotifyCompiled) {
 
 void IRCompileLayer::emit(MaterializationResponsibility R,
                           ThreadSafeModule TSM) {
-  assert(TSM.getModule() && "Module must not be null");
+  assert(TSM && "Module must not be null");
 
-  if (auto Obj = Compile(*TSM.getModule())) {
+  if (auto Obj = TSM.withModuleDo(Compile)) {
     {
       std::lock_guard<std::mutex> Lock(IRLayerMutex);
       if (NotifyCompiled)

lib/ExecutionEngine/Orc/IRTransformLayer.cpp

@@ -19,7 +19,7 @@ IRTransformLayer::IRTransformLayer(ExecutionSession &ES,
 
 void IRTransformLayer::emit(MaterializationResponsibility R,
                             ThreadSafeModule TSM) {
-  assert(TSM.getModule() && "Module must not be null");
+  assert(TSM && "Module must not be null");
 
   if (auto TransformedTSM = Transform(std::move(TSM), R))
     BaseLayer.emit(std::move(R), std::move(*TransformedTSM));

lib/ExecutionEngine/Orc/LLJIT.cpp

@@ -41,7 +41,8 @@ Error LLJIT::defineAbsolute(StringRef Name, JITEvaluatedSymbol Sym) {
 Error LLJIT::addIRModule(JITDylib &JD, ThreadSafeModule TSM) {
   assert(TSM && "Can not add null module");
 
-  if (auto Err = applyDataLayout(*TSM.getModule()))
+  if (auto Err =
+          TSM.withModuleDo([&](Module &M) { return applyDataLayout(M); }))
     return Err;
 
   return CompileLayer->add(JD, std::move(TSM), ES->allocateVModule());
@@ -166,10 +167,14 @@ Error LLLazyJITBuilderState::prepareForConstruction() {
 Error LLLazyJIT::addLazyIRModule(JITDylib &JD, ThreadSafeModule TSM) {
   assert(TSM && "Can not add null module");
 
-  if (auto Err = applyDataLayout(*TSM.getModule()))
-    return Err;
+  if (auto Err = TSM.withModuleDo([&](Module &M) -> Error {
+        if (auto Err = applyDataLayout(M))
+          return Err;
 
-  recordCtorDtors(*TSM.getModule());
+        recordCtorDtors(M);
+        return Error::success();
+      }))
+    return Err;
 
   return CODLayer->add(JD, std::move(TSM), ES->allocateVModule());
 }

lib/ExecutionEngine/Orc/Layer.cpp

@@ -29,15 +29,17 @@ IRMaterializationUnit::IRMaterializationUnit(ExecutionSession &ES,
 
   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());
-      SymbolFlags[MangledName] = JITSymbolFlags::fromGlobalValue(G);
-      SymbolToDefinition[MangledName] = &G;
+  MangleAndInterner Mangle(ES, this->TSM.getModuleUnlocked()->getDataLayout());
+  this->TSM.withModuleDo([&](Module &M) {
+    for (auto &G : M.global_values()) {
+      if (G.hasName() && !G.isDeclaration() && !G.hasLocalLinkage() &&
+          !G.hasAvailableExternallyLinkage() && !G.hasAppendingLinkage()) {
+        auto MangledName = Mangle(G.getName());
+        SymbolFlags[MangledName] = JITSymbolFlags::fromGlobalValue(G);
+        SymbolToDefinition[MangledName] = &G;
+      }
     }
-  }
+  });
 }
 
 IRMaterializationUnit::IRMaterializationUnit(
@@ -47,8 +49,9 @@ IRMaterializationUnit::IRMaterializationUnit(
       TSM(std::move(TSM)), SymbolToDefinition(std::move(SymbolToDefinition)) {}
 
 StringRef IRMaterializationUnit::getName() const {
-  if (TSM.getModule())
-    return TSM.getModule()->getModuleIdentifier();
+  if (TSM)
+    return TSM.withModuleDo(
+        [](const Module &M) { return M.getModuleIdentifier(); });
   return "<null module>";
 }
 
@@ -90,7 +93,6 @@ void BasicIRLayerMaterializationUnit::materialize(
   auto &N = R.getTargetJITDylib().getName();
 #endif // NDEBUG
 
-  auto Lock = TSM.getContextLock();
   LLVM_DEBUG(ES.runSessionLocked(
       [&]() { dbgs() << "Emitting, for " << N << ", " << *this << "\n"; }););
   L.emit(std::move(R), std::move(TSM));

lib/ExecutionEngine/Orc/ThreadSafeModule.cpp

@@ -23,41 +23,41 @@ ThreadSafeModule cloneToNewContext(ThreadSafeModule &TSM,
   if (!ShouldCloneDef)
     ShouldCloneDef = [](const GlobalValue &) { return true; };
 
-  auto Lock = TSM.getContextLock();
+  return TSM.withModuleDo([&](Module &M) {
+    SmallVector<char, 1> ClonedModuleBuffer;
 
-  SmallVector<char, 1> ClonedModuleBuffer;
+    {
+      std::set<GlobalValue *> ClonedDefsInSrc;
+      ValueToValueMapTy VMap;
+      auto Tmp = CloneModule(M, VMap, [&](const GlobalValue *GV) {
+        if (ShouldCloneDef(*GV)) {
+          ClonedDefsInSrc.insert(const_cast<GlobalValue *>(GV));
+          return true;
+        }
+        return false;
+      });
 
-  {
-    std::set<GlobalValue *> ClonedDefsInSrc;
-    ValueToValueMapTy VMap;
-    auto Tmp = CloneModule(*TSM.getModule(), VMap, [&](const GlobalValue *GV) {
-      if (ShouldCloneDef(*GV)) {
-        ClonedDefsInSrc.insert(const_cast<GlobalValue *>(GV));
-        return true;
-      }
-      return false;
-    });
+      if (UpdateClonedDefSource)
+        for (auto *GV : ClonedDefsInSrc)
+          UpdateClonedDefSource(*GV);
 
-    if (UpdateClonedDefSource)
-      for (auto *GV : ClonedDefsInSrc)
-        UpdateClonedDefSource(*GV);
+      BitcodeWriter BCWriter(ClonedModuleBuffer);
 
-    BitcodeWriter BCWriter(ClonedModuleBuffer);
+      BCWriter.writeModule(*Tmp);
+      BCWriter.writeSymtab();
+      BCWriter.writeStrtab();
+    }
 
-    BCWriter.writeModule(*Tmp);
-    BCWriter.writeSymtab();
-    BCWriter.writeStrtab();
-  }
+    MemoryBufferRef ClonedModuleBufferRef(
+        StringRef(ClonedModuleBuffer.data(), ClonedModuleBuffer.size()),
+        "cloned module buffer");
+    ThreadSafeContext NewTSCtx(llvm::make_unique<LLVMContext>());
 
-  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));
+    auto ClonedModule = cantFail(
+        parseBitcodeFile(ClonedModuleBufferRef, *NewTSCtx.getContext()));
+    ClonedModule->setModuleIdentifier(M.getName());
+    return ThreadSafeModule(std::move(ClonedModule), std::move(NewTSCtx));
+  });
 }
 
 } // end namespace orc

tools/lli/lli.cpp

@@ -695,18 +695,16 @@ int main(int argc, char **argv, char * const *envp) {
   return Result;
 }
 
-static orc::IRTransformLayer::TransformFunction createDebugDumper() {
+static std::function<void(Module &)> createDebugDumper() {
   switch (OrcDumpKind) {
   case DumpKind::NoDump:
-    return [](orc::ThreadSafeModule TSM,
-              const orc::MaterializationResponsibility &R) { return TSM; };
+    return [](Module &M) {};
 
   case DumpKind::DumpFuncsToStdOut:
-    return [](orc::ThreadSafeModule TSM,
-              const orc::MaterializationResponsibility &R) {
+    return [](Module &M) {
       printf("[ ");
 
-      for (const auto &F : *TSM.getModule()) {
+      for (const auto &F : M) {
         if (F.isDeclaration())
           continue;
 
@@ -718,31 +716,23 @@ static orc::IRTransformLayer::TransformFunction createDebugDumper() {
       }
 
       printf("]\n");
-      return TSM;
     };
 
   case DumpKind::DumpModsToStdOut:
-    return [](orc::ThreadSafeModule TSM,
-              const orc::MaterializationResponsibility &R) {
-      outs() << "----- Module Start -----\n"
-             << *TSM.getModule() << "----- Module End -----\n";
-
-      return TSM;
+    return [](Module &M) {
+      outs() << "----- Module Start -----\n" << M << "----- Module End -----\n";
     };
 
   case DumpKind::DumpModsToDisk:
-    return [](orc::ThreadSafeModule TSM,
-              const orc::MaterializationResponsibility &R) {
+    return [](Module &M) {
       std::error_code EC;
-      raw_fd_ostream Out(TSM.getModule()->getModuleIdentifier() + ".ll", EC,
-                         sys::fs::F_Text);
+      raw_fd_ostream Out(M.getModuleIdentifier() + ".ll", EC, sys::fs::F_Text);
       if (EC) {
-        errs() << "Couldn't open " << TSM.getModule()->getModuleIdentifier()
+        errs() << "Couldn't open " << M.getModuleIdentifier()
                << " for dumping.\nError:" << EC.message() << "\n";
         exit(1);
       }
-      Out << *TSM.getModule();
-      return TSM;
+      Out << M;
     };
   }
   llvm_unreachable("Unknown DumpKind");
@@ -756,12 +746,11 @@ int runOrcLazyJIT(const char *ProgName) {
   // Parse the main module.
   orc::ThreadSafeContext TSCtx(llvm::make_unique<LLVMContext>());
   SMDiagnostic Err;
-  auto MainModule = orc::ThreadSafeModule(
-      parseIRFile(InputFile, Err, *TSCtx.getContext()), TSCtx);
+  auto MainModule = parseIRFile(InputFile, Err, *TSCtx.getContext());
   if (!MainModule)
     reportError(Err, ProgName);
 
-  const auto &TT = MainModule.getModule()->getTargetTriple();
+  const auto &TT = MainModule->getTargetTriple();
   orc::LLLazyJITBuilder Builder;
 
   Builder.setJITTargetMachineBuilder(
@@ -794,11 +783,14 @@ int runOrcLazyJIT(const char *ProgName) {
 
   J->setLazyCompileTransform([&](orc::ThreadSafeModule TSM,
                                  const orc::MaterializationResponsibility &R) {
-    if (verifyModule(*TSM.getModule(), &dbgs())) {
-      dbgs() << "Bad module: " << *TSM.getModule() << "\n";
-      exit(1);
-    }
-    return Dump(std::move(TSM), R);
+    TSM.withModuleDo([&](Module &M) {
+      if (verifyModule(M, &dbgs())) {
+        dbgs() << "Bad module: " << &M << "\n";
+        exit(1);
+      }
+      Dump(M);
+    });
+    return TSM;
   });
   J->getMainJITDylib().setGenerator(
       ExitOnErr(orc::DynamicLibrarySearchGenerator::GetForCurrentProcess(
@@ -809,7 +801,8 @@ int runOrcLazyJIT(const char *ProgName) {
   ExitOnErr(CXXRuntimeOverrides.enable(J->getMainJITDylib(), Mangle));
 
   // Add the main module.
-  ExitOnErr(J->addLazyIRModule(std::move(MainModule)));
+  ExitOnErr(
+      J->addLazyIRModule(orc::ThreadSafeModule(std::move(MainModule), TSCtx)));
 
   // Create JITDylibs and add any extra modules.
   {

unittests/ExecutionEngine/Orc/ThreadSafeModuleTest.cpp

@@ -72,7 +72,7 @@ TEST(ThreadSafeModuleTest, BasicContextLockAPI) {
 
   { auto L = TSCtx.getLock(); }
 
-  { auto L = TSM.getContextLock(); }
+  { auto L = TSM.getContext().getLock(); }
 }
 
 TEST(ThreadSafeModuleTest, ContextLockPreservesContext) {