xref: /aosp_15_r20/external/pytorch/c10/util/DeadlockDetection.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1*da0073e9SAndroid Build Coastguard Worker #include <c10/util/DeadlockDetection.h>
2*da0073e9SAndroid Build Coastguard Worker 
3*da0073e9SAndroid Build Coastguard Worker #include <cstdlib>
4*da0073e9SAndroid Build Coastguard Worker 
5*da0073e9SAndroid Build Coastguard Worker namespace c10::impl {
6*da0073e9SAndroid Build Coastguard Worker 
7*da0073e9SAndroid Build Coastguard Worker namespace {
8*da0073e9SAndroid Build Coastguard Worker PythonGILHooks* python_gil_hooks = nullptr;
9*da0073e9SAndroid Build Coastguard Worker 
disable_detection()10*da0073e9SAndroid Build Coastguard Worker bool disable_detection() {
11*da0073e9SAndroid Build Coastguard Worker   return std::getenv("TORCH_DISABLE_DEADLOCK_DETECTION") != nullptr;
12*da0073e9SAndroid Build Coastguard Worker }
13*da0073e9SAndroid Build Coastguard Worker } // namespace
14*da0073e9SAndroid Build Coastguard Worker 
check_python_gil()15*da0073e9SAndroid Build Coastguard Worker bool check_python_gil() {
16*da0073e9SAndroid Build Coastguard Worker   if (!python_gil_hooks) {
17*da0073e9SAndroid Build Coastguard Worker     return false;
18*da0073e9SAndroid Build Coastguard Worker   }
19*da0073e9SAndroid Build Coastguard Worker   return python_gil_hooks->check_python_gil();
20*da0073e9SAndroid Build Coastguard Worker }
21*da0073e9SAndroid Build Coastguard Worker 
SetPythonGILHooks(PythonGILHooks * hooks)22*da0073e9SAndroid Build Coastguard Worker void SetPythonGILHooks(PythonGILHooks* hooks) {
23*da0073e9SAndroid Build Coastguard Worker   if (disable_detection()) {
24*da0073e9SAndroid Build Coastguard Worker     return;
25*da0073e9SAndroid Build Coastguard Worker   }
26*da0073e9SAndroid Build Coastguard Worker   TORCH_INTERNAL_ASSERT(!hooks || !python_gil_hooks);
27*da0073e9SAndroid Build Coastguard Worker   python_gil_hooks = hooks;
28*da0073e9SAndroid Build Coastguard Worker }
29*da0073e9SAndroid Build Coastguard Worker 
30*da0073e9SAndroid Build Coastguard Worker } // namespace c10::impl
31