Como desenvolver um novo back-end para XLA

Este guia é destinado a engenheiros de sistema que querem que o XLA gere programas voltados ao hardware de maneira eficiente. O guia não é passo a passo e requer conhecimento sobre LLVM, Bazel e XLA.

O XLA fornece uma interface abstrata que uma nova arquitetura ou acelerador pode implementar para criar um back-end e executar a saída de programas de ML pelo XLA. A nova segmentação do XLA precisa ser significativamente mais simples e mais escalonável do que a implementação de todas as operações existentes de um framework de front-end, como PyTorch ou TensorFlow, para um novo hardware.

A maioria das implementações se enquadra em um dos seguintes cenários:

  1. A arquitetura de CPU atual ainda não é oficialmente compatível com o XLA, com ou sem um back-end do LLVM.
  2. Hardware que não é semelhante à CPU com um back-end LLVM.
  3. Hardware que não é semelhante à CPU sem um back-end LLVM.

Cenário 1: a arquitetura de CPU existente ainda não é oficialmente compatível com XLA

Neste cenário, comece analisando o back-end atual da CPU do XLA. O XLA facilita o direcionamento de CPUs diferentes usando o LLVM, porque a principal diferença entre back-ends do XLA para CPUs é o código gerado pelo LLVM.

Se o fornecedor de hardware tiver um back-end LLVM para o hardware, é simples vincular o back-end ao LLVM criado com o XLA. No modo JIT, o back-end da CPU do XLA emite código para a CPU do host. Para compilação antecipada, xla::AotCompilationOptions pode fornecer um LLVM triplo para configurar a arquitetura de destino.

Se não houver um back-end LLVM, mas outro tipo de gerador de código existir, talvez seja possível reutilizar a maior parte do back-end de CPU atual.

Cenário 2: hardware não semelhante à CPU com um back-end LLVM

É possível modelar uma nova implementação xla::Compiler nas classes xla::CPUCompiler e xla::GPUCompiler existentes, já que elas já emitem IR LLVM. Dependendo da natureza do hardware, é possível que muitos aspectos da geração de IR do LLVM tenham que ser modificados, mas muitos códigos podem ser compartilhados com os back-ends existentes.

Um bom exemplo a seguir é o back-end de GPU do XLA. O back-end da GPU tem como alvo uma ISA que não é semelhante à CPU e, portanto, alguns aspectos da geração de código são exclusivos do domínio da GPU. Outros tipos de hardware, como DSPs como o Hexagon (que tem um back-end LLVM upstream), podem reutilizar partes da lógica de emissão de IR do LLVM, mas outras partes serão exclusivas.

Cenário 3: hardware semelhante à CPU sem um back-end LLVM

Se não for possível usar o LLVM, a melhor opção é implementar um novo back-end para XLA para o hardware desejado. Essa opção exige mais esforço. As classes que precisam ser implementadas são as seguintes:

  • StreamExecutor: no caso de muitos dispositivos, nem todos os métodos de StreamExecutor são necessários. Consulte as implementações atuais de StreamExecutor para mais detalhes.
  • xla::Compiler: essa classe encapsula a compilação de um cálculo de HLO em um xla::Executable.
  • xla::Executable: essa classe é usada para iniciar uma computação compilada na plataforma.
  • xla::TransferManager: essa classe permite que os back-ends forneçam mecanismos específicos da plataforma para criar dados literais do XLA usando determinados identificadores de memória do dispositivo. Em outras palavras, ele ajuda a encapsular a transferência de dados do host para o dispositivo e de volta.