/usr/local/lib64/python3.6/site-packages/torch/include/c10/util
Edit: /usr/local/lib64/python3.6/site-packages/torch/include/c10/util/ThreadLocalDebugInfo.h (2603B)
#pragma once
#include
#include
#include
#include
#include
namespace c10 {
enum class C10_API_ENUM DebugInfoKind : uint8_t {
PRODUCER_INFO = 0,
MOBILE_RUNTIME_INFO,
PROFILER_STATE,
INFERENCE_CONTEXT, // for inference usage
PARAM_COMMS_INFO,
TEST_INFO, // used only in tests
TEST_INFO_2, // used only in tests
};
class C10_API DebugInfoBase {
public:
DebugInfoBase() {}
virtual ~DebugInfoBase() = default;
};
// Thread local debug information is propagated across the forward
// (including async fork tasks) and backward passes and is supposed
// to be utilized by the user's code to pass extra information from
// the higher layers (e.g. model id) down to the lower levels
// (e.g. to the operator observers used for debugging, logging,
// profiling, etc)
class C10_API ThreadLocalDebugInfo {
public:
static DebugInfoBase* get(DebugInfoKind kind);
// Get current ThreadLocalDebugInfo
static std::shared_ptr current();
// Internal, use DebugInfoGuard/ThreadLocalStateGuard
static void _forceCurrentDebugInfo(
const std::shared_ptr& info);
// Push debug info struct of a given kind
static void _push(DebugInfoKind kind, std::shared_ptr info);
// Pop debug info, throws in case the last pushed
// debug info is not of a given kind
static std::shared_ptr _pop(DebugInfoKind kind);
// Peek debug info, throws in case the last pushed debug info is not of the
// given kind
static std::shared_ptr _peek(DebugInfoKind kind);
private:
std::shared_ptr info_;
DebugInfoKind kind_;
std::shared_ptr parent_info_;
friend class DebugInfoGuard;
};
// DebugInfoGuard is used to set debug information,
// ThreadLocalDebugInfo is semantically immutable, the values are set
// through the scope-based guard object.
// Nested DebugInfoGuard adds/overrides existing values in the scope,
// restoring the original values after exiting the scope.
// Users can access the values through the ThreadLocalDebugInfo::get() call;
class C10_API DebugInfoGuard {
public:
DebugInfoGuard(DebugInfoKind kind, std::shared_ptr info);
explicit DebugInfoGuard(std::shared_ptr info);
~DebugInfoGuard();
DebugInfoGuard(const DebugInfoGuard&) = delete;
DebugInfoGuard(DebugInfoGuard&&) = delete;
private:
bool active_ = false;
std::shared_ptr prev_info_ = nullptr;
};
} // namespace c10