Architecture XLA

Pourquoi avons-nous construit XLA ?

Nous avions plusieurs objectifs pour que XLA travaille avec TensorFlow :

  • 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 afin d'éliminer la surcharge de l'environnement d'exécution TensorFlow, fusionnez les opérations en pipeline pour réduire la surcharge de mémoire et spécialisez-vous dans des formes de tenseurs 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 en principe de nombreux tampons de stockage intermédiaires.

  • Réduisez la dépendance aux opérations personnalisées. Supprimez le besoin de nombreuses opérations personnalisées en améliorant les performances des opérations de bas niveau fusionnées automatiquement pour correspondre aux performances des opérations personnalisées fusionnées à la main.

  • Réduisez l'empreinte mobile. Éliminez l'environnement d'exécution TensorFlow en compilant à l'avance le sous-graphe et en émettant une paire de fichiers objet/en-tête pouvant être directement liée à une autre application. Les résultats peuvent réduire l'empreinte pour l'inférence mobile de plusieurs ordres de grandeur.

  • Améliorez la portabilité. Facilitez l'écriture d'un nouveau backend pour un nouveau matériel, auquel cas une grande partie des programmes TensorFlow s'exécutera sans modification sur ce matériel. Cela contraste avec l'approche consistant à spécialiser les opérations monolithiques individuelles pour le nouveau matériel, qui nécessite la réécriture des programmes TensorFlow pour utiliser ces opérations.

Comment fonctionne XLA ?

La langue d'entrée de XLA est appelée "HLO IR", ou simplement HLO (High Level Operations). La sémantique de HLO sont décrites sur la Sémantique Opération page. Il est plus commode de penser à HLO comme IR du compilateur .

XLA prend des graphes ("calculs") définis dans HLO et les compile en instructions machine pour diverses architectures. XLA est modulaire dans le sens où il est facile de fente dans un back - end alternative à cibler une nouvelle architecture de HW . Le backend CPU pour x64 et ARM64 ainsi que le backend GPU NVIDIA se trouvent dans l'arborescence des sources TensorFlow.

Le schéma suivant montre le processus de compilation dans XLA :

XLA est livré avec plusieurs optimisations et passes d'analyse qui sont indépendants de la cible, tels que le CST , la fusion de fonctionnement cible indépendante, et l' analyse du tampon d'allocation de mémoire d'exécution pour le calcul.

Après l'étape indépendante de la cible, XLA envoie le calcul HLO à un backend. Le backend peut effectuer d'autres optimisations au niveau HLO, cette fois en tenant compte des informations et des besoins spécifiques de la cible. Par exemple, le backend GPU XLA peut effectuer une fusion d'opérations bénéfique spécifiquement pour le modèle de programmation GPU et déterminer comment partitionner le calcul en flux. À ce stade, les backends peuvent également faire correspondre certaines opérations ou combinaisons de celles-ci avec des appels de bibliothèque optimisés.

L'étape suivante est la génération de code spécifique à la cible. Les CPU et GPU backends inclus avec XLA utilisation LLVM pour IR à faible niveau, l' optimisation et génération de code. Ces backends émettent l'IR LLVM nécessaire pour représenter le calcul XLA HLO de manière efficace, puis invoquent LLVM pour émettre du code natif à partir de cet IR LLVM.

Le backend GPU prend actuellement en charge les GPU NVIDIA via le backend LLVM NVPTX ; le backend CPU prend en charge plusieurs ISA CPU.