Gekacheltes Layout

Abbildung 1

Abbildung 1 zeigt, wie ein Array F32[3,5] mit 2x2-Kacheln im Speicher angeordnet wird. Eine Form mit diesem Layout wird als F32[3,5]{1,0:T(2;2)} geschrieben,wobei 1,0 sich auf die physikalische Reihenfolge der Dimensionen bezieht (Feld minor_to_major im Layout), während (2,2) nach dem Doppelpunkt die Kachelung der physischen Abmessungen durch eine 2x2-Kachel angibt.

Intuitive Kacheln werden so angeordnet, dass sie die Form abdecken. Dann werden in jeder Kachel dann Elemente ohne Kacheln angeordnet, wie im obigen Beispiel, wo der rechte Teil des Beispiels das Layout im Speicher zeigt, einschließlich der Elemente mit weißem Innenrand, die hinzugefügt werden, um vollständige 2 x 2-Kacheln zu erhalten, obwohl die ursprünglichen Arraygrenzen nicht gleichmäßig sind.

Die zusätzlichen Elemente im Abstand müssen keinen bestimmten Wert enthalten.

Lineare Indexformeln für Kacheln bei einer Form und einer Kachel

Ohne Kachel wird ein Element e=(e_n, e_n-1, ... , e₁) in einem Array mit den Array-Grenzen d=(d_n, d_n-1, ... , d₁) (d1 ist die kleinste Dimension) von der Haupt- bis zur kleineren Reihenfolge an der Position angeordnet:

linear_index(e, d)
= linear_index((e_n, e_n-1, ... , e₁), (d_n, d_n-1, ... , d₁))
= e_nd_n-1...d₁ + d_n-1...d_1-1 ... d

Zur Vereinfachung der Notation in diesem Dokument gehen wir davon aus, dass eine Kachel die gleiche Anzahl von Dimensionen wie das Array hat. Bei der Tile-Implementierung von XLA wird dies auf Kacheln mit weniger Abmessungen verallgemeinert, indem die anfänglichen größten Abmessungen unverändert bleiben und die Kachel nur auf die geringsten Abmessungen angewendet wird, sodass die angegebene Kachel ein Suffix der physischen Abmessungen der zu kachelnden Form enthält.

Bei Verwendung der Kachelgröße (t_n, t_n-1, ... , t₁) wird ein Element im Array mit Indexen (e_n, e_n-1, ..., e₁) dieser Position im endgültigen Layout zugeordnet:

_n/1_n/_nnd_nn(1)_n

Das Layout besteht aus zwei Teilen: (⌊e_n/t_n⌋, ... , ⌊e₁/t₁⌋), die einem Kachelindex in einem Array von Kacheln der Größe (⌈d_n/t_n⌉, ... , ⌋⌉, ... , ⌋⌉, ... , p _{Die Funktion ceil wird in ⌈di/ti⌉ angezeigt, da, wenn Kacheln die Grenzen des größeren Arrays überschreiten, wird das Padding wie in Abbildung 1 eingefügt. Sowohl die Kacheln als auch die Elemente innerhalb der Kacheln werden rekursiv ohne Kacheln angeordnet.}

Im Beispiel in Abbildung 1 hat das Element (2,3) einen Kachelindex (1,1) und einen Index innerhalb der Kacheln (0,1) für einen kombinierten Koordinatenvektor von (1, 1, 0, 1). Die Kachelindexe haben Grenzen (2,3) und die Kachel selbst ist (2,2) für einen kombinierten Vektor von (2,3,2,2). Der lineare Index mit der Kachel für das Element mit dem Index (2,3) in der logischen Form ist dann

linear_index_with_tile((2,3), (3,5), (2,2))
= linear_index((1,1,0,1), (2,3,2,2))
= linear_index((1,1), (2,3)) ∙ 2 ∙ 2 + linear_index((0,1) ∙ 2: (1) ∙ 2 = 1

Tiling als Pad-reshape-transpose

Das kachelbasierte Layout funktioniert so:
Betrachten Sie ein Array von Dimensionen (d_n, d_n-1, ... , d1) (d1 ist die kleinste Dimension). Wenn sie mit Kacheln der Größe (t_n, t_n-1, ... , t₁) angelegt ist (t₁ ist die kleinste Dimension), kann diese Kachel auf folgende Weise beschrieben werden:

1. Das Array wird aufgefüllt auf (⌈d_n/t_n⌉∙t_n, ... , ⌈d₁/t₁⌉∙t₁).
2. Jede Dimension i wird in (⌈d_i/ti⌉, t_i) unterteilt, d.h. das Array wird umgeformt in
   (⌈d_n/t_n⌉, t_n, ... , ⌈d₁/t₁⌉, t₁).
   Bei dieser Umformung selbst gibt es keine physische Layoutänderung. Diese Umformung ist also eine Bitcast-Änderung. Wenn nicht explizit an eine Kachel gedacht wird, könnte diese Umformung jede Form mit der gleichen Anzahl von Elementen wie die aufgefüllte Form ausdrücken. Das Beispiel hier zeigt, wie eine Kachel auf diese Weise ausgedrückt werden kann.
3. Bei einem Transponieren wird t_n, ... , t₁ in die kleinsten Dimensionen verschoben, dabei aber ihre relative Reihenfolge beibehalten, sodass die Reihenfolge der Dimensionen vom größten zum kleinsten Teil zu
   wird (⌈d_n/t_n⌉, ... , ⌈d₁/t₁⌉, t_n).

Die letzte Form hat das Präfix
(⌈d_n/t_n⌉, ... , ⌈d₁/t₁⌉), das die Anzahl der Kacheln in jeder Dimension beschreibt. Ein Element im Array (e_n, ... , e₁) wird diesem Element in der endgültigen Form zugeordnet:
(⌊e_n/t_n⌋, ... , ⌊e₀/t₀⌋, e_n mod t_n, ... , e₁). _{Es ist leicht erkennbar, dass der lineare Index des Elements wie erwartet der obigen Formel folgt.}

Wiederholtes Tiling

Die Tile-Funktion von XLA wird noch flexibler, da sie wiederholt angewendet wird.

Abbildung 2

Abbildung 2 zeigt, wie ein Array der Größe 4 x 8 durch zwei Kacheln gekachelt wird (zuerst 2 x 4, dann 2 x 1). Diese wiederholte Kachel wird als (2;4)(2;1) dargestellt. Jede Farbe steht für eine 2x4-Kachel und jedes rote Rahmenfeld ist eine 2x1-Kachel. Die Zahlen geben den linearen Index für dieses Element im gekachelten Format an. Dieses Format entspricht dem Format, das für BF16 auf TPU verwendet wird, mit der Ausnahme, dass die anfängliche Kachel größer ist, nämlich (8.128)(2,1), wobei der Zweck der zweiten Kachelung durch 2x1 darin besteht, zwei 16-Bit-Werte zu einem 32-Bit-Wert in einer Weise zusammenzuführen, die der Architektur einer TPU entspricht.

Beachten Sie, dass sich eine zweite oder spätere Kachel sowohl auf die Nebenabmessungen innerhalb der Kachel beziehen kann, wodurch die Daten innerhalb der Kachel neu angeordnet werden, wie in diesem Beispiel mit (8,128)(2,1), aber auch auf die Hauptabmessungen innerhalb der Kachel.

Dimensionen mithilfe von Kacheln kombinieren

Mit der Kachel von XLA wird auch die Kombination von Dimensionen unterstützt. Beispielsweise kann sie die Dimensionen in F32[2, 7,8,11,10]{4,3,2,1,0} zuerst in F32[112,110]{1,0} kombinieren,bevor sie mit (2,3) gekachelt wird. Die verwendete Kachel ist (∗,∗,2,∗,3). Hier impliziert ein Sternchen in einer Kachel, dass diese Dimension angenommen und mit der nächsten Nebendimension kombiniert wird. Mehrere benachbarte Dimensionen können in eine Dimension unterteilt werden. Eine subsummierte Dimension wird durch einen Kachelwert von -1 in dieser Dimension der Kachel dargestellt, der ansonsten bei einer Kachel als Dimensionsgröße nicht gültig ist.

Genauer gesagt: Wenn die Dimension i-1 der Form durch ein Sternchen im Kachel entfernt wird, wird diese Dimension, bevor die vorherige Definition der Kacheln angewendet wird, sowohl aus der gekachelten Form als auch aus dem Kachelvektor entfernt. Bei der Dimension i-1 der Form wird die Array-Grenze von d_i-1 auf d_id_i-1 erhöht. Dieser Schritt wird für jedes Sternchen im Kachelvektor wiederholt.