Aide à protéger la Grande barrière de corail avec tensorflow sur Kaggle Rejoignez Défi

Architecture XLA

Pourquoi avons-nous construit XLA?

Nous avions plusieurs objectifs pour que XLA fonctionne avec TensorFlow:

  • Améliorez la vitesse d'exécution. Compilez des sous-graphiques pour réduire le temps d'exécution des opérations de courte durée afin d'éliminer la surcharge de l'exécution de TensorFlow, fusionnez les opérations en pipeline pour réduire la charge mémoire et spécialisez-vous sur les formes de tenseurs 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 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 manuellement.

  • Réduisez l'empreinte mobile. Éliminez le runtime TensorFlow en compilant à l'avance le sous-graphique et en émettant une paire objet / fichier d'en-tête qui peut être liée directement dans une autre application. Les résultats peuvent réduire l'empreinte de l'inférence mobile de plusieurs ordres de grandeur.

  • Améliorez la portabilité. Rendez relativement facile 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 de spécialisation des 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?

Le langage d'entrée de XLA est appelé "HLO IR", ou simplement HLO (opérations de haut niveau). La sémantique de HLO est décrite sur la page Opération sémantique . Il est plus pratique de considérer HLO comme un compilateur IR .

XLA prend des graphes ("calculs") définis dans HLO et les compile en instructions machine pour différentes architectures. XLA est modulaire en ce sens qu'il est facile de s'insérer dans un backend alternatif pour cibler une nouvelle architecture matérielle . Le backend CPU pour x64 et ARM64 ainsi que le backend GPU NVIDIA se trouvent dans l'arborescence des sources de TensorFlow.

Le diagramme suivant montre le processus de compilation dans XLA:

XLA est livré avec plusieurs optimisations et passes d'analyse indépendantes de la cible, telles que CSE , la fusion d'opérations indépendante de la cible et l'analyse de la mémoire tampon pour l'allocation de la 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 gardant à l'esprit les informations et les besoins spécifiques à 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 mettre en correspondance certaines opérations ou combinaisons de celles-ci avec des appels de bibliothèque optimisés.

La prochaine étape est la génération de code spécifique à la cible. Les backends CPU et GPU inclus avec XLA utilisent LLVM pour IR de bas niveau, l'optimisation et la 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.