1 | #pragma once |
2 | #include <cstring> |
3 | #include <map> |
4 | #include <memory> |
5 | #include <ostream> |
6 | #include <sstream> |
7 | #include <unordered_map> |
8 | |
9 | #include <c10/core/impl/LocalDispatchKeySet.h> |
10 | |
11 | namespace at { |
12 | namespace vitals { |
13 | |
14 | TORCH_API bool torchVitalEnabled(); |
15 | |
16 | struct TORCH_API TorchVitalAttr { |
17 | // always initialized to empty |
18 | std::string value = "" ; |
19 | template <typename T> |
20 | TorchVitalAttr& operator<<(const T& t) { |
21 | if (torchVitalEnabled()) { |
22 | std::stringstream ss; |
23 | ss << t; |
24 | value += ss.str(); |
25 | } |
26 | return *this; |
27 | } |
28 | |
29 | template <typename T> |
30 | void write(const T& t, bool force) { |
31 | if (force || torchVitalEnabled()) { |
32 | std::stringstream ss; |
33 | ss << t; |
34 | value = ss.str(); |
35 | } |
36 | } |
37 | }; |
38 | |
39 | struct TORCH_API TorchVital { |
40 | std::string name; |
41 | std::unordered_map<std::string, TorchVitalAttr> attrs; |
42 | |
43 | explicit TorchVital(std::string n) : name(std::move(n)) {} |
44 | TorchVital() = delete; |
45 | |
46 | TorchVitalAttr& create(const std::string& attr); |
47 | TorchVitalAttr& create(const std::string& attr, bool force); |
48 | friend std::ostream& operator<<(std::ostream& os, const TorchVital& dt); |
49 | |
50 | ~TorchVital(); |
51 | }; |
52 | |
53 | std::ostream& operator<<(std::ostream& os, TorchVital const& tv); |
54 | |
55 | // A way to access vitals by string names instead of by global reference. |
56 | // This enables access to vitals from the PythonAPI. |
57 | class TORCH_API APIVitals { |
58 | public: |
59 | bool vitals_enabled; |
60 | |
61 | // Set any vital sign that was added to the map. |
62 | bool setVital( |
63 | const std::string& vital_name, |
64 | const std::string& attr_name, |
65 | const std::string& value, |
66 | bool force = false); |
67 | std::string readVitals(); |
68 | |
69 | APIVitals(); |
70 | |
71 | // Ensure this stays a singleton |
72 | APIVitals(APIVitals const& other) = delete; |
73 | APIVitals(APIVitals&& other) = delete; |
74 | APIVitals& operator=(const APIVitals&) = delete; |
75 | APIVitals& operator=(APIVitals&&) = delete; |
76 | |
77 | private: |
78 | std::unordered_map<std::string, TorchVital> name_map_; |
79 | }; |
80 | |
81 | extern TORCH_API APIVitals VitalsAPI; |
82 | |
83 | } // namespace vitals |
84 | } // namespace at |
85 | |
86 | #define TORCH_VITAL_DECLARE(name) \ |
87 | TORCH_API at::vitals::TorchVital TorchVital_##name; |
88 | |
89 | #define TORCH_VITAL_DEFINE(name) \ |
90 | TORCH_API at::vitals::TorchVital TorchVital_##name(#name); |
91 | |
92 | #define TORCH_VITAL_BASE(name) TorchVital_##name |
93 | |
94 | #define TORCH_VITAL(name, attr) TORCH_VITAL_BASE(name).create(#attr) |
95 | |