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 "tensorflow/core/platform/crash_analysis.h"
30 #include "tensorflow/core/platform/env.h"
31 #include "tensorflow/core/platform/logging.h"
32 #include "tensorflow/core/platform/path.h"
33
34 using llvm::raw_ostream;
35
36 namespace tensorflow {
37 namespace {
38
39 struct NameCounts {
40 mutex counts_mutex;
41 llvm::StringMap<int64_t> counts;
42 };
43
MakeUniqueFilename(string name)44 std::string MakeUniqueFilename(string name) {
45 static NameCounts& instance = *new NameCounts;
46
47 // Remove illegal characters from `name`.
48 for (int i = 0, e = name.size(); i < e; ++i) {
49 char ch = name[i];
50 if (ch == '/' || ch == '[' || ch == ']' || ch == '*' || ch == '?' ||
51 ch == '\\') {
52 name[i] = '_';
53 }
54 }
55
56 int count;
57 {
58 mutex_lock lock(instance.counts_mutex);
59 count = instance.counts[name]++;
60 }
61
62 std::string filename = name;
63 if (count > 0) {
64 filename = llvm::formatv("{0}_{1}", filename, count).str();
65 }
66 filename = llvm::Twine(filename).concat(".mlir").str();
67 return filename;
68 }
69
70 // Simple raw_ostream that prints to stderr.
71 struct LogInfoRawStream : public llvm::raw_ostream {
LogInfoRawStreamtensorflow::__anon72a0c6aa0111::LogInfoRawStream72 LogInfoRawStream() { SetUnbuffered(); }
73 ~LogInfoRawStream() override = default;
current_postensorflow::__anon72a0c6aa0111::LogInfoRawStream74 uint64_t current_pos() const override { return 0; }
75
write_impltensorflow::__anon72a0c6aa0111::LogInfoRawStream76 void write_impl(const char* ptr, size_t size) override {
77 fprintf(stderr, "%.*s", static_cast<int>(size), ptr);
78 }
79 };
80
81 // Simple raw_ostream that prints to a file.
82 struct WritableFileRawStream : public llvm::raw_ostream {
WritableFileRawStreamtensorflow::__anon72a0c6aa0111::WritableFileRawStream83 explicit WritableFileRawStream(std::unique_ptr<WritableFile> file)
84 : file(std::move(file)) {
85 SetUnbuffered();
86 }
87 ~WritableFileRawStream() override = default;
current_postensorflow::__anon72a0c6aa0111::WritableFileRawStream88 uint64_t current_pos() const override { return 0; }
89
write_impltensorflow::__anon72a0c6aa0111::WritableFileRawStream90 void write_impl(const char* ptr, size_t size) override {
91 // Write the file if it is still valid. If the write fails, null out the
92 // file to avoid encountering another error.
93 if (file && !file->Append(StringPiece(ptr, size)).ok()) {
94 file = nullptr;
95 }
96 }
97
98 // The file being written to.
99 std::unique_ptr<WritableFile> file;
100 };
101
102 struct CrashReproducerStream : public mlir::PassManager::ReproducerStream {
CrashReproducerStreamtensorflow::__anon72a0c6aa0111::CrashReproducerStream103 CrashReproducerStream(llvm::StringRef name,
104 std::unique_ptr<llvm::raw_ostream> file)
105 : name(name), ostream(std::move(file)) {}
106
descriptiontensorflow::__anon72a0c6aa0111::CrashReproducerStream107 llvm::StringRef description() override { return name; }
ostensorflow::__anon72a0c6aa0111::CrashReproducerStream108 raw_ostream& os() override { return *ostream; }
109
110 private:
111 std::string name;
112 std::unique_ptr<llvm::raw_ostream> ostream;
113 };
114
115 // MLIR crash reproducer which reports failures to the crash analysis system.
116 struct CrashAnalysisCrashReproducerStream
117 : public mlir::PassManager::ReproducerStream {
118 public:
CrashAnalysisCrashReproducerStreamtensorflow::__anon72a0c6aa0111::CrashAnalysisCrashReproducerStream119 CrashAnalysisCrashReproducerStream()
120 : internal_str(""), string_stream(internal_str) {}
121
~CrashAnalysisCrashReproducerStreamtensorflow::__anon72a0c6aa0111::CrashAnalysisCrashReproducerStream122 ~CrashAnalysisCrashReproducerStream() override {
123 crash_analysis::ReportEvent(
124 "mlir_crash_reproducer.mlir",
125 "Pass pipeline failure; crash reproducer attached",
126 string_stream.str());
127 }
128
descriptiontensorflow::__anon72a0c6aa0111::CrashAnalysisCrashReproducerStream129 llvm::StringRef description() override { return "mlir_crash_reproducer"; }
ostensorflow::__anon72a0c6aa0111::CrashAnalysisCrashReproducerStream130 raw_ostream& os() override { return string_stream; }
131
132 private:
133 std::string internal_str;
134 llvm::raw_string_ostream string_stream;
135 };
136
137 } // namespace
138
139 const char kCrashReproducerStdErr[] = "-";
140 const char kCrashReproducerCrashAnalysis[] = "crash_analysis";
141
CreateFileForDumping(llvm::StringRef name,std::unique_ptr<raw_ostream> * os,std::string * filepath,llvm::StringRef dirname)142 Status CreateFileForDumping(llvm::StringRef name,
143 std::unique_ptr<raw_ostream>* os,
144 std::string* filepath, llvm::StringRef dirname) {
145 std::string dir;
146 if (!dirname.empty())
147 dir = std::string(dirname);
148 else
149 dir = GetDumpDirFromEnvVar();
150
151 if (dir.empty()) {
152 return Status(error::Code::INVALID_ARGUMENT,
153 "(TF_DUMP_GRAPH_PREFIX not specified)");
154 }
155
156 if (dir == kCrashReproducerStdErr) {
157 *os = std::make_unique<LogInfoRawStream>();
158 *filepath = "(stderr)";
159 return Status();
160 }
161
162 // Get a valid file path to dump with.
163 Env* env = Env::Default();
164 Status status = env->RecursivelyCreateDir(dir);
165 if (!status.ok()) {
166 LOG(WARNING) << "Failed to create '" << dir
167 << "' directory for dumping: " << status;
168 return Status(error::Code::UNAVAILABLE, "(unavailable)");
169 }
170 *filepath = io::JoinPath(dir, MakeUniqueFilename(std::string(name)));
171
172 // Try to open the file and generate a raw_ostream.
173 std::unique_ptr<WritableFile> file;
174 status = env->NewWritableFile(*filepath, &file);
175 if (!status.ok()) {
176 LOG(WARNING) << "Failed to create file '" << filepath << "': " << status;
177 return Status(error::Code::UNAVAILABLE, "(unavailable)");
178 }
179 *os = std::make_unique<WritableFileRawStream>(std::move(file));
180 return Status();
181 }
182
DumpMlirOpToFile(llvm::StringRef name,mlir::Operation * op,llvm::StringRef dirname)183 std::string DumpMlirOpToFile(llvm::StringRef name, mlir::Operation* op,
184 llvm::StringRef dirname) {
185 std::unique_ptr<raw_ostream> os;
186 std::string filepath;
187 Status result = CreateFileForDumping(name, &os, &filepath, dirname);
188 if (!result.ok()) return result.error_message();
189
190 op->print(*os, mlir::OpPrintingFlags().useLocalScope().printGenericOpForm());
191 LOG(INFO) << "Dumped MLIR operation '" << op->getName().getStringRef().str()
192 << "' to '" << filepath << "'";
193 return filepath;
194 }
195
GetDumpDirFromEnvVar()196 std::string GetDumpDirFromEnvVar() {
197 const char* prefix_env = getenv("TF_DUMP_GRAPH_PREFIX");
198 if (!prefix_env) {
199 LOG(WARNING)
200 << "Failed to dump MLIR module because dump location is not "
201 << "specified through TF_DUMP_GRAPH_PREFIX environment variable.";
202 return "";
203 }
204
205 std::string result = prefix_env;
206
207 if (absl::EqualsIgnoreCase(result, "sponge") &&
208 !io::GetTestUndeclaredOutputsDir(&result)) {
209 LOG(WARNING) << "TF_DUMP_GRAPH_PREFIX=sponge but "
210 "TEST_UNDECLARED_OUTPUT_DIRS is not set";
211 return "";
212 }
213 return result;
214 }
215
DumpRawStringToFile(llvm::StringRef name,llvm::StringRef content,llvm::StringRef dirname)216 std::string DumpRawStringToFile(llvm::StringRef name, llvm::StringRef content,
217 llvm::StringRef dirname) {
218 std::unique_ptr<raw_ostream> os;
219 std::string filepath;
220 Status result = CreateFileForDumping(name, &os, &filepath, dirname);
221 if (!result.ok()) return result.error_message();
222
223 (*os) << content;
224 LOG(INFO) << "Outputted requested string to '" << filepath << "'";
225 return filepath;
226 }
227
SetCrashReproducer(mlir::PassManager & pm,llvm::StringRef dir_path)228 void SetCrashReproducer(mlir::PassManager& pm, llvm::StringRef dir_path) {
229 std::string path = dir_path.str();
230 if (path.empty() || path == kCrashReproducerCrashAnalysis) {
231 if (getenv("MLIR_CRASH_REPRODUCER_DIRECTORY"))
232 path = getenv("MLIR_CRASH_REPRODUCER_DIRECTORY");
233 else if (getenv("TEST_UNDECLARED_OUTPUTS_DIR"))
234 path = "sponge";
235 }
236 if (path.empty()) {
237 LOG_FIRST_N(INFO, 1) << "disabling MLIR crash reproducer, set env var "
238 "`MLIR_CRASH_REPRODUCER_DIRECTORY` to enable.";
239 return;
240 }
241
242 // Output dirs "sponge" (case-insensitive) have a special meaning: Dump into
243 // the directory specified by the environment variable
244 // TEST_UNDECLARED_OUTPUTS_DIR.
245 string lower_path = absl::AsciiStrToLower(path);
246 if (lower_path == "sponge") {
247 if (!tensorflow::io::GetTestUndeclaredOutputsDir(&path)) {
248 LOG(ERROR) << "MLIR crash reproducer is set to '" << dir_path.str()
249 << "', but environment variable TEST_UNDECLARED_OUTPUTS_DIR "
250 "is not set, so cannot dump anywhere.";
251 return;
252 }
253 }
254
255 // kCrashReproducerStdErr and kCrashReproducerCrashAnalysis settings do not
256 // require explicit file creation.
257 if (path != kCrashReproducerStdErr && path != kCrashReproducerCrashAnalysis) {
258 auto* env = tensorflow::Env::Default();
259 auto status = env->RecursivelyCreateDir(path);
260 if (!status.ok()) {
261 LOG(WARNING) << "cannot create directory '" + path +
262 "': " + status.error_message();
263 return;
264 }
265
266 path += "/mlir_reproducer_";
267
268 if (!tensorflow::Env::Default()->CreateUniqueFileName(&path, ".mlir")) {
269 LOG(WARNING) << "cannot create unique filename, won't enable MLIR crash "
270 "reproducer.";
271 return;
272 }
273 }
274
275 mlir::PassManager::ReproducerStreamFactory factory =
276 [path](std::string& error)
277 -> std::unique_ptr<mlir::PassManager::ReproducerStream> {
278 if (path == kCrashReproducerStdErr)
279 return std::make_unique<CrashReproducerStream>(
280 "(stderr)", std::make_unique<LogInfoRawStream>());
281 if (path == kCrashReproducerCrashAnalysis) {
282 return std::make_unique<CrashAnalysisCrashReproducerStream>();
283 }
284
285 // Try to open the file and generate a raw_ostream.
286 std::unique_ptr<WritableFile> file;
287 Status status = tensorflow::Env::Default()->NewWritableFile(path, &file);
288 if (!status.ok()) {
289 error = absl::StrCat("Failed to create file '", path,
290 "': ", status.error_message());
291 return nullptr;
292 }
293 return std::make_unique<CrashReproducerStream>(
294 path, std::make_unique<WritableFileRawStream>(std::move(file)));
295 };
296 pm.enableCrashReproducerGeneration(factory, /*genLocalReproducer=*/false);
297 }
298
applyTensorflowAndCLOptions(mlir::PassManager & pm,llvm::StringRef dir_path)299 void applyTensorflowAndCLOptions(mlir::PassManager& pm,
300 llvm::StringRef dir_path) {
301 mlir::applyPassManagerCLOptions(pm);
302 SetCrashReproducer(pm, dir_path);
303 }
304
305 } // namespace tensorflow
306