1 | #pragma once |
---|---|
2 | |
3 | #include <c10/macros/Export.h> |
4 | |
5 | namespace c10 { |
6 | |
7 | // Structure used to pack all the thread local boolean |
8 | // flags used by autograd |
9 | struct 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 |