Disposition en mosaïque

Figure 1

La figure 1 montre la disposition en mémoire d'un tableau F32[3,5] avec une mosaïque de 2x2. Une forme avec cette mise en page s'écrit F32[3,5]{1,0:T(2,2)}, où 1,0 correspond à l'ordre physique des dimensions (champ minor_to_major dans Layout), tandis que (2,2) après le deux-points indique le tuilage des dimensions physiques par une tuile 2x2.

Intuitivement, les tuiles sont disposées pour couvrir la forme, puis, dans chaque tuile, les éléments sont ensuite disposés sans mosaïque, comme dans l'exemple ci-dessus, où la partie droite de l'exemple montre la mise en page en mémoire, y compris les éléments de marge intérieure blancs ajoutés pour obtenir des tuiles 2x2 complètes, même si les limites du tableau d'origine ne sont pas égales.

Les éléments supplémentaires de la marge intérieure ne doivent pas nécessairement contenir une valeur particulière.

Formules d'index linéaires pour l'emplacement en fonction d'une forme et d'une tuile

Sans mosaïque, un élément e=(e_n, e_n-1, ... , e₁) dans un tableau avec les limites de tableau d=(d_n, d_n-1, ... , d₁) (d1 est la dimension la plus mineure) est réparti par ordre majeur/mineur en position:

linear_index(e, d)
= linear_index((e_n, e_n-1, ... , e₁), (d_n, d_n-1, ... , d₁))
= e_nd_n-1...d-₁ + e_n-1 + d_{n +n}

Pour simplifier la notation dans ce document, nous supposons qu'une carte a le même nombre de dimensions que le tableau. Dans l'implémentation du tuilage de XLA, cela est généralisé aux mosaïques ayant moins de dimensions, en laissant les dimensions initiales les plus principales inchangées et en appliquant le tuilage uniquement aux dimensions les plus mineures, de sorte que le tuilage spécifié mentionne un suffixe des dimensions physiques de la forme en mosaïque.

Lorsque vous utilisez un tuilage de taille (t_n, t_n-1, ... , t₁), un élément du tableau comportant des index (_e, e_n-1, ..., e₁) est mappé à cette position dans la mise en page finale:

linear_index_with_tile(e, d, t) 1, t /t


_n

La mise en page peut être considérée comme ayant deux parties : (⌊e_n/t_n⌋, ... , ⌊e₁/t₁⌋), qui correspond à un index de carte dans un tableau de tuiles de taille (⌈d_n/t_n⌋, ... ,⌈d₁/t), et La fonction "ceil" apparaît dans ⌈d_i/t_i⌉, car si les tuiles dépassent les limites du plus grand tableau, une marge intérieure est insérée, comme illustré dans la figure 1. Les tuiles et les éléments à l'intérieur des tuiles sont disposés de manière récursive, sans mosaïque.

Pour l'exemple de la figure 1, l'élément (2,3) a l'index de tuile (1,1) et l'index dans la tuile (0,1), pour un vecteur de coordonnées combiné de (1,1,0,1). Les index de carte ont des limites (2,3) et la carte elle-même correspond à (2,2) pour un vecteur combiné de (2,3,2,2). L'index linéaire avec tuile pour l'élément dont l'indice (2,3) est dans la forme logique est alors

linear_index_with_tile((2,3), (3,5), (2,2))
= linear_index((1,1,0,1), (2,3,2,2))
= index_linéaire((1,1), (2,3)) ∙ 2 ∙ 2 + linear_index(0,1) + .

Mosaïque en tant que remodelage de pad

La mise en page basée sur des emplacements fonctionne comme suit:
Considérez un tableau de dimensions (d_n, d_n-1, ... , d1) (d1 est la dimension la plus mineure). Lorsqu'il est disposé avec un tuilage de taille (t_n, t_n-1, ... , t₁) (t₁ est la dimension la plus mineure), ce tuilage peut être décrit en termes de remodelage-transposition de pad de la manière suivante.

1. Le tableau est complété jusqu'à (⌈d_n/t_n⌉∙t_n, ... , ⌈d₁/t₁⌉∙t₁).
2. Chaque dimension i est divisée en (⌈d_i/ti⌉, t_i), c'est-à-dire que le tableau est remodelé en
   (⌈d_n/t_n⌉, t_n, ... , ⌈d₁/t₁⌉, t₁).
   Cette remodelage n'entraîne aucune modification physique de la mise en page. Il s'agit donc d'un bitcast. Si l'on ne pense pas explicitement à un tuilage, cette remodelage peut exprimer toute forme avec le même nombre d'éléments que la forme remplie. Cet exemple montre comment exprimer une carte de cette manière.
3. Une transposition se produit en déplaçant t_n, ... , t₁ vers les dimensions les plus mineures, tout en conservant leur ordre relatif, de sorte que l'ordre des dimensions, de la plus majeure à la plus mineure, devienne
   (⌈d_n/t_n⌉, ... , ⌈d₁/t₁⌉, 1,_n).

La forme finale comporte le préfixe
(⌈d_n/t_n⌉, ... , ⌈d₁/t₁⌉), qui décrit le nombre de tuiles dans chaque dimension. Un élément du tableau (e_n, ... , e₁) est mappé à cet élément dans la forme finale:
(⌊e_n/t_n⌋, ... , ⌊e₀/t₀⌋, e_n mod t_n, ... , e₁). Il est facile de voir que l'index linéaire de l'élément suit comme prévu la formule ci-dessus.

Mosaïque répétée

Le tuilage de XLA devient encore plus flexible en l'appliquant à plusieurs reprises.

Figure 2

La figure 2 montre comment un tableau de taille 4x8 est tuilé selon deux niveaux de mosaïque (d'abord 2x4, puis 2x1). Nous représentons ce carrelage répété comme suit : (2,4)(2,1). Chaque couleur représente une tuile de 2x4, et chaque zone de bordure rouge est une tuile de 2x1. Les nombres indiquent l'index linéaire en mémoire de cet élément au format en mosaïque. Ce format correspond au format utilisé pour BF16 sur TPU, sauf que la tuile initiale est plus grande, à savoir (8,128)(2,1), où l'objectif de la deuxième mise en mosaïque par 2x1 est de collecter ensemble deux valeurs 16 bits pour former une valeur 32 bits de manière à s'aligner sur l'architecture d'un TPU.

Notez qu'une deuxième tuile ou ultérieure peut faire référence aux dimensions mineures dans la carte, ce qui réorganise simplement les données au sein de la carte, comme dans cet exemple avec (8,128)(2,1), mais peut également faire référence aux principales dimensions croisées de la tuile par rapport à la tuile précédente.

Combiner des dimensions à l'aide de tuiles

Le tuilage de XLA permet également de combiner des dimensions. Par exemple, il peut combiner les dimensions de F32[2,7,8,11,10]{4,3,2,1,0} dans F32[112,110]{1,0} avant de la tuiler avec (2,3). La vignette utilisée est (∗,∗,2,∗,3). Ici, un astérisque dans une carte implique de prendre cette dimension et de la combiner avec la dimension suivante, plus mineure. Plusieurs dimensions adjacentes peuvent être regroupées en une seule dimension. Une dimension en sous-somme est représentée par une valeur de tuile de -1 dans cette dimension de la carte, qui n'est autrement pas valide dans une carte en tant que taille de dimension.

Plus précisément, si la dimension i de la forme est éliminée via un astérisque dans la tuile, avant l'application de la définition précédente du mosaïque, cette dimension est supprimée à la fois de la forme en mosaïque et du vecteur de tuile, et la limite du tableau de la dimension i-1 de la forme est augmentée de d_i-1 à d_id_i-1. Cette étape est répétée pour chaque astérisque du vecteur de tuiles.