XLA (Accelerated Linear Algebra) est un compilateur de machine learning (ML) qui optimise l'algèbre linéaire, améliorant ainsi la vitesse d'exécution et l'utilisation de la mémoire. Cette page présente brièvement les objectifs et l'architecture du compilateur XLA.
Objectifs
Aujourd'hui, XLA accepte plusieurs interfaces de frameworks de ML (y compris PyTorch, TensorFlow et JAX) et fait partie du projet OpenXLA, un écosystème de technologies de compilateur Open Source pour le ML développé en collaboration par les principales organisations de matériel et de logiciels de ML. Avant la création du projet OpenXLA, XLA était développé dans le projet TensorFlow, mais les objectifs fondamentaux restent les mêmes:
Améliorer la vitesse d'exécution Compilez des sous-graphes pour réduire le temps d'exécution des opérations de courte durée et éliminer les frais généraux liés à l'environnement d'exécution, fusionner les opérations en pipeline pour réduire la surcharge de mémoire, et spécialiser les formes de Tensor connues pour permettre une propagation constante plus agressive.
Améliorez l'utilisation de la mémoire. Analysez et planifiez l'utilisation de la mémoire, en éliminant de nombreux tampons de stockage intermédiaires.
Réduisez la dépendance aux opérations personnalisées. Éliminez le besoin de nombreuses opérations personnalisées en améliorant les performances des opérations de bas niveau fusionnées automatiquement pour qu'elles correspondent aux performances des opérations personnalisées qui étaient initialement fusionnées manuellement.
Améliorez la portabilité. Facilitez l'écriture d'un nouveau backend pour du matériel novateur, afin qu'une grande partie des modèles de ML puisse s'exécuter sans modification sur ce matériel. Cette approche diffère de l'approche qui consiste à spécialiser des opérations monolithiques individuelles pour le nouveau matériel, qui nécessite de réécrire les modèles pour pouvoir les utiliser.
Comment ça marche ?
Le compilateur XLA extrait des graphes de modèle issus de frameworks de ML définis dans StableHLO et les compile dans des instructions machine pour différentes architectures. StableHLO définit un ensemble d'opérations avec gestion des versions (HLO = opérations de haut niveau) qui fournit une couche de portabilité entre les frameworks de ML et le compilateur.
En général, le processus de compilation qui convertit le graphe de modèle en un exécutable optimisé pour la cible comprend les étapes suivantes:
XLA effectue plusieurs passes d'analyse et d'optimisation intégrées sur le graphe StableHLO, qui sont indépendantes de la cible, telles que le CSE, la fusion d'opérations indépendantes de la cible et l'analyse du tampon pour allouer de la mémoire d'exécution pour le calcul. Au cours de cette étape d'optimisation, XLA convertit également le dialecte StableHLO en un dialecte HLO interne.
XLA envoie le calcul HLO à un backend pour de nouvelles optimisations au niveau de l'objectif HLO, cette fois en tenant compte des informations et des besoins spécifiques à la cible. Par exemple, le backend GPU peut effectuer des fusions d'opérations qui sont bénéfiques spécifiquement pour le modèle de programmation GPU et qui déterminent comment partitionner le calcul en flux. À ce stade, les backends peuvent également mettre en correspondance certaines opérations ou combinaisons de celles-ci avec des appels de bibliothèque optimisés.
Le backend génère ensuite le code spécifique à la cible. Les backends de processeur et de GPU inclus dans XLA utilisent LLVM pour la reconnaissance d'images, l'optimisation et la génération de code de bas niveau. Ces backends émettent l'IR LLVM nécessaire pour représenter efficacement le calcul HLO, puis appellent LLVM pour émettre du code natif à partir de cet IR LLVM.
Au cours de ce processus, le compilateur XLA est modulaire, dans le sens où il est facile de l'insérer dans un autre backend afin de cibler une nouvelle architecture matérielle. Le backend de GPU est actuellement compatible avec les GPU NVIDIA via le backend LLVM NVPTX. Le backend de processeur accepte plusieurs ISA à l'aide de processeurs.