Cette page a été traduite par l'API Cloud Translation.
Switch to English

MLIR CodeGen pour XLA

XLA fonctionne sur HloInstruction et effectue de nombreuses optimisations sur cette représentation, en partageant beaucoup entre les appareils ciblés. À un moment donné, une planification linéaire est calculée et la mémoire tampon est affectée à chaque valeur de manière statique. Le codegen spécifique à l'appareil fonctionne en parcourant cette séquence et en appelant des «émetteurs» pour générer une représentation adaptée à l'appareil (par exemple une seule fonction LLVM par calcul XLA sur le CPU, ou une séquence de «thunks» encapsulant des opérations GPU et éventuellement généré PTX lorsque ciblage GPU).

En tant qu'étape de préparation, nous sommes actuellement en train d'intercepter le processus juste après que XLA ait terminé la phase d'attribution de tampon et d'émettre à la place un module MLIR dans le dialecte lhlo . À partir de là, nous effectuons le codegen à l'aide de composants MLIR (dialecte Linalg, affine et GPU principalement) en fonction de l'appareil.

Vous trouverez ci-dessous le plan d'enregistrement pour migrer de manière incrémentielle XLA / GPU en utilisant lhlo comme entrée de codegen.

Tâches

Hôte Dispositif
Format d'entrée HloInstruction * (Tâche 1) HloInstruction * (Tâche 1)
Format de sortie xla :: Thunk (tâche 2) LLVM IR (tâche 3)
  • La tâche 1 change le format d'entrée de l'hôte et du périphérique de HloInstruction * à LHLO.
  • La tâche 2 change le format de sortie de l'hôte de thunks en «une aire d'atterrissage pour l'hôte» (voir ci-dessous).
  • La tâche 3 migre la sortie de périphérique de LLVM IR vers une forme de MLIR. Il est facultatif pour ce projet, et voir la section "Migration de périphérique LLVM IR" pour plus de détails.

Ce projet donne la priorité à des modèles exécutables de bout en bout avec les émetteurs LHLO activés autant que possible. Cela implique que l'ordre suivant liste des objectifs par priorité:

  • Rendre XLA / GPU exécutable avec les émetteurs LHLO, avec les Thunks et les émetteurs existants non modifiés.
  • Éliminez les références à HloInstruction * dans LHLO, au cas par cas:
    • Commutez un émetteur hérité vers un émetteur basé sur MLIR (par exemple Linalg), ou
    • Traduisez mécaniquement l'émetteur existant pour prendre une représentation MLIR (migrez vers Standard avec GPU Dialect).

Migration de Thunks (tâche 2)

xla :: gpu :: Thunk est une structure de données qui:

  • Peut être appelé depuis l'hôte (xla :: gpu :: Thunk :: ExecuteOnStream ()).
  • Transporte diverses données dans ses sous-classes.
  • Interagit avec BufferAllocation :: Slice et StreamExecutor.
  • Lance les noyaux
  • Appels dans toutes les bibliothèques d'exécution.

Le coût de cela comprend:

  • Représentant les données de configuration spécifiques à l'opération (par exemple, les configurations de convolution).
  • Migration de formes d'opérations et d'opérandes.
  • Représentant un arbre de thunks (while, condition, etc.).

Le travail de migration est indépendant de la migration LHLO / émetteur. Avec des ressources limitées, il est priorisé derrière la migration LHLO / émetteur.

Nous avons plusieurs choix pour abaisser la partie côté hôte de LHLO:

  • TFRT
    • (Pro) excellents emballages CUDA et HIP à utiliser.
    • (Pro) facile à implémenter les appels de bibliothèque (cuDNN, cuBLAS, cuFFT, etc.), car les opérations TFRT sont interprétées par du code C ++.
    • Le côté hôte (Con) est en cours de développement et n'est pas testé.
  • Code CPU jitted
    • (Pro) grande capacité inférieure. Créez quelques boucles et conditions et c'est fait.
    • (Con) GPUDialect ne modélise pas encore les chaînes / flux / asynchronisme / allocation d'appareils.
    • (Con) La prise en charge du runtime CUDA / HIP est minimale (chemin du toolkit, version, chargement dynamique, etc.).
  • Runtime XLA existant (interprétation)

Décision: adopter TFRT, mais également prendre en charge le code CPU jitting dans TFRT.

Migration du périphérique LLVM IR (tâche 3)

Un émetteur élémentaire génère une opération cible en la remplissant élément par élément. Chaque élément de sortie dépend d'un ensemble d'éléments des opérandes. Tous les éléments sont décrits en combinant le tampon avec des indices dynamiques. Il suffit de décrire presque toutes les opérations "mathématiques", mais pour des raisons de performances, seul un grand sous-ensemble d'opérations "mathématiques" est implémenté directement dans (Cpu | Gpu) ElementalIrEmitter.

ElementalIrEmitter est unique en ce sens que:

  • Une grande partie du code est partagée entre XLA / GPU et CPU.
  • Il représente une grande partie des opérations vues dans les modèles, y compris toutes les opérations par élément.
  • La plupart des fusions dépendent uniquement d'ElementalIrEmitter.
  • C'est structurellement simple, car il décrit un DAG de dépendance de données entre des éléments op et des éléments opérandes.
  • Il est principalement portable et de haut niveau (par exemple contrairement au GPU kReduce et au GPU kCopy).
  • La prise en charge de la forme dynamique est facile pour au moins les opérations élémentaires.

Maintenant, pour toutes les opérations, émises de manière élémentaire ou non, il existe plusieurs saveurs de l'état final de chaque opération XLA:

  1. Le code de périphérique reste comme LLVM IR.
  2. Refactoriser l'ancien émetteur pour qu'il ressemble à LHLO -> MLIR LLVM Dialect:
    • (Coût) Ce sera un travail jetable si nous voulons finalement migrer vers Standard.
    • (Avantage) C'est facile et mécanique. Peut être fait en peu de temps.
    • (Avantage) Cela ne profite pas plus que (1).
  3. Refactoriser les anciens émetteurs pour qu'ils ressemblent à LHLO -> GPU MLIR + Standard + Boucles:
    • (Coût) L'élévation des émetteurs existants vers la norme présente certains défis. Les pointeurs et les GEP doivent être convertis en MemRefs et SubViews. Assurer l'exhaustivité de amdgpu en est une autre.
    • (Coût) XLA / GPU repose fortement sur les métadonnées LLVM:
      • range pour les index de bloc / thread.
      • align , dereferenceable , invariant.load , alias.scope , noalias pour load / stores.
      • llvm.loop.unroll.disable , llvm.loop.unroll.full , llvm.loop.vectorize.enable pour les boucles séquentielles.
    • (Avantage) Peut être à long terme. Plus portable.
  4. Refactoriser les anciens émetteurs en LHLO -> Linalg, et écrire de nouveaux émetteurs Linalg
    • (Coût) C'est au cas par cas. Par rapport aux options précédentes, une nouvelle implémentation qui correspond aux performances de XLA doit passer par le benchmark <-> optimiser le flux de travail, ce qui peut représenter un coût important pour certaines opérations.
    • (Avantage) pile unifiée; soutien communautaire; portabilité; plus de potentiels d'optimisation.

Conclusions:

  • N'allez pas pour (2). (1) ou (3) sont juste mieux que (2). (2) coûte plus cher que (1), car il nécessite beaucoup de refactoring mécanique. Avec (1), nous pouvons toujours atteindre l'objectif de permettre à XLA de capter les émetteurs MLIR. C'est en faisant LHLO -> LLVM IR -> exécuter les émetteurs de périphériques hérités.
  • Les opérations ElementalIrEmitter optent pour (4), mais pas de manière incrémentielle. Il n'y a aucun moyen de le faire op par op, car toutes les opérations émises de manière élémentaire sont connectées dans le même graphe. Ce travail peut également servir de point d'unification de plusieurs forces en cours (xla / service / mlir_gpu, le générateur du noyau, Linalg).
  • Toutes les autres opérations vont pour (1). En tant qu'objectif étendu, ils peuvent être migrés vers (3) ou (4).

Priorisation

Bien que les trois tâches mentionnées ci-dessus soient parallélisables, avec des ressources limitées, elles doivent être sérialisées. La hiérarchisation se concentre sur les résultats visibles pour l'achèvement de chaque tâche.

La hiérarchisation est: Tâche1 (LHLO pour les émetteurs hérités)> Tâche 2 (Thunks)> Tâche 3 (émetteurs MLIR).

À la fin de la tâche 1, les utilisateurs de XLA peuvent générer un LHLO (par exemple un générateur de noyau) et les exécuter. Le format de compilation ne sera pas sérialisable MLIR.

À la fin de la tâche 2, LHLO passe à un MLIR sérialisable approprié. Cela permet la compilation hors ligne.

À la fin de la tâche 3, tous les émetteurs XLA sont basés sur le MLIR dans sa mise en œuvre.

Conception détaillée

Étape 1: (Tâche 1) Terminez LHLO et faites en sorte que les émetteurs hérités prennent LHLO

Cette étape permet à tous les émetteurs XLA / GPU existants d'interagir avec les opérations MLIR. Cette étape est un pur refactoring et NFC.

Cette étape est principalement mécanique, mais il convient de noter les écarts suivants entre un HloComputation non incrusté et un LHLO:

  • Chaque HloInstruction a un accès direct à ses opérandes (un DAG de flux de données). Au contraire, chaque opération LHLO n'a accès qu'à ses tampons d'opérande (un bipartite entre les opérations et les tampons). Les opérations LHLO doivent passer par des chaînes use-def pour accéder à leurs opérations d'opérande.
  • Les émetteurs hérités non imbriqués n'accèdent presque jamais empiriquement à leurs opérandes. La seule exception est kReduce.
  • Les émetteurs hérités non imbriqués accèdent à BufferAssignment uniquement pour obtenir des tranches, pas pour accéder aux structures de données auxiliaires comme dataflow_analysis () ou alias_analysis (). llvm_ir construit son propre alias_analysis () basé sur les informations de tranche.

La conclusion est que LHLO devrait s'intégrer sans problème majeur.

Étape 2: (facultatif) prise en charge du profilage

Cette étape n'est nécessaire que si nous commençons à supprimer une partie de la logique XLA Thunk (voir l'étape suivante).

Avant d'activer réellement des émetteurs basés sur MLIR, nous avons besoin d'un profilage pour les émetteurs basés sur MLIR.

Actuellement, XLA effectue son propre profilage en appelant le minuteur StreamExecutor. La minuterie sous le capot insère deux événements avant et après un lancement du noyau et mesure le temps de synchronisation entre ces deux événements.

Il existe environ trois approches pour prendre en charge le profilage dans MLIR:

  • Exécuter un profileur de bout en bout
  • Ajoutez une opération de profil pour chaque opération dans LHLO, à l'aide d'un profileur injecté.

L'approche «de bout en bout» est transparente pour MLIR, mais souffre du même problème qui fait que XLA ne l'utilise pas en premier lieu: les appels de bibliothèque collectés par un profileur (nvprof / ...) ne peuvent pas facilement se rapporter à HLO ops. Par exemple, cuDNN lance plusieurs noyaux pour chaque HLO, et il est difficile de dire quels noyaux correspondent à quel HLO.

L'approche "profileur injecté" nécessite:

  • LHLO pour prendre un profileur comme paramètre.
  • insérer profile.start / profile.end avant et après chaque opération.
  • un passage de ce profil abaissé. {start, end} à une implémentation C ++.

Le profilage exact ne peut pas être facilement effectué pour les opérations générées par MLIR, car:

  • MLIR n'a pas de minuterie et ne dépend pas de TFRT / StreamExecutor.
  • MLIR n'appelle pas facilement des fonctions C avec des paramètres compliqués.

Étape 3: (Tâche 2) Migration des Thunks

À noter, il existe environ trois types de thunks:

  • KernelThunk, qui lance un noyau.
  • Les thunks de flux de contrôle, qui ont une logique de flux de contrôle hôte (séquence conditionnelle, while, for,) et lancent des noyaux de corps.
  • Thunks de bibliothèque: cuDNN, cuBLAS, cuFFT, NCCL, etc.

Le plan est:

  • Rendre Thunks (dé) sérialisable.
  • Aidez à améliorer TFRT à un état où il peut prendre en charge cette sémantique.
  • Au fur et à mesure que l'état s'améliore, migrez les thunks individuellement de manière incrémentielle.

Ces éléments d'action ne sont que partiellement classés. Le parallélisme réel ordre d'exécution / ingénierie doit être évalué au fur et à mesure.

Étape 4: (Tâche 3) ElementalIrEmitter migré

Une fois le profilage prêt, nous pouvons terminer et régler tous les émetteurs basés sur ElementalIrEmitter dans MLIR. Ensuite, nous les activons par défaut, en supposant que tous ces émetteurs basés sur MLIR utilisent un seul flux.

Notez qu'il est également avantageux de migrer ElementalIrEmitter de XLA / CPU, car ils partagent une grande partie du code.

Une fois toutes les analyses comparatives et la recherche de performances effectuées (TODO: définir la parité des performances), nous activons le nouvel émetteur élémentaire basé sur MLIR et supprimons l'ancien ElementalIrEmitter.

Cette étape fournit également des transitions de fusion faciles (opérations imbriquées) pour la migration ultérieure.

Étape 5: Prise en charge ou suppression du multi-flux

Nous ne pouvons pas supprimer certains des émetteurs tant que nous ne l'avons pas pris en charge dans MLIR, ou que nous supprimons la fonctionnalité. C'est une quantité de travail relativement importante dans MLIR et une petite quantité de gain pour XLA. Nous devrions enquêter sur les utilisateurs actuels d'utilisateurs XLA / GPU multi-flux et essayer de supprimer cette fonctionnalité si cela est raisonnable.

Étape 6: (Tâche 3) Opérations de périphérique migrées

Cette étape migre toutes les opérations non imbriquées, puis nous pouvons supprimer tous les émetteurs non imbriqués.

Cela fait appel à une réécriture / refactorisation pour kCopy et kReduce. kReduce est déjà beaucoup travaillé, donc la quantité réelle de travail à faire reste à voir.