Membangun dari sumber

Dokumen ini menjelaskan cara membuat komponen XLA.

Jika Anda belum meng-clone repositori XLA atau menginstal Bazel, lihat bagian awal Panduan Developer XLA.

Linux

Konfigurasi

Build XLA dikonfigurasi oleh file .bazelrc di direktori root repositori. Skrip ./configure.py dapat digunakan untuk menyesuaikan setelan umum.

Jika Anda perlu mengubah konfigurasi, jalankan skrip ./configure.py dari direktori root repositori. Skrip ini memiliki flag untuk lokasi dependensi XLA dan opsi konfigurasi build tambahan (misalnya, flag compiler). Lihat bagian Contoh sesi untuk mengetahui detailnya.

Dukungan CPU

Sebaiknya gunakan image Docker yang sesuai - seperti ml-build , yang juga digunakan dalam alur kerja CI XLA di GitHub - untuk membangun dan menguji XLA. Image ml-build dilengkapi dengan Clang 18 yang telah diinstal sebelumnya.

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

Dengan menggunakan container Docker, Anda dapat membangun XLA dengan dukungan CPU dengan menjalankan perintah berikut:

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

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

Jika ingin membuat target XLA dengan dukungan CPU tanpa menggunakan Docker, Anda harus menginstal Clang. XLA saat ini dibangun dengan Clang 18 di CI, tetapi versi sebelumnya juga akan berfungsi.

Untuk mengonfigurasi dan membangun target, jalankan perintah berikut:

./configure.py --backend=CPU

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

Dukungan GPU

Sebaiknya gunakan container Docker yang sama seperti yang disebutkan di atas untuk mem-build XLA dengan dukungan GPU.

Untuk memulai container Docker dengan akses ke semua GPU, jalankan perintah berikut:

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

Untuk membangun XLA dengan dukungan GPU, jalankan perintah berikut:

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

docker exec xla_gpu bazel build \
  --spawn_strategy=sandboxed \
  --test_output=all \
  //xla/...
  • Jangan gunakan flag --gpus all saat memulai container Docker.
  • Tentukan kemampuan komputasi CUDA secara manual, misalnya:
docker exec xla_gpu ./configure.py --backend=CUDA \
  --cuda_compute_capabilities="9.0"

Untuk mengetahui detail selengkapnya terkait image Docker GPU TensorFlow, Anda dapat melihat dokumen ini.

Anda juga dapat membuat target XLA dengan dukungan GPU tanpa Docker. Konfigurasi dan bangun target menggunakan perintah berikut:

./configure.py --backend=CUDA

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

Untuk mengetahui detail selengkapnya terkait CUDA hermetik, Anda dapat melihat dokumen ini.

Membangun XLA dengan Dukungan CUDA/cuDNN Menggunakan Container CI/Rilis JAX

XLA adalah compiler yang digunakan secara internal oleh JAX. JAX didistribusikan melalui wheel PyPI. Dokumentasi Integrasi Berkelanjutan JAX menjelaskan cara membuat wheel JAX menggunakan container Docker tensorflow/ml-build:latest.

Kita juga dapat memperluas petunjuk ini untuk membuat target XLA dalam container JAX. Hal ini memastikan bahwa konfigurasi build target XLA konsisten dengan konfigurasi build JAX/XLA, yang dapat berguna jika kita ingin mereproduksi hasil workload menggunakan alat XLA yang awalnya dibuat di JAX.

Membangun Target XLA di Container CI JAX

  1. Buat clone repositori JAX dan buka direktori 'jax'
git clone <a href="https://github.com/jax-ml/jax.git">https://github.com/jax-ml/jax.git</a>

cd jax
  1. Mulai container Docker CI/Rilis JAX dengan menjalankan:
./ci/utilities/run_docker_container.sh

Tindakan ini akan memulai container Docker bernama 'jax'.

  1. Bangun target jax-cuda-plugin di dalam container menggunakan:
docker exec jax ./ci/build_artifacts.sh jax-cuda-plugin

Tindakan ini akan membuat file .jax_configure.bazelrc dengan konfigurasi build yang diperlukan, termasuk dukungan CUDA/cuDNN

  1. Akses shell interaktif di dalam container:
docker exec -ti jax /bin/bash

Sekarang Anda akan berada di direktori /jax dalam container

  1. Bangun target XLA dengan perintah berikut, misalnya:
/usr/local/bin/bazel build \
  --config=cuda_libraries_from_stubs \
  --verbose_failures=true \
  @xla//xla/tools/multihost_hlo_runner:hlo_runner_main

Secara opsional, Anda dapat menimpa HERMETIC env, misalnya:

--repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_90"
  1. Salin artefak yang dihasilkan ke /jax/dist untuk mengaksesnya dari OS host jika diperlukan
cp bazel-bin/external/xla/xla/tools/multihost_hlo_runner/hlo_runner_main \
  ./dist/
  1. Keluar dari shell interaktif:
exit