Ta strona została przetłumaczona przez Cloud Translation API.
Switch to English

Układ kafelkowy

Ryc.1

Rysunek 1 pokazuje, jak tablica F32 [3,5] jest ułożona w pamięci z kafelkami 2x2. Kształt z tym układem jest zapisywany jako F32 [3,5] {1,0: T (2,2)}, gdzie 1,0 odnosi się do fizycznej kolejności wymiarów (pole minor_to_major w Layout), podczas gdy (2,2) po dwukropku oznacza kafelkowanie wymiarów fizycznych płytką 2x2.

Płytki są układane intuicyjnie, aby zakryć kształt, a następnie w ramach każdego kafelka elementy są układane bez układania płytek, jak w powyższym przykładzie, gdzie prawa część przykładu przedstawia układ w pamięci, w tym białe elementy wyściółki, które są dodawane aby mieć pełne płytki 2x2, nawet jeśli pierwotne granice tablicy nie są równe.

Dodatkowe elementy wypełnienia nie muszą zawierać żadnej określonej wartości.

Wzory indeksu liniowego dla podziału na płytki dla danego kształtu i płytki

Bez kafelkowania element e = (e n , e n-1 , ..., e 1 ) w tablicy z granicami tablicy d = (d n , d n-1 , ..., d 1 ) (d1 to najmniejszy wymiar) jest ułożony w kolejności od głównego do podrzędnego na pozycji:

linear_index (e, d)
= indeks_liniowy ((e n , e n-1 , ..., e 1 ), (d n , d n-1 , ..., d 1 ))
= e n d n-1 ... d 1 + e n-1 d n-2 ... d 1 + ... + e 1

Dla uproszczenia notacji w tym dokumencie zakładamy, że kafelek ma taką samą liczbę wymiarów jak tablica. W implementacji kafelkowania XLA jest to uogólnione na nachylenia o mniejszych wymiarach, pozostawiając początkowe największe wymiary niezmienione i stosując kafelkowanie tylko do najmniejszych wymiarów, tak że określone kafelkowanie wspomina przyrostek fizycznych wymiarów kształt jest wyłożony kafelkami.

Gdy stosuje się kafelkowanie rozmiaru (t n , t n-1 , ..., t 1 ), element tablicy z indeksami (e n , e n-1 , ..., e 1 ) jest mapowany na to pozycja w finalnym układzie:

linear_index_with_tile (e, d, t)
= linear_index ((⌊e / t⌋, e mod t), (⌈d / t⌉, t)) (arytmetyka jest elementarna, (a, b) jest konkatenacją)
= indeks_liniowy ((⌊e n / t n ⌋, ..., ⌊e 1 / t 1 ⌋, e n mod t n , ..., e 1 mod t 1 ), (⌈d n / t n ⌉, ..., ⌈d 1 / t 1 ⌉, t n , t n-1 , ..., t 1 ))
= indeks_liniowy ((⌊e n / t n ⌋, ..., ⌊e 1 / t 1 ⌋), (⌈d n / t n ⌉, ..., ⌈d 1 / t 1 ⌉)) ∙ t n t n-1 ... t 1 + linear_index ((e n mod t n , ..., e 1 mod t 1 ), (t n , t n-1 , ..., t 1 ))

Układ można traktować jako składający się z dwóch części: (⌊e n / t n ⌋, ..., ⌊e 1 / t 1 ⌋), co odpowiada indeksowi kafelka w tablicy o rozmiarze (⌈d n / t n ⌉, ..., ⌈d 1 / t 1 ⌉) i (e n mod t n , ..., e 1 mod t 1 ), co odpowiada indeksowi wewnątrz kafelka. Funkcja sufitu pojawia się w ⌈d i / t i ⌉, ponieważ jeśli płytki przekraczają granice większej tablicy, wypełnienie jest wstawiane tak, jak na rysunku 1. Zarówno płytki, jak i elementy w płytkach są układane rekurencyjnie bez kafelkowania.

Na przykład na rysunku 1 element (2,3) ma indeks kafelka (1,1) i indeks wewnątrz kafelka (0,1) dla połączonego wektora współrzędnych (1, 1, 0, 1). Indeksy kafli mają granice (2, 3), a sam kafelek to (2, 2) dla połączonego wektora (2, 3, 2, 2). Jest to więc indeks liniowy z płytką dla elementu o indeksie (2, 3) w kształcie logicznym

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,2))
= (1 ∙ 3 + 1) ∙ 2 ∙ 2 + (0 ∙ 2 + 1)
= 17.

Kafelkowanie jako transpozycja pad-reshape

Układ oparty na kafelkach działa w następujący sposób:
Rozważ tablicę wymiarów (d n , d n-1 , ..., d1) (d1 to najmniejszy wymiar). Kiedy jest ułożony z kafelkami o rozmiarze (t n , t n-1 , ..., t 1 ) (t 1 jest najmniejszym wymiarem), to kafelkowanie można opisać w kategoriach transpozycji pad-reshape w następujący sposób sposób.

  1. Tablica jest dopełniana do (⌈d n / t n ⌉ ∙ t n , ..., ⌈d 1 / t 1 ⌉ ∙ t 1 ).
  2. Każdy wymiar i jest podzielony na (⌈d i / t i ⌉, t i ), tj. Tablica jest przekształcana do
    (⌈D n / t n ⌉, t n , ..., ⌈d 1 / t 1 ⌉, t 1 ).
    Nie ma fizycznej zmiany układu w samej tej zmianie kształtu, więc ta zmiana jest bzdurą. Jeśli ktoś nie myśli bezpośrednio o kafelkach, to przekształcenie może wyrazić dowolny kształt z taką samą liczbą elementów, jak kształt wyściełany - przykładem tutaj jest, jak wyrazić kafelek w ten sposób.
  3. Transpozycja odbywa się poprzez przesunięcie t n , ..., t 1 do najmniejszych wymiarów przy zachowaniu ich względnej kolejności, tak że kolejność wymiarów od największego do najmniejszego staje się
    (⌈D n / t n ⌉, ..., ⌈d 1 / t 1 ⌉, t n , ..., t 1 ).

Ostateczny kształt ma przedrostek
(⌈D n / t n ⌉, ..., ⌈d 1 / t 1 ⌉), który opisuje liczbę płytek w każdym wymiarze. Element tablicy (e n , ..., e 1 ) jest mapowany na ten element w ostatecznym kształcie:
(⌊E n / t n ⌋, ..., ⌊e 0 / t 0 ⌋, e n mod t n , ..., e 1 mod t 1 ). Łatwo zauważyć, że wskaźnik liniowy elementu jest zgodny z powyższym wzorem, zgodnie z oczekiwaniami.

Wielokrotne układanie płytek

Płytki XLA stają się jeszcze bardziej elastyczne dzięki wielokrotnemu nakładaniu.

Rysunek 2

Rysunek 2 pokazuje, jak tablica o rozmiarze 4x8 jest podzielona na dwa poziomy kafelków (najpierw 2x4, a następnie 2x1). Przedstawiamy to powtarzające się kafelki jako (2,4) (2,1). Każdy kolor oznacza kafelek 2x4, a każde czerwone pole obramowania to kafelek 2x1. Liczby wskazują liniowy indeks w pamięci tego elementu w formacie kafelkowym. Ten format jest zgodny z formatem używanym dla BF16 na TPU, z wyjątkiem tego, że początkowy kafelek jest większy, a mianowicie kafelkowanie to (8,128) (2,1), gdzie celem drugiego podziału na 2x1 jest zebranie razem dwóch 16-bitowych wartości do tworzą jedną 32-bitową wartość w sposób zgodny z architekturą TPU.

Zwróć uwagę, że drugi lub późniejszy kafelek może odnosić się zarówno do mniejszych wymiarów wewnątrz kafelka, który po prostu zmienia układ danych w kafelku, jak w tym przykładzie z (8,128) (2,1), ale może również odnosić się do głównego elementu krzyżowego wymiary z poprzedniego ułożenia płytek.

Łączenie wymiarów za pomocą płytek

Płytki XLA obsługują również łączenie wymiarów. Na przykład, może najpierw łączyć wymiary w F32 [2,7,8,11,10] {4,3,2,1,0} w F32 [112,110] {1,0} przed umieszczeniem go na kafelkach za pomocą (2,3 ). Użyty kafelek to (∗, ∗, 2, ∗, 3). Tutaj gwiazdka na kafelku oznacza przyjęcie tego wymiaru i połączenie go z kolejnym mniejszym wymiarem. Wiele sąsiednich wymiarów można podzielić na jeden wymiar. Wymiar podliczony jest reprezentowany przez wartość płytki -1 w tym wymiarze płytki, która w innym przypadku nie jest poprawna w płytce jako rozmiar wymiaru.

Dokładniej, jeśli wymiar i kształtu zostanie wyeliminowany za pomocą gwiazdki w kafelku, to przed zastosowaniem wcześniejszej definicji kafelkowania wymiar ten zostanie usunięty zarówno z kształtu, który ma zostać podzielony, jak i z wektora kafelka, a jaki był wymiar i-1 kształtu ma przypisaną tablicę zwiększoną z d i-1 do d i d i-1 . Ten krok jest powtarzany dla każdej gwiazdki w wektorze kafelkowym.