Se usó la API de Cloud Translation para traducir esta página.
Switch to English

Desarrollando un nuevo backend para XLA

Esta guía preliminar es para los primeros usuarios que desean reajustar fácilmente TensorFlow a su hardware de manera eficiente. La guía no es paso a paso y supone el conocimiento de LLVM , Bazel y TensorFlow.

XLA proporciona una interfaz abstracta que una nueva arquitectura o acelerador puede implementar para crear un backend para ejecutar gráficos TensorFlow. Retargeting XLA debería ser significativamente más simple y escalable que implementar cada TensorFlow Op existente para hardware nuevo.

La mayoría de las implementaciones caerán en uno de los siguientes escenarios:

  1. La arquitectura de CPU existente aún no es oficialmente compatible con XLA, con o sin un back-end LLVM existente.
  2. Hardware no similar a la CPU con un back-end LLVM existente.
  3. Hardware no similar a la CPU sin un back-end LLVM existente.

Escenario 1: arquitectura de CPU existente aún no admitida oficialmente por XLA

En este escenario, comience mirando el backend de CPU XLA existente. XLA facilita el reajuste de TensorFlow a diferentes CPU mediante el uso de LLVM, ya que la principal diferencia entre los backends XLA para CPU es el código generado por LLVM. Google prueba XLA para arquitecturas x64 y ARM64.

Si el proveedor de hardware tiene un backend LLVM para su hardware, es simple vincular el backend con el LLVM construido con XLA. En modo JIT, el backend de la CPU XLA emite código para la CPU del host. Para una compilación xla::AotCompilationOptions , xla::AotCompilationOptions puede proporcionar un triple LLVM para configurar la arquitectura de destino.

Si no hay un back-end LLVM existente pero existe otro tipo de generador de código, debería ser posible reutilizar la mayor parte del back-end de la CPU existente.

Escenario 2: hardware no similar a la CPU con un back-end LLVM existente

Es posible modelar una nueva implementación de xla::Compiler en las clases xla::CPUCompiler y xla::GPUCompiler , ya que estas ya emiten LLVM IR. Dependiendo de la naturaleza del hardware, es posible que muchos de los aspectos de generación de IR LLVM tengan que cambiarse, pero se puede compartir una gran cantidad de código con los backends existentes.

Un buen ejemplo a seguir es el backend GPU de XLA. El backend de GPU apunta a un ISA que no es similar a la CPU, y por lo tanto, algunos aspectos de su generación de código son exclusivos del dominio de GPU. Otros tipos de hardware, por ejemplo, DSP como Hexagon (que tiene un back-end LLVM aguas arriba), pueden reutilizar partes de la lógica de emisión IR LLVM, pero otras partes serán únicas.

Escenario 3: hardware no similar a la CPU sin un back-end LLVM existente

Si no es posible utilizar LLVM, la mejor opción es implementar un nuevo back-end para XLA para el hardware deseado. Esta opción requiere el mayor esfuerzo. Las clases que deben implementarse son las siguientes:

  • StreamExecutor : para muchos dispositivos no se necesitan todos los métodos de StreamExecutor . Consulte las implementaciones de StreamExecutor existentes para obtener más detalles.
  • xla::Compiler : esta clase encapsula la compilación de un cálculo HLO en un xla::Executable .
  • xla::Executable : esta clase se utiliza para iniciar un cálculo compilado en la plataforma.
  • xla::TransferManager : esta clase permite que los backends proporcionen mecanismos específicos de la plataforma para construir datos literales XLA a partir de identificadores de memoria de dispositivo dados. En otras palabras, ayuda a encapsular la transferencia de datos desde el host al dispositivo y viceversa.