Créer à partir de la source

Ce document explique comment créer des composants XLA.

Si vous n'avez pas cloné le dépôt XLA ni installé Bazel, consultez les sections initiales du Guide du développeur XLA.

Linux

Configurer

Les compilations XLA sont configurées par le fichier .bazelrc dans le répertoire racine du dépôt. Le script ./configure.py peut être utilisé pour ajuster les paramètres courants.

Si vous devez modifier la configuration, exécutez le script ./configure.py à partir du répertoire racine du dépôt. Ce script comporte des indicateurs pour l'emplacement des dépendances XLA et des options de configuration de compilation supplémentaires (indicateurs de compilation, par exemple). Pour en savoir plus, consultez la section Exemple de session.

Compatibilité avec les processeurs

Nous vous recommandons d'utiliser une image Docker appropriée, telle que ml-build, qui est également utilisée dans les workflows CI de XLA sur GitHub, pour compiler et tester XLA. L'image ml-build est fournie avec Clang 18 préinstallé.

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

À l'aide d'un conteneur Docker, vous pouvez créer XLA avec la prise en charge du processeur en exécutant les commandes suivantes :

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

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

Si vous souhaitez créer des cibles XLA avec prise en charge du processeur sans utiliser Docker, vous devez installer Clang. XLA est actuellement conçu avec Clang 18 dans l'intégration continue, mais les versions antérieures devraient également fonctionner.

Pour configurer et compiler les cibles, exécutez les commandes suivantes :

./configure.py --backend=CPU

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

Compatibilité avec les GPU

Nous vous recommandons d'utiliser le même conteneur Docker que celui mentionné ci-dessus pour compiler XLA avec la prise en charge du GPU.

Pour démarrer le conteneur Docker avec accès à tous les GPU, exécutez la commande suivante :

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

Pour compiler XLA avec la prise en charge des GPU, exécutez les commandes suivantes :

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

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

N'utilisez pas l'indicateur --gpus all lorsque vous démarrez le conteneur Docker.
Spécifiez manuellement les capacités de calcul CUDA. Par exemple :

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

Pour en savoir plus sur les images Docker de TensorFlow pour GPU, consultez ce document.

Vous pouvez également créer des cibles XLA compatibles avec les GPU sans Docker. Configurez et créez des cibles à l'aide des commandes suivantes :

./configure.py --backend=CUDA

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

Pour en savoir plus sur CUDA hermétique, consultez ce document.

Créer XLA avec la prise en charge de CUDA/cuDNN à l'aide du conteneur JAX CI/Release

XLA est un compilateur utilisé en interne par JAX. JAX est distribué via des roues PyPI. La documentation sur l'intégration continue de JAX explique comment créer des roues JAX à l'aide du conteneur Docker tensorflow/ml-build:latest.

Nous pouvons également étendre ces instructions pour créer des cibles XLA dans le conteneur JAX. Cela garantit que la configuration de compilation des cibles XLA est cohérente avec la configuration de compilation JAX/XLA, ce qui peut être utile si nous souhaitons reproduire les résultats de la charge de travail à l'aide d'outils XLA initialement créés dans JAX.

Créer des cibles XLA dans le conteneur JAX CI

Clonez le dépôt JAX et accédez au répertoire "jax".

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

cd jax

Démarrez le conteneur Docker CI/Release JAX en exécutant la commande suivante :

./ci/utilities/run_docker_container.sh

Cela démarrera un conteneur Docker nommé "jax".

Créez la cible jax-cuda-plugin dans le conteneur à l'aide de la commande suivante :

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

Cela créera le fichier .jax_configure.bazelrc avec la configuration de compilation requise, y compris la prise en charge de CUDA/cuDNN.

Accédez à une interface système interactive à l'intérieur du conteneur :

docker exec -ti jax /bin/bash

Vous devriez maintenant vous trouver dans le répertoire /jax du conteneur.

Créez la cible XLA avec la commande suivante, par exemple :

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

Vous pouvez éventuellement écraser les environnements HERMETIC, par exemple :

--repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_90"

Copiez les artefacts obtenus dans /jax/dist pour y accéder depuis l'OS hôte si nécessaire.

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

Quittez le shell interactif :

exit