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
|
#include <c10/util/ThreadLocal.h>
#include <c10/util/ThreadLocalDebugInfo.h>
namespace c10 {
C10_DEFINE_TLS_static(std::shared_ptr<ThreadLocalDebugInfo>, tls_debug_info);
#define debug_info (tls_debug_info.get())
/* static */
DebugInfoBase* ThreadLocalDebugInfo::get(DebugInfoKind kind) {
ThreadLocalDebugInfo* cur = debug_info.get();
while (cur) {
if (cur->kind_ == kind) {
return cur->info_.get();
}
cur = cur->parent_info_.get();
}
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
|
