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 | |
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 | |
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 | |
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 | |
52 | std::vector<SourceLocation> NoPythonFrames() { |
53 | SourceLocation dummy_loc; |
54 | dummy_loc.file = "No Python Frames" ; |
55 | return {dummy_loc}; |
56 | } |
57 | |
58 | std::function<std::vector<SourceLocation>()>& GetPythonFramesFunction() { |
59 | static std::function<std::vector<SourceLocation>()> func_ = NoPythonFrames; |
60 | return func_; |
61 | } |
62 | |
63 | DebugUtil::GraphFormat DebugUtil::GetDefaultGraphFormat() { |
64 | static GraphFormat format = DefaultGraphFormat(); |
65 | return format; |
66 | } |
67 | |
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 | |
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 | |
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 | |
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 | |