Crea dall'origine

Questo documento descrive come creare componenti XLA.

Se non hai clonato il repository XLA o installato Bazel, consulta le sezioni iniziali della Guida per gli sviluppatori XLA.

Linux

Configura

Le build XLA sono configurate dal file .bazelrc nella directory root del repository. Lo script ./configure.py può essere utilizzato per modificare le impostazioni comuni.

Se devi modificare la configurazione, esegui lo script ./configure.py dalla directory root del repository. Questo script contiene flag per la posizione delle dipendenze XLA e opzioni di configurazione della build aggiuntive (ad esempio, flag del compilatore). Per i dettagli, consulta la sezione Sessione di esempio.

Supporto CPU

Ti consigliamo di utilizzare un'immagine Docker adatta, ad esempio ml-build , utilizzata anche nei flussi di lavoro CI di XLA su GitHub, per la creazione e il test di XLA. L'immagine ml-build è dotata di Clang 18 preinstallato.

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

Utilizzando un container Docker, puoi creare XLA con il supporto della CPU eseguendo i seguenti comandi:

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

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

Se vuoi creare target XLA con supporto CPU senza utilizzare Docker, devi installare Clang. XLA è attualmente creato con Clang 18 in CI, ma dovrebbero funzionare anche le versioni precedenti.

Per configurare e creare i target, esegui questi comandi:

./configure.py --backend=CPU

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

Supporto GPU

Ti consigliamo di utilizzare lo stesso container Docker menzionato in precedenza per creare XLA con supporto GPU.

Per avviare il container Docker con accesso a tutte le GPU, esegui questo comando:

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

Per creare XLA con il supporto della GPU, esegui questi comandi:

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

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

Non utilizzare il flag --gpus all all'avvio del container Docker.
Specifica manualmente le funzionalità di calcolo CUDA. Ad esempio:

docker exec xla_gpu ./configure.py --backend=CUDA \
  --cuda_compute_capabilities="9.0"

Per maggiori dettagli sulle immagini Docker della GPU di TensorFlow, consulta questo documento.

Puoi creare target XLA con supporto GPU anche senza Docker. Configura e crea le destinazioni utilizzando i seguenti comandi:

./configure.py --backend=CUDA

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

Per maggiori dettagli su CUDA ermetico, consulta questo documento.

Compila XLA con il supporto di CUDA/cuDNN utilizzando il container CI/Release di JAX

XLA è un compilatore utilizzato internamente da JAX. JAX viene distribuito tramite i wheel di PyPI. La documentazione sull'integrazione continua di JAX spiega come creare ruote JAX utilizzando il container Docker tensorflow/ml-build:latest.

Possiamo estendere queste istruzioni per creare target XLA anche all'interno del contenitore JAX. In questo modo, la configurazione di compilazione dei target XLA è coerente con la configurazione di compilazione JAX/XLA, il che può essere utile se vogliamo riprodurre i risultati del carico di lavoro utilizzando strumenti XLA originariamente creati in JAX.

Crea target XLA nel container JAX CI

Clona il repository JAX e vai alla directory "jax"

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

cd jax

Avvia il container Docker JAX CI/Release eseguendo:

./ci/utilities/run_docker_container.sh

Verrà avviato un container Docker denominato "jax".

Crea il target jax-cuda-plugin all'interno del container utilizzando:

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

Verrà creato il file .jax_configure.bazelrc con la configurazione di compilazione richiesta, incluso il supporto di CUDA/cuDNN

Accedi a una shell interattiva all'interno del container:

docker exec -ti jax /bin/bash

Ora dovresti trovarti nella directory /jax all'interno del container

Crea la destinazione XLA con il seguente comando, ad esempio:

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

(Facoltativo) Puoi sovrascrivere le variabili di ambiente HERMETIC, ad esempio:

--repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_90"

Copia gli artefatti risultanti in /jax/dist per accedervi dal sistema operativo host se necessario

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

Esci dalla shell interattiva:

exit