Este guia é destinado a engenheiros de sistema que querem que o XLA gere programas destinados ao hardware de forma eficiente. O guia não é detalhado e pressupõe conhecimento de 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 de XLA precisa ser significativamente mais simples e 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:
- Arquitetura de CPU existente que ainda não é oficialmente aceita pelo XLA, com ou sem um back-end do LLVM.
- Hardware que não é semelhante à CPU com um back-end LLVM já existente.
- Hardware semelhante a uma CPU sem um back-end LLVM já existente.
Cenário 1: a arquitetura de CPU existente ainda não é oficialmente compatível com XLA
Neste cenário, comece analisando o back-end da CPU do XLA. O XLA facilita o direcionamento a diferentes CPUs 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 de LLVM para o hardware, será simples
vincular o back-end ao LLVM criado com XLA. No modo JIT, o back-end da CPU do XLA emite código para a CPU do host. Para compilação antecipada,
o xla::AotCompilationOptions
pode fornecer um triplo LLVM para configurar a arquitetura de destino.
Se não houver um back-end do LLVM, mas houver outro tipo de gerador de código, será possível reutilizar a maior parte do back-end de CPU existente.
Cenário 2: hardware que não é semelhante à CPU com um back-end LLVM já existente
É possível modelar uma nova implementação
xla::Compiler
nas classes
xla::CPUCompiler
e
xla::GPUCompiler, já que elas já emitem IR LLVM. Dependendo da natureza do
hardware, é possível que muitos aspectos da geração de IR do LLVM precisem
ser alterados, mas muito código pode ser compartilhado com os back-ends atuais.
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, por exemplo, 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 que não é semelhante à CPU e 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 em questão. Essa opção exige o maior 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 deStreamExecutoratuais para mais detalhes.xla::Compiler: essa classe encapsula a compilação de um cálculo de HLO em umxla::Executable.xla::Executable: essa 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 construir dados literais do XLA a partir de 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.