xref: /aosp_15_r20/external/pytorch/c10/util/DeadlockDetection.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1*da0073e9SAndroid Build Coastguard Worker #pragma once
2*da0073e9SAndroid Build Coastguard Worker 
3*da0073e9SAndroid Build Coastguard Worker #include <c10/macros/Export.h>
4*da0073e9SAndroid Build Coastguard Worker #include <c10/util/Exception.h>
5*da0073e9SAndroid Build Coastguard Worker 
6*da0073e9SAndroid Build Coastguard Worker /// This file provides some simple utilities for detecting common deadlocks in
7*da0073e9SAndroid Build Coastguard Worker /// PyTorch.  For now, we focus exclusively on detecting Python GIL deadlocks,
8*da0073e9SAndroid Build Coastguard Worker /// as the GIL is a wide ranging lock that is taken out in many situations.
9*da0073e9SAndroid Build Coastguard Worker /// The basic strategy is before performing an operation that may block, you
10*da0073e9SAndroid Build Coastguard Worker /// can use TORCH_ASSERT_NO_GIL_WITHOUT_PYTHON_DEP() to assert that the GIL is
11*da0073e9SAndroid Build Coastguard Worker /// not held.  This macro is to be used in contexts where no static dependency
12*da0073e9SAndroid Build Coastguard Worker /// on Python is available (we will handle indirecting a virtual call for you).
13*da0073e9SAndroid Build Coastguard Worker ///
14*da0073e9SAndroid Build Coastguard Worker /// If the GIL is held by a torchdeploy interpreter, we always report false.
15*da0073e9SAndroid Build Coastguard Worker /// If you are in a context where Python bindings are available, it's better
16*da0073e9SAndroid Build Coastguard Worker /// to directly assert on PyGILState_Check (as it avoids a vcall and also
17*da0073e9SAndroid Build Coastguard Worker /// works correctly with torchdeploy.)
18*da0073e9SAndroid Build Coastguard Worker 
19*da0073e9SAndroid Build Coastguard Worker #define TORCH_ASSERT_NO_GIL_WITHOUT_PYTHON_DEP() \
20*da0073e9SAndroid Build Coastguard Worker   TORCH_INTERNAL_ASSERT(                         \
21*da0073e9SAndroid Build Coastguard Worker       !c10::impl::check_python_gil(),            \
22*da0073e9SAndroid Build Coastguard Worker       "Holding GIL before a blocking operation!  Please release the GIL before blocking, or see https://github.com/pytorch/pytorch/issues/56297 for how to release the GIL for destructors of objects")
23*da0073e9SAndroid Build Coastguard Worker 
24*da0073e9SAndroid Build Coastguard Worker namespace c10::impl {
25*da0073e9SAndroid Build Coastguard Worker 
26*da0073e9SAndroid Build Coastguard Worker C10_API bool check_python_gil();
27*da0073e9SAndroid Build Coastguard Worker 
28*da0073e9SAndroid Build Coastguard Worker struct C10_API PythonGILHooks {
29*da0073e9SAndroid Build Coastguard Worker   virtual ~PythonGILHooks() = default;
30*da0073e9SAndroid Build Coastguard Worker   // Returns true if we hold the GIL.  If not linked against Python we
31*da0073e9SAndroid Build Coastguard Worker   // always return false.
32*da0073e9SAndroid Build Coastguard Worker   virtual bool check_python_gil() const = 0;
33*da0073e9SAndroid Build Coastguard Worker };
34*da0073e9SAndroid Build Coastguard Worker 
35*da0073e9SAndroid Build Coastguard Worker C10_API void SetPythonGILHooks(PythonGILHooks* factory);
36*da0073e9SAndroid Build Coastguard Worker 
37*da0073e9SAndroid Build Coastguard Worker // DO NOT call this registerer from a torch deploy instance!  You will clobber
38*da0073e9SAndroid Build Coastguard Worker // other registrations
39*da0073e9SAndroid Build Coastguard Worker struct C10_API PythonGILHooksRegisterer {
PythonGILHooksRegistererPythonGILHooksRegisterer40*da0073e9SAndroid Build Coastguard Worker   explicit PythonGILHooksRegisterer(PythonGILHooks* factory) {
41*da0073e9SAndroid Build Coastguard Worker     SetPythonGILHooks(factory);
42*da0073e9SAndroid Build Coastguard Worker   }
~PythonGILHooksRegistererPythonGILHooksRegisterer43*da0073e9SAndroid Build Coastguard Worker   ~PythonGILHooksRegisterer() {
44*da0073e9SAndroid Build Coastguard Worker     SetPythonGILHooks(nullptr);
45*da0073e9SAndroid Build Coastguard Worker   }
46*da0073e9SAndroid Build Coastguard Worker };
47*da0073e9SAndroid Build Coastguard Worker 
48*da0073e9SAndroid Build Coastguard Worker } // namespace c10::impl
49