1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
|
#include <c10/core/DispatchKeySet.h>
#include <c10/core/SafePyObject.h>
#include <c10/core/impl/LocalDispatchKeySet.h>
#include <c10/core/impl/PythonDispatcherTLS.h>
namespace c10 {
namespace impl {
thread_local PyInterpreter* pythonDispatcherState;
void PythonDispatcherTLS::set_state(PyInterpreter* state) {
if (state) {
c10::impl::tls_set_dispatch_key_included(
DispatchKey::PythonDispatcher, true);
} else {
PythonDispatcherTLS::reset_state();
}
pythonDispatcherState = state;
}
PyInterpreter* PythonDispatcherTLS::get_state() {
return pythonDispatcherState;
}
void PythonDispatcherTLS::reset_state() {
pythonDispatcherState = nullptr;
c10::impl::tls_set_dispatch_key_included(
DispatchKey::PythonDispatcher, false);
}
} // namespace impl
} // namespace c10
|
