Compila a partir del código fuente

En este documento, se describe cómo compilar componentes de XLA.

Si no clonaste el repositorio de XLA ni instalaste Bazel, consulta las secciones iniciales de la Guía para desarrolladores de XLA.

Linux

Configurar

Las compilaciones de XLA se configuran con el archivo .bazelrc en el directorio raíz del repositorio. La secuencia de comandos ./configure.py se puede usar para ajustar la configuración común.

Si necesitas cambiar la configuración, ejecuta la secuencia de comandos ./configure.py desde el directorio raíz del repositorio. Esta secuencia de comandos tiene marcas para la ubicación de las dependencias de XLA y opciones de configuración de compilación adicionales (por ejemplo, marcas del compilador). Consulta la sección Sesión de muestra para obtener más detalles.

Compatibilidad con la CPU

Te recomendamos que uses una imagen de Docker adecuada, como ml-build, que también se usa en los flujos de trabajo de CI de XLA en GitHub, para compilar y probar XLA. La imagen ml-build incluye Clang 18 preinstalado.

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

Con un contenedor de Docker, puedes compilar XLA con compatibilidad para CPU ejecutando los siguientes comandos:

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

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

Si deseas compilar destinos de XLA con compatibilidad con CPU sin usar Docker, deberás instalar Clang. Actualmente, XLA se compila con Clang 18 en la CI, pero las versiones anteriores también deberían funcionar.

Para configurar y compilar los destinos, ejecuta los siguientes comandos:

./configure.py --backend=CPU

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

Asistencia de GPU

Te recomendamos que uses el mismo contenedor de Docker que se mencionó anteriormente para compilar XLA con compatibilidad con GPU.

Para iniciar el contenedor de Docker con acceso a todas las GPUs, ejecuta el siguiente 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

Para compilar XLA con compatibilidad con GPU, ejecuta los siguientes comandos:

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

docker exec xla_gpu bazel build \
  --spawn_strategy=sandboxed \
  --test_output=all \
  //xla/...
  • No uses la marca --gpus all cuando inicies el contenedor de Docker.
  • Especifica las capacidades de procesamiento de CUDA de forma manual. Por ejemplo:
docker exec xla_gpu ./configure.py --backend=CUDA \
  --cuda_compute_capabilities="9.0"

Para obtener más detalles sobre las imágenes de Docker de la GPU de TensorFlow, consulta este documento.

También puedes compilar destinos de XLA con compatibilidad con GPU sin Docker. Configura y compila los destinos con los siguientes comandos:

./configure.py --backend=CUDA

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

Para obtener más detalles sobre CUDA hermético, consulta este documento.

Compila XLA con compatibilidad con CUDA/cuDNN usando el contenedor de CI/versión de JAX

XLA es un compilador que JAX usa de forma interna. JAX se distribuye a través de ruedas de PyPI. En la documentación de integración continua de JAX, se explica cómo compilar ruedas de JAX con el contenedor de Docker tensorflow/ml-build:latest.

También podemos extender estas instrucciones para compilar destinos de XLA dentro del contenedor de JAX. Esto garantiza que la configuración de compilación de los destinos de XLA sea coherente con la configuración de compilación de JAX/XLA, lo que puede ser útil si queremos reproducir los resultados de la carga de trabajo con herramientas de XLA que se crearon originalmente en JAX.

Compila objetivos de XLA en el contenedor de CI de JAX

  1. Clona el repositorio de JAX y navega al directorio "jax".
git clone <a href="https://github.com/jax-ml/jax.git">https://github.com/jax-ml/jax.git</a>

cd jax
  1. Para iniciar el contenedor de Docker de CI/versión de JAX, ejecuta el siguiente comando:
./ci/utilities/run_docker_container.sh

Esto iniciará un contenedor de Docker llamado "jax".

  1. Compila el destino jax-cuda-plugin dentro del contenedor con el siguiente comando:
docker exec jax ./ci/build_artifacts.sh jax-cuda-plugin

Esto creará el archivo .jax_configure.bazelrc con la configuración de compilación requerida, incluida la compatibilidad con CUDA/cuDNN.

  1. Accede a una shell interactiva dentro del contenedor:
docker exec -ti jax /bin/bash

Ahora deberías estar en el directorio /jax dentro del contenedor.

  1. Compila el destino de XLA con el siguiente comando, p.ej.:
/usr/local/bin/bazel build \
  --config=cuda_libraries_from_stubs \
  --verbose_failures=true \
  @xla//xla/tools/multihost_hlo_runner:hlo_runner_main

De manera opcional, puedes reemplazar las variables de entorno HERMETIC, p.ej.:

--repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_90"
  1. Copia los artefactos resultantes en /jax/dist para acceder a ellos desde el SO host si es necesario.
cp bazel-bin/external/xla/xla/tools/multihost_hlo_runner/hlo_runner_main \
  ./dist/
  1. Sal de la shell interactiva:
exit