ThreadSafeModule.h 5.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172
  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. #ifndef LLVM_EXECUTIONENGINE_ORC_THREADSAFEMODULE_H
  13. #define LLVM_EXECUTIONENGINE_ORC_THREADSAFEMODULE_H
  14. #include "llvm/IR/LLVMContext.h"
  15. #include "llvm/IR/Module.h"
  16. #include "llvm/Support/Compiler.h"
  17. #include <functional>
  18. #include <memory>
  19. #include <mutex>
  20. namespace llvm {
  21. namespace orc {
  22. /// An LLVMContext together with an associated mutex that can be used to lock
  23. /// the context to prevent concurrent access by other threads.
  24. class ThreadSafeContext {
  25. private:
  26. struct State {
  27. State(std::unique_ptr<LLVMContext> Ctx) : Ctx(std::move(Ctx)) {}
  28. std::unique_ptr<LLVMContext> Ctx;
  29. std::recursive_mutex Mutex;
  30. };
  31. public:
  32. // RAII based lock for ThreadSafeContext.
  33. class LLVM_NODISCARD Lock {
  34. public:
  35. Lock(std::shared_ptr<State> S) : S(std::move(S)), L(this->S->Mutex) {}
  36. private:
  37. std::shared_ptr<State> S;
  38. std::unique_lock<std::recursive_mutex> L;
  39. };
  40. /// Construct a null context.
  41. ThreadSafeContext() = default;
  42. /// Construct a ThreadSafeContext from the given LLVMContext.
  43. ThreadSafeContext(std::unique_ptr<LLVMContext> NewCtx)
  44. : S(std::make_shared<State>(std::move(NewCtx))) {
  45. assert(S->Ctx != nullptr &&
  46. "Can not construct a ThreadSafeContext from a nullptr");
  47. }
  48. /// Returns a pointer to the LLVMContext that was used to construct this
  49. /// instance, or null if the instance was default constructed.
  50. LLVMContext *getContext() { return S ? S->Ctx.get() : nullptr; }
  51. /// Returns a pointer to the LLVMContext that was used to construct this
  52. /// instance, or null if the instance was default constructed.
  53. const LLVMContext *getContext() const { return S ? S->Ctx.get() : nullptr; }
  54. Lock getLock() const {
  55. assert(S && "Can not lock an empty ThreadSafeContext");
  56. return Lock(S);
  57. }
  58. private:
  59. std::shared_ptr<State> S;
  60. };
  61. /// An LLVM Module together with a shared ThreadSafeContext.
  62. class ThreadSafeModule {
  63. public:
  64. /// Default construct a ThreadSafeModule. This results in a null module and
  65. /// null context.
  66. ThreadSafeModule() = default;
  67. ThreadSafeModule(ThreadSafeModule &&Other) = default;
  68. ThreadSafeModule &operator=(ThreadSafeModule &&Other) {
  69. // We have to explicitly define this move operator to copy the fields in
  70. // reverse order (i.e. module first) to ensure the dependencies are
  71. // protected: The old module that is being overwritten must be destroyed
  72. // *before* the context that it depends on.
  73. // We also need to lock the context to make sure the module tear-down
  74. // does not overlap any other work on the context.
  75. if (M) {
  76. auto L = TSCtx.getLock();
  77. M = nullptr;
  78. }
  79. M = std::move(Other.M);
  80. TSCtx = std::move(Other.TSCtx);
  81. return *this;
  82. }
  83. /// Construct a ThreadSafeModule from a unique_ptr<Module> and a
  84. /// unique_ptr<LLVMContext>. This creates a new ThreadSafeContext from the
  85. /// given context.
  86. ThreadSafeModule(std::unique_ptr<Module> M, std::unique_ptr<LLVMContext> Ctx)
  87. : M(std::move(M)), TSCtx(std::move(Ctx)) {}
  88. /// Construct a ThreadSafeModule from a unique_ptr<Module> and an
  89. /// existing ThreadSafeContext.
  90. ThreadSafeModule(std::unique_ptr<Module> M, ThreadSafeContext TSCtx)
  91. : M(std::move(M)), TSCtx(std::move(TSCtx)) {}
  92. ~ThreadSafeModule() {
  93. // We need to lock the context while we destruct the module.
  94. if (M) {
  95. auto L = TSCtx.getLock();
  96. M = nullptr;
  97. }
  98. }
  99. /// Boolean conversion: This ThreadSafeModule will evaluate to true if it
  100. /// wraps a non-null module.
  101. explicit operator bool() const {
  102. if (M) {
  103. assert(TSCtx.getContext() &&
  104. "Non-null module must have non-null context");
  105. return true;
  106. }
  107. return false;
  108. }
  109. /// Locks the associated ThreadSafeContext and calls the given function
  110. /// on the contained Module.
  111. template <typename Func> decltype(auto) withModuleDo(Func &&F) {
  112. assert(M && "Can not call on null module");
  113. auto Lock = TSCtx.getLock();
  114. return F(*M);
  115. }
  116. /// Locks the associated ThreadSafeContext and calls the given function
  117. /// on the contained Module.
  118. template <typename Func> decltype(auto) withModuleDo(Func &&F) const {
  119. auto Lock = TSCtx.getLock();
  120. return F(*M);
  121. }
  122. /// Get a raw pointer to the contained module without locking the context.
  123. Module *getModuleUnlocked() { return M.get(); }
  124. /// Get a raw pointer to the contained module without locking the context.
  125. const Module *getModuleUnlocked() const { return M.get(); }
  126. /// Returns the context for this ThreadSafeModule.
  127. ThreadSafeContext getContext() const { return TSCtx; }
  128. private:
  129. std::unique_ptr<Module> M;
  130. ThreadSafeContext TSCtx;
  131. };
  132. using GVPredicate = std::function<bool(const GlobalValue &)>;
  133. using GVModifier = std::function<void(GlobalValue &)>;
  134. /// Clones the given module on to a new context.
  135. ThreadSafeModule
  136. cloneToNewContext(const ThreadSafeModule &TSMW,
  137. GVPredicate ShouldCloneDef = GVPredicate(),
  138. GVModifier UpdateClonedDefSource = GVModifier());
  139. } // End namespace orc
  140. } // End namespace llvm
  141. #endif // LLVM_EXECUTIONENGINE_ORC_THREADSAFEMODULE_H