OpenXLA 是一种针对特定领域的线性代数编译器，能够加快 TensorFlow 模型的运行速度，而且可能不需要更改源代码。
简介
运行 TensorFlow 程序后，所有操作均由 TensorFlow 执行程序单独执行。每个 TensorFlow 操作都有一个预编译的 GPU 内核实现，可以将执行程序分派给该实现。
XLA 提供了一种运行模型的替代模式：它会将 TensorFlow 图编译成一系列专门为给定模型生成的计算内核。由于这些内核是模型特有的，因此它们可以利用模型专属信息进行优化。以 XLA 在简单的 TensorFlow 计算环境中进行的优化为例：
def model_fn(x, y, z):
return tf.reduce_sum(x + y * z)
如果在不使用 XLA 的情况下运行，图会启动三个内核：分别对应于乘法、加法和减法运算。但是，XLA 可以优化该图，使其启动一次内核就能计算结果。它通过将加法、乘法和减法“融合”到一个 GPU 内核中来实现这一点。此外，这种融合操作不会将由 y*z 和 x+y*z 生成的中间值写出到内存中；而是直接将这些中间计算的结果“流式传输”给用户，同时将它们完全保留在 GPU 寄存器中。融合是 XLA 采用的最重要的一项优化措施。
内存带宽通常是硬件加速器上最稀缺的资源，因此消除内存操作是提高性能的最佳方法之一。
为 TensorFlow 模型启用 XLA
使用 tf.function(jit_compile=True) 进行显式编译
借助 Explicit compilation API，您可以精细地控制应选择哪些函数进行编译。例如，以下执行 MNIST 训练的 TensorFlow 函数使用 XLA 进行编译：
@tf.function(jit_compile=True)
def train_mnist(images, labels):
images, labels = cast(images, labels)
with tf.GradientTape() as tape:
predicted_labels = layer(images)
loss = tf.reduce_mean(tf.nn.sparse_softmax_cross_entropy_with_logits(
logits=predicted_labels, labels=labels
))
layer_variables = layer.trainable_variables
grads = tape.gradient(loss, layer_variables)
optimizer.apply_gradients(zip(grads, layer_variables))
jit_compile API 具有必须编译语义：要么使用 XLA 编译整个函数，要么抛出 errors.InvalidArgumentError 异常。目前，如果维度无法推断出来，也就是说，无法在不运行完整计算的情况下推断所有张量的维度，那么 XLA 无法编译包含此维度的函数。例如，以下函数将无法编译：
@tf.function
def not_compilable(x):
return tf.unique(x)
但形状可能会在不同的运行会话中有所不同：
@tf.function(jit_compile=True)
def recompiled_on_launch(a, b):
return a + b
recompiled_on_launch(tf.ones([1, 10]), tf.ones([1, 10]))
recompiled_on_launch(tf.ones([1, 100]), tf.ones([1, 100]))
如需查看更详细的用法示例，请参阅教程 Colab；如需了解 jit_compile=True 的使用方法，请观看教程视频。
与 Keras 搭配使用
对于 Keras 模型，jit_compile=True 可以设置为 model.compile 的参数：
model.compile(optimizer="adam", jit_compile=True)
使用分布式策略
您可以通过为步骤函数添加 jit_compile=True 注解，将 XLA:GPU 与 TF 分布式策略（MirroredStrategy 或 MultiWorkerMirroredStrategy）搭配使用：
@tf.function(jit_compile=True)
def step_fn():
t = tf.ones(shape=[100], dtype=tf.float32)
ctx = tf.distribute.get_replica_context()
return ctx.all_reduce(tf.distribute.ReduceOp.SUM, t)
@tf.function
def run_fn():
return strategy.run(step_fn)
自动聚类
若要在 TensorFlow 模型中开始使用 XLA 并且无需做出任何更改，最简单的方法是启用“自动聚类”，此功能会自动在可以使用 XLA 进行编译和执行的 TensorFlow 函数中查找聚类（连通的子图）。您可以通过设置 TF_XLA_FLAGS 环境变量，在 GPU 上启用自动聚类功能：
$ TF_XLA_FLAGS=--tf_xla_auto_jit=2 path/to/your/tf/program
自动聚类目前已针对 GPU 工作负载进行了优化，但您也可以通过另外使用 --tf_xla_cpu_global_jit 标记在 CPU 上启用它：
$ TF_XLA_FLAGS="--tf_xla_auto_jit=2 --tf_xla_cpu_global_jit" path/to/your/program
如需查看详细的用法示例，请参阅自动聚类教程 Colab。
使用 tfcompile 针对 CPU 进行 AOT（预先）编译
您还可以使用独立的 tfcompile 工具，此工具会将 TensorFlow 图转换为可执行代码（仅适用于 x86-64 CPU）。
检查已编译的程序
XLA 提供了自省工具，您可以利用这些工具检查生成的程序。如需转储生成的程序，请使用环境变量 XLA_FLAGS：
$ XLA_FLAGS="--xla_dump_to=/tmp/generated" TF_XLA_FLAGS="--tf_xla_auto_jit=2" my/tensorflow/program
执行转储后，您可以在 /tmp/generated 中找到以下文件：
module_XXXX.*_optimizations.txt：生成的 XLA 程序，每个已编译的聚类对应一个。在提交 XLA 错误报告时附加这些文件非常有帮助！module_XXXX.ptx：生成的 PTX 文件。
您还可以使用以下命令，对可视化 TensorFlow 图中 XLA 聚类的嵌入方式的图进行转储：
$ TF_DUMP_GRAPH_PREFIX=/tmp/generated TF_XLA_FLAGS="--tf_xla_clustering_debug"
可重现的 bug 报告
如果错误报告包含已生成 XLA 程序和已用自动聚类嵌入的转储文件，则更容易重现。如需为使用自动聚类运行的 TensorFlow 程序生成此类报告，请启动以下命令：
$ TF_DUMP_GRAPH_PREFIX=/tmp/generated \
TF_XLA_FLAGS="--tf_xla_clustering_debug --tf_xla_auto_jit=2" \
XLA_FLAGS="--xla_dump_hlo_as_text --xla_dump_to=/tmp/generated" \
my/tensorflow/program"
提交 bug 时，请附上 /tmp/generated 目录（如上所述）的内容。
如果可能，请尝试使用 run_hlo_module 并在生成的程序上以迭代方式运行它，使错误具体到单个 XLA 程序中。
深入阅读
- OpenXLA 文档 OpenXLA 文档
- 已知问题：关于 XLA+TF 的已知问题列表
- XLA - TensorFlow，已编译：阅读 Google Developers 博客
- 查看 GitHub 上的 XLA 源代码！
XLA 前端
除了 TensorFlow，还可以通过以下工具生成 XLA 程序：