tensorflow xla tfcompile命令和interactive_graphviz工具
tfcompile
tfcompile有两种使用方式,一是直接编译成命令,二是用tf_library规则。
命令
编译:
bazel build tensorflow/compiler/aot:tfcompile
使用:
bazel-bin/tensorflow/compiler/aot/tfcompile --graph=./tensorflow/compiler/aot/tests/test_graph_tfadd.pb --config=./tensorflow/compiler/aot/tests/test_graph_tfadd.config.pbtxt --cpp_class="myns::test"
如果需要在编译图的时候导出hlo中间图,需要添加环境变量:
export XLA_FLAGS="--xla_hlo_profile --xla_dump_to=/tmp/foo --xla_dump_hlo_as_text"
环境变量说明:
https://github.com/tensorflow/tensorflow/commit/39587aaeb7760c6c963d8668bd6f614360e914f8
规则
参考官方文档:
https://www.tensorflow.org/xla/tfcompile
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/xla/g3doc/tfcompile.md
interactive_graphviz
interactive_graphviz是一个debug工具,可以查看hlo图。代码提交:
https://github.com/tensorflow/tensorflow/commit/4817dabb5aba45547c6c1f5ef79d2b0c2040d51f
从这个源码可以很容易看出hlo的结构:
module->computation->instruction
以及,其实可以手动构造hlo module,有两种方式,一是从字符串生成,二是用testing::opcode_matchers
编译的时候,如果不编译GPU,需要把BUILD文件中对gpu的依赖注释掉。
编译规则:
tf_cc_binary(
name = "interactive_graphviz",
srcs = ["interactive_graphviz.cc"],
deps = [
":hlo_extractor",
"//tensorflow/compiler/xla/client:client_library",
"//tensorflow/compiler/xla/client:local_client",
"//tensorflow/compiler/xla/service:compiler",
"//tensorflow/compiler/xla/service:cpu_plugin",
#"//tensorflow/compiler/xla/service:gpu_plugin",
"//tensorflow/compiler/xla/service:hlo_graph_dumper",
"//tensorflow/compiler/xla/service:hlo_proto",
"//tensorflow/compiler/xla/service:hlo_runner",
"//tensorflow/compiler/xla/service:local_service",
"//tensorflow/compiler/xla/service:platform_util",
"//tensorflow/core:framework_internal",
"//tensorflow/core:lib",
"//tensorflow/core:protos_all_cc",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/strings",
],
)
编译:
bazel build tensorflow/compiler/xla/tools:interactive_graphviz
运行:
bazel-bin/tensorflow/compiler/xla/tools/interactive_graphviz --hlo_text=/tmp/foo/module_0000.before_optimizations.txt --platform=cpu
Comments