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
11namespace at {
12namespace vitals {
13
14TORCH_API bool torchVitalEnabled();
15
16struct 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
39struct 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
53std::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.
57class 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
81extern 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