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
11namespace torch {
12namespace lazy {
13
14c10::optional<SourceLocation> GetPythonFrameTop() {
15 if (!Py_IsInitialized()) {
16 return c10::nullopt;
17 }
18 pybind11::gil_scoped_acquire gil;
19 PyFrameObject* frame = PyEval_GetFrame();
20 if (frame == nullptr) {
21 return c10::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
31std::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