ExecutionEngine/Orc/ThreadSafeModule.h

0001 //===----------- ThreadSafeModule.h -- Layer interfaces ---------*- C++ -*-===//
0002 //
0003 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
0004 // See https://llvm.org/LICENSE.txt for license information.
0005 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
0006 //
0007 //===----------------------------------------------------------------------===//
0008 //
0009 // Thread safe wrappers and utilities for Module and LLVMContext.
0010 //
0011 //===----------------------------------------------------------------------===//
0012
0013 #ifndef LLVM_EXECUTIONENGINE_ORC_THREADSAFEMODULE_H
0014 #define LLVM_EXECUTIONENGINE_ORC_THREADSAFEMODULE_H
0015
0016 #include "llvm/IR/LLVMContext.h"
0017 #include "llvm/IR/Module.h"
0018 #include "llvm/Support/Compiler.h"
0019
0020 #include <functional>
0021 #include <memory>
0022 #include <mutex>
0023
0024 namespace llvm {
0025 namespace orc {
0026
0027 /// An LLVMContext together with an associated mutex that can be used to lock
0028 /// the context to prevent concurrent access by other threads.
0029 class ThreadSafeContext {
0030 private:
0031   struct State {
0032     State(std::unique_ptr<LLVMContext> Ctx) : Ctx(std::move(Ctx)) {}
0033
0034     std::unique_ptr<LLVMContext> Ctx;
0035     std::recursive_mutex Mutex;
0036   };
0037
0038 public:
0039   // RAII based lock for ThreadSafeContext.
0040   class [[nodiscard]] Lock {
0041   public:
0042     Lock(std::shared_ptr<State> S) : S(std::move(S)), L(this->S->Mutex) {}
0043
0044   private:
0045     std::shared_ptr<State> S;
0046     std::unique_lock<std::recursive_mutex> L;
0047   };
0048
0049   /// Construct a null context.
0050   ThreadSafeContext() = default;
0051
0052   /// Construct a ThreadSafeContext from the given LLVMContext.
0053   ThreadSafeContext(std::unique_ptr<LLVMContext> NewCtx)
0054       : S(std::make_shared<State>(std::move(NewCtx))) {
0055     assert(S->Ctx != nullptr &&
0056            "Can not construct a ThreadSafeContext from a nullptr");
0057   }
0058
0059   /// Returns a pointer to the LLVMContext that was used to construct this
0060   /// instance, or null if the instance was default constructed.
0061   LLVMContext *getContext() { return S ? S->Ctx.get() : nullptr; }
0062
0063   /// Returns a pointer to the LLVMContext that was used to construct this
0064   /// instance, or null if the instance was default constructed.
0065   const LLVMContext *getContext() const { return S ? S->Ctx.get() : nullptr; }
0066
0067   Lock getLock() const {
0068     assert(S && "Can not lock an empty ThreadSafeContext");
0069     return Lock(S);
0070   }
0071
0072 private:
0073   std::shared_ptr<State> S;
0074 };
0075
0076 /// An LLVM Module together with a shared ThreadSafeContext.
0077 class ThreadSafeModule {
0078 public:
0079   /// Default construct a ThreadSafeModule. This results in a null module and
0080   /// null context.
0081   ThreadSafeModule() = default;
0082
0083   ThreadSafeModule(ThreadSafeModule &&Other) = default;
0084
0085   ThreadSafeModule &operator=(ThreadSafeModule &&Other) {
0086     // We have to explicitly define this move operator to copy the fields in
0087     // reverse order (i.e. module first) to ensure the dependencies are
0088     // protected: The old module that is being overwritten must be destroyed
0089     // *before* the context that it depends on.
0090     // We also need to lock the context to make sure the module tear-down
0091     // does not overlap any other work on the context.
0092     if (M) {
0093       auto L = TSCtx.getLock();
0094       M = nullptr;
0095     }
0096     M = std::move(Other.M);
0097     TSCtx = std::move(Other.TSCtx);
0098     return *this;
0099   }
0100
0101   /// Construct a ThreadSafeModule from a unique_ptr<Module> and a
0102   /// unique_ptr<LLVMContext>. This creates a new ThreadSafeContext from the
0103   /// given context.
0104   ThreadSafeModule(std::unique_ptr<Module> M, std::unique_ptr<LLVMContext> Ctx)
0105       : M(std::move(M)), TSCtx(std::move(Ctx)) {}
0106
0107   /// Construct a ThreadSafeModule from a unique_ptr<Module> and an
0108   /// existing ThreadSafeContext.
0109   ThreadSafeModule(std::unique_ptr<Module> M, ThreadSafeContext TSCtx)
0110       : M(std::move(M)), TSCtx(std::move(TSCtx)) {}
0111
0112   ~ThreadSafeModule() {
0113     // We need to lock the context while we destruct the module.
0114     if (M) {
0115       auto L = TSCtx.getLock();
0116       M = nullptr;
0117     }
0118   }
0119
0120   /// Boolean conversion: This ThreadSafeModule will evaluate to true if it
0121   /// wraps a non-null module.
0122   explicit operator bool() const {
0123     if (M) {
0124       assert(TSCtx.getContext() &&
0125              "Non-null module must have non-null context");
0126       return true;
0127     }
0128     return false;
0129   }
0130
0131   /// Locks the associated ThreadSafeContext and calls the given function
0132   /// on the contained Module.
0133   template <typename Func> decltype(auto) withModuleDo(Func &&F) {
0134     assert(M && "Can not call on null module");
0135     auto Lock = TSCtx.getLock();
0136     return F(*M);
0137   }
0138
0139   /// Locks the associated ThreadSafeContext and calls the given function
0140   /// on the contained Module.
0141   template <typename Func> decltype(auto) withModuleDo(Func &&F) const {
0142     assert(M && "Can not call on null module");
0143     auto Lock = TSCtx.getLock();
0144     return F(*M);
0145   }
0146
0147   /// Locks the associated ThreadSafeContext and calls the given function,
0148   /// passing the contained std::unique_ptr<Module>. The given function should
0149   /// consume the Module.
0150   template <typename Func> decltype(auto) consumingModuleDo(Func &&F) {
0151     auto Lock = TSCtx.getLock();
0152     return F(std::move(M));
0153   }
0154
0155   /// Get a raw pointer to the contained module without locking the context.
0156   Module *getModuleUnlocked() { return M.get(); }
0157
0158   /// Get a raw pointer to the contained module without locking the context.
0159   const Module *getModuleUnlocked() const { return M.get(); }
0160
0161   /// Returns the context for this ThreadSafeModule.
0162   ThreadSafeContext getContext() const { return TSCtx; }
0163
0164 private:
0165   std::unique_ptr<Module> M;
0166   ThreadSafeContext TSCtx;
0167 };
0168
0169 using GVPredicate = std::function<bool(const GlobalValue &)>;
0170 using GVModifier = std::function<void(GlobalValue &)>;
0171
0172 /// Clones the given module on to a new context.
0173 ThreadSafeModule
0174 cloneToNewContext(const ThreadSafeModule &TSMW,
0175                   GVPredicate ShouldCloneDef = GVPredicate(),
0176                   GVModifier UpdateClonedDefSource = GVModifier());
0177
0178 } // End namespace orc
0179 } // End namespace llvm
0180
0181 #endif // LLVM_EXECUTIONENGINE_ORC_THREADSAFEMODULE_H