Desarrolla un nuevo backend para XLA

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

XLA proporciona una interfaz abstracta que una arquitectura o un acelerador nuevos pueden implementar a fin de crear un backend para ejecutar los programas de AA salientes de XLA. La resegmentación de XLA debería ser mucho más simple y más escalable que implementar todas las operaciones existentes de un framework de frontend, como PyTorch o TensorFlow, para el hardware nuevo.

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

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

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

En esta situación, comienza por observar el backend de CPU de XLA existente. XLA facilita la segmentación a diferentes CPU mediante 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 de tipo 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 cambiar muchos aspectos de la generación de IR de LLVM, pero se puede compartir una gran cantidad de código con los backends existentes.

Un buen ejemplo a continuación es el backend de GPU de XLA. El backend de la GPU se orienta a un ISA que no es del tipo 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., 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 similar a una 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 existentes de StreamExecutor 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 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 de XLA a partir de controladores de memoria del dispositivo determinados. En otras palabras, ayuda a encapsular la transferencia de datos del host al dispositivo y viceversa.