Criar a partir da fonte

Este documento descreve como criar componentes XLA.

Se você não clonou o repositório XLA nem instalou o Bazel, confira as seções iniciais do Guia do desenvolvedor do XLA.

Linux

Configurar

Os builds de XLA são configurados pelo arquivo .bazelrc no diretório raiz do repositório. O script ./configure.py pode ser usado para ajustar configurações comuns.

Se você precisar mudar a configuração, execute o script ./configure.py no diretório raiz do repositório. Esse script tem flags para o local das dependências do XLA e outras opções de configuração de build (flags do compilador, por exemplo). Consulte a seção Exemplo de sessão para mais detalhes.

Suporte a CPU

Recomendamos usar uma imagem do Docker adequada, como ml-build, que também é usada nos fluxos de trabalho de CI do XLA no GitHub, para criar e testar o XLA. A imagem ml-build vem com o Clang 18 pré-instalado.

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

Usando um contêiner do Docker, é possível criar o XLA com suporte à CPU executando os seguintes comandos:

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

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

Se você quiser criar destinos XLA com suporte a CPU sem usar o Docker, instale o Clang. Atualmente, o XLA é criado com o Clang 18 em CI, mas versões anteriores também devem funcionar.

Para configurar e criar os destinos, execute os seguintes comandos:

./configure.py --backend=CPU

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

Suporte a GPUs

Recomendamos usar o mesmo contêiner do Docker mencionado acima para criar o XLA com suporte a GPU.

Para iniciar o contêiner do Docker com acesso a todas as GPUs, execute o seguinte 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 criar o XLA com suporte a GPU, execute os seguintes comandos:

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

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

Não use a flag --gpus all ao iniciar o contêiner do Docker.
Especifique manualmente as capacidades de computação CUDA. Por exemplo:

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

Para mais detalhes sobre imagens do Docker de GPU do TensorFlow, consulte este documento.

Também é possível criar destinos XLA com suporte a GPU sem o Docker. Configure e crie destinos usando os seguintes comandos:

./configure.py --backend=CUDA

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

Para mais detalhes sobre CUDA hermético, consulte este documento.

Criar XLA com suporte a CUDA/cuDNN usando o contêiner de CI/lançamento do JAX

O XLA é um compilador usado internamente pelo JAX. O JAX é distribuído por pacotes wheel do PyPI. A documentação de integração contínua do JAX explica como criar wheels do JAX usando o contêiner Docker tensorflow/ml-build:latest.

Podemos estender essas instruções para criar destinos XLA também no contêiner JAX. Isso garante que a configuração de build dos destinos do XLA seja consistente com a configuração de build do JAX/XLA, o que pode ser útil se quisermos reproduzir resultados de carga de trabalho usando ferramentas do XLA originalmente criadas no JAX.

Criar destinos do XLA no contêiner de CI do JAX

Clone o repositório JAX e navegue até o diretório "jax".

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

cd jax

Inicie o contêiner do Docker de CI/lançamento do JAX executando:

./ci/utilities/run_docker_container.sh

Isso vai iniciar um contêiner do Docker chamado "jax".

Crie o destino jax-cuda-plugin dentro do contêiner usando:

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

Isso vai criar o arquivo .jax_configure.bazelrc com a configuração de build necessária, incluindo suporte a CUDA/cuDNN.

Acesse um shell interativo dentro do contêiner:

docker exec -ti jax /bin/bash

Agora você está no diretório /jax dentro do contêiner.

Crie o destino XLA com o seguinte comando, por exemplo:

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

Se quiser, substitua os ambientes HERMETIC. Por exemplo:

--repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_90"

Copie os artefatos resultantes para /jax/dist e acesse-os no SO do host se necessário.

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

Saia do shell interativo:

exit