1 #include <c10/util/irange.h>
2 #include <torch/csrc/lazy/core/debug_util.h>
3
4 #include <torch/csrc/lazy/backend/backend_device.h>
5 #include <torch/csrc/lazy/core/helpers.h>
6 #include <torch/csrc/lazy/core/ir.h>
7 #include <torch/csrc/lazy/core/ir_dump_util.h>
8 #include <torch/csrc/lazy/core/ir_util.h>
9 #include <torch/csrc/lazy/core/unique.h>
10
11 #include <fstream>
12 #include <mutex>
13 #include <sstream>
14 #include <unordered_set>
15
16 namespace torch {
17 namespace lazy {
18 namespace {
19
GetEnvString(const char * name,const std::string & defval)20 std::string GetEnvString(const char* name, const std::string& defval) {
21 const char* env = std::getenv(name);
22 return env != nullptr ? env : defval;
23 }
24
DefaultGraphFormat()25 DebugUtil::GraphFormat DefaultGraphFormat() {
26 std::string fmt_str = GetEnvString("LTC_SAVE_TENSORS_FMT", "text");
27 if (fmt_str == "text") {
28 return DebugUtil::GraphFormat::kText;
29 } else if (fmt_str == "backend") {
30 return DebugUtil::GraphFormat::kBackend;
31 } else if (fmt_str == "dot") {
32 return DebugUtil::GraphFormat::kDot;
33 }
34 LOG(ERROR) << "Invalid save graph format: " << fmt_str;
35 return DebugUtil::GraphFormat::kText;
36 }
37
LoadExperiments()38 std::unordered_set<std::string>* LoadExperiments() {
39 std::unique_ptr<std::unordered_set<std::string>> xset =
40 std::make_unique<std::unordered_set<std::string>>();
41 std::string experiments = GetEnvString("LTC_EXPERIMENTAL", "");
42 std::vector<std::string> experiment_list =
43 torch::lazy::StrSplit(experiments, ':');
44 for (auto& name : experiment_list) {
45 xset->insert(name);
46 }
47 return xset.release();
48 }
49
50 } // namespace
51
NoPythonFrames()52 static std::vector<SourceLocation> NoPythonFrames() {
53 SourceLocation dummy_loc;
54 dummy_loc.file = "No Python Frames";
55 return {dummy_loc};
56 }
57
GetPythonFramesFunction()58 std::function<std::vector<SourceLocation>()>& GetPythonFramesFunction() {
59 static std::function<std::vector<SourceLocation>()> func_ = NoPythonFrames;
60 return func_;
61 }
62
GetDefaultGraphFormat()63 DebugUtil::GraphFormat DebugUtil::GetDefaultGraphFormat() {
64 static GraphFormat format = DefaultGraphFormat();
65 return format;
66 }
67
GetFirstUserFrameInPython()68 std::string GetFirstUserFrameInPython() {
69 std::string empty;
70 if (!torch::lazy::GetPythonFramesFunction()) {
71 return empty;
72 }
73
74 auto frames = torch::lazy::GetPythonFramesFunction()();
75
76 for (auto i = frames.size(); i > 0; i--) {
77 auto& loc = frames[i - 1];
78 if (loc.file.find("site-packages") == std::string::npos) {
79 std::stringstream ss;
80 ss << loc.file << " " << loc.function << " " << loc.line;
81 return ss.str();
82 }
83 }
84 return empty;
85 }
86
GetTensorsGraphInfo(c10::ArrayRef<torch::lazy::LazyTensorPtr> tensors,const std::vector<size_t> * indices,GraphFormat format)87 std::string DebugUtil::GetTensorsGraphInfo(
88 c10::ArrayRef<torch::lazy::LazyTensorPtr> tensors,
89 const std::vector<size_t>* indices,
90 GraphFormat format) {
91 std::vector<const torch::lazy::Node*> root_nodes;
92 std::vector<torch::lazy::Value> root_values;
93 std::vector<torch::lazy::hash_t> root_hashes;
94 torch::lazy::Unique<torch::lazy::BackendDevice> unique_device;
95 if (indices != nullptr) {
96 for (auto index : *indices) {
97 const torch::lazy::LazyTensorPtr& tensor = tensors[index];
98 torch::lazy::Value ir_value = tensor->CurrentIrValue();
99 if (ir_value) {
100 root_nodes.push_back(ir_value.node.get());
101 root_hashes.push_back(ir_value.hash());
102 root_values.push_back(std::move(ir_value));
103 unique_device.set(tensor->GetDevice());
104 }
105 }
106 } else {
107 for (auto& tensor : tensors) {
108 torch::lazy::Value ir_value = tensor->CurrentIrValue();
109 if (ir_value) {
110 root_nodes.push_back(ir_value.node.get());
111 root_hashes.push_back(ir_value.hash());
112 root_values.push_back(std::move(ir_value));
113 unique_device.set(tensor->GetDevice());
114 }
115 }
116 }
117 std::stringstream ss;
118 // Call into a function pointer that may backed by python or empty depending
119 // on runtime
120 std::vector<SourceLocation> frames = GetPythonFramesFunction()();
121 ss << "Python Stacktrace:\n";
122 for (auto& location : frames) {
123 ss << " " << location.function << " (" << location.file << ":"
124 << location.line << ")\n";
125 }
126 ss << "\nHashes: (";
127 for (const auto i : c10::irange(root_hashes.size())) {
128 if (i > 0) {
129 ss << ", ";
130 }
131 ss << torch::lazy::HashToString(root_hashes[i]);
132 }
133 ss << ")\n";
134
135 std::string graph_str;
136 if (format == GraphFormat::kText) {
137 graph_str = torch::lazy::DumpUtil::ToText(root_nodes);
138 } else if (format == GraphFormat::kDot) {
139 graph_str = torch::lazy::DumpUtil::ToDot(root_nodes);
140 } else if (format == GraphFormat::kBackend) {
141 graph_str = torch::lazy::DumpUtil::ToBackend(
142 root_values,
143 unique_device ? *unique_device : torch::lazy::BackendDevice());
144 } else {
145 LOG(ERROR) << "Invalid graph format: " << format;
146 }
147 ss << "\n## BEGIN_GRAPH\n" << graph_str << "\n## END_GRAPH\n\n";
148 return ss.str();
149 }
150
SaveTensorsGraphInfo(const char * name,c10::ArrayRef<torch::lazy::LazyTensorPtr> tensors,const std::vector<size_t> * indices,GraphFormat format)151 void DebugUtil::SaveTensorsGraphInfo(
152 const char* name,
153 c10::ArrayRef<torch::lazy::LazyTensorPtr> tensors,
154 const std::vector<size_t>* indices,
155 GraphFormat format) {
156 static const std::string save_file =
157 GetEnvString("LTC_SAVE_TENSORS_FILE", "");
158 if (!save_file.empty()) {
159 static std::mutex lock;
160 std::string info = GetTensorsGraphInfo(tensors, indices, format);
161 std::lock_guard<std::mutex> guard(lock);
162 std::ofstream graph_file(save_file, std::ios_base::app);
163 graph_file << "[" << name << "]\n" << info << "\n";
164 }
165 }
166
ExperimentEnabled(const std::string & name)167 bool DebugUtil::ExperimentEnabled(const std::string& name) {
168 static const std::unordered_set<std::string>* xset = LoadExperiments();
169 return xset->find(name) != xset->end();
170 }
171
172 } // namespace lazy
173 } // namespace torch
174