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