從原始碼開始建構

本文說明如何建構 XLA 元件。

如果您未複製 XLA 存放區或安裝 Bazel,請參閱 XLA 開發人員指南的初始章節。

Linux

設定

XLA 建構作業是透過存放區根目錄中的 .bazelrc 檔案設定。您可以使用 ./configure.py 指令碼調整常見設定。

如要變更設定,請從存放區的根目錄執行 ./configure.py 指令碼。這個指令碼包含 XLA 依附元件位置的旗標,以及其他建構設定選項 (例如編譯器旗標)。詳情請參閱「範例工作階段」一節。

CPU 支援

建議使用合適的 Docker 映像檔 (例如 ml-build),這類映像檔也用於 GitHub 上的 XLA CI 工作流程,可建構及測試 XLA。ml-build 映像檔已預先安裝 Clang 18。

docker run -itd --rm \
--name xla \
-w /xla \
-v $PWD:/xla \
us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest \
bash

使用 Docker 容器時,您可以執行下列指令,建構支援 CPU 的 XLA:

docker exec xla ./configure.py --backend=CPU

docker exec xla bazel build \
  --spawn_strategy=sandboxed \
  --test_output=all \
  //xla/...

如要不使用 Docker 建構支援 CPU 的 XLA 目標,您需要安裝 Clang。XLA 目前在 CI 中以 Clang 18 建構,但舊版也應可運作。

如要設定及建構目標,請執行下列指令:

./configure.py --backend=CPU

bazel build \
  --spawn_strategy=sandboxed \
  --test_output=all \
  //xla/...

GPU 支援

建議您使用上述 Docker 容器,建構支援 GPU 的 XLA。

如要啟動 Docker 容器並存取所有 GPU,請執行下列指令:

docker run -itd --rm \
  --gpus all \
  --name xla_gpu \
  -w /xla \
  -v $PWD:/xla \
  us-docker.pkg.dev/ml-oss-artifacts-published/ml-public-container/ml-build:latest \
  bash

如要建構支援 GPU 的 XLA,請執行下列指令:

docker exec xla_gpu ./configure.py --backend=CUDA

docker exec xla_gpu bazel build \
  --spawn_strategy=sandboxed \
  --test_output=all \
  //xla/...
  • 啟動 Docker 容器時,請不要使用 --gpus all 旗標。
  • 手動指定 CUDA 運算功能,例如:
docker exec xla_gpu ./configure.py --backend=CUDA \
  --cuda_compute_capabilities="9.0"

如要進一步瞭解 TensorFlow 的 GPU Docker 映像檔,請參閱這份文件。

您也可以在不使用 Docker 的情況下,建構支援 GPU 的 XLA 目標。使用下列指令設定及建構目標:

./configure.py --backend=CUDA

bazel build \
  --spawn_strategy=sandboxed \
  --test_output=all \
  //xla/...

如要進一步瞭解密封式 CUDA,請參閱這份文件

使用 JAX CI/發布容器建構支援 CUDA/cuDNN 的 XLA

XLA 是 JAX 內部使用的編譯器。JAX 是透過 PyPI Wheel 發布。 JAX 持續整合說明文件說明如何使用 tensorflow/ml-build:latest Docker 容器建構 JAX 輪子。

我們也可以擴充這些指令,在 JAX 容器中建構 XLA 目標。這可確保 XLA 目標的建構設定與 JAX/XLA 建構設定一致。如果我們想使用最初在 JAX 中建立的 XLA 工具重現工作負載結果,這就非常實用。

在 JAX CI 容器中建構 XLA 目標

  1. 複製 JAX 存放區,然後前往「jax」目錄
git clone <a href="https://github.com/jax-ml/jax.git">https://github.com/jax-ml/jax.git</a>

cd jax
  1. 執行下列指令,啟動 JAX CI/Release Docker 容器:
./ci/utilities/run_docker_container.sh

這會啟動名為「jax」的 Docker 容器。

  1. 在容器內使用下列指令建構 jax-cuda-plugin 目標:
docker exec jax ./ci/build_artifacts.sh jax-cuda-plugin

這會建立 .jax_configure.bazelrc 檔案,其中包含必要的建構設定,包括 CUDA/cuDNN 支援

  1. 存取容器內的互動式殼層:
docker exec -ti jax /bin/bash

您現在應該位於容器內的 /jax 目錄

  1. 使用下列指令建構 XLA 目標,例如:
/usr/local/bin/bazel build \
  --config=cuda_libraries_from_stubs \
  --verbose_failures=true \
  @xla//xla/tools/multihost_hlo_runner:hlo_runner_main

您也可以視需要覆寫 HERMETIC envs,例如:

--repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_90"
  1. 將產生的構件複製到 /jax/dist,以便從主機 OS 存取 (如有需要)
cp bazel-bin/external/xla/xla/tools/multihost_hlo_runner/hlo_runner_main \
  ./dist/
  1. 結束互動式殼層:
exit