Arquitetura XLA

Por que criamos o XLA?

Tínhamos vários objetivos para o XLA funcionar com o TensorFlow:

  • Melhore a velocidade de execução. Compile subgráficos para reduzir o tempo de execução de operações de curta duração para eliminar a sobrecarga do tempo de execução do TensorFlow, fundir operações em pipeline para reduzir a sobrecarga de memória e se especializar em formas tensoras conhecidas para permitir uma propagação constante mais agressiva.

  • Melhore o uso de memória. Analise e programe o uso da memória, em princípio eliminando muitos buffers de armazenamento intermediários.

  • Reduza a dependência de operações personalizadas. Elimine a necessidade de muitas operações personalizadas, melhorando o desempenho das Ops de baixo nível fundidas automaticamente para corresponder ao desempenho das Ops personalizadas que foram fundidas manualmente.

  • Reduza a pegada móvel. Elimine o tempo de execução do TensorFlow compilando antecipadamente o subgráfico e emitindo um par de arquivo de objeto / cabeçalho que pode ser vinculado diretamente a outro aplicativo. Os resultados podem reduzir a pegada para inferência móvel em várias ordens de magnitude.

  • Melhore a portabilidade. Torne relativamente fácil escrever um novo back-end para um novo hardware, ponto em que uma grande fração dos programas TensorFlow será executada sem modificações nesse hardware. Isso contrasta com a abordagem de especializar operações monolíticas individuais para novo hardware, que exige que os programas do TensorFlow sejam reescritos para fazer uso dessas operações.

Como funciona o XLA?

O idioma de entrada para o XLA é denominado "HLO IR" ou apenas HLO (Operações de alto nível). A semântica do HLO é descrita na página Operation Semantics . É mais conveniente pensar no HLO como um IR do compilador .

O XLA pega gráficos ("cálculos") definidos no HLO e os compila em instruções de máquina para várias arquiteturas. O XLA é modular no sentido de que é fácil encaixar em um back-end alternativo para atingir alguma nova arquitetura de HW . O back-end da CPU para x64 e ARM64, bem como o back-end da GPU NVIDIA, estão na árvore de origem do TensorFlow.

O diagrama a seguir mostra o processo de compilação no XLA:

O XLA vem com várias otimizações e passos de análise que são independentes do destino, como CSE , fusão de operação independente do destino e análise de buffer para alocar memória de tempo de execução para a computação.

Após a etapa independente do destino, o XLA envia a computação HLO para um back-end. O back-end pode realizar mais otimizações de nível HLO, desta vez com informações e necessidades específicas do destino em mente. Por exemplo, o back-end da GPU XLA pode realizar fusão de operação benéfica especificamente para o modelo de programação da GPU e determinar como particionar a computação em fluxos. Nesse estágio, os back-ends também podem corresponder a certas operações ou combinações de padrões para chamadas de biblioteca otimizadas.

A próxima etapa é a geração de código específico para o destino. Os back-ends de CPU e GPU incluídos no XLA usam LLVM para IR de baixo nível, otimização e geração de código. Esses back-ends emitem o LLVM IR necessário para representar a computação XLA HLO de maneira eficiente e, em seguida, invocam o LLVM para emitir o código nativo desse LLVM IR.

O backend de GPU atualmente suporta GPUs NVIDIA por meio do backend LLVM NVPTX; o backend da CPU oferece suporte a vários ISAs de CPU.