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 | |
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 | |
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 | |
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 | |
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 | |
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 | |
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 | |
84 | ScopePusher::ScopePusher(const std::string& name) { |
85 | PushScope(name); |
86 | } |
87 | |
88 | ScopePusher::~ScopePusher() { |
89 | PopScope(); |
90 | } |
91 | |
92 | void ScopePusher::ResetScopes() { |
93 | ResetScopeContext(); |
94 | } |
95 | |
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 |