1 /* Copyright 2019 The TensorFlow Authors. All Rights Reserved.
2
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6
7 http://www.apache.org/licenses/LICENSE-2.0
8
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15
16 #include "tensorflow/compiler/mlir/tensorflow/utils/dump_mlir_util.h"
17
18 #include <cstdint>
19 #include <cstring>
20 #include <memory>
21 #include <string>
22
23 #include "llvm/ADT/StringMap.h"
24 #include "llvm/ADT/StringRef.h"
25 #include "llvm/ADT/Twine.h"
26 #include "llvm/Support/FormatVariadic.h"
27 #include "llvm/Support/raw_ostream.h"
28 #include "mlir/IR/Operation.h" // from @llvm-project
29 #include "mlir/Pass/Pass.h" // from @llvm-project
30 #include "tensorflow/core/platform/crash_analysis.h"
31 #include "tensorflow/core/platform/env.h"
32 #include "tensorflow/core/platform/logging.h"
33 #include "tensorflow/core/platform/path.h"
34
35 using llvm::raw_ostream;
36
37 namespace tensorflow {
38 namespace {
39
40 struct NameCounts {
41 mutex counts_mutex;
42 llvm::StringMap<int64_t> counts;
43 };
44
MakeUniqueFilename(string name)45 std::string MakeUniqueFilename(string name) {
46 static NameCounts& instance = *new NameCounts;
47
48 // Remove illegal characters from `name`.
49 for (int i = 0, e = name.size(); i < e; ++i) {
50 char ch = name[i];
51 if (ch == '/' || ch == '[' || ch == ']' || ch == '*' || ch == '?' ||
52 ch == '\\') {
53 name[i] = '_';
54 }
55 }
56
57 int count;
58 {
59 mutex_lock lock(instance.counts_mutex);
60 count = instance.counts[name]++;
61 }
62
63 std::string filename = name;
64 if (count > 0) {
65 filename = llvm::formatv("{0}_{1}", filename, count).str();
66 }
67 filename = llvm::Twine(filename).concat(".mlir").str();
68 return filename;
69 }
70
71 // Simple raw_ostream that prints to stderr.
72 struct LogInfoRawStream : public llvm::raw_ostream {
LogInfoRawStreamtensorflow::__anondcab91740111::LogInfoRawStream73 LogInfoRawStream() { SetUnbuffered(); }
74 ~LogInfoRawStream() override = default;
current_postensorflow::__anondcab91740111::LogInfoRawStream75 uint64_t current_pos() const override { return 0; }
76
write_impltensorflow::__anondcab91740111::LogInfoRawStream77 void write_impl(const char* ptr, size_t size) override {
78 fprintf(stderr, "%.*s", static_cast<int>(size), ptr);
79 }
80 };
81
82 // Simple raw_ostream that prints to a file.
83 struct WritableFileRawStream : public llvm::raw_ostream {
WritableFileRawStreamtensorflow::__anondcab91740111::WritableFileRawStream84 explicit WritableFileRawStream(std::unique_ptr<WritableFile> file)
85 : file(std::move(file)) {
86 SetUnbuffered();
87 }
88 ~WritableFileRawStream() override = default;
current_postensorflow::__anondcab91740111::WritableFileRawStream89 uint64_t current_pos() const override { return 0; }
90
write_impltensorflow::__anondcab91740111::WritableFileRawStream91 void write_impl(const char* ptr, size_t size) override {
92 // Write the file if it is still valid. If the write fails, null out the
93 // file to avoid encountering another error.
94 if (file && !file->Append(StringPiece(ptr, size)).ok()) {
95 file = nullptr;
96 }
97 }
98
99 // The file being written to.
100 std::unique_ptr<WritableFile> file;
101 };
102
103 struct CrashReproducerStream : public mlir::PassManager::ReproducerStream {
CrashReproducerStreamtensorflow::__anondcab91740111::CrashReproducerStream104 CrashReproducerStream(llvm::StringRef name,
105 std::unique_ptr<llvm::raw_ostream> file)
106 : name(name), ostream(std::move(file)) {}
107
descriptiontensorflow::__anondcab91740111::CrashReproducerStream108 llvm::StringRef description() override { return name; }
ostensorflow::__anondcab91740111::CrashReproducerStream109 raw_ostream& os() override { return *ostream; }
110
111 private:
112 std::string name;
113 std::unique_ptr<llvm::raw_ostream> ostream;
114 };
115
116 // MLIR crash reproducer which reports failures to the crash analysis system.
117 struct CrashAnalysisCrashReproducerStream
118 : public mlir::PassManager::ReproducerStream {
119 public:
CrashAnalysisCrashReproducerStreamtensorflow::__anondcab91740111::CrashAnalysisCrashReproducerStream120 CrashAnalysisCrashReproducerStream()
121 : internal_str(""), string_stream(internal_str) {}
122
~CrashAnalysisCrashReproducerStreamtensorflow::__anondcab91740111::CrashAnalysisCrashReproducerStream123 ~CrashAnalysisCrashReproducerStream() override {
124 crash_analysis::ReportEvent(
125 "mlir_crash_reproducer.mlir",
126 "Pass pipeline failure; crash reproducer attached",
127 string_stream.str());
128 }
129
descriptiontensorflow::__anondcab91740111::CrashAnalysisCrashReproducerStream130 llvm::StringRef description() override { return "mlir_crash_reproducer"; }
ostensorflow::__anondcab91740111::CrashAnalysisCrashReproducerStream131 raw_ostream& os() override { return string_stream; }
132
133 private:
134 std::string internal_str;
135 llvm::raw_string_ostream string_stream;
136 };
137
138 } // namespace
139
CreateFileForDumping(llvm::StringRef name,std::unique_ptr<raw_ostream> * os,std::string * filepath,llvm::StringRef dirname)140 Status CreateFileForDumping(llvm::StringRef name,
141 std::unique_ptr<raw_ostream>* os,
142 std::string* filepath, llvm::StringRef dirname) {
143 std::string dir;
144 if (!dirname.empty())
145 dir = std::string(dirname);
146 else
147 dir = GetDumpDirFromEnvVar();
148
149 if (dir.empty()) {
150 return Status(error::Code::INVALID_ARGUMENT,
151 "(TF_DUMP_GRAPH_PREFIX not specified)");
152 }
153
154 if (dir == kCrashReproducerStdErr) {
155 *os = std::make_unique<LogInfoRawStream>();
156 *filepath = "(stderr)";
157 return Status();
158 }
159
160 // Get a valid file path to dump with.
161 Env* env = Env::Default();
162 Status status = env->RecursivelyCreateDir(dir);
163 if (!status.ok()) {
164 LOG(WARNING) << "Failed to create '" << dir
165 << "' directory for dumping: " << status;
166 return Status(error::Code::UNAVAILABLE, "(unavailable)");
167 }
168 *filepath = io::JoinPath(dir, MakeUniqueFilename(std::string(name)));
169
170 // Try to open the file and generate a raw_ostream.
171 std::unique_ptr<WritableFile> file;
172 status = env->NewWritableFile(*filepath, &file);
173 if (!status.ok()) {
174 LOG(WARNING) << "Failed to create file '" << filepath << "': " << status;
175 return Status(error::Code::UNAVAILABLE, "(unavailable)");
176 }
177 *os = std::make_unique<WritableFileRawStream>(std::move(file));
178 return Status();
179 }
180
181 // Prints the pass pipeline of `pass_manager` to `os`.
PrintPassPipeline(const mlir::PassManager & pass_manager,mlir::Operation * op,llvm::raw_ostream & os)182 void PrintPassPipeline(const mlir::PassManager& pass_manager,
183 mlir::Operation* op, llvm::raw_ostream& os) {
184 std::string str;
185 llvm::raw_string_ostream passOS(str);
186 llvm::interleaveComma(
187 pass_manager.getPasses(), passOS,
188 [&](mlir::Pass& pass) { pass.printAsTextualPipeline(passOS); });
189 os << "{-# external_resources: { mlir_reproducer: { pipeline: \""
190 << passOS.str() << "\", ";
191 os << "disable_threading: true, ";
192 os << "verify_each: true } } #-}";
193 os << "\n\n";
194 }
195
DumpCrashReproducerToFile(llvm::StringRef name,const mlir::PassManager & pm,mlir::Operation * op,llvm::StringRef dirname)196 std::string DumpCrashReproducerToFile(llvm::StringRef name,
197 const mlir::PassManager& pm,
198 mlir::Operation* op,
199 llvm::StringRef dirname) {
200 std::unique_ptr<llvm::raw_ostream> os;
201 std::string filepath;
202 Status result = CreateFileForDumping(name, &os, &filepath, dirname);
203 if (!result.ok()) return result.error_message();
204
205 PrintPassPipeline(pm, op, *os);
206 op->print(*os, mlir::OpPrintingFlags().useLocalScope());
207 LOG(INFO) << "Dumped MLIR operation '" << op->getName().getStringRef().str()
208 << "' to '" << filepath << "'";
209 return filepath;
210 }
211
DumpMlirOpToFile(llvm::StringRef name,mlir::Operation * op,llvm::StringRef dirname,const mlir::PassManager * pass_manager)212 std::string DumpMlirOpToFile(llvm::StringRef name, mlir::Operation* op,
213 llvm::StringRef dirname,
214 const mlir::PassManager* pass_manager) {
215 std::unique_ptr<raw_ostream> os;
216 std::string filepath;
217 Status result = CreateFileForDumping(name, &os, &filepath, dirname);
218 if (!result.ok()) return result.error_message();
219
220 if (pass_manager) PrintPassPipeline(*pass_manager, op, *os);
221 op->print(*os, mlir::OpPrintingFlags().useLocalScope());
222 LOG(INFO) << "Dumped MLIR operation '" << op->getName().getStringRef().str()
223 << "' to '" << filepath << "'";
224 return filepath;
225 }
226
GetDumpDirFromEnvVar()227 std::string GetDumpDirFromEnvVar() {
228 const char* prefix_env = getenv("TF_DUMP_GRAPH_PREFIX");
229 if (!prefix_env) {
230 LOG(WARNING)
231 << "Failed to dump MLIR module because dump location is not "
232 << "specified through TF_DUMP_GRAPH_PREFIX environment variable.";
233 return "";
234 }
235
236 std::string result = prefix_env;
237
238 if (absl::EqualsIgnoreCase(result, "sponge") &&
239 !io::GetTestUndeclaredOutputsDir(&result)) {
240 LOG(WARNING) << "TF_DUMP_GRAPH_PREFIX=sponge but "
241 "TEST_UNDECLARED_OUTPUT_DIRS is not set";
242 return "";
243 }
244 return result;
245 }
246
DumpRawStringToFile(llvm::StringRef name,llvm::StringRef content,llvm::StringRef dirname)247 std::string DumpRawStringToFile(llvm::StringRef name, llvm::StringRef content,
248 llvm::StringRef dirname) {
249 std::unique_ptr<raw_ostream> os;
250 std::string filepath;
251 Status result = CreateFileForDumping(name, &os, &filepath, dirname);
252 if (!result.ok()) return result.error_message();
253
254 (*os) << content;
255 LOG(INFO) << "Outputted requested string to '" << filepath << "'";
256 return filepath;
257 }
258
SetCrashReproducer(mlir::PassManager & pm,llvm::StringRef dir_path)259 void SetCrashReproducer(mlir::PassManager& pm, llvm::StringRef dir_path) {
260 std::string path = dir_path.str();
261 if (path.empty() || path == kCrashReproducerCrashAnalysis) {
262 if (getenv("MLIR_CRASH_REPRODUCER_DIRECTORY"))
263 path = getenv("MLIR_CRASH_REPRODUCER_DIRECTORY");
264 else if (getenv("TEST_UNDECLARED_OUTPUTS_DIR"))
265 path = "sponge";
266 }
267 if (path.empty()) {
268 LOG_FIRST_N(INFO, 1) << "disabling MLIR crash reproducer, set env var "
269 "`MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.";
270 return;
271 }
272
273 // Output dirs "sponge" (case-insensitive) have a special meaning: Dump into
274 // the directory specified by the environment variable
275 // TEST_UNDECLARED_OUTPUTS_DIR.
276 string lower_path = absl::AsciiStrToLower(path);
277 if (lower_path == "sponge") {
278 if (!tensorflow::io::GetTestUndeclaredOutputsDir(&path)) {
279 LOG(ERROR) << "MLIR crash reproducer is set to '" << dir_path.str()
280 << "', but environment variable TEST_UNDECLARED_OUTPUTS_DIR "
281 "is not set, so cannot dump anywhere.";
282 return;
283 }
284 }
285
286 // kCrashReproducerStdErr and kCrashReproducerCrashAnalysis settings do not
287 // require explicit file creation.
288 if (path != kCrashReproducerStdErr && path != kCrashReproducerCrashAnalysis) {
289 auto* env = tensorflow::Env::Default();
290 auto status = env->RecursivelyCreateDir(path);
291 if (!status.ok()) {
292 LOG(WARNING) << "cannot create directory '" + path +
293 "': " + status.error_message();
294 return;
295 }
296
297 path += "/mlir_reproducer_";
298
299 if (!tensorflow::Env::Default()->CreateUniqueFileName(&path, ".mlir")) {
300 LOG(WARNING) << "cannot create unique filename, won't enable MLIR crash "
301 "reproducer.";
302 return;
303 }
304 }
305
306 mlir::PassManager::ReproducerStreamFactory factory =
307 [path](std::string& error)
308 -> std::unique_ptr<mlir::PassManager::ReproducerStream> {
309 if (path == kCrashReproducerStdErr)
310 return std::make_unique<CrashReproducerStream>(
311 "(stderr)", std::make_unique<LogInfoRawStream>());
312 if (path == kCrashReproducerCrashAnalysis) {
313 return std::make_unique<CrashAnalysisCrashReproducerStream>();
314 }
315
316 // Try to open the file and generate a raw_ostream.
317 std::unique_ptr<WritableFile> file;
318 Status status = tensorflow::Env::Default()->NewWritableFile(path, &file);
319 if (!status.ok()) {
320 error = absl::StrCat("Failed to create file '", path,
321 "': ", status.error_message());
322 return nullptr;
323 }
324 return std::make_unique<CrashReproducerStream>(
325 path, std::make_unique<WritableFileRawStream>(std::move(file)));
326 };
327 pm.enableCrashReproducerGeneration(factory, /*genLocalReproducer=*/false);
328 }
329
applyTensorflowAndCLOptions(mlir::PassManager & pm,llvm::StringRef dir_path)330 void applyTensorflowAndCLOptions(mlir::PassManager& pm,
331 llvm::StringRef dir_path) {
332 mlir::applyPassManagerCLOptions(pm);
333 SetCrashReproducer(pm, dir_path);
334 }
335
336 } // namespace tensorflow
337