Desarrolla un nuevo backend para XLA

Esta guía está destinada a ingenieros de sistemas que desean que XLA genere programas que se orienten a su hardware de manera eficiente. La guía no es paso a paso y se da por sentado que conoces LLVM, Bazel y XLA.

XLA proporciona una interfaz abstracta que una arquitectura o un acelerador nuevos pueden implementar para crear un backend y así ejecutar los resultados de los programas de AA mediante XLA. La resegmentación de XLA debería ser mucho más simple y escalable que implementar cada op existente desde un framework de frontend, como PyTorch o TensorFlow, para hardware nuevo.

La mayoría de las implementaciones se encontrarán en una de las siguientes situaciones:

  1. Arquitectura de CPU existente aún no compatible oficialmente con XLA, con o sin un backend LLVM existente.
  2. Hardware que no es similar a la CPU con un backend de LLVM existente.
  3. Hardware que no es similar a la CPU sin un backend de LLVM existente.

Situación 1: Arquitectura de CPU existente que XLA aún no es compatible oficialmente

En esta situación, comienza por observar el backend de CPU de XLA existente. XLA facilita la orientación a diferentes CPU mediante el uso de LLVM, ya que la diferencia principal entre los backends de XLA para CPU es el código que genera LLVM.

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

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

Situación 2: Hardware que no es similar a la CPU con un backend de LLVM existente

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

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

Situación 3: Hardware que no es similar a la CPU sin un backend de LLVM existente

Si no es posible usar LLVM, la mejor opción es implementar un backend nuevo para XLA en el hardware deseado. Esta opción requiere el mayor esfuerzo. Las clases que se deben implementar son las siguientes:

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