Ce document décrit l'analyse de l'indexation des HLO, qui vous permet de calculer symboliquement des mappages d'indexation pour les opérations HLO. Le mappage d'indexation est une fonction qui mappe les indices d'un Tensor aux index d'un autre (par exemple, les indices d'une sortie d'instruction HLO aux index d'entrées d'instructions HLO, et inversement).
Exemple
Pour une diffusion du tensor<20xf32> au tensor<10x20x30xf32>
p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}
le mappage d'indexation de la sortie à l'entrée est \((i, j, k) \mapsto (j)\) pour $i \in [0, 10]\(, \)j \in [0, 20]\( and \)k \in [0, 30]$.
Motivation
Le GPU XLA utilise plusieurs solutions sur mesure pour interpréter la coalisation, l'utilisation des opérandes et le tuilage (plus de détails ci-dessous). L'objectif de l'analyse d'indexation est de fournir un composant réutilisable pour de tels cas d'utilisation. L'analyse d'indexation est basée sur l'infrastructure de carte Affine de MLIR et ajoute la sémantique HLO.
Coalescent
Le raisonnement sur la coalisation de la mémoire devient réalisable pour les cas complexes, lorsque nous savons quels éléments/tranches d'entrée sont lus pour calculer un élément de la sortie.
Utilisation de l'opérande
L'utilisation de l'opérande dans XLA indique la quantité d'utilisation de chaque entrée de l'instruction, en supposant que sa sortie est entièrement utilisée. Actuellement, l'utilisation n'est pas non plus calculée pour un cas générique. L'analyse d'indexation permet de calculer l'utilisation avec précision.
Mosaïque
Une tuile/tranche est un sous-ensemble hyperrectangulaire d'un Tensor paramétré par des décalages, des tailles et des pas. La propagation de tuiles est un moyen de calculer les paramètres de tuile du producteur/consommateur de l'opération à l'aide des paramètres de tuile de l'opération elle-même. Il existe déjà une bibliothèque qui effectue cette opération pour softmax et .dot. La propagation des tuiles peut être rendue plus générique et plus robuste si elle est exprimée via des cartes d'indexation.
Fonction et domaine
Le mappage d'indexation est une fonction \(\boldsymbol{f}(\boldsymbol{d}, \boldsymbol{s})\)qui mappe le multi-index \(\boldsymbol{d}\) d'un Tensor \(A\) à des éléments/plages d'un Tensor \(B\). Le paramètre \(\boldsymbol{s}\) fait référence aux plages d'indices des dimensions présentes dans le Tensor \(B\), mais pas dans le Tensor \(A\).
Par exemple, si nous avons une réduction de tensor<2x4x8x16xf32> à tensor<4x8xf32>, le mappage d'indexation de la sortie 2D vers l'entrée 4D est\((d_0, d_1) \mapsto (s_0, d_0, d_1, s_1)\), où \(d_i\) sont les paramètres de dimension qui correspondent aux index du Tensor de sortie. Les paramètres \(s_j\)encodent plusieurs valeurs. Par exemple, pour calculer un élément \((d_0, d_1)\) de la sortie, nous avons besoin \((s_0, d_0, d_1, s_1)\) d'éléments d'entrée, où \(s_0 \in [0, 2)\) et\(s_1 \in [0, 16)\).
Ce mappage peut être créé à partir des attributs des instructions HLO, ou les mappages des instructions non fusionnées peuvent être composés pour obtenir l'indexation d'une fusion. Le mappage comporte également un domaine, qui spécifie les éléments du Tensor dont il existe le mappage.
\[ \begin{eqnarray} \boldsymbol{f}(\boldsymbol{d}, \boldsymbol{s})\; &s.t.& \\ \boldsymbol{lb}_d &\leq& \boldsymbol{d} \leq \boldsymbol{ub}_d \\ \boldsymbol{lb}_s &\leq& \boldsymbol{s} \leq \boldsymbol{ub}_s \\ \boldsymbol{lb}_g &\leq& \boldsymbol{g}(\boldsymbol{d}, \boldsymbol{s}) \leq \boldsymbol{ub}_g \end{eqnarray} \]
Comme nous voulons minimiser le nombre de nouveaux calculs, nous avons besoin d'une bibliothèque pour les calculs symboliques. XLA dépend déjà de MLIR. Nous utilisons donc mlir::AffineMap au lieu d'écrire une bibliothèque arithmétique symbolique.
Voici à quoi ressemble un AffineMap type
(d0)[s0, s1] -> (s0 + 5, d0 * 2, s1 * 3 + 50)
AffineMap comporte deux types de paramètres: les dimensions et les symboles que nous pouvons utiliser respectivement pour \(\boldsymbol d\) et \(\boldsymbol s\) .
AffineMap ne contient pas de métadonnées sur les plages des dimensions. Nous devons donc fournir ces données nous-mêmes.
struct Range {
int64_t lower_bound;
int64_t upper_bound;
};
struct IndexingMap {
mlir::AffineMap affine_map;
std::vector<Range> dimension_ranges;
std::vector<Range> symbol_ranges;
llvm::DenseMap<mlir::AffineExpr, Range> expr_ranges;
};
dim_ranges encode les contraintes de boîte incluses pour les paramètres de dimension \(\boldsymbol{d}\) du mappage d'indexation, qui coïncident généralement avec la forme du Tensor de sortie pour les opérations telles que la transposition, la réduction, l'élément par élément, les points, mais il existe des exceptions telles que HloConcatenateInstruction.
symbol_ranges encode les valeurs possibles que les paramètres \(\boldsymbol {s}\) peuvent prendre.
Étudions-les par exemple pour comprendre ce que signifie réellement tous les éléments ci-dessus.
Indexer des cartes pour les opérations non fusionnées
Par élément
Pour une opération élément par élément, le mappage d'indexation est une identité.
p0 = f32[10, 20] parameter(0)
p1 = f32[10, 20] parameter(1)
add = f32[10, 20] add(p0, p1)
La sortie en entrée est mappée comme suit:
- sortie -> input_0: \((d_0, d_1) \mapsto (d_0, d_1)\) pour $\boldsymbol{d} \in [0,9] \times [0, 19]\(, i.e. \)\boldsymbol{d} \in {\rm Dom}(sortie)$
- sortie -> input_1: \((d_0, d_1) \mapsto (d_0, d_1)\) pour $\boldsymbol{d} \in {\rm Dom} (sortie)$
Les mappages d'entrée et de sortie
- input_i -> output: \((d_0, d_1) \mapsto (d_0, d_1)\) pour $\boldsymbol{d} \in {\rm Dom}(output)$
Diffusion
Le broadcasting signifie que certaines dimensions sont supprimées lorsque nous mappons la sortie à l'entrée et ajoutées lorsque nous mappons l'entrée à la sortie.
p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}
Mapper le résultat en entrée:
- sortie -> entrée: \((d_0, d_1, d_2) \mapsto (d_1)\) pour $\boldsymbol{d} \in {\rm Dom}(output)$
Le mappage entre l'entrée et la sortie
- entrée -> sortie: \((d_0) \mapsto (s_0, d_1, s_1)\) pour $\boldsymbol{d} \in {\rm Dom}(sortie)\( and \)\boldsymbol{s} \in [0, 9] \times [0, 29]$.
Notez que nous avons maintenant \(\boldsymbol s\) à droite pour le mappage d'entrée/sortie. Ce sont les symboles qui représentent des plages de valeurs. Par exemple, dans ce cas particulier, chaque élément d'entrée ayant l'indice \(d_0\) est mappé à une tranche 10 x 1 x 30 de la sortie.
Constant et Iota
Pour être utiles, elles ne comportent aucun paramètre d'entrée. Il n'y a donc aucun calcul pour l'indexation.
Transpose
Le mappage d'indexation pour la transposition est une permutation de dimensions d'entrée/sortie.
p0 = f32[3, 12288, 6, 128] parameter(0)
transpose = f32[3, 6, 128, 12288] transpose(p0), dimensions={0, 2, 3, 1}
Mapper le résultat en entrée:
- sortie -> entrée: \((d_0, d_1, d_2, d_3) \mapsto (d_0, d_3, d_1, d_2)\) pour \(\boldsymbol{d} \in {\rm Dom}(output)\)
Le mappage entre l'entrée et la sortie:
- entrée -> sortie: \((d_0, d_1, d_2, d_3) \mapsto (d_0, d_2, d_3, d_1)\) pour \(\boldsymbol{d} \in {\rm Dom}(input)\)
Inverser
L'indexation de la carte inversée remplace les dimensions inversées par $upper_bound(d_i) - d_i$:
p0 = f32[1, 17, 9, 9] parameter(0)
reverse = f32[1, 17, 9, 9] reverse(p0), dimensions={1, 2}
Mapper le résultat en entrée:
- sortie -> entrée: $(d_0, d_1, d_2, d_3) \mapsto (d_0, -d_1 + 16, -d_2 + 8, d_3)\( for \)\boldsymbol{d} \in {\rm Dom}(output)$
Le mappage entre l'entrée et la sortie:
- entrée -> sortie: $(d_0, d_1, d_2, d_3) \mapsto (d_0, -d_1 + 16, -d_2 + 8, d_3)\( for \)\boldsymbol{d} \in {\rm Dom}(input)$
(Variadique)Réduire
La réduction variadique comporte plusieurs entrées et plusieurs intents. Le mappage de la sortie à l'entrée ajoute les dimensions réduites. Cela se comporte comme l'inverse d'une diffusion, en quelque sorte.
p0 = f32[256,10] parameter(0)
p0_init = f32[] constant(-inf)
p1 = s32[256,10] parameter(1)
p1_init = s32[] constant(0)
reduce = (f32[10], s32[10]) reduce(p0, p1, p0_init, p1_init),
dimensions={0}, to_apply=min
La sortie en entrée est mappée comme suit:
- sortie -> input_j: \((d_0) \mapsto (s_0, d_0)\) pour $\boldsymbol{d} \in {\rm Dom}(sortie)\( and \)\boldsymbol{s} \in [0, 9]$
- sortie -> init_j: \((d_0) \mapsto ()\) pour $\boldsymbol{d} \in {\rm Dom}(output)$
Les mappages d'entrée et de sortie sont les suivants:
- input_i -> output_j: \((d_0, d_1) \mapsto (d_1)\) pour $\boldsymbol{d} \in {\rm Dom}(entrée)$
- init_i -> output_j: \(() \mapsto (s_0)\) pour \(\boldsymbol{s} \in [0, 9]\)
pour \(i, j = 0, \ldots, INPUT\\_COUNT\).
Tranche
L'indexation de la sortie en entrée pour les tranches génère un mappage d'index étendu valide pour chaque élément de la sortie. Le mappage de l'entrée à la sortie est limité à une plage échelonnée des éléments de l'entrée.
p0 = f32[10, 20, 50] parameter(0)
slice = f32[5, 3, 25] slice(f32[10, 20, 50] p0),
slice={[5:10:1], [3:20:7], [0:50:2]}
Mapper le résultat en entrée:
- sortie -> entrée: \((d_0, d_1, d_2) \mapsto (d_0 + 5, 7d_1 + 3, 2d_2)\) pour \(\boldsymbol{d} \in {\rm Dom}(output)\)
Le mappage entre l'entrée et la sortie:
- input -> output: \((d_0, d_1, d_2) \mapsto (d_0, d_1 / 7, d_2 / 2)\) for \(\boldsymbol{d} \in [5, 9] \times [3, 19] \times [0, 49]\) with foules $[1, 7, 2]$.
À définir: indexation des entrées-sorties
Remodeler
Les remodelages peuvent prendre différentes saveurs.
Réduire la forme
Il s'agit d'une restructuration "linéaire" de N-D à 1D.
p0 = f32[4,8] parameter(0)
reshape = f32[32] reshape(p0)
Mapper le résultat en entrée:
- sortie -> entrée: \((d_0) \mapsto (d_0 / 8, d_0 \mod 8)\) pour $\boldsymbol{d} \in {\rm Dom}(output)$
Le mappage entre l'entrée et la sortie:
- entrée -> sortie: \((d_0, d_1) \mapsto (8 d_0 + d_1)\) pour $\boldsymbol{d} \in {\rm Dom}(input)$.
Développer la forme
Il s'agit d'une opération de "réduction de forme" inverse, qui transforme une entrée unidimensionnelle en sortie N-D.
p0 = f32[32] parameter(0)
reshape = f32[4, 8] reshape(p0)
Mapper le résultat en entrée:
- sortie -> entrée: \((d_0, d_1) \mapsto (8 d_0 + d_1)\) pour $\boldsymbol{d} \in {\rm Dom}(output)$
Le mappage entre l'entrée et la sortie:
- entrée -> sortie: \((d_0) \mapsto (d_0 / 8, d_0 \mod 8)\) pour $\boldsymbol{d} \in {\rm Dom}(input)$.
Remodelage générique
Il s'agit des opérations de remodelage qui ne peuvent pas être représentées sous la forme d'une seule forme de développement ou de réduction. Elles ne peuvent être représentées que sous la forme d'une composition de deux ou plusieurs formes d'agrandissement ou de réduction.
Exemple 1: linéarisation-délinéarisation
p0 = f32[4,8] parameter(0)
reshape = f32[2, 4, 4] reshape(p0)
Cette remodelage peut être constitué d'une forme réduite de tensor<4x8xf32> à tensor<32xf32>, puis d'une expansion de forme vers tensor<2x4x4xf32>.
Mapper le résultat en entrée:
- sortie -> entrée: $(d_0, d_1, d_2) \mapsto (2d_0 + (4d_1 + d_2) / 8, 4d_1 + d_2) \mod 8)$
pour \(\boldsymbol{d} \in {\rm Dom}(output)\)
Le mappage entre l'entrée et la sortie:
- entrée -> sortie: $(d_0, d_1) \mapsto ((8d_0 + d_1) / 16, ((8d_0 + d_1) \mod 16) / 4, d_1 \mod 4)$
pour \(\boldsymbol{d} \in {\rm Dom}(input)\).
Exemple 2: Sous-formes développées et réduites
p0 = f32[4, 8, 12] parameter(0)
reshape = f32[32, 3, 4] reshape(p0)
Cette remodelation peut être représentée par la combinaison de deux remodelages. La première réduit les dimensions de plus haut niveau tensor<4x8x12xf32> en tensor<32x12xf32> et la seconde développe la dimension la plus interne tensor<32x12xf32> en tensor<32x3x4xf32>.
Mapper le résultat en entrée:
- sortie -> entrée: \((d_0, d_1, d_2) \mapsto (d_0 / 8, d_0 \mod 8, 4d_1 + d_2)\) pour \(\boldsymbol{d} \in {\rm Dom}(output)\)
Le mappage entre l'entrée et la sortie:
- entrée -> sortie: \((d_0, d_1, d_2) \mapsto (8d_0 + d_1, d_2 / 4, d_2 \mod 4)\)pour \(\boldsymbol{d} \in {\rm Dom}(input)\).
Bitcast
Une opération de bitcast peut être représentée comme une séquence de transpose-remodelage-transpose. Par conséquent, ses mappages d'indexation ne constituent qu'une composition de mappages d'indexation pour cette séquence.
Concatenate
Le mappage de sortie/entrée pour la concaténation est défini pour toutes les entrées, mais avec des domaines qui ne se chevauchent pas, c'est-à-dire qu'une seule des entrées sera utilisée à la fois.
p0 = f32[3,50] parameter(0)
p1 = f32[3,30] parameter(1)
concat = f32[3,80] concatenate(f32[3,50] p0, f32[3,30] p1),
dimensions={1}
Mapper le résultat en entrée:
- sortie -> entrée 1:
\((d_0, d_1) \mapsto (d_0, d_1)\) pour \(\boldsymbol{d} \in [0, 2] \times [0, 49]\)
- sortie -> entrée 2:
\((d_0, d_1) \mapsto (d_0, d_1 - 50)\) pour $\boldsymbol{d} \in [0, 2] \times [50, 79]$
Les entrées du mappage de sortie:
- entrée 1 -> sortie: \((d_0, d_1) \mapsto (d_0, d_1)\) pour $\boldsymbol{d} \in {\rm Dom}(input_1)$.
- entrée 2 -> sortie: \((d_0, d_1) \mapsto (d_0, d_1 + 50)\) pour $\boldsymbol{d} \in {\rm Dom}(input_2)$.
Point (mise en œuvre du processus de sortie en entrée)
L'indexation des cartes par point est très similaire à celle de la méthode "reduce".
p0 = f32[4, 128, 256] parameter(0)
p1 = f32[4, 256, 64] parameter(1)
dot = f32[4, 128, 64] dot(p0, p1),
lhs_batch_dims={0}, rhs_batch_dims={0},
lhs_contracting_dims={2}, rhs_contracting_dims={1}
La sortie vers les entrées correspond aux mappages suivants:
- sortie -> input_1: \((d_0, d_1, d_2) \mapsto (d_0, d_1, s_0)\) pour\(\boldsymbol{d} \in {\rm Dom}(output)\) et \(\boldsymbol{s} \in [0, 255]\)
- sortie -> input_2: \((d_0, d_1, d_2) \mapsto (d_0, s_0, d_2)\) pour\(\boldsymbol{d} \in {\rm Dom}(output)\) et \(\boldsymbol{s} \in [0, 255]\)
Les entrées des mappages de sortie:
- input_1 -> output: \((d_0, d_1, d_2) \mapsto (d_0, d_1, s_0)\) pour \(\boldsymbol{d} \in {\rm Dom}(input_1)\) et \(\boldsymbol{s} \in [0, 63]\)
- input_2 -> output: \((d_0, d_1, d_2) \mapsto (d_0, s_0, d_1)\) pour \(\boldsymbol{d} \in {\rm Dom}(input_2)\) et \(\boldsymbol{s} \in [0, 127]\)
Réduire la fenêtre (à déterminer)
Pad (à définir)
Indexer Maps pour Fusion
L'indexage de la carte pour l'opération de fusion consiste en une composition d'indexage des mappages pour chaque opération dans le cluster. Il peut arriver que certaines entrées soient lues plusieurs fois avec des modèles d'accès différents.
Une entrée, plusieurs mappages d'indexation
Voici un exemple pour \(p_0 + p_0^T\)
f {
p0 = f32[1000, 1000] parameter(0)
transpose_p0 = f32[1000, 1000]{0, 1} transpose(p0), dimensions={1, 0}
ROOT a0 = f32[1000, 1000] add(p0, transpose_p0)
}
Les mappages d'indexation de sortie/entrée pour p0 seront $(d_0, d_1) \mapsto (d_0,
d_1)\( and \)(d_0, d_1) \mapsto (d_1, d_0)$. Cela signifie que pour calculer un élément de la sortie, nous devrons peut-être lire le paramètre d'entrée deux fois.
Une entrée, carte d'indexation dédupliquée
Dans certains cas, les mappages d'indexation sont en fait identiques, même si cela n'est pas immédiatement évident.
f {
p0 = f32[20, 10, 50] parameter(0)
lhs_transpose_1 = f32[10, 20, 50] transpose(p0), dimensions={1, 0, 2}
lhs_e = f32[10, 20, 50] exponential(lhs_transpose_1)
lhs_transpose_2 = f32[10, 50, 20] transpose(lhs_e), dimensions={0, 2, 1}
rhs_transpose_1 = f32[50, 10, 20] transpose(p0), dimensions={2, 1, 0}
rhs_log = f32[50, 10, 20] exponential(rhs_transpose_1)
rhs_transpose_2 = f32[10, 50, 20] transpose(rhs_log), dimensions={1, 0, 2}
ROOT add = f32[10, 50, 20] add(lhs_transpose_2, rhs_transpose_2)
}
Dans ce cas, le mappage d'indexation de sortie/entrée pour p0 est simplement $(d_0, d_1, d_2)\mapsto (d_2, d_0, d_1)$.
Softmax
L'indexation de sortie/entrée est mappée pour parameter 0 pour softmax:
- \((d_0, d_1, d_2) \mapsto (d_0, d_1, d_2)\)
- \((d_0, d_1, d_2)[s_0] \mapsto (d_0, d_1, s_0)\)
pour \(\boldsymbol{d} \in {\rm Dom}(output)\) et \(\boldsymbol{s} \in [0, 124]\)fait référence à la dimension la plus interne de l'entrée.
Simplificateur de carte d'indexation
Le simplificateur par défaut pour mlir::AffineMap en amont ne peut faire aucune hypothèse concernant les plages de dimensions/symboles. Par conséquent, elle ne peut pas simplifier efficacement les expressions avec mod et div.
Nous pouvons exploiter les connaissances sur les limites inférieure et supérieure des sous-expressions dans les cartes affine pour les simplifier encore plus.
Le simplificateur peut réécrire les expressions suivantes.
- \((d_0, d_1) \mapsto (d_0 + d1 / 16, d1 \mod 16)\) pour $\boldsymbol{d} \in [0, 6] \times [0, 14]\( becomes \)(d_0, d_1) \mapsto (d_0, d_1)$
- $(d_0, d_1, d_2) \mapsto ((100d_0 + 10d_1 + d_2) /100, ((100d_0 + 10d_1 + d_2) \mod 100) / 10, d_2 \mod 10)\( for \)d_i, d_2\( becomes \)
- $(d_0, d_1, d_2) \mapsto ((16d_0 + 4d_1 + d_2) /8, (16d_0 + 4d_1 + d_2) \mod 8)\( for \)d_i \in [0, 9]\( becomes \)(d_0, d_1, d_2) \0(d_2d) +
- \((d_0, d_1) \mapsto (-(-11d_0 - d_1 + 109) / 11 + 9)\) pour $\boldsymbol{d} \in [0, 9] \times [0, 10]\( becomes \)(d_0, d_1) \mapsto (d_0)$.
Le simplificateur de carte d'indexation nous permet de comprendre que certaines reformes en chaîne dans HLO s'annulent les unes les autres.
p0 = f32[10, 10, 10] parameter(0)
reshape1 = f32[50, 20] reshape(p0)
reshape2 = f32[10, 10, 10] reshape(reshape1)
Après la composition des cartes d'indexation et leur simplification, nous obtiendrons
\((d_0, d_1, d_2) \mapsto (d_0, d_1, d_2)\).