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

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

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

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

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

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

В этом сценарии начните с просмотра существующей серверной части ЦП XLA . XLA упрощает перенацеливание TensorFlow на разные ЦП с помощью LLVM, поскольку основное различие между бэкендами XLA для ЦП заключается в коде, сгенерированном LLVM. Google тестирует XLA для архитектур x64 и ARM64.

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

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

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

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

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

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

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

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