W tym dokumencie opisano analizę indeksowania HLO, która umożliwia symboliczne obliczanie map indeksowania dla operacji HLO. Mapa indeksowania to funkcja, która mapuje indeksy jednego tensora na indeksy innego tensora, np. indeksy wyników instrukcji HLO na indeksy danych wejściowych instrukcji HLO lub odwrotnie.
Przykład
Transmisja od: tensor<20xf32> do: tensor<10x20x30xf32>
p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}
mapa indeksowania od danych wyjściowych do danych wejściowych wynosi \((i, j, k) \mapsto (j)\) dla $i \in [0, 10]\(, \)j \in [0, 20]\( and \)k \in [0, 30]$.
Motywacja
GPU XLA wykorzystuje kilka spersonalizowanych rozwiązań do łączenia, wykorzystywania argumentów i układów kafelkowych (więcej informacji poniżej). Celem analizy indeksowania jest udostępnienie komponentu wielokrotnego użytku na potrzeby takich zastosowań. Analiza indeksowania jest oparta na infrastrukturze Affine Map MLIR i dodaje semantykę HLO.
Połączenie
Rozumowanie dotyczące łączenia pamięci staje się możliwe w nieskomplikowanych przypadkach, gdy wiemy, które elementy lub wycinki danych wejściowych są odczytywane, aby obliczyć element danych wyjściowych.
Wykorzystanie argumentów
Wykorzystanie argumentu w XLA wskazuje, w jakim stopniu poszczególne dane wejściowe instrukcji są używane przy założeniu, że dane wyjściowe zostały w pełni wykorzystane. Obecnie wykorzystanie nie jest też obliczane dla przypadku ogólnego. Analiza indeksowania umożliwia dokładne obliczenie wykorzystania.
Segmentacja:
Kafelek/wycinek to hiperprostokątny podzbiór tensora z parametrami przesunięcia, rozmiaru i kroków. Propagacja kafelków to sposób obliczania parametrów kafelków producenta/konsumenta operacji za pomocą parametrów kafelków samej operacji. Istnieje już biblioteka, która obsługuje funkcje softmax i kropka. Propagacja kafelków może być bardziej ogólna i skuteczna, jeśli odzwierciedla się ją za pomocą map indeksowania.
Funkcja i domena
Mapa indeksowania to funkcja \(\boldsymbol{f}(\boldsymbol{d}, \boldsymbol{s})\) mapująca wiele indeksu \(\boldsymbol{d}\) tensora \(A\) na elementy/zakresy tensora \(B\). Parametr \(\boldsymbol{s}\) odwołuje się do zakresów indeksów wymiarów, które występują w tensorze \(B\), ale nie w tensorze. \(A\)
Jeśli np. zmniejszymy liczbę z tensor<2x4x8x16xf32> do tensor<4x8xf32>, mapa indeksowania z danych wyjściowych 2D na dane wejściowe 4D będzie wyglądać tak:\((d_0, d_1) \mapsto (s_0, d_0, d_1, s_1)\), gdzie \(d_i\) są parametry wymiarów odpowiadające indeksom tensora wyjściowego. Parametry \(s_j\)kodują wiele wartości, np. aby obliczyć \((d_0, d_1)\) element danych wyjściowych, potrzebujemy \((s_0, d_0, d_1, s_1)\) elementów danych wejściowych, gdzie \(s_0 \in [0, 2)\) i\(s_1 \in [0, 16)\).
Mapowanie to można utworzyć na podstawie atrybutów instrukcji HLO. Można też utworzyć mapowania nieskonsolidowanych instrukcji w celu uzyskania indeksowania na potrzeby fuzji. Mapowanie zawiera też domenę, która określa, dla jakich elementów tensora występuje mapowanie.
\[ \begin{eqnarray} \boldsymbol{f}(\boldsymbol{d}, \boldsymbol{s})\; &s.t.& \\ \boldsymbol{lb}_d &\leq& \boldsymbol{d} \leq \boldsymbol{ub}_d \\ \boldsymbol{lb}_s &\leq& \boldsymbol{s} \leq \boldsymbol{ub}_s \\ \boldsymbol{lb}_g &\leq& \boldsymbol{g}(\boldsymbol{d}, \boldsymbol{s}) \leq \boldsymbol{ub}_g \end{eqnarray} \]
Chcemy zminimalizować ponowne obliczanie, dlatego potrzebujemy biblioteki do obliczeń symbolicznych. XLA zależy już od MLIR, więc zamiast tworzyć bibliotekę symboliczną arytmetyczną, używamy mlir::AffineMap.
Typowy AffineMap wygląda tak
(d0)[s0, s1] -> (s0 + 5, d0 * 2, s1 * 3 + 50)
AffineMap zawiera 2 rodzaje parametrów: wymiary i symbole, których można użyć \(\boldsymbol d\) i \(\boldsymbol s\) odpowiadając.
Parametr AffineMap nie zawiera żadnych metadanych dotyczących zakresów wymiarów, więc dane muszą zostać dostarczone sami.
struct Range {
int64_t lower_bound;
int64_t upper_bound;
};
struct IndexingMap {
mlir::AffineMap affine_map;
std::vector<Range> dimension_ranges;
std::vector<Range> symbol_ranges;
llvm::DenseMap<mlir::AffineExpr, Range> expr_ranges;
};
dim_ranges koduje ograniczenia pola włącznie dla parametrów wymiarów \(\boldsymbol{d}\) mapy indeksowania, które zwykle zbiegają się z kształtem tensora danych wyjściowych w przypadku operacji takich jak transponowanie, redukcja, element elementu i kropka. Są jednak pewne wyjątki, np. HloConcatenateInstruction.
symbol_ranges koduje możliwe wartości, które \(\boldsymbol {s}\) mogą przyjmować parametry.
Przeanalizujmy każdy przykład, aby zrozumieć, co faktycznie oznacza wszystkie powyższe informacje.
Indeksowanie map dla nieujednoliconych operacji
Elementwise
W przypadku operacji związanych z elementami mapa indeksowania jest tożsamością.
p0 = f32[10, 20] parameter(0)
p1 = f32[10, 20] parameter(1)
add = f32[10, 20] add(p0, p1)
Dane wyjściowe do wprowadzania map:
- wyjściowe ->input_0: \((d_0, d_1) \mapsto (d_0, d_1)\) dla $\boldsymbol{d} \in [0,9] \times [0, 19]\(, i.e. \)\boldsymbol{d} \in {\rm Dom}(output)$
- wyjściowe ->input_1: \((d_0, d_1) \mapsto (d_0, d_1)\) dla $\boldsymbol{d} \in {\rm Dom} (output)$
Dane wyjściowe map wyjściowych
- wejściowe_i -> dane wyjściowe: \((d_0, d_1) \mapsto (d_0, d_1)\) dla $\boldsymbol{d} \in {\rm Dom}(output)$
Komunikaty
Przesyłanie oznacza, że niektóre wymiary zostaną usunięte, gdy przypiszemy dane wyjściowe do danych wejściowych i dodamy je podczas mapowania danych wejściowych na dane wyjściowe.
p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}
Dane wyjściowe do wpisania mapy:
- wyjściowe -> wejściowe: \((d_0, d_1, d_2) \mapsto (d_1)\) dla $\boldsymbol{d} \in {\rm Dom}(output)$
Dane wejściowe mapy wyjściowej
- wejściowe -> dane wyjściowe: \((d_0) \mapsto (s_0, d_1, s_1)\) for $\boldsymbol{d} \in {\rm Dom}(output)\( and \)\boldsymbol{s} \in [0, 9] \times [0, 29]$.
Zwróć uwagę, że teraz mamy \(\boldsymbol s\) po prawej stronie mapowania danych wejściowych i wyjściowych. Te symbole reprezentują zakresy wartości. Na przykład w tym przypadku każdy element danych wejściowych z indeksem \(d_0\) jest zmapowany na wycinek danych wyjściowych o wymiarach 10 × 1 × 30.
Stała i Iota
Nie mają żadnych parametrów wejściowych, więc nie ma potrzeby ich obliczania.
Transponuj
Mapa indeksowania na potrzeby transponowania to permutacja wymiarów wejściowych i wyjściowych.
p0 = f32[3, 12288, 6, 128] parameter(0)
transpose = f32[3, 6, 128, 12288] transpose(p0), dimensions={0, 2, 3, 1}
Dane wyjściowe do wpisania mapy:
- dane wyjściowe -> dane wejściowe: \((d_0, d_1, d_2, d_3) \mapsto (d_0, d_3, d_1, d_2)\) dla\(\boldsymbol{d} \in {\rm Dom}(output)\)
Dane wejściowe mapy wyjściowej:
- wejście -> dane wyjściowe: \((d_0, d_1, d_2, d_3) \mapsto (d_0, d_2, d_3, d_1)\) dla\(\boldsymbol{d} \in {\rm Dom}(input)\)
Wsad za plecami
Mapa indeksowania dla odwrotnej zmiany zmienia przywrócone wymiary na $upper_bound(d_i) – d_i$:
p0 = f32[1, 17, 9, 9] parameter(0)
reverse = f32[1, 17, 9, 9] reverse(p0), dimensions={1, 2}
Dane wyjściowe do wpisania mapy:
- dane wyjściowe -> wejściowe: $(d_0, d_1, d_2, d_3) \mapsto (d_0, -d_1 + 16, -d_2 + 8, d_3)\( for \)\boldsymbol{d} \in {\rm Dom}(output)$
Dane wejściowe mapy wyjściowej:
- wejściowe -> dane wyjściowe: $(d_0, d_1, d_2, d_3) \mapsto (d_0, -d_1 + 16, -d_2 + 8, d_3)\( for \)\boldsymbol{d} \in {\rm Dom}(input)$
(Variadic)Reduce
Redukcja variadyczna ma kilka wejść i kilka inicjacji, a mapa między danymi wyjściowymi a danymi wejściowymi dodaje mniejsze wymiary. Dlatego w pewnym sensie działa ona jak odwrotność transmisji.
p0 = f32[256,10] parameter(0)
p0_init = f32[] constant(-inf)
p1 = s32[256,10] parameter(1)
p1_init = s32[] constant(0)
reduce = (f32[10], s32[10]) reduce(p0, p1, p0_init, p1_init),
dimensions={0}, to_apply=min
Dane wyjściowe do wprowadzania map:
- wyjściowe -> insert_j: \((d_0) \mapsto (s_0, d_0)\) dla $\boldsymbol{d} \in {\rm Dom}(output)\( and \)\boldsymbol{s} \in [0, 9]$
- wyjściowe -> init_j: \((d_0) \mapsto ()\) dla $\boldsymbol{d} \in {\rm Dom}(output)$
Dane wejściowe map wyjściowych:
- wejściowe_i -> wyjściowe_j: \((d_0, d_1) \mapsto (d_1)\) dla $\boldsymbol{d} \in {\rm Dom}(input)$
- init_i -> wyjściowy_j: \(() \mapsto (s_0)\) dla \(\boldsymbol{s} \in [0, 9]\)
dla domeny \(i, j = 0, \ldots, INPUT\\_COUNT\).
Wycinek
Indeksowanie od danych wyjściowych do danych wejściowych na potrzeby wycinka oznacza skróconą mapę indeksowania, która jest poprawna w przypadku każdego elementu danych wyjściowych. Mapowanie danych wejściowych na dane wyjściowe jest ograniczone do ograniczonego zakresu elementów danych wejściowych.
p0 = f32[10, 20, 50] parameter(0)
slice = f32[5, 3, 25] slice(f32[10, 20, 50] p0),
slice={[5:10:1], [3:20:7], [0:50:2]}
Dane wyjściowe do wpisania mapy:
- dane wyjściowe -> dane wejściowe: \((d_0, d_1, d_2) \mapsto (d_0 + 5, 7d_1 + 3, 2d_2)\) dla\(\boldsymbol{d} \in {\rm Dom}(output)\)
Dane wejściowe mapy wyjściowej:
- wejściowe -> dane wyjściowe: \((d_0, d_1, d_2) \mapsto (d_0, d_1 / 7, d_2 / 2)\) dla\(\boldsymbol{d} \in [5, 9] \times [3, 19] \times [0, 49]\) z krokami $[1, 7, 2]$.
Do ustalenia: indeksowanie między danymi wejściowymi
Zmień kształt
Modyfikacje mają różne style.
Zwiń kształt
To jest „liniowość” przekształcenie z N-D na 1D.
p0 = f32[4,8] parameter(0)
reshape = f32[32] reshape(p0)
Dane wyjściowe do wpisania mapy:
- dane wyjściowe -> wejściowe: \((d_0) \mapsto (d_0 / 8, d_0 \mod 8)\) dla $\boldsymbol{d} \in {\rm Dom}(output)$
Dane wejściowe mapy wyjściowej:
- wejściowe -> dane wyjściowe: \((d_0, d_1) \mapsto (8 d_0 + d_1)\) dla $\boldsymbol{d} \in {\rm Dom}(input)$.
Rozwiń kształt
Jest to odwrotna operacja „zwinięcia kształtu”. Przekształca dane wejściowe 1D w dane wyjściowe N-D.
p0 = f32[32] parameter(0)
reshape = f32[4, 8] reshape(p0)
Dane wyjściowe do wpisania mapy:
- wyjściowe -> wejściowe: \((d_0, d_1) \mapsto (8 d_0 + d_1)\) dla $\boldsymbol{d} \in {\rm Dom}(output)$
Dane wejściowe mapy wyjściowej:
- wejściowe -> dane wyjściowe: \((d_0) \mapsto (d_0 / 8, d_0 \mod 8)\) dla $\boldsymbol{d} \in {\rm Dom}(input)$.
Ogólna zmiana
Są to operacje zmiany kształtu, których nie można reprezentować jako pojedynczego kształtu rozwinięcia lub zwinięcia. Mogą być reprezentowane jako kompozycja z co najmniej 2 kształtów rozwijania lub zwijania.
Przykład 1. Linearizacja-delinearizacja.
p0 = f32[4,8] parameter(0)
reshape = f32[2, 4, 4] reshape(p0)
Zmiana kształtu może być przedstawiona jako kompozycja zwiniętego kształtu o wymiarach tensor<4x8xf32> do tensor<32xf32>, a następnie rozwinięcie kształtu do tensor<2x4x4xf32>.
Dane wyjściowe do wpisania mapy:
- dane wyjściowe -> dane wejściowe: $(d_0, d_1, d_2) \mapsto (2d_0 + (4d_1 + d_2) / 8, 4d_1 + d_2) \mod 8)$
przez \(\boldsymbol{d} \in {\rm Dom}(output)\)
Dane wejściowe mapy wyjściowej:
- wejściowe -> dane wyjściowe: $(d_0, d_1) \mapsto ((8d_0 + d_1) / 16, ((8d_0 + d_1) \mod 16) / 4, d_1 \mod 4)$
dla domeny \(\boldsymbol{d} \in {\rm Dom}(input)\).
Przykład 2. Rozwinięte i zwinięte podkształty
p0 = f32[4, 8, 12] parameter(0)
reshape = f32[32, 3, 4] reshape(p0)
Takie zmiany można przedstawić jako kompozycję 2 kształtów. Pierwszy zwija
skrajne wymiary tensor<4x8x12xf32> do tensor<32x12xf32>, a drugi tensor<32x12xf32>
do wartości tensor<32x3x4xf32>.
Dane wyjściowe do wpisania mapy:
- dane wyjściowe -> dane wejściowe: \((d_0, d_1, d_2) \mapsto (d_0 / 8, d_0 \mod 8, 4d_1 + d_2)\) dla \(\boldsymbol{d} \in {\rm Dom}(output)\)
Dane wejściowe mapy wyjściowej:
- wejściowe -> dane wyjściowe: \((d_0, d_1, d_2) \mapsto (8d_0 + d_1, d_2 / 4, d_2 \mod 4)\) dla \(\boldsymbol{d} \in {\rm Dom}(input)\).
Bitcast
operacja bitcastu może być reprezentowana jako sekwencja transponowania-zmiany kształtu. Dlatego mapy indeksowania są tylko kompozycją map indeksowania dla tej sekwencji.
Połącz
Mapowanie danych wyjściowych na dane wejściowe na potrzeby concat jest zdefiniowane dla wszystkich danych wejściowych, ale w przypadku niepokrywających się domen, tj. w danym momencie będzie używana tylko jedna z tych wartości.
p0 = f32[3,50] parameter(0)
p1 = f32[3,30] parameter(1)
concat = f32[3,80] concatenate(f32[3,50] p0, f32[3,30] p1),
dimensions={1}
Dane wyjściowe do wpisania mapy:
- dane wyjściowe -> dane wejściowe 1:
\((d_0, d_1) \mapsto (d_0, d_1)\) przez \(\boldsymbol{d} \in [0, 2] \times [0, 49]\)
- dane wyjściowe -> dane wejściowe 2:
\((d_0, d_1) \mapsto (d_0, d_1 - 50)\) for $\boldsymbol{d} \in [0, 2] \times [50, 79]$
Dane wejściowe do utworzenia mapy wyjściowej:
- wejściowe 1 -> dane wyjściowe: \((d_0, d_1) \mapsto (d_0, d_1)\) dla $\boldsymbol{d} \in {\rm Dom}(input_1)$.
- wejściowe 2 -> dane wyjściowe: \((d_0, d_1) \mapsto (d_0, d_1 + 50)\) dla $\boldsymbol{d} \in {\rm Dom}(input_2)$.
Kropka (implementacja danych wyjściowych na dane wejściowe
Mapy indeksowania dla punktów są bardzo podobne do map redukcji.
p0 = f32[4, 128, 256] parameter(0)
p1 = f32[4, 256, 64] parameter(1)
dot = f32[4, 128, 64] dot(p0, p1),
lhs_batch_dims={0}, rhs_batch_dims={0},
lhs_contracting_dims={2}, rhs_contracting_dims={1}
Dane wyjściowe map są mapowane:
- dane wyjściowe ->input_1: \((d_0, d_1, d_2) \mapsto (d_0, d_1, s_0)\) dla\(\boldsymbol{d} \in {\rm Dom}(output)\) i \(\boldsymbol{s} \in [0, 255]\)
- dane wyjściowe -> data_2: \((d_0, d_1, d_2) \mapsto (d_0, s_0, d_2)\) dla\(\boldsymbol{d} \in {\rm Dom}(output)\) i \(\boldsymbol{s} \in [0, 255]\)
Dane wejściowe map wyjściowych:
- wejściowe_1 -> dane wyjściowe: \((d_0, d_1, d_2) \mapsto (d_0, d_1, s_0)\) dla\(\boldsymbol{d} \in {\rm Dom}(input_1)\) i \(\boldsymbol{s} \in [0, 63]\)
- wejściowe_2 -> dane wyjściowe: \((d_0, d_1, d_2) \mapsto (d_0, s_0, d_1)\) dla\(\boldsymbol{d} \in {\rm Dom}(input_2)\) i \(\boldsymbol{s} \in [0, 127]\)
Zmniejsz okno (do ustalenia)
Pad (do ustalenia)
Indeksowanie map dla Fusion
Mapa indeksowania dla operacji fusion to kompozycja map indeksowania dla każdej operacji w klastrze. Może się zdarzyć, że niektóre dane wejściowe będą odczytywane kilka razy z zastosowaniem różnych wzorców dostępu.
Jedno dane wejściowe, kilka map indeksowania
Oto przykład dla: \(p_0 + p_0^T\)
f {
p0 = f32[1000, 1000] parameter(0)
transpose_p0 = f32[1000, 1000]{0, 1} transpose(p0), dimensions={1, 0}
ROOT a0 = f32[1000, 1000] add(p0, transpose_p0)
}
Mapy indeksowania danych wyjściowych na dane wejściowe w przypadku p0 mają postać $(d_0, d_1) \mapsto (d_0, d_1)\( and \)(d_0, d_1) \mapsto (d_1, d_0)$. Oznacza to, że podczas obliczenia 1 elementu danych wyjściowych będziemy musieli dwukrotnie odczytać parametr wejściowy.
1 dane wejściowe – mapa indeksowania bez duplikatów
W niektórych przypadkach mapy indeksowania są w rzeczywistości takie same, chociaż nie są to od razu oczywiste.
f {
p0 = f32[20, 10, 50] parameter(0)
lhs_transpose_1 = f32[10, 20, 50] transpose(p0), dimensions={1, 0, 2}
lhs_e = f32[10, 20, 50] exponential(lhs_transpose_1)
lhs_transpose_2 = f32[10, 50, 20] transpose(lhs_e), dimensions={0, 2, 1}
rhs_transpose_1 = f32[50, 10, 20] transpose(p0), dimensions={2, 1, 0}
rhs_log = f32[50, 10, 20] exponential(rhs_transpose_1)
rhs_transpose_2 = f32[10, 50, 20] transpose(rhs_log), dimensions={1, 0, 2}
ROOT add = f32[10, 50, 20] add(lhs_transpose_2, rhs_transpose_2)
}
Mapa indeksowania danych wyjściowych na dane wejściowe dla funkcji p0 w tym przypadku to tylko $(d_0, d_1, d_2)
\mapsto (d_2, d_0, d_1)$.
Softmax
Mapy indeksowania danych wyjściowych na dane wejściowe dla parameter 0 na potrzeby softmax:
- \((d_0, d_1, d_2) \mapsto (d_0, d_1, d_2)\)
- \((d_0, d_1, d_2)[s_0] \mapsto (d_0, d_1, s_0)\)
\(\boldsymbol{d} \in {\rm Dom}(output)\) i \(\boldsymbol{s} \in [0, 124]\)odwołują się do najbardziej wewnętrznego wymiaru danych wejściowych.
Upraszczanie map indeksowania
Domyślny upraszczający dla zapytania mlir::AffineMap na wyższym poziomie nie może przyjmować żadnych założeń dotyczących zakresów wymiarów/symboli. Dlatego nie można skutecznie uprościć wyrażeń za pomocą mod i div.
Możemy wykorzystać wiedzę o dolnych i górnych granicach podwyrażeń w mapach afinicznych, aby je jeszcze bardziej uprościć.
Upraszczanie może przeredagować poniższe wyrażenia.
- \((d_0, d_1) \mapsto (d_0 + d1 / 16, d1 \mod 16)\) for $\boldsymbol{d} \in [0, 6] \times [0, 14]\( becomes \)(d_0, d_1) \mapsto (d_0, d_1)$
- $(d_0, d_1, d_2) \mapsto ((100d_0 + 10d_1 + d_2) /100, ((100d_0 + 10d_1 + d_2) \mod 100) / 10, d_2 \mod 10) [d_mapy, d_1, 0, d_1, 0, d_2, 0, d_2, 0,\( for \)\( becomes \)
- $(d_0, d_1, d_2) \mapsto ((16d_0 + 4d_1 + d_2) /8, (16d_0 + 4d_1 + d_2) \mod 8)\( for \)d_i \in [0, 9]\( becomes \)(d_0, d_mod, d_2/2)
- \((d_0, d_1) \mapsto (-(-11d_0 - d_1 + 109) / 11 + 9)\) for $\boldsymbol{d} \in [0, 9] \times [0, 10]\( becomes \)(d_0, d_1) \mapsto (d_0)$.
Dzięki uproszczeniu map indeksowania możemy zrozumieć, że niektóre łańcuchowe zmiany kształtów w HLO nakładają się na siebie.
p0 = f32[10, 10, 10] parameter(0)
reshape1 = f32[50, 20] reshape(p0)
reshape2 = f32[10, 10, 10] reshape(reshape1)
Po stworzeniu map indeksowania i ich uproszczeniu uzyskamy
\((d_0, d_1, d_2) \mapsto (d_0, d_1, d_2)\).