xref: /aosp_15_r20/external/pytorch/torch/csrc/profiler/combined_traceback.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 
3 #include <torch/csrc/jit/runtime/interpreter.h>
4 #include <torch/csrc/profiler/unwind/unwind.h>
5 
6 namespace torch {
7 
8 // struct that holds the result of symbolizing multiple tracebacks
9 // each traceback is a list of indices into all_frames
10 // (lots of Frames get duplicated across traces)
11 struct TORCH_API SymbolizedTracebacks {
12   std::vector<unwind::Frame> all_frames;
13   // index into all_frames, so that
14   // it is possible to dedupe frame objects in
15   // construction of python objects
16   std::vector<std::vector<uint64_t>> tracebacks;
17 };
18 
19 struct TORCH_API CapturedTraceback : public c10::GatheredContext {
20   struct PyFrame {
21     void* code; // PyCodeObject*, but python headers not present
22     int lasti;
23   };
24 
25   static std::shared_ptr<CapturedTraceback> gather(
26       bool python,
27       bool script,
28       bool cpp);
29   CapturedTraceback() = default;
30   CapturedTraceback(const CapturedTraceback&) = delete;
31   CapturedTraceback& operator=(const CapturedTraceback&) = delete;
32   CapturedTraceback(CapturedTraceback&&) noexcept = default;
33   CapturedTraceback& operator=(CapturedTraceback&&) noexcept = delete;
34   ~CapturedTraceback() override;
35 
36   using visitproc = int (*)(void* self, void* arg);
37 
38   struct Python {
39     virtual std::vector<PyFrame> gather() = 0;
40     virtual void release(std::vector<PyFrame>& frames) = 0;
41     virtual void appendSymbolized(
42         const std::vector<PyFrame>& to_symbolize,
43         SymbolizedTracebacks& st) = 0;
44     // tp_traverse/tp_clear implementations
45     virtual int traverse(
46         std::vector<PyFrame>& frames,
47         visitproc visit,
48         void* arg) = 0;
49     virtual int clear(std::vector<PyFrame>& frames) = 0;
50     virtual ~Python() = default;
51     Python* next_ = nullptr;
52   };
53   // called once by each python interpreter to
54   // register python stack recording functionality
55   // p cannot be deleted once added.
56   static void addPythonUnwinder(Python* p);
57 
58   int traversePython(visitproc visit, void* arg);
59   int clearPython();
60 
61  private:
62   std::vector<PyFrame> frames_;
63   std::vector<void*> cpp_frames_;
64   std::vector<jit::StackEntry> script_frames_;
65   friend TORCH_API SymbolizedTracebacks
66   symbolize(const std::vector<CapturedTraceback*>& to_symbolize);
67 
68   // non-owning reference to one of the immortal Python* objects
69   // registered above.
70   Python* python_ = nullptr;
71 };
72 
73 TORCH_API SymbolizedTracebacks
74 symbolize(const std::vector<CapturedTraceback*>& to_symbolize);
75 
76 } // namespace torch
77