From e384e28a97bc9f1da1402c93087e0394e5a0168c Mon Sep 17 00:00:00 2001
From: Justin Lebar <jlebar@google.com>
Date: Tue, 24 Oct 2017 18:55:42 -0700
Subject: [PATCH] Re-land: Dump the computation's SessionModule as part of the
tf_compile rule.
Nondeterminism in the SessionModule proto dumped by aot/compile.cc was
causing problems for some users. Re-landed with the SessionModule proto
being generated in a different genrule (so as not to disturb existing
users), and with more determinism in the dumped proto.
PiperOrigin-RevId: 173344189
---
tensorflow/compiler/aot/compile.cc | 10 +++++++---
tensorflow/compiler/aot/flags.cc | 5 ++---
tensorflow/compiler/aot/flags.h | 2 +-
tensorflow/compiler/aot/tfcompile.bzl | 28 +++++++++++++++++++++++++++
4 files changed, 38 insertions(+), 7 deletions(-)
diff --git a/tensorflow/compiler/aot/compile.cc b/tensorflow/compiler/aot/compile.cc
index eac8da0ab1b..2b8cc6024cb 100644
--- a/tensorflow/compiler/aot/compile.cc
+++ b/tensorflow/compiler/aot/compile.cc
@@ -97,11 +97,15 @@ Status CompileGraph(const GraphDef& graph_def, const tf2xla::Config& config,
TF_RETURN_IF_ERROR(ConvertGraphDefToXla(graph_def, config, client,
&computation,
&compile_result->has_context_arg));
- if (!flags.debug_dir.empty()) {
+ if (!flags.out_session_module.empty()) {
TF_ASSIGN_OR_RETURN(std::unique_ptr<xla::SessionModule> module,
computation.Snapshot());
- string file = io::JoinPath(flags.debug_dir, "tfcompile_xla_module.pb");
- TF_RETURN_IF_ERROR(WriteBinaryProto(Env::Default(), file, *module));
+ // Serialize the SessionModule deterministically so that all the outputs of
+ // a tf_library genrule are deterministic.
+ string proto;
+ TF_RET_CHECK(SerializeToStringDeterministic(*module, &proto));
+ TF_RETURN_IF_ERROR(
+ WriteStringToFile(Env::Default(), flags.out_session_module, proto));
}
xla::cpu::CpuAotCompilationOptions aot_opts(
flags.target_triple, flags.target_cpu, flags.target_features,
diff --git a/tensorflow/compiler/aot/flags.cc b/tensorflow/compiler/aot/flags.cc
index 5aff10346fa..7c2f27e550d 100644
--- a/tensorflow/compiler/aot/flags.cc
+++ b/tensorflow/compiler/aot/flags.cc
@@ -33,9 +33,6 @@ void AppendMainFlags(std::vector<Flag>* flag_list, MainFlags* flags) {
"fetch nodes will be dumped to stdout in a comma-separated list. "
"Typically used to format arguments for other tools, e.g. "
"freeze_graph."},
- {"debug_dir", &flags->debug_dir,
- "Specifies a directory to dump debugging information, including "
- "rewritten graphs and the XLA HLO module."},
// Flags controlling the XLA ahead-of-time compilation, that correspond to
// the fields of xla::cpu::CpuAotCompilationOptions.
//
@@ -64,6 +61,8 @@ void AppendMainFlags(std::vector<Flag>* flag_list, MainFlags* flags) {
"namespaces are given, within the global namespace."},
{"out_object", &flags->out_object, "Output object file name."},
{"out_header", &flags->out_header, "Output header file name."},
+ {"out_session_module", &flags->out_session_module,
+ "Output session module proto."},
{"gen_name_to_index", &flags->gen_name_to_index,
"Generate name-to-index data for Lookup{Arg,Result}Index methods."},
{"gen_program_shape", &flags->gen_program_shape,
diff --git a/tensorflow/compiler/aot/flags.h b/tensorflow/compiler/aot/flags.h
index 3246dbf95c8..3519659e3af 100644
--- a/tensorflow/compiler/aot/flags.h
+++ b/tensorflow/compiler/aot/flags.h
@@ -29,7 +29,6 @@ struct MainFlags {
string graph;
string config;
bool dump_fetch_nodes = false;
- string debug_dir;
string target_triple;
string target_cpu;
string target_features;
@@ -37,6 +36,7 @@ struct MainFlags {
string cpp_class;
string out_object;
string out_header;
+ string out_session_module;
// C++ codegen options
bool gen_name_to_index = false;
diff --git a/tensorflow/compiler/aot/tfcompile.bzl b/tensorflow/compiler/aot/tfcompile.bzl
index 4888760acd4..2adb1dc65ed 100644
--- a/tensorflow/compiler/aot/tfcompile.bzl
+++ b/tensorflow/compiler/aot/tfcompile.bzl
@@ -165,6 +165,34 @@ def tf_library(name, graph, config,
tags=tags,
)
+ # Rule that runs tfcompile to produce the SessionModule proto, useful for
+ # debugging. TODO(b/64813587): Once the SessionModule proto is
+ # deterministic, move this into the main rule above.
+ session_module_pb = name + "_session_module.pb"
+ native.genrule(
+ name=(name + "_session_module"),
+ srcs=[
+ tfcompile_graph,
+ config,
+ ],
+ outs=[
+ session_module_pb,
+ ],
+ cmd=("$(location " + tfcompile_tool + ")" +
+ " --graph=$(location " + tfcompile_graph + ")" +
+ " --config=$(location " + config + ")" +
+ " --entry_point=" + ep +
+ " --cpp_class=" + cpp_class +
+ " --target_triple=" + target_llvm_triple() +
+ " --out_session_module=$(@D)/" + session_module_pb +
+ " " + (tfcompile_flags or "")),
+ tools=[tfcompile_tool],
+ visibility=visibility,
+ testonly=testonly,
+ local=1,
+ tags=tags,
+ )
+
# The cc_library rule packaging up the header and object file, and needed
# kernel implementations.
need_xla_data_proto = (tfcompile_flags and