Layout a mosaico

Figura 1

La figura 1 mostra come un array F32[3,5] è disposto in memoria con tiling 2x2. Una forma con questo layout è scritta come F32[3,5]{1,0:T(2,2)}, dove 1,0 si riferisce all'ordine fisico delle dimensioni (campo minor_to_major in Layout) mentre (2,2) dopo i due punti indica la suddivisione delle dimensioni fisiche in un riquadro 2 x 2.

Intuitivamente, i riquadri vengono disposti in modo da coprire la forma e, all'interno di ogni riquadro, gli elementi vengono poi disposti senza riquadro, come nell'esempio precedente, dove la parte destra dell'esempio mostra il layout in memoria, inclusi gli elementi di spaziatura interna bianca aggiunti per ottenere riquadri 2 x 2 completi anche se i limiti dell'array originali non sono uniformi.

Gli elementi aggiuntivi nella spaziatura interna non devono contenere alcun valore specifico.

Formule dell'indice lineare per la creazione di riquadri in base a una forma e a un riquadro

Senza tiling, un elemento e=(e_n, e_n-1, ... , e₁) in un array con limiti dell'array d=(d_n, d_n-1, ... , d₁) (d1 è la dimensione minore) viene disposto in ordine maggiore-minore nella posizione:

indice_lineare(e, d)
= indice_lineare((e_n, e_n-1, ... , e₁), (d_n, d_n-1, ... , d₁))
= e_nd_n-1...d₁ + e_n-...d_n-1 + e_n-2d +n-1

Per semplicità di notazione in questo documento, supponiamo che un riquadro abbia lo stesso numero di dimensioni dell'array. Nell'implementazione dei riquadri di XLA, questo viene generalizzato ai riquadri con meno dimensioni lasciando invariate le dimensioni iniziali più importanti e applicando i riquadri solo alle dimensioni più minori, in modo che il riquadro specificato menzioni un suffisso delle dimensioni fisiche della forma affiancata.

Quando viene utilizzato il tiling di dimensioni (t_n, t_n-1, ... , t₁), un elemento nell'array con indici (e_n, e_n-1, ..., e₁) viene mappato in questa posizione nel layout finale:

n_nnnnn

Il layout può essere considerato come costituito da due parti: (⌊e_n/t_n⌋, ... , ⌊e₁/t₁⌋), che corrisponde a un indice di riquadro in un array di riquadri di dimensioni (⌈d_n/t_n⌉, ... , ⌈_d/1/ La funzione ceil viene visualizzata in ⌈d_i/t_i⌉ perché se i riquadri superano i limiti dell'array più grande, la spaziatura interna viene inserita come nella Figura 1. Sia i riquadri che gli elementi al loro interno sono disposti in modo ricorsivo senza piastrelle.

Nell'esempio nella Figura 1, l'elemento (2,3) ha l'indice dei riquadri (1,1) e l'indice all'interno del riquadro (0,1), per un vettore di coordinate combinato di (1,1,0,1). Gli indici dei riquadri hanno limiti (2,3) e il riquadro stesso è (2,2) per un vettore combinato di (2,3,2,2). L'indice lineare con riquadro per l'elemento con indice (2,3) nella forma logica viene quindi

indice_lineare_con_tile((2,3), (3,5), (2,2))
= indice_lineare((1,1,0,1), (2,3,2,2))
= indice_lineare((1,1), (2,3)) ∙ 2 ∙ 2 + indice_lineare((0,1), = +1)

Affiancare come pad-rimodellare-trasporre

Il layout basato su riquadri funziona come segue:
Considera un array di dimensioni (d_n, d_n-1, ... , d1) (d1 è la dimensione minore). Quando è strutturato con riquadri di dimensioni (t_n, t_n-1, ... , t₁) (t₁ è la dimensione più minore), questo riquadro può essere descritto in termini di pad-riforma-trasposizione nel seguente modo.

1. L'array viene riempito in (⌈d_n/t_n⌉∙t_n, ... , ⌈d₁/t₁⌉∙t₁).
2. Ogni dimensione i viene suddivisa in (⌈d_i/ti⌉, t_i), ovvero l'array viene rimodellato in
   (⌈d_n/t_n⌉, t_n, ... , ⌈d₁/t₁⌉, t₁).
   Poiché questa rimodellazione non ha subito alcuna modifica al layout fisico, questa modifica è leggermente distorta. Se non si pensa esplicitamente a un riquadro, questa rimodella può esprimere qualsiasi forma con lo stesso numero di elementi della forma imbottita; l'esempio qui è di come esprimere un riquadro in questo modo.
3. Una trasposizione avviene spostando t_n, ... , t₁ alle dimensioni minori, mantenendone l'ordine relativo, in modo che l'ordine delle dimensioni dalla maggiore alla maggiore minore diventi
   (⌈d_n/t_n⌉, ... , ⌈d₁/t₁⌉, t_n, t_n).

La forma finale ha il prefisso
(⌈d_n/t_n⌉, ... , ⌈d₁/t₁⌉), che descrive il numero di riquadri in ogni dimensione. Un elemento dell'array (e_n, ... , e₁) viene mappato a questo elemento nella forma finale:
(⌊e_n/t_n⌋, ... , ⌊e₀/t₀⌋, e_n mod t_n, ... , e₁ mod). È facile vedere che l'indice lineare dell'elemento segue la formula di cui sopra come previsto.

Riquadri ripetuti

Applicando ripetutamente i riquadri XLA, i riquadri diventano ancora più flessibili.

Figura 2

La Figura 2 mostra come un array di dimensioni 4 x 8 venga affiancato da due livelli di tiling (prima 2 x 4 poi 2 x 1). Rappresentiamo questo tiling ripetuto come (2,4)(2,1). Ogni colore indica un riquadro 2 x 4 e ogni casella con bordo rosso è un riquadro 2 x 1. I numeri indicano l'indice lineare in memoria dell'elemento nel formato a mosaico. Questo formato corrisponde al formato utilizzato per BF16 su TPU, tranne per il fatto che il riquadro iniziale è più grande, ovvero (8,128)(2,1), dove lo scopo del secondo tiling di 2x1 è raccogliere insieme due valori a 16 bit per formare un valore a 32 bit in modo da allinearsi all'architettura di una TPU.

Nota che un secondo riquadro o un riquadro successivo possono fare riferimento a entrambe le dimensioni secondarie all'interno del riquadro, che riorganizzano semplicemente i dati all'interno del riquadro, come in questo esempio con (8,128)(2,1), ma possono anche fare riferimento alle dimensioni incrociate principali dei riquadri precedenti.

Combinazione di dimensioni con riquadri

I riquadri di XLA supportano anche la combinazione di dimensioni. Ad esempio, può combinare le dimensioni in F32[2,7,8,11,10]{4,3,2,1,0} in F32[112,110]{1,0} prima di affiancarlo a (2,3). La tessera usata è (∗,∗,2,∗,3). In questo caso, un asterisco in un riquadro implica di prendere quella dimensione e di combinarla con la dimensione minore successiva. Più dimensioni adiacenti possono essere sommate in un'unica dimensione. Una dimensione dedotta è rappresentata da un valore di riquadro pari a -1 nella dimensione del riquadro, che non sarebbe altrimenti valido in un riquadro come dimensione.

Più precisamente, se la dimensione i della forma viene eliminata tramite un asterisco nel riquadro, prima che venga applicata la definizione precedente di riquadro, questa dimensione viene rimossa sia dalla forma che viene affiancata dal vettore di riquadri, e per la dimensione i-1 della forma il limite dell'array aumenta da d_i-1 a d_id_i-1. Questo passaggio viene ripetuto per ogni asterisco nel vettore di riquadro.