Kształty i układ

Struktura operacji XLA

Rozważmy przykładowy HLO:

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

Składa się z tych komponentów:

Nazwa operacji: add.936
- Jest to unikalna nazwa operacji.
Kształt: bf16[8,1,1280,16384]
- Jest to kształt danych wyjściowych operacji. W tym przypadku typ danych to bf16, a kształt to [8,1,1280,16384].
Układ (z kafelkowaniem): 3,2,0,1:T(8,128)(2,1)
- Opisuje sposób przechowywania tablicy w pamięci. Symbol 3,2,0,1 oznacza kolejność osi w pamięci (np. kolejność kolumn, wierszy itp.), a symbol T(8,128)(2,1) oznacza użyte kafelkowanie i wypełnienie.
- Układ jest opcjonalny. Jeśli nie zostanie określona, nie ma kafelkowania, a wymiary są uporządkowane od najważniejszego do najmniej ważnego.
Działanie: add
- Operacja, która jest wykonywana. W tym przypadku jest to Add, co jest również wymienione w nazwie operacji.
Argumenty: exponential.183, broadcast.3115
- Ta operacja przyjmuje 2 argumenty, które są określane za pomocą unikalnych nazw.

Rozważmy inny przykład, operację fuzji:

%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

Oprócz opisanych wcześniej komponentów obejmuje ona:

Atrybuty: kind i calls
- Zawierają one więcej informacji o wykonywanej operacji, w tym przypadku: fuzji.
Lokalizacja w pamięci (identyfikator miejsca w pamięci): S(1)
- Oznacza to miejsce w pamięci, w którym jest przechowywana tablica. S(1) oznacza, że ta tablica znajduje się w pamięci VMEM (na TPU).
Szczegóły kształtu i układu argumentu wejściowego %fusion.32

W sekcjach poniżej opisujemy kształty, układ i identyfikatory przestrzeni pamięci. Więcej informacji o kafelkowaniu znajdziesz w artykule Układ kafelkowy.

Kształty

Protokół XLA ShapeProto proto (xla_data.proto) określa liczbę wymiarów, rozmiar i typ danych tablicy N-wymiarowej (w skrócie tablicy).

Terminologia, notacja i konwencje

Prawdziwa liczba wymiarów tablicy to liczba wymiarów o rozmiarze większym niż 1.
Wymiary są numerowane od 0 do N-1 w przypadku tablicy N-wymiarowej. Rozmiar wymiaru jest nieujemną liczbą całkowitą. W szczególności rozmiar 0 jest prawidłowy. Numery wymiarów to dowolne etykiety, które ułatwiają korzystanie z tej funkcji. Kolejność tych liczb wymiarów nie oznacza konkretnego porządku mniejszego/większego w układzie kształtu. Układ jest określany przez protokół LayoutProto.
Zgodnie z konwencją wymiary są wymienione w kolejności rosnącej numeru wymiaru. Na przykład w przypadku 3-wymiarowej tablicy o rozmiarze [A x B x C] wymiar 0 ma rozmiar A, wymiar 1 ma rozmiar B, a wymiar 2 ma rozmiar C.

Niektóre narzędzia w XLA obsługują też indeksowanie ujemne podobne do tego w Pythonie: Dimension-1 to ostatni wymiar (odpowiednik N-1 w przypadku tablicy N-wymiarowej). Na przykład w przypadku 3-wymiarowej tablicy opisanej powyżej wymiar –1 ma rozmiar C, wymiar –2 ma rozmiar B itd.
Tablice dwu-, trzy- i czterowymiarowe często mają określone litery powiązane z wymiarami. Na przykład w przypadku tablicy dwuwymiarowej:
- wymiar 0: y
- Wymiar 1: x
W przypadku tablicy 3D:
- wymiar 0: z
- Wymiar 1: y
- wymiar 2: x
W przypadku tablicy 4D:
- wymiar 0: p
- Wymiar 1: z
- wymiar 2: y
- wymiar 3: x
Funkcje w interfejsie XLA API, które przyjmują wymiary, robią to w kolejności rosnącej numeru wymiaru. Odpowiada to kolejności używanej podczas przekazywania wymiarów jako initializer_list, np.

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

utworzy kształt, którego tablica rozmiarów wymiarów będzie zawierać sekwencję [A, B, C, D].

Układ

Protokół LayoutProto opisuje, jak tablica jest reprezentowana w pamięci. Zawiera te pola:

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

Kolejność wymiarów od najmniejszego do największego

Jedynym wymaganym polem jest minor_to_major. To pole opisuje kolejność wymiarów w kształcie od najmniejszego do największego. Wartości w minor_to_major to kolejność wymiarów tablicy (od 0 do N-1 w przypadku tablicy N-wymiarowej), przy czym pierwsza wartość to najmniejszy wymiar, a ostatnia to największy wymiar. Najmniej znaczący wymiar to wymiar, który zmienia się najszybciej podczas przechodzenia przez elementy tablicy ułożonej w pamięci liniowej.

Rozważmy na przykład tę dwuwymiarową tablicę o rozmiarze [2 x 3]:

a b c
d e f

W tym przypadku wymiar 0 ma rozmiar 2, a wymiar 1 ma rozmiar 3. Jeśli pole minor_to_major w układzie ma wartość [0, 1], wymiar 0 jest najmniej znaczącym wymiarem, a wymiar 1 jest najbardziej znaczącym wymiarem. Odpowiada to następującemu układowi w pamięci liniowej:

a d b e c f

Kolejność wymiarów od najmniejszego do największego, czyli od 0 do N-1, jest podobna do kolejności głównej kolumny (w przypadku 2-wymiarowych). Zakładając monotoniczne uporządkowanie wymiarów, możemy odwoływać się do tego układu w kodzie po prostu jako „wymiar 0 jest mniejszy”.

Z drugiej strony, jeśli pole minor_to_major w układzie ma wartość [1, 0], układ w pamięci liniowej wygląda tak:

a b c d e f

Kolejność wymiarów od najmniej do najbardziej znaczącego, czyli od N-1 do 0 w przypadku tablicy N-wymiarowej, jest podobna do kolejności wierszowej (w przypadku tablic 2-wymiarowych). Zakładając monotoniczne uporządkowanie wymiarów, możemy też odwoływać się do tego układu w kodzie jako „wymiar 0 jest główny”.

Domyślne sortowanie od najmniejszej do największej wersji

Domyślny układ nowo utworzonych kształtów to „dimension order is major-to-minor” (czyli [N-1, ..., 0]).

Dopełnienie

Pole tail_padding_alignment_in_elements określa wyrównanie tablicy kafelkowej pod względem liczby elementów. Po zastosowaniu kafelkowania na końcu układu zostaną dodane elementy dopełniające, aż łączna liczba elementów będzie wielokrotnością tej wartości.

Indeksowanie tablic

Klasa IndexUtil w pliku index_util.h zawiera narzędzia do konwertowania indeksów wielowymiarowych na indeksy liniowe na podstawie kształtu i układu. Indeksy wielowymiarowe zawierają indeks int64 dla każdego wymiaru. Indeksy liniowe to pojedyncza wartość int64, która indeksuje bufor zawierający tablicę. W tym samym katalogu znajdziesz narzędzia shape_util.h i layout_util.h, które ułatwiają tworzenie i manipulowanie kształtami oraz układami.

Identyfikatory przestrzeni pamięci

W HLO każda tablica może być opatrzona identyfikatorem przestrzeni pamięci, zapisanym jako S(n).

S(0) (często pomijana) oznacza pamięć o dużej przepustowości (HBM).
S(1) oznacza pamięć wirtualną na urządzeniu (VMEM).
S(2), S(3) itd. odpowiadają dodatkowym obszarom pamięci na urządzeniu.
S(5) – pamięć hosta.