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