1#pragma once
2
3#include <c10/macros/Export.h>
4
5namespace c10 {
6
7// Structure used to pack all the thread local boolean
8// flags used by autograd
9struct C10_API AutogradState {
10 static AutogradState& get_tls_state();
11 static void set_tls_state(AutogradState state);
12
13 AutogradState(
14 bool grad_mode,
15 bool inference_mode,
16 bool fw_grad_mode,
17 bool multithreading_enabled)
18 : grad_mode_(grad_mode),
19 inference_mode_(inference_mode),
20 fw_grad_mode_(fw_grad_mode),
21 mulithreading_enabled_(multithreading_enabled) {}
22
23 void set_grad_mode(bool enabled) {
24 grad_mode_ = enabled;
25 }
26
27 void set_fw_grad_mode(bool enabled) {
28 fw_grad_mode_ = enabled;
29 }
30
31 void set_inference_mode(bool enabled) {
32 inference_mode_ = enabled;
33 }
34
35 void set_multithreading_enabled(bool mulithreading_enabled) {
36 mulithreading_enabled_ = mulithreading_enabled;
37 }
38
39 void set_view_replay_enabled(bool view_replay_enabled) {
40 view_replay_enabled_ = view_replay_enabled;
41 }
42
43 bool get_grad_mode() const {
44 return grad_mode_;
45 }
46
47 bool get_fw_grad_mode() const {
48 return fw_grad_mode_;
49 }
50
51 bool get_inference_mode() const {
52 return inference_mode_;
53 }
54
55 bool get_multithreading_enabled() const {
56 return mulithreading_enabled_;
57 }
58
59 bool get_view_replay_enabled() const {
60 return view_replay_enabled_;
61 }
62
63 private:
64 bool grad_mode_ : 1;
65 bool inference_mode_ : 1;
66 bool fw_grad_mode_ : 1;
67 bool mulithreading_enabled_ : 1;
68 bool view_replay_enabled_ : 1;
69};
70
71} // namespace c10
72