Ce document décrit le parcours d'un module High Level Optimizer (HLO) XLA, de son état initial à un exécutable final. Parfois, nous omettons le terme "module" et parlons simplement de "HLO".
HLO pré-optimisé
Nous commençons par le module HLO de pré-optimisation. Le HLO pré-optimisé ne contient pas d'opérations (ops) considérées comme internes à XLA, telles que fusion ou bitcast. Les opérations n'ont pas de mise en page à ce stade, ou si elles en ont une, elle sera ignorée. Le HLO de pré-optimisation est généralement produit par des frameworks de niveau supérieur tels que TensorFlow et JAX. Lorsque vous utilisez l'indicateur XLA -xla_dump_to, le HLO de pré-optimisation est transféré dans un fichier dont le nom se termine par "before_optimizations.txt".
Optimiser le module HLO
Le pipeline XLA:GPU transforme le HLO pré-optimisé en HLO optimisé en exécutant une séquence de passes. Les passes peuvent être regroupées de manière sémantique et exécutées dans l'ordre suivant :
Cartes associées au sharding
Cela inclut les passes telles que Shardy Partitioner ou celles pour le partitionnement SPMD.
Passes d'optimisation
Cela peut inclure à la fois des passes de légalisation et de simplification.
Passes d'optimisation collective
Semblable aux passes d'optimisation, mais se concentre sur les opérations collectives.
Passes d'attribution de la mise en page
Chaque opération HLO se voit attribuer une mise en page qui fait partie de la forme de l'instruction. La mise en page contrôle la façon dont le Tensor est physiquement disposé en mémoire.
Exemple de forme avec une mise en page :
f32[10,20,30]{2,0,1}
Après le type d'élément, vous trouverez les dimensions logiques de la forme, suivies de la permutation de mise en page dans l'ordre croissant. Dans cet exemple, la dimension la plus mineure est 30, la deuxième dimension la plus mineure est 10 et la dimension majeure est 20.
L'objectif de l'attribution de la mise en page est de minimiser le nombre de transpositions physiques requises à l'aide d'une stratégie gourmande. Il commence par certaines contraintes de mise en page (par exemple, les bibliothèques cuDNN/cuBLAS s'attendent à des dimensions consécutives) et propage les mises en page "vers le bas" puis "vers le haut" du graphique HLO. À la fin de la propagation de la mise en page, certaines instructions peuvent avoir des mises en page conflictuelles, l'une propagée à partir d'un opérande, l'autre propagée à partir d'un utilisateur. Pour résoudre ce conflit, une instruction HLO copy est insérée. Elle modifie la mise en page de la mise en page de l'opérande en mise en page de l'instruction.
Passes de normalisation de la mise en page
Étant donné qu'il est quelque peu difficile de déterminer la forme physique, la normalisation de la mise en page tente de réécrire la forme de manière à ce qu'elle utilise la mise en page par défaut {rank-1, rank-2, …, 0}. Dans l'exemple ci-dessus, la forme normalisée serait f32[20,10,30]{2,1,0}. Les opérations de copie qui modifient les mises en page sont réécrites sous la forme d'une combinaison de transpose et bitcast. Étant donné que nous ne pouvons pas normaliser toutes les opérations pour le moment, certaines opérations peuvent encore avoir des mises en page non par défaut, notamment gather et dot. Aux limites entre les opérations normalisées et non normalisées, il y aura des opérations bitcast qui représentent une transposition, c'est-à-dire une transposition avec une mise en page attribuée qui en fait une opération sans effet physique.
La normalisation de la mise en page rend également explicites certaines transpositions implicites, ce qui est important, car le générateur de code peut gérer les transpositions explicites avec un émetteur dédié. Par exemple, une opération de remodelage est techniquement autorisée à avoir une disposition physique différente entre l'opérande et le résultat (par exemple, en raison d'un rang différent). Le pass ReshapeDecomposer qui s'exécute dans le cadre des passes de normalisation de la mise en page transforme une mise en forme en une séquence de transpose, de mise en forme bitcast et de transpose.
Passes d'optimisation de l'attribution de mise en page
Les passes les plus importantes ici sont les fusions Triton (fusions GEMM + fusions Softmax/Layernorm) ou les réécritures des appels de bibliothèque. L'optimisation automatique s'exécute également à cette étape, où XLA choisit entre différents émetteurs, sélectionne le meilleur algorithme pour les convolutions ou les points, trouve le meilleur pavage pour les fusions gérées par l'émetteur Triton, etc.
Cartes Fusion
Les deux principaux sont la fusion PriorityFusion et Multi-Output.
Dans PriorityFusion, nous formons des fusions guidées par le modèle de coût. Lors de la fusion, nous autorisons la duplication des opérations avec plusieurs utilisateurs si l'opération peut être fusionnée pour tous les utilisateurs. Nous autoriserions également l'extension des fusions Triton Softmax existantes si possible.
La fusion Multi-Output est un pass distinct qui permet de fusionner des opérations/fusions qui partagent un opérande. Il peut également fusionner des opérandes/fusions d'opérandes dans des utilisateurs sans duplication en ajoutant des sorties supplémentaires, de sorte que les autres utilisateurs de l'opération à fusionner peuvent être redirigés vers ces sorties. Cette passe doit veiller à ne pas introduire de cycles dans le graphique HLO.
Après la fusion multi-sorties, l'élimination des sous-expressions communes (pass HloCSE) s'exécute, ce qui peut fusionner les opérations précédemment dupliquées si elles se retrouvent dans la même fusion.
Plusieurs passes post-fusion
Plusieurs passes liées aux collectifs (par exemple, les transformer en asynchrones ou appliquer un certain ordre relatif des collectifs).
Enfin, nous exécutons CopyInsertion où des copies sont ajoutées pour nous assurer que les opérations sur place n'écrasent pas les données qui sont encore nécessaires ailleurs.
À la fin de l'optimisation, le HLO optimisé est exporté si l'indicateur -xla_dump_to est utilisé dans un fichier dont le nom se termine par "after_optimizations.txt". Si vous souhaitez vider le HLO après les passes intermédiaires qui modifient réellement le HloModule, vous pouvez utiliser l'indicateur -xla_dump_hlo_pass_re=.* (ou une expression régulière spécifique pour le limiter à certaines passes).
Planification
Un module HLO sans programmation dispose toujours d'une certaine marge de manœuvre dans l'ordre de traitement des opérations. Tout tri topologique respectant les relations opérande/résultat et les dépendances de contrôle est valide. La planification détermine l'ordre spécifique à utiliser. À ce stade, la principale préoccupation est la consommation maximale de mémoire, qui dépend de la durée de vie des Tensors. Dans une première étape, nous essayons différents algorithmes de planification et choisissons celui qui devrait minimiser la consommation de mémoire maximale. Notez qu'à ce stade, nous ne travaillons pas encore avec des tampons physiques (cela se fera dans "Buffer Assignment") et que nous simulons l'utilisation de la mémoire.
Ensuite, les passes LatencyHidingScheduler s'exécutent et tentent de maximiser le chevauchement entre le calcul et la communication. Mais cela peut à nouveau augmenter l'utilisation de la mémoire.
Enfin, si la consommation de mémoire maximale est supérieure à la quantité de mémoire dont nous disposons, nous exécutons HloRematerialization. Ce pass tente de réduire l'utilisation de la mémoire au détriment des performances, car certaines fusions peuvent être fractionnées et certaines opérations peuvent être dupliquées pour réduire la durée de vie des tampons. Si une rematérialisation se produit, il peut être utile de rechercher des moyens de réduire les besoins en mémoire côté modèle (par exemple, en utilisant des tailles de lot plus petites).
Attribution de tampon
Juste avant de passer à LLVM IR, nous exécutons les passes d'attribution de tampon qui attribuent des tranches de tampon à chaque instruction du graphique HLO. L'attribution du tampon se déroule en plusieurs étapes :
HloDataflowAnalysisattribueHloValues(essentiellement des tampons logiques) aux instructions. Pour les opérations sur place, leHloValued'un opérande peut être réutilisé. Une opération peut définir plusieursHloValue(par exemple, avec une forme de résultat de tuple).HloAliasAnalysistente de combiner des tampons pour les opérations d'alias et calcule un mappage deHloValueàHloBuffer.BufferAssignmentcalcule un mappage deHloBufferssur des tranches de tampon à l'intérieur d'un grand tampon de manière à ce que la même tranche de tampon ne soit pas utilisée pour différentsHloBuffersavec des durées de vie qui se chevauchent. Pour les opérations qui peuvent être des alias, il est acceptable qu'il y ait un léger chevauchement (l'heure de fin d'unHloBufferpeut coïncider avec l'heure de début de l'autreHloBuffer). Lorsque vous utilisez le flag-xla_dump_to, certaines informations sur l'attribution des tampons sont enregistrées dans un fichier dont le nom se termine par "after_optimizations-buffer-assignment.txt".
Thunks
Une fois qu'un graphique HLO est optimisé et planifié, il est abaissé en une séquence linéaire de thunks pour un backend spécifique (CPU ou GPU).
Dans XLA, un Thunk est une abstraction d'une unité de travail autonome que le runtime exécute. Il peut s'agir d'un lancement de noyau compilé, d'une opération spécifique, d'un appel de bibliothèque, d'une construction de flux de contrôle, d'une communication collective, etc. Une séquence Thunk représente l'intégralité de l'exécutable pour un backend spécifique.
Émission de thunk
Le processus de conversion d'un calcul HLO planifié en séquence de thunks est appelé "émission de thunks". Cela est géré par une classe d'émetteur dédiée dans chaque backend.
Pour le backend GPU, cela est géré par IrEmitterUnnested.
EmitHloComputation itère sur la liste planifiée des instructions HLO dans un calcul et les distribue à une méthode Emit... spécialisée (par exemple,
EmitFusion, EmitConvolutionThunk, EmitWhile). Chacune de ces méthodes construit les objets Thunk appropriés et les ajoute à la séquence Thunk.
Pour le backend du processeur, ThunkEmitter remplit ce rôle et est organisé de manière similaire. Le ThunkSequence final est intégré au CpuExecutable.
Notez que chaque instruction dans le calcul d'entrée d'un module HLO peut correspondre à zéro (kTuple, kConstant, etc.), un ou plusieurs thunks (par exemple, une instruction de tri) dans la séquence de thunks finale.
Tampons de commandes : optimiser l'exécution sur le GPU
Le matériel GPU moderne permet d'enregistrer une séquence d'opérations GPU (lancements de noyaux, copies de mémoire, etc.) une seule fois, puis de la relire plusieurs fois avec une surcharge de processeur minimale. Il s'agit d'une optimisation des performances essentielle, en particulier pour les charges de travail comportant de nombreux petits kernels à lancement rapide. XLA utilise Command Buffer comme abstraction des graphiques CUDA ou HIP. L'interface principale est définie dans GpuCommandBuffer.
Un tampon de commande est représenté dans une séquence thunk par CommandBufferThunk.
L'émetteur ne produit pas ce thunk directement à partir des instructions HLO. Au lieu de cela, cela est fait par CommandBufferConversionPass qui s'exécute sur ThunkSequence lui-même.
Le pass identifie les sous-séquences contiguës de thunks compatibles (par exemple, une série de KernelThunk et de GemmThunk). Il remplace ensuite la sous-séquence trouvée par un seul CommandBufferThunk. Le nouveau thunk encapsule la logique des thunks d'origine sous la forme d'une liste d'objets CommandBufferCmd légers.
Lorsqu'un CommandBufferThunk s'exécute pour la première fois sur un flux GPU donné, il "enregistre" sa séquence de commandes dans un tampon de commandes matériel. Lors de toutes les exécutions suivantes, il émet simplement une seule commande au GPU pour "rejouer" la séquence enregistrée. Cela évite la surcharge du processeur liée au lancement de chaque noyau individuel.
Exécutable
Le produit final du pipeline de compilation XLA est un exécutable autonome et spécifique à la plate-forme. Cet objet encapsule toutes les informations nécessaires pour exécuter le programme compilé sur un appareil cible, tel qu'un processeur ou un GPU. Il s'agit du pont entre le compilateur et l'environnement d'exécution. Les environnements d'exécution modernes tels que PJRT utilisent des abstractions de niveau légèrement supérieur (voir PjRtExecutable), mais celles-ci encapsulent en fin de compte un exécutable spécifique au backend.
Un Executable contient plusieurs informations clés générées lors de la compilation. Bien que le contenu exact varie selon le backend, il inclut généralement les éléments suivants :
Code compilé : il s'agit du code machine de bas niveau qui s'exécutera sur l'appareil. Pour les processeurs, il s'agit généralement d'un ou de plusieurs fichiers objets. Pour les GPU, il s'agit du code de périphérique compilé au format PTX ou HSACO, qui est chargé sur le GPU au moment de l'exécution.
Plan d'exécution (ThunkSequence) : le cœur de la logique d'exécution. Il s'agit d'une séquence linéaire d'objets Thunk. Chaque thunk représente une seule unité de travail, comme le lancement d'un noyau, l'appel d'une fonction de bibliothèque (par exemple, cuBLAS) ou la gestion du flux de contrôle. Le runtime exécute le programme en itérant sur cette séquence.
Disposition de la mémoire (BufferAssignment) : cette métadonnée essentielle, produite par BufferAssigner, décrit la disposition complète de la mémoire pour le calcul. Il spécifie la taille de chaque tampon et la façon dont la mémoire est allouée et réutilisée pour les paramètres, les sorties et les valeurs temporaires. L'environnement d'exécution l'utilise pour allouer de la mémoire à l'appareil et transmettre les pointeurs appropriés à chaque thunk.
(facultatif) Module HLO : pour le débogage et le profilage, l'exécutable conserve souvent une référence au HloModule final et optimisé à partir duquel il a été compilé.
La création de l'exécutable final est orchestrée par le compilateur pour chaque backend spécifique. La méthode RunBackend d'une implémentation de compilateur est la dernière étape du processus de compilation, qui regroupe tous les artefacts compilés dans un objet Executable.
GpuCompiler et CpuCompiler ciblent respectivement le GPU et le CPU.
Lorsqu'un utilisateur appelle Execute... sur un exécutable, le runtime utilise BufferAssignment pour allouer de la mémoire, puis appelle ThunkSequence pour lancer les opérations sur l'appareil à l'aide du code compilé.