在优化前后可视化 XLA 图答案

【问题标题】：Visualizing XLA graphs before and after optimizations在优化前后可视化 XLA 图
【发布时间】：2019-10-31 19:34:52
【问题描述】：

我正在使用 XLA，并希望直观地了解 XLA 执行的优化类型，特别是内核融合。我一直在使用以下标志转储图表并将它们导入 tensorboard，但这些似乎是优化之前的图表。我确认 XLA 正在融合同样被转储的 HLA IR 的内核。

我正在使用的环境变量及其值

TF_DUMP_GRAPH_PREFIX=hlo TF_XLA_FLAGS="--tf_xla_clustering_debug --tf_xla_auto_jit=2" XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=./hlo

【问题讨论】：

XLA 优化中存在用于节点可视化的悬挂末端，但我不知道能够将它们可视化的工具。然而，这是我所知道的最好的单源视频来解释 XLA 的作用和原因：youtube.com/watch?v=u3cWOd99xX0

标签： tensorflow tensorboard

【解决方案1】：

将以下代码 sn-p 添加到您要检查的优化通道中。对于任何 hlo 模块，DumpHloModuleIfEnabled 会将 HLO 图发送到 ./hlo。

DumpHloModuleIfEnabled(*module, "your_hlo_module_to_print");

// use case:
StatusOr<bool> CustomGraphRewriter::Run(HloModule* module) {
  // this will dump the graph before the optimization
  DumpHloModuleIfEnabled(*module, "before_graph_rewriting");
  bool is_rewriting_done = false;
  TF_ASSIGN_OR_RETURN(is_rewriting_done,
                    RunGraphRewriting(module));
  // this will dump the graph after the optimization
  DumpHloModuleIfEnabled(*module, "after_graph_rewriting");
  return is_rewriting_done;
}

对于可视化工具，使用 graphviz 会有所帮助。安装 graphviz 后，您应该将图形发送到 *.dot 文件，并带有以下标志。

XLA_FLAGS="--xla_dump_hlo_as_dot"

为了可视化点文件，这会起作用：

dot -Tsvg your_hlo_graph.dot > visualized_output.svg

在那里，这将发出如图所示的结果。 sample HLO graph visualized

【讨论】：