Разработка нового бэкэнда для XLA

Это руководство предназначено для системных инженеров, которые хотят, чтобы XLA эффективно выводил программы, ориентированные на их оборудование. Руководство не пошаговое и предполагает знание LLVM , Bazel и XLA.

XLA предоставляет абстрактный интерфейс, который может реализовать новая архитектура или ускоритель для создания серверной части для запуска программ ML, выводимых XLA. Перенацеливание XLA должно быть значительно проще и масштабируемее, чем реализация каждой существующей операции из внешней среды, такой как PyTorch или TensorFlow, для нового оборудования.

Большинство реализаций попадают в один из следующих сценариев:

  1. Существующая архитектура ЦП еще официально не поддерживается XLA, с существующим бэкэндом LLVM или без него.
  2. Аппаратное обеспечение, не похожее на ЦП, с существующим бэкэндом LLVM.
  3. Аппаратное обеспечение, не похожее на ЦП, без существующего бэкэнда LLVM.

Сценарий 1: Существующая архитектура ЦП еще официально не поддерживается XLA.

В этом сценарии начните с рассмотрения существующей серверной части процессора XLA . XLA позволяет легко ориентироваться на разные процессоры с помощью LLVM, поскольку основное различие между серверными модулями XLA для процессоров — это код, генерируемый LLVM.

Если у поставщика оборудования есть серверная часть LLVM для своего оборудования, ее можно легко связать с LLVM, созданной с помощью XLA. В режиме JIT серверная часть ЦП XLA генерирует код для ЦП хоста. Для предварительной компиляции xla::AotCompilationOptions может предоставить тройку LLVM для настройки целевой архитектуры.

Если существующего бэкэнда LLVM нет, но существует другой тип генератора кода, должна быть возможность повторно использовать большую часть существующего бэкенда ЦП.

Сценарий 2: Аппаратное обеспечение, не похожее на ЦП, с существующим серверным компонентом LLVM.

Можно смоделировать новую реализацию xla::Compiler на основе существующих классов xla::CPUCompiler и xla::GPUCompiler , поскольку они уже излучают LLVM IR. В зависимости от характера оборудования, возможно, придется изменить многие аспекты генерации LLVM IR, но большая часть кода может использоваться совместно с существующими бэкэндами.

Хорошим примером для подражания является серверная часть XLA с графическим процессором . Серверная часть графического процессора ориентирована на ISA, не похожую на процессор, и поэтому некоторые аспекты генерации кода уникальны для домена графического процессора. Другие виды аппаратного обеспечения, например, DSP, такие как Hexagon (который имеет входящую часть LLVM), могут повторно использовать части логики IR-излучения LLVM, но другие части будут уникальными.

Сценарий 3: Аппаратное обеспечение, не похожее на ЦП, без существующего бэкэнда LLVM.

Если использовать LLVM невозможно, то лучшим вариантом будет реализация нового бэкэнда для XLA для нужного оборудования. Этот вариант требует наибольших усилий. Классы, которые необходимо реализовать, следующие:

  • StreamExecutor : для многих устройств необходимы не все методы StreamExecutor . Подробности см. в существующих реализациях StreamExecutor .
  • xla::Compiler : этот класс инкапсулирует компиляцию вычислений HLO в xla::Executable .
  • xla::Executable : этот класс используется для запуска скомпилированных вычислений на платформе.
  • xla::TransferManager : этот класс позволяет серверным модулям предоставлять специфичные для платформы механизмы для создания литеральных данных XLA из заданных дескрипторов памяти устройства. Другими словами, он помогает инкапсулировать передачу данных от хоста к устройству и обратно.