message: Dump out an annotated TF graph from mark_for_compilation_pass for better visualization
diff: @@ -41,6 +41,7 @@ limitations under the License.
  #include "tensorflow/core/framework/types.h"
  #include "tensorflow/core/graph/algorithm.h"
  #include "tensorflow/core/graph/control_flow.h"
  +#include "tensorflow/core/graph/graph_constructor.h"
  #include "tensorflow/core/kernels/bounds_check.h"
  #include "tensorflow/core/lib/gtl/cleanup.h"
  #include "tensorflow/core/lib/strings/stringprintf.h"
  @@ -1150,6 +1151,27 @@ Status MarkForCompilationPass::RunImpl(
  if (flags->tf_xla_clustering_debug) {
  dump_graph::DumpGraphToFile("mark_for_compilation", **options.graph,
  options.flib_def);
  +
  + // We also dump out an annoated version of the TF graph where the nodes
  + // names are prefixed with the cluster names. This can help visualizing the
  + // clustering decisions on TensorBoard.
  + Graph new_graph((*options.graph)->op_registry());
  + CopyGraph(**options.graph, &new_graph);
  +
  + for (Node* n : new_graph.nodes()) {
  + if (absl::optional<absl::string_view> cluster_name =
  + GetXlaClusterForNode(*n)) {
  + n->set_name(absl::StrCat(*cluster_name, "/", n->name()));
  + } else {
  + // There is room for improvement here. In particular, it may help to
  + // split these unclustered nodes into classes where every node in a
  + // specific class has edges to and from the same set of clusters.
  + n->set_name(absl::StrCat("unclustered/", n->name()));
  + }
  + }
  +
  + dump_graph::DumpGraphToFile("mark_for_compilation_annotated", new_graph,
  + options.flib_def);
  }
  VLogClusteringSummary(*graph);
  