Сборка из исходного кода

В этом документе описывается, как создавать компоненты XLA.

Если вы не клонировали репозиторий XLA или не установили Bazel, ознакомьтесь с начальными разделами Руководства разработчика XLA .

Linux

Настройка

Настройка сборок XLA осуществляется с помощью файла .bazelrc , расположенного в корневом каталоге репозитория. Для изменения общих параметров можно использовать скрипт ./configure.py .

Если вам необходимо изменить конфигурацию, запустите скрипт ./configure.py из корневого каталога репозитория. Этот скрипт содержит флаги для указания расположения зависимостей XLA и дополнительные параметры конфигурации сборки (например, флаги компилятора). Подробности см. в разделе «Пример сессии» .

поддержка ЦП

Для сборки и тестирования XLA мы рекомендуем использовать подходящий образ Docker, например ml-build , который также используется в рабочих процессах CI XLA на GitHub. Образ ml-build поставляется с предустановленным Clang 18.

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

Используя контейнер Docker, вы можете собрать XLA с поддержкой ЦП, выполнив следующие команды:

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

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

Если вы хотите собрать целевые объекты XLA с поддержкой ЦП без использования Docker , вам потребуется установить Clang. В настоящее время XLA собирается с помощью Clang 18 в CI, но более ранние версии также должны работать.

Для настройки и сборки целевых объектов выполните следующие команды:

./configure.py --backend=CPU

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

поддержка графического процессора

Мы рекомендуем использовать тот же контейнер Docker, что и выше, для сборки XLA с поддержкой графического процессора.

Чтобы запустить контейнер Docker с доступом ко всем графическим процессорам, выполните следующую команду:

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

Для сборки XLA с поддержкой графического процессора выполните следующие команды:

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

docker exec xla_gpu bazel build \
  --spawn_strategy=sandboxed \
  --test_output=all \
  //xla/...
  • Не используйте флаг --gpus all при запуске контейнера Docker.
  • Укажите вычислительные возможности CUDA вручную, например:
docker exec xla_gpu ./configure.py --backend=CUDA \
  --cuda_compute_capabilities="9.0"

Более подробную информацию об образах Docker для работы с GPU в TensorFlow можно найти в этом документе.

Вы можете создавать целевые объекты XLA с поддержкой GPU и без Docker. Настройте и соберите целевые объекты, используя следующие команды:

./configure.py --backend=CUDA

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

Для получения более подробной информации о герметичной архитектуре CUDA вы можете ознакомиться с этим документом.

Сборка XLA с поддержкой CUDA/cuDNN с использованием контейнера JAX CI/Release

XLA — это компилятор, используемый внутри JAX. JAX распространяется через пакеты PyPI. В документации по непрерывной интеграции JAX объясняется, как собирать пакеты JAX с помощью контейнера Docker tensorflow/ml-build:latest .

Мы можем расширить эти инструкции, чтобы собирать цели XLA и внутри контейнера JAX. Это гарантирует, что конфигурация сборки целей XLA будет соответствовать конфигурации сборки JAX/XLA, что может быть полезно, если мы хотим воспроизвести результаты рабочей нагрузки с использованием инструментов XLA, которые изначально были созданы в JAX.

Создание целей XLA в контейнере JAX CI

  1. Клонируйте репозиторий JAX и перейдите в каталог 'jax'.
git clone https://github.com/jax-ml/jax.git

cd jax
  1. Запустите контейнер Docker JAX CI/Release, выполнив команду:
./ci/utilities/run_docker_container.sh

Это запустит контейнер Docker с именем 'jax'.

  1. Создайте целевой объект jax-cuda-plugin внутри контейнера, используя:
docker exec jax ./ci/build_artifacts.sh jax-cuda-plugin

Это создаст файл .jax_configure.bazelrc с необходимой конфигурацией сборки, включая поддержку CUDA/cuDNN.

  1. Получите доступ к интерактивной оболочке внутри контейнера:
docker exec -ti jax /bin/bash

Теперь вы должны находиться в каталоге /jax внутри контейнера.

  1. Создайте целевой объект XLA с помощью следующей команды, например:
/usr/local/bin/bazel build \
  --config=cuda_libraries_from_stubs \
  --verbose_failures=true \
  @xla//xla/tools/multihost_hlo_runner:hlo_runner_main

При желании вы можете перезаписать HERMETIC среды, например:

--repo_env=HERMETIC_CUDA_COMPUTE_CAPABILITIES="sm_90"
  1. Скопируйте полученные артефакты в /jax/dist , чтобы при необходимости получить к ним доступ из операционной системы хоста.
cp bazel-bin/external/xla/xla/tools/multihost_hlo_runner/hlo_runner_main \
  ./dist/
  1. Выйти из интерактивной оболочки:
exit

Windows

Сборка XLA в Windows нативно — это процесс, задействующий только процессор, поскольку CUDA напрямую не поддерживается в Windows; для поддержки CUDA необходимо использовать WSL2. Также требуются определенные настройки среды, включая оболочку Bash, компилятор Clang и Visual Studio.

Предварительные требования

  1. Visual Studio: Для настройки набора инструментов C++ (который предоставляет необходимые системные заголовочные файлы и библиотеки) необходимо установить Visual Studio 2019 версии 16.5 или новее.
  2. Среда Bash: Для сборки XLA в Windows необходимо использовать оболочку Bash (например, MSYS2 или Git Bash).
  3. Компилятор Clang: В сборках XLA для Windows используется clang-cl вместо стандартного компилятора MSVC. Убедитесь, что LLVM/Clang установлен.
  4. Python: Python 3 должен быть установлен и доступен в PATH вашей системы.
  5. Режим разработчика: Необходимо включить режим разработчика в Windows или запустить оболочку Bash от имени администратора. Это необходимо, поскольку Bazel использует создание символических ссылок для дерева исполняемых файлов в процессе сборки.

Создание с нуля

Для компиляции бэкенда XLA для ЦП в Windows выполните следующую команду в терминале Bash. Мы используем конфигурацию xla_windows_x86_cpu_2022 для автоматической настройки инструментария clang-cl .

Обратите внимание, что некоторые целевые платформы и теги необходимо явно исключить из сборки Windows для ЦП, например, целевые платформы для графических процессоров и функции, которые еще не поддерживаются в Windows:

# Note: Ensure your PATH includes Python and Bazel before running this
bazel build \
  --config=xla_windows_x86_cpu_2022 \
  --keep_going \
  --build_tag_filters=-no_oss,-oss_excluded,-gpu,-no_windows,-windows_excluded \
  -- //xla/... -//xla/hlo/experimental/... -//xla/python_api/... -//xla/python/...