Formen und Layout

Struktur eines XLA-Vorgangs

Hier ein Beispiel für ein HLO:

add.936 = bf16[8,1,1280,16384]{3,2,0,1:T(8,128)(2,1)}
          add(exponential.183, broadcast.3115)

Sie besteht aus den folgenden Komponenten:

Name des Vorgangs: add.936
- Dies ist der eindeutige Name für den Vorgang.
Form: bf16[8,1,1280,16384]
- Dies ist die Ausgabedimension des Vorgangs. Hier ist der Datentyp bf16 und die Dimension [8,1,1280,16384].
Layout (mit Kacheln): 3,2,0,1:T(8,128)(2,1)
- Dies beschreibt, wie das Array im Speicher abgelegt wird. 3,2,0,1 gibt die Reihenfolge der Achsen im Arbeitsspeicher an (z. B. spaltenweise, zeilenweise) und T(8,128)(2,1) gibt die verwendete Kachelung und Auffüllung an.
- Das Layout ist optional. Wenn nicht angegeben, erfolgt keine Kachelung und die Dimensionen werden in der Reihenfolge von der wichtigsten zur unwichtigsten angenommen.
Vorgang: add
- Der ausgeführte Vorgang. Hier ist es Add, was auch im Op-Namen erwähnt wird.
Argumente: exponential.183, broadcast.3115
- Dieser Vorgang akzeptiert zwei Argumente, die mit ihren eindeutigen Namen angegeben werden.

Sehen wir uns ein weiteres Beispiel an, einen Fusionsvorgang:

%fusion.3 = bf16[32,32,4096]{2,1,0:T(8,128)(2,1)S(1)}
            fusion(bf16[32,32,8192]{2,1,0:T(8,128)(2,1)S(1)} %fusion.32),
            kind=kCustom, calls=%all-reduce-scatter.3

Zusätzlich zu den zuvor beschriebenen Komponenten besteht sie aus:

Attribute: kind und calls
- Sie enthalten weitere Informationen zum ausgeführten Vorgang, in diesem Fall zur Zusammenführung.
Speicherort (Speicherplatz-ID): S(1)
- Dies gibt den Speicherplatz bzw. den Speicherort an, an dem das Array gespeichert ist. S(1) bedeutet hier, dass sich dieses Array im VMEM (auf einer TPU) befindet.
Details zu Form und Layout für das Eingabeargument %fusion.32

In den folgenden Abschnitten werden Formen, Layout und Speicherplatz-IDs beschrieben. Weitere Informationen zum Kacheln finden Sie unter Gekacheltes Layout.

Formen

Das XLA-ShapeProto-Proto (xla_data.proto) beschreibt die Anzahl der Dimensionen, die Größe und den Datentyp eines N-dimensionalen Arrays (kurz array).

Terminologie, Notation und Konventionen

Die tatsächliche Anzahl der Dimensionen eines Arrays ist die Anzahl der Dimensionen mit einer Größe von mehr als 1.
Dimensionen werden für ein N-dimensionales Array von 0 bis N-1 nummeriert. Die Größe einer Dimension ist eine nicht negative Ganzzahl. Insbesondere ist die Größe 0 gültig. Die Dimensionsnummern sind beliebige Labels zur besseren Übersicht. Die Reihenfolge dieser Dimensionsnummern impliziert keine bestimmte Reihenfolge von untergeordneten und übergeordneten Dimensionen im Layout der Form. Das Layout wird durch das LayoutProto-Proto bestimmt.
Konventionsgemäß werden Dimensionen in aufsteigender Reihenfolge der Dimensionsnummer aufgeführt. Bei einem dreidimensionalen Array der Größe [A x B x C] hat Dimension 0 beispielsweise die Größe A, Dimension 1 die Größe B und Dimension 2 die Größe C.

Einige Dienstprogramme in XLA unterstützen auch die Python-ähnliche negative Indexierung: Dimension -1 ist die letzte Dimension (entspricht N-1 für ein N-dimensionales Array). Für das oben beschriebene dreidimensionale Array hat die Dimension -1 beispielsweise die Größe C, die Dimension -2 die Größe B usw.
Zwei-, drei- und vierdimensionale Arrays haben oft bestimmte Buchstaben, die mit den Dimensionen verknüpft sind. Beispiel für ein 2D-Array:
- Dimension 0: y
- Dimension 1: x
Für ein 3D-Array:
- Dimension 0: z
- Dimension 1: y
- Dimension 2: x
Für ein 4D-Array:
- Dimension 0: p
- Dimension 1: z
- Dimension 2: y
- Dimension 3: x
Funktionen in der XLA API, die Dimensionen verwenden, tun dies in aufsteigender Reihenfolge der Dimensionsnummer. Dies entspricht der Reihenfolge, die beim Übergeben von Dimensionen als initializer_list verwendet wird, z.B.

ShapeUtil::MakeShape(F32, {A, B, C, D})

erstellt eine Form, deren Dimensionsgrößen-Array aus der Sequenz [A, B, C, D] besteht.

Layout

Das LayoutProto-Proto beschreibt, wie ein Array im Arbeitsspeicher dargestellt wird. Sie enthält die folgenden Felder:

message LayoutProto {
  repeated int64 minor_to_major;
  int64 tail_padding_alignment_in_elements;
  ...
}

Reihenfolge von untergeordneten zu übergeordneten Dimensionen

Das einzige Pflichtfeld ist minor_to_major. In diesem Feld wird die Reihenfolge der Dimensionen innerhalb einer Form von der untergeordneten zur übergeordneten Dimension beschrieben. Die Werte in minor_to_major sind eine Reihenfolge der Dimensionen des Arrays (0 bis N-1 für ein N-dimensionales Array), wobei der erste Wert die am wenigsten wichtige Dimension und der letzte Wert die wichtigste Dimension ist. Die am wenigsten wichtige Dimension ist die Dimension, die sich am schnellsten ändert, wenn die Elemente des Arrays im linearen Speicher durchlaufen werden.

Betrachten Sie beispielsweise das folgende 2D-Array der Größe [2 x 3]:

a b c
d e f

Hier ist die Dimension 0 die Größe 2 und die Dimension 1 die Größe 3. Wenn das Feld minor_to_major im Layout [0, 1] ist, ist die Dimension 0 die am wenigsten wichtige Dimension und die Dimension 1 die wichtigste Dimension. Dies entspricht dem folgenden Layout im linearen Speicher:

a d b e c f

Diese Reihenfolge von untergeordneten zu übergeordneten Dimensionen von 0 bis N-1 entspricht der spaltenweisen Reihenfolge (für 2-dimensionale Daten). Unter der Annahme einer monotonen Anordnung der Dimensionen können wir dieses Layout im Code auch einfach als „dim 0 is minor“ bezeichnen.

Wenn das Feld minor_to_major im Layout hingegen [1, 0] ist, sieht das Layout im linearen Speicher so aus:

a b c d e f

Eine Neben-zu-Haupt-Dimensionsreihenfolge von N-1 bis 0 für ein N-dimensionales Array entspricht zeilenweise (für 2-dimensionale Arrays). Unter der Annahme einer monotonen Anordnung der Dimensionen kann dieses Layout im Code auch einfach als „dim 0 is major“ bezeichnet werden.

Standardmäßige Reihenfolge von untergeordneten zu übergeordneten Elementen

Das Standardlayout für neu erstellte Formen ist „Dimension order is major-to-minor“ ([N-1, ..., 0]).

Abstand

Das Feld tail_padding_alignment_in_elements definiert die Ausrichtung des tiled-Arrays in Bezug auf die Anzahl der Elemente. Nach dem Kacheln werden am Ende des Layouts gepolsterte Elemente hinzugefügt, bis die Gesamtzahl der Elemente ein Vielfaches dieses Werts ist.

Auf Arrays zugreifen

Die Klasse IndexUtil in index_util.h bietet Dienstprogramme zum Konvertieren zwischen mehrdimensionalen und linearen Indexen für eine bestimmte Form und ein bestimmtes Layout. Multidimensionale Indexe enthalten für jede Dimension einen int64-Index. Lineare Indexe sind ein einzelner int64-Wert, mit dem auf den Puffer zugegriffen wird, der das Array enthält. Im selben Verzeichnis finden Sie unter shape_util.h und layout_util.h Hilfsprogramme, die das Erstellen und Bearbeiten von Formen und Layouts vereinfachen.

Speicherplatz-IDs

In HLO kann jedes Array mit einer Speicherplatz-ID versehen werden, die als S(n) geschrieben wird.

S(0) (wird oft weggelassen) steht für den HBM (High Bandwidth Memory) des Geräts.
S(1) steht für den virtuellen Speicher (VMEM) auf dem Gerät.
S(2), S(3) usw. entsprechen zusätzlichen gerätespezifischen Speicherbereichen.
S(5) gibt den Hostspeicher an.