ソースからビルドする

このドキュメントでは、XLA コンポーネントをビルドする方法について説明します。

XLA リポジトリのクローンを作成していない場合や、Bazel をインストールしていない場合は、XLA デベロッパー ガイドの最初のセクションをご覧ください。

Linux

設定

XLA ビルドは、リポジトリのルート ディレクトリにある .bazelrc ファイルで構成されます。./configure.py スクリプトを使用して、一般的な設定を調整できます。

構成を変更する必要がある場合は、リポジトリのルート ディレクトリから ./configure.py スクリプトを実行します。このスクリプトには、XLA の依存関係の場所と追加のビルド構成オプション（コンパイラ フラグなど）のフラグがあります。詳細については、セッションの例をご覧ください。

CPU のサポート

XLA のビルドとテストには、適切な Docker イメージ（GitHub の XLA の CI ワークフローでも使用されている ml-build など）を使用することをおすすめします。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 ホイールをビルドする方法について説明しています。

これらの手順を拡張して、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/リリース Docker コンテナを起動します。
./ci/utilities/run_docker_container.sh

これにより、「jax」という名前の Docker コンテナが起動します。

  1. コンテナ内で jax-cuda-plugin ターゲットをビルドします。
docker exec jax ./ci/build_artifacts.sh jax-cuda-plugin

これにより、CUDA/cuDNN のサポートなど、必要なビルド構成を含む .jax_configure.bazelrc ファイルが作成されます。

  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 環境を上書きできます。例:

--repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_90"
  1. 必要に応じてホスト OS からアクセスできるように、結果のアーティファクトを /jax/dist にコピーします。
cp bazel-bin/external/xla/xla/tools/multihost_hlo_runner/hlo_runner_main \
  ./dist/
  1. インタラクティブ シェルを終了します。
exit