XLA: компиляция с опережением обработки на графическом процессоре (AOT)

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

В отличие от компиляции «на лету» (JIT), которая выполняет вычисления в запущенном процессе при первом выполнении, AOT исключает компиляцию из пути обслуживания, что уменьшает задержку при запуске и позволяет выполнять компиляцию в выделенной среде сборки.

Рабочие процессы и API

Стандартный конвейер AOT состоит из этапов компиляции, сериализации и загрузки/десериализации:

• Точки входа PjRt: В интерфейсе PjRt на C++ ( xla/pjrt/pjrt_client.h ) компиляция инициируется с помощью PjRtClient::Compile (или CompileAndLoad ). Полученный исполняемый файл сериализуется в строку с помощью PjRtExecutable::SerializeExecutable . На хосте, обслуживающем приложение, среда выполнения восстанавливает исполняемый файл с помощью PjRtClient::DeserializeExecutable или загружает его непосредственно для выполнения с помощью PjRtClient::LoadSerializedExecutable .
• Артефакты на уровне компилятора: На уровне компилятора предварительная компиляция создает объект GpuAotCompilationResult ( CompiledModule ; AotCompilationResult — устаревший псевдоним). Этот объект инкапсулирует сериализованный GpuExecutableProto и предоставляет доступ к SerializeAsString() и LoadExecutable() .

• Инструменты командной строки: Инструмент xla_compile компилирует модули HLO или StableHLO в сериализованные исполняемые файлы:

  xla_compile --module_file=module.hlo --output_file=output.bin \
      --platform=gpu --gpu_target_config=gpu_target_config.pbtxt
  

  Компиляция может выполняться на хосте без графического процессора при условии указания конфигурации целевого устройства (см. файл README xla/backends/gpu/target_config ).

Что такое сериализованный

Контейнером верхнего уровня для сериализованных программ для графического процессора является GpuExecutableProto :

• thunks : Последовательность выполнения thunk-функций. Kernel thunks (например, CustomKernelThunk ) встраивают собственный скомпилированный двоичный файл ядра ( cubin для NVIDIA, hsaco для ROCm), поэтому большая часть кода устройства находится здесь.
• binary : Двоичный файл устройства на уровне модуля. Сегодня он содержит только буферы констант модуля; ядра больше здесь не хранятся.
• buffer_allocations : Размеры буферов и срезы выделения памяти.
• gpu_compute_capability : Возможности аппаратной архитектуры целевого объекта компиляции.
• executable_abi_version : Версии ABI инструментария и среды выполнения.
• hlo_module_with_config : Встроенный сериализованный HloModuleProto и отладочные метаданные.

Каждое протокольное сообщение и поле, транзитивно достижимое из GpuExecutableProto является частью контракта совместимости, описанного в руководстве по совместимости GPU AOT .

Как загружается и выполняется артефакт

1. Восстановите граф thunk-функций: десериализуйте записи ThunkProto в последовательность thunk-функций в памяти, разрешая при этом пользовательские символы вызовов, FFI и ядра. Встроенный модуль HLO также анализируется, поэтому код операции HLO, неизвестный загружаемому двоичному файлу, приводит к сбою загрузки.
2. Выполнение проходов преобразования графа thunk: Выполните оптимизацию во время загрузки и проходы преобразования для thunk-объектов, находящихся в памяти (например, создание CommandBufferThunk во время выполнения для отправки работы через CUDA Graphs). CommandBufferThunk динамически перестраивается во время загрузки и никогда не сериализуется напрямую.
3. Выделение буферов: инициализация областей памяти и выделение буферов в соответствии с последовательно распределенными фрагментами назначения буферов.
4. Выполнить: Отправить подготовленную последовательность thunk-функций в потоки данных устройства.

Изменения в параметрах Pass, Fusion, Codegen и Autotuning, сохраняющие целостность сериализованной схемы и семантики времени выполнения, не влияют на совместимость с AOT.

Устранение неполадок: сбои загрузки, не связанные с ошибками версионирования.

Не все сбои загрузки исполняемых файлов вызваны ошибками сериализации. Два распространенных несоответствия среды приводят к ошибкам во время загрузки:

• Несоответствие вычислительных возможностей: GpuExecutable::FromProto проверяет, соответствует ли сериализованная gpu_compute_capability целевому устройству. Исполняемый файл, скомпилированный для NVIDIA Hopper, не может работать на NVIDIA Ampere.
• Несоответствие ABI инструментария CUDA: CudaRuntimeAbiVersion::IsCompatibleWith проверяет, соответствуют ли версии инструментария CUDA, cuDNN и CUB среды выполнения минимальным версиям, указанным в ExecutableAbiVersionProto , или превышают их. Если среда драйвера или библиотеки хоста старше, чем целевая среда компиляции, загрузка завершается с ошибкой FailedPreconditionError .