tensorflow - 可视化优化前后的 XLA 图

Question

我正在研究 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

Answer 1

将以下代码片段添加到您要检查的优化过程中。 DumpHloModuleIfEnabled 会将 HLO 图发送到 ./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