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
6namespace torch {
7namespace lazy {
8
9void 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
25std::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
36namespace {
37
38struct ScopeEntry {
39 std::string name;
40 size_t saved_next_id = 1;
41};
42
43struct ScopeContext {
44 std::vector<ScopeEntry> scopes;
45 size_t next_id = 1;
46};
47
48thread_local ScopeContext g_scope_context;
49
50std::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
62void 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
69void 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
75void 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
84ScopePusher::ScopePusher(const std::string& name) {
85 PushScope(name);
86}
87
88ScopePusher::~ScopePusher() {
89 PopScope();
90}
91
92void ScopePusher::ResetScopes() {
93 ResetScopeContext();
94}
95
96MetaData 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