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