include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h

   1 //===----------- ThreadSafeModule.h -- Layer interfaces ---------*- C++ -*-===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // Thread safe wrappers and utilities for Module and LLVMContext.
  10 //
  11 //===----------------------------------------------------------------------===//
  12
  13 #ifndef LLVM_EXECUTIONENGINE_ORC_THREADSAFEMODULEWRAPPER_H
  14 #define LLVM_EXECUTIONENGINE_ORC_THREADSAFEMODULEWRAPPER_H
  15
  16 #include "llvm/IR/LLVMContext.h"
  17 #include "llvm/IR/Module.h"
  18 #include "llvm/Support/Compiler.h"
  19
  20 #include <functional>
  21 #include <memory>
  22 #include <mutex>
  23
  24 namespace llvm {
  25 namespace orc {
  26
  27 /// An LLVMContext together with an associated mutex that can be used to lock
  28 /// the context to prevent concurrent access by other threads.
  29 class ThreadSafeContext {
  30 private:
  31   struct State {
  32     State(std::unique_ptr<LLVMContext> Ctx) : Ctx(std::move(Ctx)) {}
  33
  34     std::unique_ptr<LLVMContext> Ctx;
  35     std::recursive_mutex Mutex;
  36   };
  37
  38 public:
  39   // RAII based lock for ThreadSafeContext.
  40   class LLVM_NODISCARD Lock {
  41   public:
  42     Lock(std::shared_ptr<State> S) : S(std::move(S)), L(this->S->Mutex) {}
  43
  44   private:
  45     std::shared_ptr<State> S;
  46     std::unique_lock<std::recursive_mutex> L;
  47   };
  48
  49   /// Construct a null context.
  50   ThreadSafeContext() = default;
  51
  52   /// Construct a ThreadSafeContext from the given LLVMContext.
  53   ThreadSafeContext(std::unique_ptr<LLVMContext> NewCtx)
  54       : S(std::make_shared<State>(std::move(NewCtx))) {
  55     assert(S->Ctx != nullptr &&
  56            "Can not construct a ThreadSafeContext from a nullptr");
  57   }
  58
  59   /// Returns a pointer to the LLVMContext that was used to construct this
  60   /// instance, or null if the instance was default constructed.
  61   LLVMContext *getContext() { return S ? S->Ctx.get() : nullptr; }
  62
  63   /// Returns a pointer to the LLVMContext that was used to construct this
  64   /// instance, or null if the instance was default constructed.
  65   const LLVMContext *getContext() const { return S ? S->Ctx.get() : nullptr; }
  66
  67   Lock getLock() const {
  68     assert(S && "Can not lock an empty ThreadSafeContext");
  69     return Lock(S);
  70   }
  71
  72 private:
  73   std::shared_ptr<State> S;
  74 };
  75
  76 /// An LLVM Module together with a shared ThreadSafeContext.
  77 class ThreadSafeModule {
  78 public:
  79   /// Default construct a ThreadSafeModule. This results in a null module and
  80   /// null context.
  81   ThreadSafeModule() = default;
  82
  83   ThreadSafeModule(ThreadSafeModule &&Other) = default;
  84
  85   ThreadSafeModule &operator=(ThreadSafeModule &&Other) {
  86     // We have to explicitly define this move operator to copy the fields in
  87     // reverse order (i.e. module first) to ensure the dependencies are
  88     // protected: The old module that is being overwritten must be destroyed
  89     // *before* the context that it depends on.
  90     // We also need to lock the context to make sure the module tear-down
  91     // does not overlap any other work on the context.
  92     if (M) {
  93       auto L = TSCtx.getLock();
  94       M = nullptr;
  95     }
  96     M = std::move(Other.M);
  97     TSCtx = std::move(Other.TSCtx);
  98     return *this;
  99   }
 100
 101   /// Construct a ThreadSafeModule from a unique_ptr<Module> and a
 102   /// unique_ptr<LLVMContext>. This creates a new ThreadSafeContext from the
 103   /// given context.
 104   ThreadSafeModule(std::unique_ptr<Module> M, std::unique_ptr<LLVMContext> Ctx)
 105       : M(std::move(M)), TSCtx(std::move(Ctx)) {}
 106
 107   /// Construct a ThreadSafeModule from a unique_ptr<Module> and an
 108   /// existing ThreadSafeContext.
 109   ThreadSafeModule(std::unique_ptr<Module> M, ThreadSafeContext TSCtx)
 110       : M(std::move(M)), TSCtx(std::move(TSCtx)) {}
 111
 112   ~ThreadSafeModule() {
 113     // We need to lock the context while we destruct the module.
 114     if (M) {
 115       auto L = TSCtx.getLock();
 116       M = nullptr;
 117     }
 118   }
 119
 120   /// Boolean conversion: This ThreadSafeModule will evaluate to true if it
 121   /// wraps a non-null module.
 122   explicit operator bool() const {
 123     if (M) {
 124       assert(TSCtx.getContext() &&
 125              "Non-null module must have non-null context");
 126       return true;
 127     }
 128     return false;
 129   }
 130
 131   /// Locks the associated ThreadSafeContext and calls the given function
 132   /// on the contained Module.
 133   template <typename Func>
 134   auto withModuleDo(Func &&F) -> decltype(F(std::declval<Module &>())) {
 135     assert(M && "Can not call on null module");
 136     auto Lock = TSCtx.getLock();
 137     return F(*M);
 138   }
 139
 140   /// Locks the associated ThreadSafeContext and calls the given function
 141   /// on the contained Module.
 142   template <typename Func>
 143   auto withModuleDo(Func &&F) const
 144       -> decltype(F(std::declval<const Module &>())) {
 145     auto Lock = TSCtx.getLock();
 146     return F(*M);
 147   }
 148
 149   /// Get a raw pointer to the contained module without locking the context.
 150   Module *getModuleUnlocked() { return M.get(); }
 151
 152   /// Get a raw pointer to the contained module without locking the context.
 153   const Module *getModuleUnlocked() const { return M.get(); }
 154
 155   /// Returns the context for this ThreadSafeModule.
 156   ThreadSafeContext getContext() const { return TSCtx; }
 157
 158 private:
 159   std::unique_ptr<Module> M;
 160   ThreadSafeContext TSCtx;
 161 };
 162
 163 using GVPredicate = std::function<bool(const GlobalValue &)>;
 164 using GVModifier = std::function<void(GlobalValue &)>;
 165
 166 /// Clones the given module on to a new context.
 167 ThreadSafeModule
 168 cloneToNewContext(ThreadSafeModule &TSMW,
 169                   GVPredicate ShouldCloneDef = GVPredicate(),
 170                   GVModifier UpdateClonedDefSource = GVModifier());
 171
 172 } // End namespace orc
 173 } // End namespace llvm
 174
 175 #endif // LLVM_EXECUTIONENGINE_ORC_THREADSAFEMODULEWRAPPER_H