gil.h 6.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204
  1. #if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
  2. /*
  3. pybind11/gil.h: RAII helpers for managing the GIL
  4. Copyright (c) 2016 Wenzel Jakob <wenzel.jakob@epfl.ch>
  5. All rights reserved. Use of this source code is governed by a
  6. BSD-style license that can be found in the LICENSE file.
  7. */
  8. #pragma once
  9. #if defined(PYBIND11_SIMPLE_GIL_MANAGEMENT)
  10. # include "detail/common.h"
  11. # include "gil_simple.h"
  12. PYBIND11_NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
  13. using gil_scoped_acquire = gil_scoped_acquire_simple;
  14. using gil_scoped_release = gil_scoped_release_simple;
  15. PYBIND11_NAMESPACE_END(PYBIND11_NAMESPACE)
  16. #else
  17. # include "detail/common.h"
  18. # include "detail/internals.h"
  19. # include <cassert>
  20. PYBIND11_NAMESPACE_BEGIN(PYBIND11_NAMESPACE)
  21. PYBIND11_NAMESPACE_BEGIN(detail)
  22. PYBIND11_WARNING_PUSH
  23. PYBIND11_WARNING_DISABLE_GCC("-Wredundant-decls")
  24. // forward declarations
  25. PyThreadState *get_thread_state_unchecked();
  26. PYBIND11_WARNING_POP
  27. PYBIND11_NAMESPACE_END(detail)
  28. /* The functions below essentially reproduce the PyGILState_* API using a RAII
  29. * pattern, but there are a few important differences:
  30. *
  31. * 1. When acquiring the GIL from an non-main thread during the finalization
  32. * phase, the GILState API blindly terminates the calling thread, which
  33. * is often not what is wanted. This API does not do this.
  34. *
  35. * 2. The gil_scoped_release function can optionally cut the relationship
  36. * of a PyThreadState and its associated thread, which allows moving it to
  37. * another thread (this is a fairly rare/advanced use case).
  38. *
  39. * 3. The reference count of an acquired thread state can be controlled. This
  40. * can be handy to prevent cases where callbacks issued from an external
  41. * thread would otherwise constantly construct and destroy thread state data
  42. * structures.
  43. *
  44. * See the Python bindings of NanoGUI (http://github.com/wjakob/nanogui) for an
  45. * example which uses features 2 and 3 to migrate the Python thread of
  46. * execution to another thread (to run the event loop on the original thread,
  47. * in this case).
  48. */
  49. class gil_scoped_acquire {
  50. public:
  51. PYBIND11_NOINLINE gil_scoped_acquire() {
  52. auto &internals = detail::get_internals();
  53. tstate = internals.tstate.get();
  54. if (!tstate) {
  55. /* Check if the GIL was acquired using the PyGILState_* API instead (e.g. if
  56. calling from a Python thread). Since we use a different key, this ensures
  57. we don't create a new thread state and deadlock in PyEval_AcquireThread
  58. below. Note we don't save this state with internals.tstate, since we don't
  59. create it we would fail to clear it (its reference count should be > 0). */
  60. tstate = PyGILState_GetThisThreadState();
  61. }
  62. if (!tstate) {
  63. tstate = PyThreadState_New(internals.istate);
  64. # if defined(PYBIND11_DETAILED_ERROR_MESSAGES)
  65. if (!tstate) {
  66. pybind11_fail("scoped_acquire: could not create thread state!");
  67. }
  68. # endif
  69. tstate->gilstate_counter = 0;
  70. internals.tstate = tstate;
  71. } else {
  72. release = detail::get_thread_state_unchecked() != tstate;
  73. }
  74. if (release) {
  75. PyEval_AcquireThread(tstate);
  76. }
  77. inc_ref();
  78. }
  79. gil_scoped_acquire(const gil_scoped_acquire &) = delete;
  80. gil_scoped_acquire &operator=(const gil_scoped_acquire &) = delete;
  81. void inc_ref() { ++tstate->gilstate_counter; }
  82. PYBIND11_NOINLINE void dec_ref() {
  83. --tstate->gilstate_counter;
  84. # if defined(PYBIND11_DETAILED_ERROR_MESSAGES)
  85. if (detail::get_thread_state_unchecked() != tstate) {
  86. pybind11_fail("scoped_acquire::dec_ref(): thread state must be current!");
  87. }
  88. if (tstate->gilstate_counter < 0) {
  89. pybind11_fail("scoped_acquire::dec_ref(): reference count underflow!");
  90. }
  91. # endif
  92. if (tstate->gilstate_counter == 0) {
  93. # if defined(PYBIND11_DETAILED_ERROR_MESSAGES)
  94. if (!release) {
  95. pybind11_fail("scoped_acquire::dec_ref(): internal error!");
  96. }
  97. # endif
  98. PyThreadState_Clear(tstate);
  99. if (active) {
  100. PyThreadState_DeleteCurrent();
  101. }
  102. detail::get_internals().tstate.reset();
  103. release = false;
  104. }
  105. }
  106. /// This method will disable the PyThreadState_DeleteCurrent call and the
  107. /// GIL won't be released. This method should be used if the interpreter
  108. /// could be shutting down when this is called, as thread deletion is not
  109. /// allowed during shutdown. Check _Py_IsFinalizing() on Python 3.7+, and
  110. /// protect subsequent code.
  111. PYBIND11_NOINLINE void disarm() { active = false; }
  112. PYBIND11_NOINLINE ~gil_scoped_acquire() {
  113. dec_ref();
  114. if (release) {
  115. PyEval_SaveThread();
  116. }
  117. }
  118. private:
  119. PyThreadState *tstate = nullptr;
  120. bool release = true;
  121. bool active = true;
  122. };
  123. class gil_scoped_release {
  124. public:
  125. // PRECONDITION: The GIL must be held when this constructor is called.
  126. explicit gil_scoped_release(bool disassoc = false) : disassoc(disassoc) {
  127. assert(PyGILState_Check());
  128. // `get_internals()` must be called here unconditionally in order to initialize
  129. // `internals.tstate` for subsequent `gil_scoped_acquire` calls. Otherwise, an
  130. // initialization race could occur as multiple threads try `gil_scoped_acquire`.
  131. auto &internals = detail::get_internals();
  132. // NOLINTNEXTLINE(cppcoreguidelines-prefer-member-initializer)
  133. tstate = PyEval_SaveThread();
  134. if (disassoc) {
  135. internals.tstate.reset();
  136. }
  137. }
  138. gil_scoped_release(const gil_scoped_release &) = delete;
  139. gil_scoped_release &operator=(const gil_scoped_release &) = delete;
  140. /// This method will disable the PyThreadState_DeleteCurrent call and the
  141. /// GIL won't be acquired. This method should be used if the interpreter
  142. /// could be shutting down when this is called, as thread deletion is not
  143. /// allowed during shutdown. Check _Py_IsFinalizing() on Python 3.7+, and
  144. /// protect subsequent code.
  145. PYBIND11_NOINLINE void disarm() { active = false; }
  146. ~gil_scoped_release() {
  147. if (!tstate) {
  148. return;
  149. }
  150. // `PyEval_RestoreThread()` should not be called if runtime is finalizing
  151. if (active) {
  152. PyEval_RestoreThread(tstate);
  153. }
  154. if (disassoc) {
  155. detail::get_internals().tstate = tstate;
  156. }
  157. }
  158. private:
  159. PyThreadState *tstate;
  160. bool disassoc;
  161. bool active = true;
  162. };
  163. PYBIND11_NAMESPACE_END(PYBIND11_NAMESPACE)
  164. #endif // !PYBIND11_SIMPLE_GIL_MANAGEMENT
  165. #else
  166. #error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
  167. #endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)