从源代码构建

本文档介绍了如何构建 XLA 组件。

如果您未克隆 XLA 代码库或安装 Bazel，请查看 XLA 开发者指南的初始部分。

Linux

配置

XLA build 通过代码库根目录中的 .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。

如需启动可访问所有 GPU 的 Docker 容器，请运行以下命令：

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 软件包分发。 JAX 持续集成文档介绍了如何使用 tensorflow/ml-build:latest Docker 容器构建 JAX build。

我们还可以扩展这些说明，以在 JAX 容器内构建 XLA 目标。这样可确保 XLA 目标的 build 配置与 JAX/XLA build 配置保持一致，如果我们想使用最初在 JAX 中创建的 XLA 工具重现工作负载结果，这会非常有用。

在 JAX CI 容器中构建 XLA 目标

git clone <a href="https://github.com/jax-ml/jax.git">https://github.com/jax-ml/jax.git</a>

cd jax

运行以下命令，启动 JAX CI/Release Docker 容器：

./ci/utilities/run_docker_container.sh

此命令将启动名为“jax”的 Docker 容器。

使用以下命令在容器内构建 jax-cuda-plugin 目标：

docker exec jax ./ci/build_artifacts.sh jax-cuda-plugin

此命令将创建包含所需 build 配置（包括 CUDA/cuDNN 支持）的 .jax_configure.bazelrc 文件

访问容器内的交互式 shell：

docker exec -ti jax /bin/bash

您现在应该位于容器内的 /jax 目录中

使用以下命令构建 XLA 目标，例如：

/usr/local/bin/bazel build \
  --config=cuda_libraries_from_stubs \
  --verbose_failures=true \
  @xla//xla/tools/multihost_hlo_runner:hlo_runner_main

您可以选择性地覆盖 HERMETIC 环境，例如：

--repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_90"

将生成的制品复制到 /jax/dist，以便在需要时从主机操作系统访问它们

cp bazel-bin/external/xla/xla/tools/multihost_hlo_runner/hlo_runner_main \
  ./dist/

退出交互式 shell：

exit