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

@@ -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,

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,

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;

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