A álgebra linear acelerada (XA, na sigla em inglês) é 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 de maneira colaborativa 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:
Melhorar 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 em pipeline para reduzir a sobrecarga de memória e especializar formas de tensor conhecidas para permitir uma propagação constante mais agressiva.
Melhore o uso da memória. Analisar e programar o uso da memória, eliminando muitos buffers de armazenamento intermediários.
Confiar 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 que foram originalmente fundidas manualmente.
Melhore a portabilidade. Facilitar a programação de um novo back-end para um novo hardware, de modo que uma grande parte dos modelos de ML possa ser executada sem modificações nesse hardware. Isso é em contraste com a abordagem de especializar operações monolíticas individuais para novos hardwares, que exigem que os modelos sejam reescritos para usar essas operações.
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 = operações de alto nível) 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 de otimização e análise integradas 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 outras 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 fazer a correspondência de padrão a determinadas operações ou combinações delas para chamadas de biblioteca otimizadas.
Em seguida, o back-end executa a geração de códigos específicos do destino. Os back-ends de CPU e GPU incluídos no XLA usam o 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 HLO de maneira eficiente e, em seguida, invocam o LLVM para emitir código nativo a partir desse IR LLVM.
Nesse processo, o compilador XLA é modular porque é fácil encaixar um back-end alternativo para direcionar a uma nova arquitetura HW. No momento, o back-end da GPU oferece suporte a GPUs NVIDIA pelo back-end LLVM NVPTX. O back-end da CPU é compatível com várias ISAs de CPU.