ThreadSafeModule.h 5.7 KB

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