Arquitetura XLA

Por que construímos 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 especializar-se em formas de tensor 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 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 de HLO são descritos na Semântica Operação página. É mais conveniente para pensar em HLO como IR compilador .

O XLA pega gráficos ("cálculos") definidos no HLO e os compila em instruções de máquina para várias arquiteturas. XLA é modular no sentido de que é fácil de fenda em uma infra-estrutura alternativa para combater alguns romance arquitetura 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:

XLA vem com várias optimizações e passagens de análise que são independentes-alvo, tais como CSE , fusão operação independente do alvo, e o tampão de análise para a atribuição de memória de tempo de execução para a computação.

Após a etapa independente do destino, o XLA envia o cálculo 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. Neste 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. As CPU e GPU backends incluído com XLA uso 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.