A álgebra linear acelerada (XLA) é um compilador de machine learning (ML) que otimiza a álgebra linear, oferecendo melhorias na velocidade de execução e no uso de memória. Nesta página, você encontra uma breve visão geral dos objetivos e da arquitetura do compilador XLA.
Objetivos
Atualmente, o XLA oferece suporte a vários front-ends de framework de ML (incluindo PyTorch, TensorFlow e JAX) e faz parte do projeto OpenXLA, um ecossistema de tecnologias de compilador de código aberto para ML desenvolvido em colaboração pelas principais organizações de hardware e software de ML. Antes da criação do projeto OpenXLA, o XLA foi desenvolvido dentro do projeto do TensorFlow, mas os objetivos fundamentais permanecem os mesmos:
Melhore a velocidade de execução. Compile subgráficos para reduzir o tempo de execução de operações de curta duração e eliminar a sobrecarga do ambiente de execução, fundir operações com pipeline para reduzir a sobrecarga de memória e especializar formas de tensor conhecidas para permitir uma propagação constante mais agressiva.
Melhorar o uso da memória. Analisar e programar o uso da memória, eliminando muitos buffers de armazenamento intermediários.
Confie menos em operações personalizadas. Elimine a necessidade de muitas operações personalizadas, melhorando o desempenho das operações de baixo nível fundidas automaticamente para corresponder ao desempenho das operações personalizadas originalmente fundidas manualmente.
Melhore a portabilidade. Facilitar a criação de um novo back-end para hardware novo, de modo que uma grande fração dos modelos de ML possa ser executada sem modificações nesse hardware. Isso é diferente da abordagem de especialização de operações monolíticas individuais para novo hardware, que exige a reescrita de modelos para que as operações sejam usadas.
Como funciona
O compilador XLA usa gráficos de modelo de frameworks de ML definidos em StableHLO e os compila em instruções de máquina para várias arquiteturas. StableHLO define um conjunto de operações com controle de versões (HLO, na sigla em inglês) que fornece uma camada de portabilidade entre os frameworks de ML e o compilador.
Em geral, o processo de compilação que converte o grafo do modelo em um executável otimizado para destino inclui estas etapas:
O XLA executa várias passagens integradas de otimização e análise no gráfico StableHLO que são independentes de destino, como CSE, fusão de operações independentes de destino e análise de buffer para alocar memória de tempo de execução para a computação. Durante esse estágio de otimização, o XLA também converte o dialeto StableHLO em um dialeto HLO interno.
O XLA envia a computação do HLO a um back-end para mais otimizações no nível do HLO, desta vez com informações e necessidades específicas do destino em mente. Por exemplo, o back-end da GPU pode realizar fusões de operações que são benéficas especificamente para o modelo de programação da GPU e determinar como particionar a computação em streams. Nessa fase, os back-ends também podem corresponder a certas operações ou combinações delas com as chamadas de biblioteca otimizadas.
Em seguida, o back-end executa a geração de código específico do destino. Os back-ends de CPU e GPU incluídos no XLA usam LLVM para IR, otimização e geração de código de baixo nível. Esses back-ends emitem o IR do LLVM necessário para representar a computação do HLO de maneira eficiente e, em seguida, invocam o LLVM para emitir código nativo a partir desse IR do LLVM.
Nesse processo, o compilador XLA é modular porque é fácil encaixar um back-end alternativo para direcionar a uma nova arquitetura HW. Atualmente, o back-end da GPU oferece suporte a GPUs NVIDIA pelo back-end LLVM NVPTX. O back-end da CPU oferece suporte a vários ISAs de CPU.