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
|
#pragma once
#include <c10/macros/Export.h>
namespace c10 {
// Structure used to pack all the thread local boolean
// flags used by autograd
struct C10_API AutogradState {
static AutogradState& get_tls_state();
static void set_tls_state(AutogradState state);
AutogradState(
bool grad_mode,
bool inference_mode,
bool fw_grad_mode,
bool multithreading_enabled)
: grad_mode_(grad_mode),
inference_mode_(inference_mode),
fw_grad_mode_(fw_grad_mode),
multithreading_enabled_(multithreading_enabled),
view_replay_enabled_(false) {}
void set_grad_mode(bool enabled) {
grad_mode_ = enabled;
}
void set_fw_grad_mode(bool enabled) {
fw_grad_mode_ = enabled;
}
void set_inference_mode(bool enabled) {
inference_mode_ = enabled;
}
void set_multithreading_enabled(bool multithreading_enabled) {
multithreading_enabled_ = multithreading_enabled;
}
void set_view_replay_enabled(bool view_replay_enabled) {
view_replay_enabled_ = view_replay_enabled;
}
bool get_grad_mode() const {
return grad_mode_;
}
bool get_fw_grad_mode() const {
return fw_grad_mode_;
}
bool get_inference_mode() const {
return inference_mode_;
}
bool get_multithreading_enabled() const {
return multithreading_enabled_;
}
bool get_view_replay_enabled() const {
return view_replay_enabled_;
}
private:
bool grad_mode_ : 1;
bool inference_mode_ : 1;
bool fw_grad_mode_ : 1;
bool multithreading_enabled_ : 1;
// NOLINTNEXTLINE(cppcoreguidelines-use-default-member-init)
bool view_replay_enabled_ : 1;
};
} // namespace c10
|
