Questa pagina è stata tradotta dall'API Cloud Translation.

Analisi dell'indicizzazione

Questo documento descrive l'analisi di indicizzazione HLO, che consente di calcolare simbolicamente le mappe di indicizzazione per le operazioni HLO. La mappa di indicizzazione è una funzione che mappa gli indici di un tensore agli indici di un altro, ad esempio gli indici di un'istruzione HLO a indici di input dell'istruzione HLO o viceversa.

Esempio

Per una trasmissione da tensor<20xf32> a tensor<10x20x30xf32>

p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}

la mappa di indicizzazione dall'output all'input è $(i, j, k) \mapsto (j)$ per $i \in [0, 10]$, $j \in [0, 20]$ and $k \in [0, 30]$.

Motivazione

La GPU XLA utilizza diverse soluzioni su misura per motivi di coalescenza, utilizzo degli operandi e schemi di tiling (maggiori dettagli di seguito). L'obiettivo dell'analisi dell'indicizzazione è fornire un componente riutilizzabile per questi casi d'uso. L'analisi dell'indicizzazione si basa sull'infrastruttura Affine Map di MLIR e aggiunge la semantica HLO.

Coalescenza

Ragionare sulla coalescenza della memoria diventa utilizzabile per casi non banali, quando sappiamo quali elementi/sezioni degli input vengono letti per calcolare un elemento dell'output.

Utilizzo degli operandi

L'utilizzo dell'operando in XLA indica quanto viene utilizzato ciascun input dell'istruzione, supponendo che il relativo output sia completamente utilizzato. Attualmente, l'utilizzo non viene calcolato per un caso generico. L'analisi dell'indicizzazione consente di calcolare con precisione l'utilizzo.

Riquadri:

Un riquadro/sezione è un sottoinsieme iper-rettangolare di un tensore parametrizzato da offset, dimensioni e passi. La propagazione dei riquadri è un modo per calcolare i parametri dei riquadri del produttore/consumatore dell'operazione utilizzando i parametri di tiling dell'operazione stessa. Esiste già una libreria che lo fa per softmax e dot. La propagazione dei riquadri può essere resa più generica e efficace se viene espressa tramite mappe di indicizzazione.

Funzione e dominio

La mappa di indicizzazione è una funzione $\boldsymbol{f}(\boldsymbol{d}, \boldsymbol{s})$ che mappa un multi-indice $\boldsymbol{d}$ di un tensore $A$ a elementi/intervalli di tensore $B$. Il parametro $\boldsymbol{s}$ fa riferimento agli intervalli di indici delle dimensioni che sono presenti nel tensore $B$, ma non nel tensore $A$.

Ad esempio, se otteniamo una riduzione da tensor<2x4x8x16xf32> a tensor<4x8xf32>, la mappa di indicizzazione dall'output 2D all'input 4D è$(d_0, d_1) \mapsto (s_0, d_0, d_1, s_1)$, dove $d_i$ sono i parametri di dimensione che corrispondono agli indici del tensore di output. I parametri $s_j$ codificano più valori, ovvero per calcolare un $(d_0, d_1)$ elemento dell'output sono necessari $(s_0, d_0, d_1, s_1)$ elementi dell'input, dove $s_0 \in [0, 2)$ e $s_1 \in [0, 16)$.

Questa mappatura può essere creata dagli attributi delle istruzioni HLO oppure puoi comporre le mappature delle istruzioni non fuse per ottenere l'indicizzazione per una fusione. La mappatura ha anche un dominio, che specifica quali elementi del tensore esistono.

\[ \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} \]

Poiché vogliamo ridurre al minimo il ricalcolo, abbiamo bisogno di una libreria per i calcoli simbolici. La XLA dipende già da MLIR, quindi utilizziamo mlir::AffineMap anziché scrivere una libreria aritmetica simbolica.

Un tipico AffineMap sembra

(d0)[s0, s1] -> (s0 + 5, d0 * 2, s1 * 3 + 50)

AffineMap ha comodamente due tipi di parametri: dimensioni e simboli per cui possiamo utilizzare $\boldsymbol d$ e $\boldsymbol s$ rispettivamente. AffineMap non contiene metadati relativi a intervalli di dimensioni, quindi dobbiamo fornire questi dati noi stessi.

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 codifica i vincoli casella inclusivi per i parametri di dimensione $\boldsymbol{d}$ della mappa di indicizzazione, che di solito coincidono con la forma del tensore di output per operazioni come traspose, reduce, elementwise, dot, ma esistono alcune eccezioni come HloConcatenateInstruction.

symbol_ranges codifica i possibili valori che i $\boldsymbol {s}$ parametri possono assumere.

Analizziamo per esempio per capire cosa significa in realtà tutto quanto riportato sopra.

Indicizzazione delle mappe per unfused Ops

Elementwise

Per un'operazione a livello di elemento, la mappa di indicizzazione è un'identità.

  p0 = f32[10, 20] parameter(0)
  p1 = f32[10, 20] parameter(1)
  add = f32[10, 20] add(p0, p1)

L'output per le mappe di input:

output -> input_0: $(d_0, d_1) \mapsto (d_0, d_1)$ per $\boldsymbol{d} \in [0,9] \times [0, 19]$, i.e. $\boldsymbol{d} \in {\rm Dom}(output)$
output -> input_1: $(d_0, d_1) \mapsto (d_0, d_1)$ per $\boldsymbol{d} \in {\rm Dom} (output)$

L'input per le mappe di output

input_i -> output: $(d_0, d_1) \mapsto (d_0, d_1)$ per $\boldsymbol{d} \in {\rm Dom}(output)$

Trasmetti un annuncio

La trasmissione significa che alcune dimensioni verranno rimosse quando mapperemo l'output all'input e verranno aggiunte quando mapperemo l'input all'output.

p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}

L'output della mappa di input:

output -> input: $(d_0, d_1, d_2) \mapsto (d_1)$ per $\boldsymbol{d} \in {\rm Dom}(output)$

La mappa di input per output

input -> output: $(d_0) \mapsto (s_0, d_1, s_1)$ per $\boldsymbol{d} \in {\rm Dom}(output)$ and $\boldsymbol{s} \in [0, 9] \times [0, 29]$.

Nota che ora abbiamo $\boldsymbol s$ sul lato destro per la mappatura input-output. Questi sono i simboli che rappresentano intervalli di valori. Ad esempio, in questo caso particolare ogni elemento di input con indice $d_0$ è mappato a una sezione 10 x 1 x 30 dell'output.

Costante e Iota

Praticamente non hanno parametri di input, quindi non c'è nulla di cui calcolare l'indicizzazione.

Transpose

La mappa di indicizzazione per la trasposizione è una permutazione delle dimensioni di input/output.

p0 = f32[3, 12288, 6, 128] parameter(0)
transpose = f32[3, 6, 128, 12288] transpose(p0), dimensions={0, 2, 3, 1}

L'output della mappa di input:

output -> input: $(d_0, d_1, d_2, d_3) \mapsto (d_0, d_3, d_1, d_2)$ per $\boldsymbol{d} \in {\rm Dom}(output)$

La mappa di input per l'output:

input -> output: $(d_0, d_1, d_2, d_3) \mapsto (d_0, d_2, d_3, d_1)$ per $\boldsymbol{d} \in {\rm Dom}(input)$

All'indietro

Mappa di indicizzazione per modifiche inverse alle dimensioni ripristinate a $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}

L'output della mappa di input:

output -> input: $(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)$

La mappa di input per l'output:

input -> output: $(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)$

(Variadica)Riduci

La riduzione variadica ha diversi input e diversi init, la mappa dall'output all'input aggiunge le dimensioni ridotte. Quindi, in un certo senso si comporta come l'opposto di una trasmissione.

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

L'output per le mappe di input:

output -> input_j: $(d_0) \mapsto (s_0, d_0)$ per $\boldsymbol{d} \in {\rm Dom}(output)$ and $\boldsymbol{s} \in [0, 9]$
output -> init_j: $(d_0) \mapsto ()$ per $\boldsymbol{d} \in {\rm Dom}(output)$

L'input per le mappe di output:

input_i -> output_j: $(d_0, d_1) \mapsto (d_1)$ per $\boldsymbol{d} \in {\rm Dom}(input)$
init_i -> output_j: $() \mapsto (s_0)$ per $\boldsymbol{s} \in [0, 9]$

per $i, j = 0, \ldots, INPUT\\_COUNT$.

Sezione

L'indicizzazione dall'output all'input per la sezione genera una mappa di indicizzazione a righe valida per ogni elemento dell'output. La mappatura dall'input all'output è limitata a un intervallo a righe di elementi nell'input.

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]}

L'output della mappa di input:

output -> input: $(d_0, d_1, d_2) \mapsto (d_0 + 5, 7d_1 + 3, 2d_2)$ per $\boldsymbol{d} \in {\rm Dom}(output)$

La mappa di input per l'output:

input -> output: $(d_0, d_1, d_2) \mapsto (d_0, d_1 / 7, d_2 / 2)$ per $\boldsymbol{d} \in [5, 9] \times [3, 19] \times [0, 49]$ con passi di $[1, 7, 2]$.

Da definire: indicizzazione input-to-output

Rimodella

Sono disponibili diversi sapori.

Comprimi forma

Si tratta di una rimodellazione "linearizzazione" da N-D a 1D.

p0 = f32[4,8] parameter(0)
reshape = f32[32] reshape(p0)

L'output della mappa di input:

output -> input: $(d_0) \mapsto (d_0 / 8, d_0 \mod 8)$ per $\boldsymbol{d} \in {\rm Dom}(output)$

La mappa di input per l'output:

input -> output: $(d_0, d_1) \mapsto (8 d_0 + d_1)$ per $\boldsymbol{d} \in {\rm Dom}(input)$.

Espandi forma

Si tratta di un'operazione inversa per la "forma di compressione", che trasforma un input 1D nell'output N-D.

p0 = f32[32] parameter(0)
reshape = f32[4, 8] reshape(p0)

L'output della mappa di input:

output -> input: $(d_0, d_1) \mapsto (8 d_0 + d_1)$ per $\boldsymbol{d} \in {\rm Dom}(output)$

La mappa di input per l'output:

input -> output: $(d_0) \mapsto (d_0 / 8, d_0 \mod 8)$ per $\boldsymbol{d} \in {\rm Dom}(input)$.

Rimodellazione generica

Si tratta delle operazioni di rimodellazione che non possono essere rappresentate come un'unica forma di espansione o compressione. Possono essere rappresentate solo come una composizione di due o più forme di espansione o compressione.

Esempio 1: Linearizzazione-delinearizzazione.

p0 = f32[4,8] parameter(0)
reshape = f32[2, 4, 4] reshape(p0)

Questa riforma può essere rappresentata come una composizione della forma di compressione da tensor<4x8xf32> a tensor<32xf32> e poi un'espansione della forma in tensor<2x4x4xf32>.

L'output della mappa di input:

output -> input: $(d_0, d_1, d_2) \mapsto (2d_0 + (4d_1 + d_2) / 8, 4d_1 + d_2) \mod 8)$

per $\boldsymbol{d} \in {\rm Dom}(output)$

La mappa di input per l'output:

input -> output: $(d_0, d_1) \mapsto ((8d_0 + d_1) / 16, ((8d_0 + d_1) \mod 16) / 4, d_1 \mod 4)$

per $\boldsymbol{d} \in {\rm Dom}(input)$.

Esempio 2: forme secondarie espanse e compresse

p0 = f32[4, 8, 12] parameter(0)
reshape = f32[32, 3, 4] reshape(p0)

Questa rimodellazione può essere rappresentata come una composizione di due riforme. La prima comprime le dimensioni più esterne tensor<4x8x12xf32> in tensor<32x12xf32> e la seconda espande la dimensione più interna tensor<32x12xf32> in tensor<32x3x4xf32>.

L'output della mappa di input:

output -> input: $(d_0, d_1, d_2) \mapsto (d_0 / 8, d_0 \mod 8, 4d_1 + d_2)$ per $\boldsymbol{d} \in {\rm Dom}(output)$

La mappa di input per l'output:

input -> output: $(d_0, d_1, d_2) \mapsto (8d_0 + d_1, d_2 / 4, d_2 \mod 4)$ per $\boldsymbol{d} \in {\rm Dom}(input)$.

Bitcast

Un'operazione bitcast può essere rappresentata come sequenza di trasposizione-rimodellazione-trasposizione. Di conseguenza, le relative mappe di indicizzazione sono solo una composizione di mappe di indicizzazione per questa sequenza.

Concatena

La mappatura da output a input per concat è definita per tutti gli input, ma con domini non sovrapposti, ovvero verrà utilizzato un solo input alla volta.

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}

L'output della mappa di input:

output -> input 1:

$(d_0, d_1) \mapsto (d_0, d_1)$ per $\boldsymbol{d} \in [0, 2] \times [0, 49]$

output -> input 2:

$(d_0, d_1) \mapsto (d_0, d_1 - 50)$ per $\boldsymbol{d} \in [0, 2] \times [50, 79]$

La mappa degli input da restituire:

input 1 -> output: $(d_0, d_1) \mapsto (d_0, d_1)$ per $\boldsymbol{d} \in {\rm Dom}(input_1)$.
input 2 -> output: $(d_0, d_1) \mapsto (d_0, d_1 + 50)$ per $\boldsymbol{d} \in {\rm Dom}(input_2)$.

Punto (da output a input implementato

Le mappe di indicizzazione per il punto sono molto simili a quelle di riduzione.

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}

L'output e gli input vengono mappati:

output -> input_1: $(d_0, d_1, d_2) \mapsto (d_0, d_1, s_0)$ per $\boldsymbol{d} \in {\rm Dom}(output)$ e $\boldsymbol{s} \in [0, 255]$
output -> input_2: $(d_0, d_1, d_2) \mapsto (d_0, s_0, d_2)$ per $\boldsymbol{d} \in {\rm Dom}(output)$ e $\boldsymbol{s} \in [0, 255]$

Gli input per le mappe di output:

input_1 -> output: $(d_0, d_1, d_2) \mapsto (d_0, d_1, s_0)$ per $\boldsymbol{d} \in {\rm Dom}(input_1)$ e $\boldsymbol{s} \in [0, 63]$
input_2 -> output: $(d_0, d_1, d_2) \mapsto (d_0, s_0, d_1)$ per $\boldsymbol{d} \in {\rm Dom}(input_2)$ e $\boldsymbol{s} \in [0, 127]$

Riduzione della finestra (da definire)

Cuscinetto (da definire)

Indicizzazione di Maps per Fusion

Mappa di indicizzazione per operazioni di fusione è una composizione di mappe di indicizzazione per ogni operazione nel cluster. Può accadere che alcuni input vengano letti più volte con pattern di accesso diversi.

Un input, diverse mappe di indicizzazione

Ecco un esempio per $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)
}

Le mappe di indicizzazione da output a input per p0 saranno $(d_0, d_1) \mapsto (d_0, d_1)$ and $(d_0, d_1) \mapsto (d_1, d_0)$. Significa che per calcolare un elemento dell'output potrebbe essere necessario leggere il parametro di input due volte.

Una mappa di indicizzazione deduplicata di input

A volte, le mappe di indicizzazione sono effettivamente le stesse, anche se non sono immediatamente evidenti.

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)
}

La mappa di indicizzazione da output a input per p0 in questo caso è solo $(d_0, d_1, d_2) \mapsto (d_2, d_0, d_1)$.

Softmax

L'indicizzazione da output a input viene mappata per parameter 0 per la funzionalità 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)$

for $\boldsymbol{d} \in {\rm Dom}(output)$ e $\boldsymbol{s} \in [0, 124]$ si riferisce alla dimensione più interna dell'input.

Semplificatore della mappa di indicizzazione

Il semplificatore predefinito per mlir::AffineMap upstream non può fare alcuna ipotesi sugli intervalli di dimensioni/simboli. Di conseguenza, non può semplificare le espressioni con mod ed divin modo efficiente.

Possiamo sfruttare la conoscenza dei limiti inferiore e superiore delle sottoespressioni nelle mappe affine per semplificarle ancora di più.

Il semplificatore è in grado di riscrivere le seguenti espressioni.

$(d_0, d_1) \mapsto (d_0 + d1 / 16, d1 \mod 16)$ per $\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 1\0]$ for $d_i \d_1$ 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_1, d_4) \
$(d_0, d_1) \mapsto (-(-11d_0 - d_1 + 109) / 11 + 9)$ per $\boldsymbol{d} \in [0, 9] \times [0, 10]$ becomes $(d_0, d_1) \mapsto (d_0)$.

La semplificazione delle mappe di indicizzazione ci consente di capire che alcune delle rimodelle concatenate in HLO si annullano a vicenda.

p0 = f32[10, 10, 10] parameter(0)
reshape1 = f32[50, 20] reshape(p0)
reshape2 = f32[10, 10, 10] reshape(reshape1)

Dopo la composizione delle mappe di indicizzazione e la loro semplificazione, otteniamo

$(d_0, d_1, d_2) \mapsto (d_0, d_1, d_2)$.