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) 進行明確編譯
明確編譯 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 具有 must-compile 語意：不是使用 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)
搭配分散式策略使用
XLA：GPU 可搭配 TF 分散式策略 (MirroredStrategy 或 MultiWorkerMirroredStrategy) 使用，方法是使用 jit_compile=True 標註步驟函式：
@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"
可重現的錯誤報告
如果錯誤報告包含所產生 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"
提報錯誤時，請附上 /tmp/generated 目錄的內容 (如上所示)。
請盡可能使用 run_hlo_module 將錯誤隔離到單一 XLA 程式，並透過產生的程式疊代執行該程式。
延伸閱讀
- OpenXLA 說明文件 OpenXLA 說明文件
- 已知問題 XLA 和 TF 的已知問題清單
- 經過編譯的 XLA - TensorFlow：前往 Google Developers 網誌閱讀
- 前往 GitHub 查看 XLA 原始碼！
XLA 前端
除了 TensorFlow 以外，您還可以透過下列方式產生 XLA 程式：