xref: /aosp_15_r20/external/pytorch/torch/csrc/lazy/python/python_util.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #include <torch/csrc/lazy/python/python_util.h>
2 
3 #include <Python.h>
4 #include <frameobject.h>
5 #include <pybind11/pybind11.h>
6 #include <torch/csrc/lazy/core/debug_util.h>
7 #include <torch/csrc/utils/pybind.h>
8 #include <torch/csrc/utils/python_compat.h>
9 #include <torch/csrc/utils/python_strings.h>
10 
11 namespace torch {
12 namespace lazy {
13 
GetPythonFrameTop()14 std::optional<SourceLocation> GetPythonFrameTop() {
15   if (!Py_IsInitialized()) {
16     return std::nullopt;
17   }
18   pybind11::gil_scoped_acquire gil;
19   PyFrameObject* frame = PyEval_GetFrame();
20   if (frame == nullptr) {
21     return std::nullopt;
22   }
23   SourceLocation loc;
24   auto code = THPCodeObjectPtr(PyFrame_GetCode(frame));
25   loc.line = PyFrame_GetLineNumber(frame);
26   loc.file = THPUtils_unpackString(code->co_filename);
27   loc.function = THPUtils_unpackString(code->co_name);
28   return loc;
29 }
30 
GetPythonFrames()31 std::vector<SourceLocation> GetPythonFrames() {
32   std::vector<SourceLocation> frames;
33   if (Py_IsInitialized()) {
34     pybind11::gil_scoped_acquire gil;
35     PyFrameObject* frame = PyEval_GetFrame();
36     if (frame != nullptr) {
37       Py_INCREF(frame);
38     }
39     while (frame != nullptr) {
40       SourceLocation loc;
41       auto code = THPCodeObjectPtr(PyFrame_GetCode(frame));
42       loc.line = PyFrame_GetLineNumber(frame);
43       loc.file = THPUtils_unpackString(code->co_filename);
44       loc.function = THPUtils_unpackString(code->co_name);
45       frames.push_back(std::move(loc));
46       auto new_frame = PyFrame_GetBack(frame);
47       Py_DECREF(frame);
48       frame = new_frame;
49     }
50   }
51   return frames;
52 }
53 
54 } // namespace lazy
55 } // namespace torch
56