Architecture XLA

XLA (Accelerated Linear Algebra) est un compilateur de machine learning (ML) qui optimise l'algèbre linéaire, et améliore 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 est compatible avec plusieurs interfaces de frameworks de ML (y compris PyTorch, TensorFlow et JAX) et fait partie du projet OpenXLA, un écosystème de technologies de compilation 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é au sein du projet TensorFlow, mais les objectifs fondamentaux restent les mêmes:

  • Améliorer la vitesse d'exécution Compilez des sous-graphes pour réduire la durée d'exécution des opérations de courte durée et éliminer la surcharge de l'environnement d'exécution, fusionner les opérations en pipeline pour réduire la surcharge de la mémoire, et spécialiser les formes de Tensor connues pour permettre une propagation constante plus agressive.

  • Améliorer 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. Supprimez la nécessité de nombreuses opérations personnalisées en améliorant les performances des opérations de bas niveau fusionnées automatiquement pour qu'elles correspondent à celles des opérations personnalisées initialement fusionnées à la main.

  • Améliorez la portabilité. facilitent l'écriture d'un nouveau backend pour un nouveau matériel, de sorte 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 des modèles pour exploiter ces opérations.

Fonctionnement

Le compilateur XLA extrait des graphes de modèle issus de frameworks de ML définis dans StableHLO et les compile en instructions machine pour différentes architectures. StableHLO définit un ensemble d'opérations avec versions gérées (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 du modèle en exécutable optimisé pour la cible comprend les étapes suivantes:

  1. XLA effectue plusieurs passes d'analyse et d'optimisation intégrées au 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 de la mémoire tampon pour l'allocation de mémoire d'exécution pour le calcul. Au cours de cette étape d'optimisation, XLA convertit également le dialecte StableHLO en dialecte HLO interne.

  2. XLA envoie le calcul HLO à un backend afin d'effectuer d'autres optimisations au niveau de l'objet 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 utiles, spécifiquement pour le modèle de programmation GPU, et déterminer comment partitionner le calcul en flux. À ce stade, les backends peuvent également appliquer des schémas à certaines opérations ou combinaisons de celles-ci pour optimiser les appels de bibliothèque.

  3. Le backend génère ensuite du code spécifique à la cible. Les backends de processeur et de GPU inclus dans XLA utilisent LLVM pour l'IA de bas niveau, l'optimisation et la génération de code. 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.

Dans ce processus, le compilateur XLA est modulaire dans le sens où il peut facilement être placé dans un backend alternatif afin de cibler une nouvelle architecture matérielle. Le backend de GPU est actuellement compatible avec les GPU NVIDIA via le backend NVPTX de LLVM. Le backend de processeur accepte plusieurs ISA de processeur.