xref: /aosp_15_r20/external/pytorch/torch/csrc/lazy/core/ir_metadata.cpp (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1 #include <torch/csrc/lazy/core/config.h>
2 #include <torch/csrc/lazy/core/debug_util.h>
3 #include <torch/csrc/lazy/core/ir_metadata.h>
4 #include <functional>
5 
6 namespace torch {
7 namespace lazy {
8 
EmitShortFrameInfo(std::ostream & stream,const std::vector<SourceLocation> & frames)9 void EmitShortFrameInfo(
10     std::ostream& stream,
11     const std::vector<SourceLocation>& frames) {
12   if (!frames.empty()) {
13     const SourceLocation& frame = frames.front();
14     std::string::size_type pos = frame.file.find_last_of('/');
15     if (pos == std::string::npos) {
16       pos = 0;
17     } else {
18       ++pos;
19     }
20     stream << ", location=" << frame.function << "@" << frame.file.substr(pos)
21            << ":" << frame.line;
22   }
23 }
24 
operator <<(std::ostream & stream,const std::vector<SourceLocation> & frames)25 std::ostream& operator<<(
26     std::ostream& stream,
27     const std::vector<SourceLocation>& frames) {
28   stream << "Frames:\n";
29   for (auto& location : frames) {
30     stream << "  " << location.function << " (" << location.file << ":"
31            << location.line << ")\n";
32   }
33   return stream;
34 }
35 
36 namespace {
37 
38 struct ScopeEntry {
39   std::string name;
40   size_t saved_next_id = 1;
41 };
42 
43 struct ScopeContext {
44   std::vector<ScopeEntry> scopes;
45   size_t next_id = 1;
46 };
47 
48 thread_local ScopeContext g_scope_context;
49 
GetCurrentScope()50 std::string GetCurrentScope() {
51   std::string scope;
52   for (auto& scope_entry : g_scope_context.scopes) {
53     if (scope.empty()) {
54       scope = scope_entry.name;
55     } else {
56       scope += "/" + scope_entry.name;
57     }
58   }
59   return scope;
60 }
61 
PushScope(const std::string & name)62 void PushScope(const std::string& name) {
63   size_t id = g_scope_context.next_id;
64   g_scope_context.scopes.push_back(
65       {c10::str(name, ".", id), g_scope_context.next_id + 1});
66   g_scope_context.next_id = 1;
67 }
68 
PopScope()69 void PopScope() {
70   TORCH_CHECK(!g_scope_context.scopes.empty());
71   g_scope_context.next_id = g_scope_context.scopes.back().saved_next_id;
72   g_scope_context.scopes.pop_back();
73 }
74 
ResetScopeContext()75 void ResetScopeContext() {
76   if (!g_scope_context.scopes.empty()) {
77     TORCH_CHECK(
78         false, "Expecting scope to be empty but it is " + GetCurrentScope());
79   }
80   g_scope_context.next_id = 1;
81 }
82 } // namespace
83 
ScopePusher(const std::string & name)84 ScopePusher::ScopePusher(const std::string& name) {
85   PushScope(name);
86 }
87 
~ScopePusher()88 ScopePusher::~ScopePusher() {
89   PopScope();
90 }
91 
ResetScopes()92 void ScopePusher::ResetScopes() {
93   ResetScopeContext();
94 }
95 
GetMetaDataIfDebugging()96 MetaData GetMetaDataIfDebugging() {
97   if (!FLAGS_torch_lazy_ir_debug) {
98     return MetaData();
99   }
100   MetaData meta;
101   meta.scope = GetCurrentScope();
102   meta.frame_info = torch::lazy::GetPythonFramesFunction()();
103   return meta;
104 }
105 
106 } // namespace lazy
107 } // namespace torch
108