compilation - 如何在 TensorFlow 中从 XLA 获取 LLVM IR 转储？

Question

我正在尝试获取 TensorFlow 中 XLA 编译器生成的 LLVM IR。我知道整个 LLVM 上下文包含在 llvm_module 对象中。然后使用文件中 Compile() 函数中的实用函数 llvm_ir::DumpModuleToString(*llvm_module) 函数将其转换为字符串: //tensorflow/compiler/xla/service/cpu.cpu_compiler.cc。

但我一直在尝试使用 tensorflow/core/logging.h 中的 VLOG(2) 来记录它。没有显示日志。但是，其他文件中剩余的 VLOG(2) 语句会记录在我的 Python 运行中。

>>> import tensorflow as tf
>>> hello = tf.constant('Hello, TensorFlow!')
>>> sess = tf.Session()
>>> print(sess.run(hello))
2017-03-10 22:36:43.226843: I tensorflow/compiler/xla/service/platform_util.cc:58] platform Host present with 8 visible devices
2017-03-10 22:36:43.227931: I tensorflow/compiler/xla/service/service.cc:183] XLA service 0x2821510 executing computations on platform Host. Devices:
2017-03-10 22:36:43.227951: I tensorflow/compiler/xla/service/service.cc:191]   StreamExecutor device (0): <undefined>, <undefined>
b'Hello, TensorFlow!'

Answer 1

[仅供引用，我无法发表评论，因为我刚刚加入，显然还没有声誉。]

首先，请务必阅读本文，包括带星号的蓝色框。特别注意，为整个 session 打开 XLA 目前仅对 GPU 执行 JIT，而不是对 CPU 执行 JIT。 https://www.tensorflow.org/performance/xla/jit

现在假设您已正确设置所有内容。您的示例中的程序不会使用 XLA 进行编译，原因有两个:

1. 正如 @mrry 所指出的，XLA 不处理字符串。
2. 即使您将字符串替换为数字，您仍然看不到任何 IR 转储，因为它只是一个常量，XLA 会将其常量折叠起来。

在评论中您提到在 mnist_softmax 上运行，大概遵循上面链接上的说明。如果您确实在 CPU 上编译和运行，那么唯一剩下的问题是使用 VLOG(2)。仅当您设置命令行标志将其打开时，VLOG 才会启用。

因此，请尝试将 VLOG(2) 替换为 LOG(INFO)，您应该会在日志中看到 IR 转储。