forked from pytorch/pytorch
-
Notifications
You must be signed in to change notification settings - Fork 0
/
callstack_debug_info_serialization.cpp
252 lines (242 loc) · 9.66 KB
/
callstack_debug_info_serialization.cpp
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
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
#include <torch/csrc/jit/api/compilation_unit.h>
#include <torch/csrc/jit/mobile/type_parser.h>
#include <torch/csrc/jit/serialization/callstack_debug_info_serialization.h>
#include <torch/csrc/jit/serialization/pickle.h>
namespace torch {
namespace jit {
namespace {
const int64_t kInvalidSourceRangeTag = -1;
} // namespace
c10::IValue InlinedCallStackSerializer::serialize(
const InlinedCallStackPtr& cs_ptr,
const SourceRangeTagMap& source_range_tags) {
if (!cs_ptr) {
return c10::IValue();
}
auto cs_it = serialized_inlined_callstack_.find(cs_ptr);
if (cs_it != serialized_inlined_callstack_.end()) {
return cs_it->second;
}
// Inlined callstack pointer is serialized as tuple of 4 elements
// {IValue(module_instance_info), source_range_tag, IValue(InlinedCallStack),
// function name} Note function name is serialized separately because Function
// is only in memory structure. It gets constructed by JIT from serialized
// Code at runtime. As such even InlinedCallStack get constructed by JIT at
// runtime during graph inlining. However, we introduce
// serialization/deserialization of it in order to generate callstack debug
// information, _when_ equivalent InlinedCallStack cannot be constructed at
// runtime. For example, in lite interpreter or delegated backend.
std::vector<c10::IValue> elements;
elements.reserve(4);
elements.emplace_back(
serialize_module_instance_info(cs_ptr->module_instance()));
int64_t source_range_tag{kInvalidSourceRangeTag};
const SourceRange& sr = cs_ptr->source_range().findSourceRangeThatGenerated()
? cs_ptr->source_range().findSourceRangeThatGenerated().value()
: cs_ptr->source_range();
auto sr_it = source_range_tags.find(sr);
if (sr_it != source_range_tags.end()) {
source_range_tag = sr_it->second;
}
elements.emplace_back(source_range_tag);
if (cs_ptr->callee()) {
elements.emplace_back(
serialize(cs_ptr->callee().value(), source_range_tags));
} else {
elements.emplace_back(c10::IValue());
}
auto fn_name = cs_ptr->function_name();
if (!fn_name.empty()) {
elements.emplace_back(fn_name);
} else {
elements.emplace_back("FunctionName_UNKNOWN");
}
c10::IValue serialized_cs = c10::ivalue::Tuple::create(elements);
serialized_inlined_callstack_[cs_ptr] = serialized_cs;
return serialized_cs;
}
c10::IValue InlinedCallStackSerializer::serialize_module_instance_info(
const c10::optional<ModuleInstanceInfo>& m) {
if (!m) {
return c10::IValue();
}
const auto& m_val = m.value();
std::string module_type_name = m_val.class_type()->name()->qualifiedName();
auto module_instance_name = m_val.instance_name();
if (m_val.class_type()) {
module_type_name = m_val.class_type()->name()->qualifiedName();
}
auto key_val = module_type_name + module_instance_name;
auto m_inst_it = serialized_module_instance_info_.find(key_val);
if (m_inst_it != serialized_module_instance_info_.end()) {
return m_inst_it->second;
}
// Module instance info is serialized as
// {type name, instance name}
serialized_module_instance_info_[key_val] =
c10::ivalue::Tuple::create({module_type_name, module_instance_name});
return serialized_module_instance_info_[key_val];
}
std::vector<char> CallStackDebugInfoPickler::pickle(
const std::unordered_map<int64_t, DebugInfoTuple>& callstack_ptrs,
const SourceRangeTagMap& source_range_tags) {
std::vector<c10::IValue> ivalues;
for (const auto& it : callstack_ptrs) {
int64_t debug_handle = it.first;
std::vector<c10::IValue> elements;
/*
* Debug handles and debug info (source range + inlinded callstack)
* are serialized as a tuple of 3 elements
* {debug_handle, source_range_tag, serialized_callstack}
*/
elements.reserve(4);
elements.emplace_back(debug_handle);
int64_t source_range_tag{kInvalidSourceRangeTag};
const auto& source_range =
std::get<kDebugInfoTupleSourceRangeIndex>(it.second);
const SourceRange& sr = source_range.findSourceRangeThatGenerated()
? source_range.findSourceRangeThatGenerated().value()
: source_range;
auto sr_it = source_range_tags.find(sr);
if (sr_it != source_range_tags.end()) {
source_range_tag = sr_it->second;
}
elements.emplace_back(source_range_tag);
elements.emplace_back(std::get<kDebugInfoTupleNodeNameIndex>(it.second));
const auto& inlined_cs_ptr =
std::get<kDebugInfoTupleInlinedCSIndex>(it.second);
elements.emplace_back(css_.serialize(inlined_cs_ptr, source_range_tags));
ivalues.emplace_back(c10::ivalue::Tuple::create(elements));
}
std::vector<at::Tensor> table;
c10::IValue ivalue = c10::ivalue::Tuple::create(std::move(ivalues));
auto result = jit::pickle(ivalue, &table);
TORCH_CHECK(table.size() == 0, "Expected 0 tensors to be written");
return result;
}
InlinedCallStackPtr InlinedCallStackDeserializer::deserialize(
const c10::IValue& iv,
const ska::flat_hash_map<int64_t, SourceRange>& source_range_map,
const std::shared_ptr<CompilationUnit>& cu) {
if (iv.isNone()) {
return c10::intrusive_ptr<InlinedCallStack>();
}
auto tup = iv.toTuple();
auto it = cached_inlined_callstacks_.find(tup);
if (it != cached_inlined_callstacks_.end()) {
return it->second;
}
const auto& tup_elems = tup->elements();
TORCH_INTERNAL_ASSERT(tup_elems.size() == 4);
// {IValue(module_instance_info), source_range_tag, IValue(InlinedCallStack),
// function name}
auto module_instance_info =
deserialize_module_instance_info(tup_elems[0], cu);
int64_t source_range_tag = tup_elems[1].toInt();
auto source_range_it = source_range_map.find(source_range_tag);
TORCH_CHECK(
source_range_tag == kInvalidSourceRangeTag ||
source_range_it != source_range_map.end(),
"Source range tag must exist in deserialized source range map."
" Not found source range tag:",
source_range_tag);
SourceRange source_range;
if (source_range_tag != kInvalidSourceRangeTag) {
source_range = source_range_it->second;
}
auto callee = deserialize(tup_elems[2], source_range_map, cu);
auto function_name = tup_elems[3].toStringRef();
InlinedCallStackPtr cs_ptr;
if (callee) {
cs_ptr = c10::make_intrusive<InlinedCallStack>(
callee, nullptr, source_range, module_instance_info);
} else {
cs_ptr = c10::make_intrusive<InlinedCallStack>(
nullptr, source_range, module_instance_info);
}
cs_ptr->set_function_name(function_name);
cached_inlined_callstacks_[tup] = cs_ptr;
// Invoking move constructor
// It is not clear if copy-ellision can happen since
// cs_ptr is copied into map above.
// This is to help avoid ref count update
return cs_ptr;
}
c10::optional<ModuleInstanceInfo> InlinedCallStackDeserializer::
deserialize_module_instance_info(
const c10::IValue& iv,
const std::shared_ptr<CompilationUnit>& cu) {
if (iv.isNone()) {
return c10::nullopt;
}
auto tup = iv.toTuple();
auto it = cached_module_instance_info_.find(tup);
if (it != cached_module_instance_info_.end()) {
return it->second;
}
const auto& tup_elems = iv.toTupleRef().elements();
TORCH_CHECK(tup_elems.size() == 2);
std::string type_name = tup_elems[0].toStringRef();
std::string instance_name = tup_elems[1].toStringRef();
// type_name might be empty string ""
// In that case type_ptr should be just nullptr
auto type_ptr = cu->get_class(type_name);
if (!type_ptr) {
// We may have lost type information. For example in lowered backends
// original class type has no relevance.
// However, to correlate ops to their original modules
// we saved both type name and instance name.
// In such cases, when module is absorbed by lowered backend
// we augment instance name with type name instead of losing it.
auto last_dot_position = type_name.find_last_of('.');
size_t substring_pos{0};
if (last_dot_position != std::string::npos) {
substring_pos = last_dot_position + 1;
}
type_name = type_name.substr(substring_pos);
instance_name = instance_name + "(" + type_name + ")";
}
cached_module_instance_info_[tup] =
ModuleInstanceInfo(type_ptr, instance_name);
return cached_module_instance_info_[tup];
}
ska::flat_hash_map<int64_t, DebugInfoTuple> CallStackDebugInfoUnpickler::
unpickle(
at::DataPtr&& data,
size_t size,
const ska::flat_hash_map<int64_t, SourceRange>& source_range_map,
const std::shared_ptr<CompilationUnit>& cu) {
auto ival = jit::unpickle(
reinterpret_cast<const char*>(data.get()),
size,
nullptr,
{},
c10::parseType);
ska::flat_hash_map<int64_t, DebugInfoTuple> callstack_ptrs;
const auto& ivalues = ival.toTupleRef().elements();
for (auto& val : ivalues) {
const auto& tup_elems = val.toTupleRef().elements();
TORCH_CHECK(
tup_elems.size() == 4,
"Pickled map must have four elements: "
"debug_handle, source_range_tag, op name, IValue(inlined_call_stack)");
int64_t debug_handle = tup_elems[0].toInt();
int64_t source_range_tag = tup_elems[1].toInt();
const std::string& node_name = tup_elems[2].toStringRef();
auto source_range_it = source_range_map.find(source_range_tag);
TORCH_CHECK(
source_range_it != source_range_map.end(),
"Source range tag must exist in deserialized source range map.");
auto source_range = source_range_it->second;
TORCH_CHECK(
callstack_ptrs.count(debug_handle) == 0,
"Debug handles should be unique.");
callstack_ptrs[debug_handle] = std::make_tuple(
source_range,
node_name,
csds_.deserialize(tup_elems[3], source_range_map, cu));
}
return callstack_ptrs;
}
} // namespace jit
} // namespace torch