experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Minor changes on dumping utility functions.

Browse Source 
PiperOrigin-RevId: 361069523
Change-Id: I270db07162393cd071226f80a0dcc99a10bd6384
This commit is contained in:
Peter Ma 2021-03-04 21:34:13 -08:00 committed by TensorFlower Gardener
parent d379154ca0
commit 2a5cbed7e9
2 changed files with 24 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

26
tensorflow/compiler/xla/service/dump.cc
View File

@ -353,10 +353,12 @@ int64 StepNumberForModule(const HloModule& module) {
tensorflow::mutex_lock lock(mu);
return module_id_to_step_number[module.unique_id()]++;
}
} // namespace
string TimestampFor(const HloModule& module) {
if (!module.config().debug_options().xla_dump_include_timestamp()) {
// Get a timestamp which we can use as a filename prefix specific to this
// module.
string TimestampFor(const HloModule& module,
const DebugOptions& debug_options) {
if (!debug_options.xla_dump_include_timestamp()) {
return "";
}
tensorflow::mutex_lock lock(mu);
@ -365,6 +367,12 @@ string TimestampFor(const HloModule& module) {
return std::to_string(timestamp_emplace.first->second);
}
string TimestampFor(const HloModule& module) {
return TimestampFor(module, module.config().debug_options());
}
} // namespace
string FilenameFor(const HloModule& module, string_view prefix,
string_view suffix) {
return StrFormat("%s%smodule_%04d.%s", prefix, prefix.empty() ? "" : ".",
@ -406,13 +414,19 @@ void DumpExecutionOptions(const ExecutionOptions& execution_options,
}
}
void DumpHloModuleIfEnabled(const HloModule& module, string_view name) {
CanonicalDebugOptions opts(module.config().debug_options());
void DumpHloModuleIfEnabled(const HloModule& module, string_view name,
const DebugOptions& debug_options) {
CanonicalDebugOptions opts(debug_options);
if (opts.should_dump_module(module.name())) {
DumpHloModuleImpl(module, /*buffer_assn=*/nullptr, /*profile=*/nullptr,
TimestampFor(module), name, opts);
TimestampFor(module, debug_options), name, opts);
}
}
void DumpHloModuleIfEnabled(const HloModule& module, string_view name) {
DumpHloModuleIfEnabled(module, name, module.config().debug_options());
}
void DumpHloModuleIfEnabled(const HloModule& module,
const BufferAssignment& buffer_assn,
string_view name) {

9
tensorflow/compiler/xla/service/dump.h
View File

@ -33,10 +33,6 @@ class BufferAssignment;
class HloExecutionProfile;
class HloSnapshot;
// Get a timestamp which we can use as a filename prefix specific to this
// module.
string TimestampFor(const HloModule& module);
// Create the filename we will use to dump in DumpToFileInDir.
string FilenameFor(const HloModule& module, absl::string_view prefix,
absl::string_view suffix);
@ -60,7 +56,7 @@ void DumpToFileInDirOrStdout(const HloModule& module,
void DumpExecutionOptions(const ExecutionOptions& execution_options,
const DebugOptions& debug_options);
// Dumps the given HLO module if dumping is enabled for the module. Exactly
// Dumps the given HLO module if dumping is enabled for the module. Exactly
// where and in what formats it's dumped is determined by the module's config.
//
// If you pass an HloExecutionProfile, note that currently only DOT-based output
@ -73,6 +69,9 @@ void DumpHloModuleIfEnabled(const HloModule& module,
void DumpHloModuleIfEnabled(const HloModule& module,
const HloExecutionProfile& profile,
absl::string_view name);
// Checks provided debug_options instead of the one from module's config.
void DumpHloModuleIfEnabled(const HloModule& module, absl::string_view name,
const DebugOptions& debug_options);
// Dumps the given HLO module after running one HLO pass and before running
// another, if that's enabled. Returns the full file paths of all dumps of the