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

Desenvolvendo um novo backend para XLA

Este guia preliminar é para os adotantes que desejam redirecionar facilmente o TensorFlow para o hardware de maneira eficiente. O guia não é passo a passo e assume o 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 TensorFlow. O redirecionamento do XLA deve ser significativamente mais simples e escalável do que a implementação de todos os TensorFlow Op existentes para o novo hardware.

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

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

Cenário 1: Arquitetura da CPU existente ainda não oficialmente suportada pelo XLA

Nesse cenário, comece examinando 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 backends do XLA para as CPUs é o código gerado pelo LLVM. O Google testa o XLA para arquiteturas x64 e ARM64.

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

Se não houver back-end LLVM, mas existir outro tipo de gerador de código, será possível reutilizar a maior parte do back-end da CPU existente.

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

É possível modelar uma nova implementação do xla::Compiler nas xla::CPUCompiler e xla::GPUCompiler , pois elas já emitem LLVM IR. Dependendo da natureza do hardware, é possível que muitos aspectos da geração de LLVM IR precisem ser alterados, mas muito código pode ser compartilhado com os back-end existentes.

Um bom exemplo a seguir é o back-end da GPU do XLA. O back-end da GPU tem como alvo 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 possui um back-end LLVM upstream), podem reutilizar partes da lógica de emissão IR LLVM, mas outras serão únicas.

Cenário 3: Hardware sem CPU, sem um back-end LLVM existente

Se não for possível utilizar o LLVM, a melhor opção é implementar um novo backend 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. Veja as implementações existentes do StreamExecutor para obter detalhes.
  • xla::Compiler : Esta classe encapsula a compilação de uma computação HLO em um xla::Executable .
  • xla::Executable : Esta classe é usada para iniciar um cálculo compilado na plataforma.
  • xla::TransferManager : essa classe permite que os back-ends forneçam mecanismos específicos da plataforma para a construção de dados literais XLA a partir de 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.