mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-21 05:34:18 +08:00
9c83e4160d
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/59017 See the comment in ThreadLocal.h for context. I used a slightly dirty preprocessor hack to minimize the number of changes. The hope is that we'll be able to revert all of these soon. Test Plan: CI. Built FB4A with gnustl and saw no references to cxa_thread_atexit in the PyTorch libraries. Reviewed By: ilia-cher Differential Revision: D28720762 fbshipit-source-id: 0f13c7ac5a108b95f8fde6dbc63c6b8bdb8599de
94 lines
2.3 KiB
C++
94 lines
2.3 KiB
C++
#include <c10/util/ThreadLocal.h>
|
|
#include <c10/util/ThreadLocalDebugInfo.h>
|
|
|
|
namespace c10 {
|
|
|
|
// NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
|
|
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
|