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
|
#include <c10/util/ThreadLocalDebugInfo.h>
namespace c10 {
namespace {
thread_local std::shared_ptr<ThreadLocalDebugInfo> debug_info = nullptr;
}
/* static */
std::shared_ptr<DebugInfoBase> ThreadLocalDebugInfo::get(
DebugInfoKind kind) {
auto cur = debug_info;
while (cur) {
if (cur->kind_ == kind) {
return cur->info_;
}
cur = cur->parent_info_;
}
return nullptr;
}
/* static */
std::shared_ptr<ThreadLocalDebugInfo> ThreadLocalDebugInfo::current() {
return debug_info;
}
/* static */
void ThreadLocalDebugInfo::_forceCurrentDebugInfo(
const std::shared_ptr<ThreadLocalDebugInfo>& info) {
debug_info = info;
}
/* static */
void ThreadLocalDebugInfo::_push(
DebugInfoKind kind,
std::shared_ptr<DebugInfoBase> info) {
auto prev_info = debug_info;
debug_info = std::make_shared<ThreadLocalDebugInfo>();
debug_info->parent_info_ = prev_info;
debug_info->kind_ = kind;
debug_info->info_ = info;
}
/* static */
std::shared_ptr<DebugInfoBase> ThreadLocalDebugInfo::_pop(DebugInfoKind kind) {
TORCH_CHECK(
debug_info && debug_info->kind_ == kind,
"Expected debug info of type ", (size_t)kind);
auto res = debug_info;
debug_info = debug_info->parent_info_;
return res->info_;
}
/* static */
std::shared_ptr<DebugInfoBase> ThreadLocalDebugInfo::_peek(DebugInfoKind kind) {
TORCH_CHECK(
debug_info && debug_info->kind_ == kind,
"Expected debug info of type ",
(size_t)kind);
return debug_info->info_;
}
DebugInfoGuard::DebugInfoGuard(
DebugInfoKind kind, std::shared_ptr<DebugInfoBase> info) {
if (!info) {
return;
}
prev_info_ = debug_info;
ThreadLocalDebugInfo::_push(kind, info);
active_ = true;
}
DebugInfoGuard::~DebugInfoGuard() {
if (active_) {
debug_info = prev_info_;
}
}
// Used only for setting a debug info after crossing the thread boundary;
// in this case we assume that thread pool's thread does not have an
// active debug info
DebugInfoGuard::DebugInfoGuard(
std::shared_ptr<ThreadLocalDebugInfo> info) {
if (!info) {
return;
}
prev_info_ = debug_info;
debug_info = info;
active_ = true;
}
} // namespace c10
|
