Este guia preliminar é para os primeiros usuários que desejam redirecionar facilmente o TensorFlow para seu 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. O redirecionamento do XLA deve ser significativamente mais simples e escalável do que implementar cada TensorFlow Op existente para um novo hardware.
A maioria das implementações cairá em um dos seguintes cenários:
- Arquitetura de CPU existente ainda não suportada oficialmente pelo XLA, com ou sem um back-end LLVM existente.
- Hardware não semelhante a CPU com um back-end LLVM existente.
- Hardware não semelhante a CPU sem um back-end LLVM existente.
Cenário 1: arquitetura de CPU existente ainda não suportada oficialmente pelo XLA
Nesse cenário, comece examinando o back- end de 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 XLA para CPUs é o código gerado pelo LLVM. O Google testa XLA para arquiteturas x64 e ARM64.
Se o fornecedor de hardware tiver um back-end LLVM para seu hardware, é simples vincular o back-end ao LLVM criado com XLA. No modo JIT, o backend da CPU XLA emite código para a CPU do host. Para compilação antecipada, xla::AotCompilationOptions pode fornecer um triplo LLVM para configurar a arquitetura de destino.
Se não houver 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 xla::Compiler nas xla::CPUCompiler e xla::GPUCompiler , pois estas já emitem LLVM IR. Dependendo da natureza do hardware, é possível que muitos dos aspectos de geração de IR do LLVM tenham que ser alterados, mas muito código pode ser compartilhado com os backends 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 o Hexagon (que possui um backend upstream LLVM), 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, 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 deStreamExecutorsão necessários. Consulte as implementações existentesStreamExecutorpara obter detalhes. -
xla::Compiler: Esta classe encapsula a compilação de uma computação HLO em umxla::Executable. -
xla::Executable: Esta 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 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.