[XLA] Dump hlo module names along with the module id
PiperOrigin-RevId: 380936785
Change-Id: Ibc535b1d260893c97b8951f1058591263350e31c
diff --git a/tensorflow/compiler/xla/service/dump.cc b/tensorflow/compiler/xla/service/dump.cc
index 663b6ef..43961f1 100644
--- a/tensorflow/compiler/xla/service/dump.cc
+++ b/tensorflow/compiler/xla/service/dump.cc
@@ -16,6 +16,7 @@
#include "tensorflow/compiler/xla/service/dump.h"
#include "absl/strings/ascii.h"
+#include "absl/strings/str_cat.h"
#include "tensorflow/compiler/xla/service/hlo_graph_dumper.h"
#include "tensorflow/compiler/xla/service/hlo_module.h"
#include "tensorflow/compiler/xla/service/hlo_proto_util.h"
@@ -410,15 +411,27 @@
return std::to_string(timestamp_emplace.first->second);
}
-static string FilenameFor(int unique_id, string_view prefix,
- string_view suffix) {
- return StrFormat("%s%smodule_%04d.%s", prefix, prefix.empty() ? "" : ".",
- unique_id, suffix);
+static string FilenameFor(int unique_id, string_view module_name,
+ string_view prefix, string_view suffix) {
+ string filename;
+ if (!prefix.empty()) {
+ absl::StrAppend(&filename, prefix, ".");
+ }
+ absl::StrAppendFormat(&filename, "module_%04d", unique_id);
+ if (!module_name.empty()) {
+ absl::StrAppend(&filename, ".", module_name);
+ }
+ absl::StrAppend(&filename, ".", suffix);
+ // Skip the module name if the resulting length is too long.
+ if (!module_name.empty() && filename.size() > 255) {
+ return FilenameFor(unique_id, "", prefix, suffix);
+ }
+ return filename;
}
string FilenameFor(const HloModule& module, string_view prefix,
string_view suffix) {
- return FilenameFor(module.unique_id(), prefix, suffix);
+ return FilenameFor(module.unique_id(), module.name(), prefix, suffix);
}
void DumpToFileInDir(const HloModule& module, string_view file_prefix,
@@ -435,10 +448,11 @@
}
void DumpToFileInDirOrStdout(const DebugOptions& debug_options, int unique_id,
- string_view file_prefix, string_view file_suffix,
- string_view contents) {
- DumpToFileInDirOrStdoutImpl(FilenameFor(unique_id, file_prefix, file_suffix),
- contents, CanonicalDebugOptions(debug_options));
+ string_view module_name, string_view file_prefix,
+ string_view file_suffix, string_view contents) {
+ DumpToFileInDirOrStdoutImpl(
+ FilenameFor(unique_id, module_name, file_prefix, file_suffix), contents,
+ CanonicalDebugOptions(debug_options));
}
void DumpExecutionOptions(const ExecutionOptions& execution_options,
diff --git a/tensorflow/compiler/xla/service/dump.h b/tensorflow/compiler/xla/service/dump.h
index 97e9715..ccd9825 100644
--- a/tensorflow/compiler/xla/service/dump.h
+++ b/tensorflow/compiler/xla/service/dump.h
@@ -59,6 +59,7 @@
// directory specified, or if that directory is equal to "-", writes to stdout
// instead.
void DumpToFileInDirOrStdout(const DebugOptions& debug_options, int unique_id,
+ absl::string_view module_name,
absl::string_view file_prefix,
absl::string_view file_suffix,
absl::string_view contents);
diff --git a/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc b/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc
index 366aa36..14651d9 100644
--- a/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc
+++ b/tensorflow/compiler/xla/service/llvm_ir/llvm_util.cc
@@ -608,16 +608,18 @@
void DumpIrIfEnabled(mlir::ModuleOp mlir_module, int unique_id,
const DebugOptions& debug_options) {
- absl::string_view module_name = "<unnamed>";
+ absl::optional<absl::string_view> module_name;
if (llvm::Optional<llvm::StringRef> mlir_module_name =
mlir_module.getName()) {
module_name = AsStringView(*mlir_module_name);
}
- if (!DumpingEnabledForHloModule(module_name, debug_options)) {
+ if (!DumpingEnabledForHloModule(module_name.value_or("<unnamed>"),
+ debug_options)) {
return;
}
- DumpToFileInDirOrStdout(debug_options, unique_id, /*file_prefix=*/"",
+ DumpToFileInDirOrStdout(debug_options, unique_id, module_name.value_or(""),
+ /*file_prefix=*/"",
/*file_suffix=*/"lmhlo", DumpToString(mlir_module));
}
