Эта страница была переведа с помощью Cloud Translation API.
Switch to English

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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