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
|
#pragma once
#include <c10/macros/Export.h>
#include <c10/util/Exception.h>
#include <memory>
#include <string>
#include <unordered_map>
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<ThreadLocalDebugInfo> current();
// Internal, use DebugInfoGuard/ThreadLocalStateGuard
static void _forceCurrentDebugInfo(
const std::shared_ptr<ThreadLocalDebugInfo>& info);
// Push debug info struct of a given kind
static void _push(DebugInfoKind kind, std::shared_ptr<DebugInfoBase> info);
// Pop debug info, throws in case the last pushed
// debug info is not of a given kind
static std::shared_ptr<DebugInfoBase> _pop(DebugInfoKind kind);
// Peek debug info, throws in case the last pushed debug info is not of the
// given kind
static std::shared_ptr<DebugInfoBase> _peek(DebugInfoKind kind);
private:
std::shared_ptr<DebugInfoBase> info_;
DebugInfoKind kind_;
std::shared_ptr<ThreadLocalDebugInfo> 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<DebugInfoBase> info);
explicit DebugInfoGuard(std::shared_ptr<ThreadLocalDebugInfo> info);
~DebugInfoGuard();
DebugInfoGuard(const DebugInfoGuard&) = delete;
DebugInfoGuard(DebugInfoGuard&&) = delete;
private:
bool active_ = false;
std::shared_ptr<ThreadLocalDebugInfo> prev_info_ = nullptr;
};
} // namespace c10
|
