[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

Author: lhames

llvm/include/llvm/ExecutionEngine/Orc/CompileOnDemandLayer.h

@@ -68,17 +68,14 @@ class CompileOnDemandLayer2 : public IRLayer {
   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,15 +84,14 @@ class CompileOnDemandLayer2 : public IRLayer {
   IndirectStubsManager &getStubsManager(const JITDylib &JD);
 
   void emitExtractedFunctionsModule(MaterializationResponsibility R,
-                                    std::unique_ptr<Module> M);
+                                    ThreadSafeModule TSM);
 
   mutable std::mutex CODLayerMutex;
 
   IRLayer &BaseLayer;
   JITCompileCallbackManager &CCMgr;
   IndirectStubsManagerBuilder BuildIndirectStubsManager;
   StubManagersMap StubsMgrs;
-  GetAvailableContextFunction GetAvailableContext;
 };
 
 /// Compile-on-demand layer.

llvm/include/llvm/ExecutionEngine/Orc/IRCompileLayer.h

@@ -34,15 +34,15 @@ class IRCompileLayer2 : public IRLayer {
       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);
 
   void setNotifyCompiled(NotifyCompiledFunction NotifyCompiled);
 
   void emit(MaterializationResponsibility R, VModuleKey K,
-            std::unique_ptr<Module> M) override;
+            ThreadSafeModule TSM) override;
 
 private:
   mutable std::mutex IRLayerMutex;

llvm/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 @@ class IRTransformLayer2 : public IRLayer {
   }
 
   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:

llvm/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 @@ class LLJIT {
   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 @@ class LLLazyJIT : public LLJIT {
   }
 
   /// 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);
 

llvm/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 @@ class IRLayer {
   /// 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 @@ class IRMaterializationUnit : public MaterializationUnit {
 
   /// 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 @@ class IRMaterializationUnit : public MaterializationUnit {
 class BasicIRLayerMaterializationUnit : public IRMaterializationUnit {
 public:
   BasicIRLayerMaterializationUnit(IRLayer &L, VModuleKey K,
-                                  std::unique_ptr<Module> M);
+                                  ThreadSafeModule TSM);
+
 private:
 
   void materialize(MaterializationResponsibility R) override;

llvm/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

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