• Home
  • Line#
  • Scopes#
  • Navigate#
  • Raw
  • Download
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/xla/service/dump.h"
17 #include "absl/strings/ascii.h"
18 #include "tensorflow/compiler/xla/service/hlo_graph_dumper.h"
19 #include "tensorflow/compiler/xla/service/hlo_module.h"
20 #include "tensorflow/compiler/xla/service/hlo_proto_util.h"
21 #include "tensorflow/compiler/xla/util.h"
22 #include "tensorflow/core/lib/io/path.h"
23 #include "tensorflow/core/lib/strings/proto_serialization.h"
24 #include "tensorflow/core/platform/env.h"
25 #include "tensorflow/core/platform/regexp.h"
26 
27 namespace xla {
28 
29 namespace {
30 
31 using absl::StrCat;
32 using absl::StrFormat;
33 using absl::string_view;
34 
35 struct CanonicalDebugOptions {
CanonicalDebugOptionsxla::__anonf8e0eea60111::CanonicalDebugOptions36   explicit CanonicalDebugOptions(const DebugOptions& opts)
37       : dump_to(opts.xla_dump_to()),
38         dump_as_text(opts.xla_dump_hlo_as_text()),
39         dump_as_proto(opts.xla_dump_hlo_as_proto()),
40         dump_as_dot(opts.xla_dump_hlo_as_dot()),
41         dump_as_html(opts.xla_dump_hlo_as_html()),
42         dump_as_url(opts.xla_dump_hlo_as_url()),
43         dump_snapshots(opts.xla_dump_hlo_snapshots()) {
44     // This constructor examines the values in `opts` and turns on other flags
45     // based on what we think is the user's intent.  To reduce confusion about
46     // what was a user-specified value versus an extrapolated value, within this
47     // function we treat this struct's members as write-only, and read only from
48     // `opts`.
49 
50     // If dump_to is empty, default to dumping to stdout.
51     if (opts.xla_dump_to().empty()) {
52       dump_to = "-";
53     }
54 
55     // Did the user specifiy an explicit format for dumping?
56     bool output_format_specified =
57         opts.xla_dump_hlo_as_text() || opts.xla_dump_hlo_as_proto() ||
58         opts.xla_dump_hlo_as_dot() || opts.xla_dump_hlo_as_html() ||
59         opts.xla_dump_hlo_as_url() || opts.xla_dump_hlo_snapshots();
60 
61     // If we haven't specified an output format, default to dumping as text.
62     if (!output_format_specified) {
63       dump_as_text = true;
64     }
65 
66     // If we specified a regular expression restricting which modules to dump,
67     // respect that.
68     //
69     // If we didn't specify which modules to dump but we passed some other flag
70     // which implies dumping modules, dump all modules.
71     //
72     // Otherwise, don't dump any HLO modules.
73     if (!opts.xla_dump_hlo_module_re().empty()) {
74       // RE2 object is not copyable, and we can't capture "by move", so we
75       // resort to this hack.
76       string pattern = opts.xla_dump_hlo_module_re();
77       should_dump_module = [pattern](string_view module_name) {
78         return RE2::PartialMatch(string(module_name), pattern);
79       };
80     } else if (!opts.xla_dump_hlo_pass_re().empty() ||
81                !opts.xla_dump_to().empty() || output_format_specified) {
82       should_dump_module = [](string_view) { return true; };
83     } else {
84       should_dump_module = [](string_view) { return false; };
85     }
86 
87     // Initialize should_dump_pass.  This one is easy: We only dump per-pass
88     // data if the user asked for it explicitly.
89     if (!opts.xla_dump_hlo_pass_re().empty()) {
90       string pattern = opts.xla_dump_hlo_pass_re();
91       should_dump_pass = [pattern](string_view pass_name) {
92         return RE2::PartialMatch(string(pass_name), pattern);
93       };
94     } else {
95       should_dump_pass = [](string_view) { return false; };
96     }
97 
98     // Output dirs "sponge" and "test_undeclared_outputs_dir" (case-insensitive)
99     // have a special meaning: Dump into the directory specified by the
100     // environment variable TEST_UNDECLARED_OUTPUTS_DIR.
101     string dump_to_lower = absl::AsciiStrToLower(opts.xla_dump_to());
102     if (dump_to_lower == "sponge" ||
103         dump_to_lower == "test_undeclared_outputs_dir") {
104       const char* dir = getenv("TEST_UNDECLARED_OUTPUTS_DIR");
105       if (dir != nullptr) {
106         dump_to = dir;
107       } else {
108         LOG(ERROR) << "--xla_dump_to=" << opts.xla_dump_to()
109                    << ", but environment variable TEST_UNDECLARED_OUTPUTS_DIR "
110                       "is not set, so cannot dump anywhere.";
111         should_dump_module = [](string_view) { return false; };
112         should_dump_pass = [](string_view) { return false; };
113       }
114     }
115   }
116 
dumping_to_stdoutxla::__anonf8e0eea60111::CanonicalDebugOptions117   bool dumping_to_stdout() const { return dump_to == "-"; }
118 
119   string dump_to;
120   std::function<bool(string_view module_name)> should_dump_module;
121   std::function<bool(string_view pass_name)> should_dump_pass;
122 
123   // dump_ir isn't present here because this file is mostly concerned with
124   // dumping HLO.
125   bool dump_as_text;
126   bool dump_as_proto;
127   bool dump_as_dot;
128   bool dump_as_html;
129   bool dump_as_url;
130   bool dump_snapshots;
131 };
132 
FilenameFor(const HloModule & module,string_view suffix)133 string FilenameFor(const HloModule& module, string_view suffix) {
134   return StrFormat("module_%04d.%s", module.unique_id(), suffix);
135 }
136 
DumpToFileInDirImpl(string_view filename,string_view contents,const CanonicalDebugOptions & opts)137 void DumpToFileInDirImpl(string_view filename, string_view contents,
138                          const CanonicalDebugOptions& opts) {
139   if (opts.dumping_to_stdout()) {
140     LOG(ERROR) << "Refusing to write " << filename
141                << " to stdout.  Pass --xla_dump_to=<path> to write to a file.";
142     return;
143   }
144 
145   const string& dir = opts.dump_to;
146   VLOG(1) << "Dumping " << filename << " to " << dir;
147 
148   tensorflow::Env* env = tensorflow::Env::Default();
149   // Two threads can race to observe the absence of the dump directory and
150   // simultaneously try to create it, causing the "losing" thread to get a
151   // "directory already exists" error.  We can work around this by checking
152   // again whether the dir exists.
153   if (!env->IsDirectory(dir).ok()) {
154     auto status = env->RecursivelyCreateDir(dir);
155     if (!status.ok() && !env->IsDirectory(dir).ok()) {
156       LOG(ERROR) << "Could not create directory " << dir
157                  << " for dumping XLA debug data: " << status;
158       return;
159     }
160   }
161 
162   string file_path =
163       tensorflow::io::JoinPath(dir, SanitizeFileName(string(filename)));
164   auto status = tensorflow::WriteStringToFile(env, file_path, contents);
165   if (!status.ok()) {
166     LOG(ERROR) << "Could not write XLA debug data to " << file_path << ": "
167                << status;
168   }
169 }
170 
DumpToFileInDirOrStdoutImpl(string_view filename,string_view contents,const CanonicalDebugOptions & opts)171 void DumpToFileInDirOrStdoutImpl(string_view filename, string_view contents,
172                                  const CanonicalDebugOptions& opts) {
173   // Dump to stdout if that's called for.
174   if (opts.dumping_to_stdout()) {
175     std::cout << "*** Begin " << filename << " ***\n"
176               << contents << "\n*** End " << filename << " ***" << std::endl;
177     return;
178   }
179 
180   // Otherwise, dump to a file.
181   DumpToFileInDirImpl(filename, contents, opts);
182 }
183 
DumpHloModuleImpl(const HloModule & module,const BufferAssignment * buffer_assn,const HloExecutionProfile * profile,string_view suffix,const CanonicalDebugOptions & opts)184 void DumpHloModuleImpl(const HloModule& module,
185                        const BufferAssignment* buffer_assn,
186                        const HloExecutionProfile* profile, string_view suffix,
187                        const CanonicalDebugOptions& opts) {
188   string filename = FilenameFor(module, suffix);
189 
190   if (opts.dump_as_text) {
191     DumpToFileInDirOrStdoutImpl(StrCat(filename, ".txt"), module.ToString(),
192                                 opts);
193   }
194 
195   if (opts.dump_as_proto) {
196     HloProto module_proto =
197         buffer_assn ? MakeHloProto(module, *buffer_assn) : MakeHloProto(module);
198     string pb;
199     if (!tensorflow::SerializeToStringDeterministic(module_proto, &pb)) {
200       pb = "Failed to serialize HLO module proto.";
201     }
202     DumpToFileInDirImpl(StrCat(filename, ".hlo.pb"), pb, opts);
203   }
204 
205   auto render_graph = [&](RenderedGraphFormat format) {
206     StatusOr<string> rendered_graph = RenderGraph(
207         *module.entry_computation(),
208         /*label=*/filename, module.config().debug_options(), format, profile);
209     if (rendered_graph.ok()) {
210       return std::move(rendered_graph).ValueOrDie();
211     }
212     return StrFormat("Error rendering graph: %s",
213                      rendered_graph.status().ToString());
214   };
215 
216   if (opts.dump_as_dot) {
217     DumpToFileInDirImpl(StrFormat("%s.dot", filename),
218                         render_graph(RenderedGraphFormat::kDot), opts);
219   }
220 
221   if (opts.dump_as_html) {
222     DumpToFileInDirImpl(StrFormat("%s.html", filename),
223                         render_graph(RenderedGraphFormat::kHtml), opts);
224   }
225 
226   // Special case for rendering graphs as URLs.  We'll dump them to a file
227   // because why not, but we always log them to stdout as well.
228   if (opts.dump_as_url) {
229     string url = render_graph(RenderedGraphFormat::kUrl);
230     std::cout << filename << " --> " << url << std::endl;
231     if (!opts.dumping_to_stdout()) {
232       DumpToFileInDirImpl(StrFormat("%s.url", filename), url, opts);
233     }
234   }
235 }
236 
237 static tensorflow::mutex mu(tensorflow::LINKER_INITIALIZED);
238 
239 // Maps a module's unique ID to a counter indicating how many times we've dumped
240 // this module during the compilation pipeline.  This lets us keep the filenames
241 // ordered nicely.
242 //
243 // Entries added here leak forever; we have no way to GC them when a module
244 // dies.  But we only add an entry if dumping is enabled for this module, and
245 // dumping a module leaks buffer space in stdout or bytes on disk *way* faster
246 // than this hashtable leaks memory.
247 static auto& module_id_to_step_number GUARDED_BY(mu) =
248     *new absl::flat_hash_map<int64, int64>();
249 
250 }  // namespace
251 
DumpToFileInDir(const HloModule & module,string_view suffix,string_view contents)252 void DumpToFileInDir(const HloModule& module, string_view suffix,
253                      string_view contents) {
254   DumpToFileInDirImpl(FilenameFor(module, suffix), contents,
255                       CanonicalDebugOptions(module.config().debug_options()));
256 }
257 
DumpToFileInDirOrStdout(const HloModule & module,string_view suffix,string_view contents)258 void DumpToFileInDirOrStdout(const HloModule& module, string_view suffix,
259                              string_view contents) {
260   DumpToFileInDirOrStdoutImpl(
261       FilenameFor(module, suffix), contents,
262       CanonicalDebugOptions(module.config().debug_options()));
263 }
264 
DumpHloModuleIfEnabled(const HloModule & module,string_view name)265 void DumpHloModuleIfEnabled(const HloModule& module, string_view name) {
266   CanonicalDebugOptions opts(module.config().debug_options());
267   if (opts.should_dump_module(module.name())) {
268     DumpHloModuleImpl(module, /*buffer_assn=*/nullptr, /*profile=*/nullptr,
269                       name, opts);
270   }
271 }
DumpHloModuleIfEnabled(const HloModule & module,const BufferAssignment & buffer_assn,string_view name)272 void DumpHloModuleIfEnabled(const HloModule& module,
273                             const BufferAssignment& buffer_assn,
274                             string_view name) {
275   CanonicalDebugOptions opts(module.config().debug_options());
276   if (opts.should_dump_module(module.name())) {
277     DumpHloModuleImpl(module, &buffer_assn, /*profile=*/nullptr, name, opts);
278   }
279 }
280 
DumpHloModuleIfEnabled(const HloModule & module,const HloExecutionProfile & profile,string_view name)281 void DumpHloModuleIfEnabled(const HloModule& module,
282                             const HloExecutionProfile& profile,
283                             string_view name) {
284   CanonicalDebugOptions opts(module.config().debug_options());
285   if (opts.should_dump_module(module.name())) {
286     DumpHloModuleImpl(module, /*buffer_assn=*/nullptr, &profile, name, opts);
287   }
288 }
289 
DumpingEnabledForHloModule(string_view hlo_module_name,const DebugOptions & opts)290 bool DumpingEnabledForHloModule(string_view hlo_module_name,
291                                 const DebugOptions& opts) {
292   return CanonicalDebugOptions(opts).should_dump_module(hlo_module_name);
293 }
294 
DumpingToStdout(const DebugOptions & opts)295 bool DumpingToStdout(const DebugOptions& opts) {
296   return CanonicalDebugOptions(opts).dumping_to_stdout();
297 }
298 
DumpHloModuleBetweenPassesIfEnabled(string_view pipeline_name,string_view before_pass_name,string_view after_pass_name,const HloModule & module)299 void DumpHloModuleBetweenPassesIfEnabled(string_view pipeline_name,
300                                          string_view before_pass_name,
301                                          string_view after_pass_name,
302                                          const HloModule& module) {
303   CanonicalDebugOptions opts(module.config().debug_options());
304   if (!opts.should_dump_module(module.name())) {
305     return;
306   }
307 
308   if (!opts.should_dump_pass(before_pass_name) &&
309       !opts.should_dump_pass(after_pass_name)) {
310     return;
311   }
312 
313   int64 step_number;
314   {
315     tensorflow::mutex_lock lock(mu);
316     step_number = module_id_to_step_number[module.unique_id()]++;
317   }
318 
319   string filename_suffix =
320       StrFormat("%04d.%s.after_%s.before_%s", step_number, pipeline_name,
321                 after_pass_name, before_pass_name);
322   DumpHloModuleImpl(module, /*buffer_assn=*/nullptr, /*profile=*/nullptr,
323                     filename_suffix, opts);
324 }
325 
DumpHloModuleDuringPassIfEnabled(string_view pass_name,string_view step_name,const HloModule & module)326 void DumpHloModuleDuringPassIfEnabled(string_view pass_name,
327                                       string_view step_name,
328                                       const HloModule& module) {
329   CanonicalDebugOptions opts(module.config().debug_options());
330   if (!opts.should_dump_module(module.name()) ||
331       !opts.should_dump_pass(pass_name)) {
332     return;
333   }
334 
335   int64 step_number;
336   {
337     tensorflow::mutex_lock lock(mu);
338     step_number = module_id_to_step_number[module.unique_id()]++;
339   }
340 
341   string filename_suffix =
342       StrFormat("%04d.%s.%s", step_number, pass_name, step_name);
343   DumpHloModuleImpl(module, /*buffer_assn=*/nullptr, /*profile=*/nullptr,
344                     filename_suffix, opts);
345 }
346 
DumpHloSnapshotIfEnabled(const HloModule & module,const HloSnapshot & snapshot)347 void DumpHloSnapshotIfEnabled(const HloModule& module,
348                               const HloSnapshot& snapshot) {
349   CanonicalDebugOptions opts(module.config().debug_options());
350   if (!opts.should_dump_module(module.name()) || !opts.dump_snapshots) {
351     return;
352   }
353   int64 execution_count;
354   {
355     static auto& module_id_to_execution_count GUARDED_BY(mu) =
356         *new absl::flat_hash_map<int64, int64>();
357     tensorflow::mutex_lock lock(mu);
358     execution_count = module_id_to_execution_count[module.unique_id()]++;
359   }
360   string filename =
361       StrCat(FilenameFor(module, StrFormat("execution_%04d", execution_count)),
362              ".hlo_snapshot.pb");
363   if (opts.dumping_to_stdout()) {
364     LOG(ERROR) << "Refusing to write HLO snapshot proto for " << filename
365                << " to stdout.  Pass --xla_dump_to=<path> to write to a file.";
366     return;
367   }
368   string pb;
369   if (!tensorflow::SerializeToStringDeterministic(snapshot, &pb)) {
370     LOG(ERROR) << "Failed to serialize HLO snapshot proto " << filename;
371   }
372   DumpToFileInDirImpl(filename, pb, opts);
373 }
374 
DumpHloSnapshotIfEnabled(const HloSnapshot & snapshot,const DebugOptions & opts)375 void DumpHloSnapshotIfEnabled(const HloSnapshot& snapshot,
376                               const DebugOptions& opts) {
377   CanonicalDebugOptions canonical_opts(opts);
378   string name = snapshot.hlo().hlo_module().name();
379   if (!canonical_opts.should_dump_module(name) ||
380       !canonical_opts.dump_snapshots) {
381     return;
382   }
383 
384   // We don't have a unique id for an HloSnapshot, so in this overload we just
385   // have to use its name.
386   int64 execution_count;
387   {
388     static auto& module_name_to_execution_count GUARDED_BY(mu) =
389         *new absl::flat_hash_map<string, int64>();
390     tensorflow::mutex_lock lock(mu);
391     execution_count = module_name_to_execution_count[name]++;
392   }
393   string filename = StrFormat("module_%s.execution_%04d.hlo_snapshot.pb", name,
394                               execution_count);
395   if (canonical_opts.dumping_to_stdout()) {
396     LOG(ERROR) << "Refusing to write HLO snapshot proto for " << filename
397                << " to stdout.  Pass --xla_dump_to=<path> to write to a file.";
398     return;
399   }
400   string pb;
401   if (!tensorflow::SerializeToStringDeterministic(snapshot, &pb)) {
402     LOG(ERROR) << "Failed to serialize HLO snapshot proto " << filename;
403   }
404   DumpToFileInDirImpl(filename, pb, canonical_opts);
405 }
406 
407 }  // namespace xla
408