Esta página foi traduzida pela API Cloud Translation.
Switch to English

Desenvolvendo um novo back-end para XLA

Este guia preliminar é para os primeiros usuários que desejam redirecionar facilmente o TensorFlow para o hardware de maneira eficiente. O guia não é passo a passo e pressupõe conhecimento de LLVM , Bazel e TensorFlow.

O XLA fornece uma interface abstrata que uma nova arquitetura ou acelerador pode implementar para criar um back-end para executar gráficos do TensorFlow. Retargeting XLA deve ser significativamente mais simples e escalonável do que implementar cada TensorFlow Op existente para novo hardware.

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

  1. A arquitetura de CPU existente ainda não é oficialmente suportada pelo XLA, com ou sem um back-end LLVM existente.
  2. Hardware diferente de CPU com um back-end LLVM existente.
  3. Hardware diferente de CPU sem um back-end LLVM existente.

Cenário 1: arquitetura de CPU existente ainda não oficialmente suportada pelo XLA

Nesse cenário, comece observando o back-end da CPU XLA existente. O XLA facilita o redirecionamento do TensorFlow para diferentes CPUs usando o LLVM, pois a principal diferença entre os back-ends do XLA para CPUs é o código gerado pelo LLVM. O Google testa o XLA para arquiteturas x64 e ARM64.

Se o fornecedor do hardware possui um backend LLVM para seu hardware, é simples vincular o backend ao LLVM construído com o XLA. No modo JIT, o backend da CPU XLA emite código para a CPU host. Para compilação xla::AotCompilationOptions , xla::AotCompilationOptions pode fornecer um triplo LLVM para configurar a arquitetura de destino.

Se não houver um backend LLVM existente, mas existir outro tipo de gerador de código, deve ser possível reutilizar a maior parte do backend de CPU existente.

Cenário 2: Hardware não semelhante a CPU com um back-end LLVM existente

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

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

Cenário 3: Hardware não semelhante a CPU sem um back-end LLVM existente

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

  • StreamExecutor : Para muitos dispositivos, nem todos os métodos do StreamExecutor são necessários. Consulte as implementações existentes do StreamExecutor para obter detalhes.
  • xla::Compiler : Esta classe encapsula a compilação de um cálculo HLO em um xla::Executable .
  • xla::Executable : esta classe é usada para lançar uma computação compilada na plataforma.
  • xla::TransferManager : esta classe permite que os back-ends forneçam mecanismos específicos da plataforma para construir dados literais XLA de determinados identificadores de memória do dispositivo. Em outras palavras, ajuda a encapsular a transferência de dados do host para o dispositivo e vice-versa.