Questa pagina è stata tradotta dall'API Cloud Translation.
Switch to English

Layout piastrellato

Figura 1

La Figura 1 mostra come un array F32 [3,5] è disposto in memoria con affiancamento 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 piastrellatura delle dimensioni fisiche da una piastrella 2x2.

Intuitivamente le tessere sono disposte per coprire la forma e quindi all'interno di ciascuna tessera, gli elementi vengono quindi disposti senza piastrelle, come nell'esempio sopra, dove la parte destra dell'esempio mostra il layout in memoria, inclusi gli elementi di imbottitura bianchi che vengono aggiunti per avere tessere 2x2 complete anche se i limiti della matrice originale non sono pari.

Gli elementi extra nel padding non devono contenere alcun valore particolare.

Formule di indice lineare per la piastrellatura data una forma e una piastrella

Senza piastrellatura, un elemento e = (e n , e n-1 , ..., e 1 ) in un array con limiti di array d = (d n , d n-1 , ..., d 1 ) (d1 è la dimensione più piccola) è disposta da ordine maggiore a minore nella posizione:

linear_index (e, d)
= lineare_indice ((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

Per semplicità di notazione in questo documento assumiamo che un riquadro abbia lo stesso numero di dimensioni dell'array. Nell'implementazione della piastrellatura da parte di XLA, questa viene generalizzata alle piastrellature con meno dimensioni lasciando invariate le dimensioni iniziali più importanti e applicando la piastrellatura solo alle dimensioni più piccole, in modo che la piastrellatura specificata menzioni un suffisso delle dimensioni fisiche del forma essendo piastrellato.

Quando si usa la piastrellatura di dimensione (t n , t n-1 , ..., t 1 ), un elemento dell'array con indici (e n , e n-1 , ..., e 1 ) viene mappato a questo posizione nel layout finale:

linear_index_with_tile (e, d, t)
= linear_index ((⌊e / t⌋, e mod t), (⌈d / t⌉, t)) (l'aritmetica è elementwise, (a, b) è la concatenazione)
= linear_index ((⌊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 ))
= linear_index ((⌊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 ))

Si può pensare che il layout abbia due parti: (⌊e n / t n ⌋, ..., ⌊e 1 / t 1 ⌋), che corrisponde a un indice di tessera in una matrice di tessere di dimensione (⌈d n / t n ⌉, ..., ⌈d 1 / t 1 ⌉) e (e n mod t n , ..., e 1 mod t 1 ), che corrisponde a un indice all'interno di una tessera. La funzione ceil appare in ⌈d i / t i ⌉ perché se le tessere superano i limiti dell'array più grande, il riempimento è inserito come in Figura 1. Sia le tessere che gli elementi all'interno delle tessere sono disposte in modo ricorsivo senza tessere.

Per l'esempio nella Figura 1, l'elemento (2,3) ha un indice di tessera (1,1) e un indice di tessera (0,1), per un vettore di coordinate combinato di (1, 1, 0, 1). Gli indici delle tessere hanno limiti (2, 3) e la tessera stessa è (2, 2) per un vettore combinato di (2, 3, 2, 2). L'indice lineare con piastrella per l'elemento con indice (2, 3) nella forma logica è quindi

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.

Piastrellatura come pad-rimodellare-trasporre

Il layout basato su piastrellatura funziona come segue:
Considera un array di dimensioni (d n , d n-1 , ..., d1) (d1 è la dimensione minore). Quando è disposto con piastrellatura di dimensione (t n , t n-1 , ..., t 1 ) (t 1 è la dimensione minore), quella piastrellatura può essere descritta in termini di pad-rimodellare-trasporre nel seguito modo.

  1. L'array viene riempito con (⌈d n / t n ⌉ ∙ t n , ..., ⌈d 1 / t 1 ⌉ ∙ t 1 ).
  2. Ogni dimensione i viene suddivisa in (⌈d i / t i ⌉, t i ), ovvero l'array viene rimodellato
    (⌈D n / t n ⌉, t n , ..., ⌈d 1 / t 1 ⌉, t 1 ).
    Non vi è alcun cambiamento di layout fisico in questo rimodellamento da solo, quindi questo rimodellamento è un po 'bitcast. Se non si pensa esplicitamente a una piastrellatura, questo rimodellamento potrebbe esprimere qualsiasi forma con lo stesso numero di elementi della forma imbottita - l'esempio qui è di come esprimere una piastrella in questo modo.
  3. Una trasposizione avviene spostando t n , ..., t 1 alle dimensioni minori mantenendo il loro ordine relativo, in modo che l'ordine delle dimensioni dalla maggiore alla minore diventi
    (⌈D n / t n ⌉, ..., ⌈d 1 / t 1 ⌉, t n , ..., t 1 ).

La forma finale ha il prefisso
(⌈D n / t n ⌉, ..., ⌈d 1 / t 1 ⌉), che descrive il numero di tessere in ogni dimensione. Un elemento dell'array (e n , ..., e 1 ) è mappato a questo elemento nella forma finale:
(⌊E n / t n ⌋, ..., ⌊e 0 / t 0 ⌋, e n mod t n , ..., e 1 mod t 1 ). È facile vedere che l'indice lineare dell'elemento segue la formula sopra come previsto.

Piastrellatura ripetuta

La piastrellatura di XLA diventa ancora più flessibile applicandola ripetutamente.

figura 2

La Figura 2 mostra come un array di dimensioni 4x8 viene affiancato da due livelli di affiancamento (prima 2x4 poi 2x1). Rappresentiamo questa piastrellatura ripetuta come (2,4) (2,1). Ogni colore indica una tessera 2x4 e ogni casella con bordo rosso è una tessera 2x1. I numeri indicano l'indice lineare in memoria di quell'elemento nel formato piastrellato. Questo formato corrisponde al formato utilizzato per BF16 su TPU, tranne per il fatto che la piastrella iniziale è più grande, ovvero la piastrellatura è (8,128) (2,1), dove lo scopo della seconda piastrellatura per 2x1 è raccogliere insieme due valori a 16 bit per formare un valore a 32 bit in un modo che si allinea con l'architettura di un TPU.

Si noti che una seconda tessera o successiva può fare riferimento sia alle dimensioni minori all'interno della tessera, che riorganizza solo i dati all'interno della tessera, come in questo esempio con (8,128) (2,1), ma può anche fare riferimento alla tessera incrociata maggiore dimensioni dalla piastrellatura precedente.

Combinare le dimensioni usando le piastrelle

La piastrellatura di XLA supporta anche la combinazione delle 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 con (2,3 ). La tessera utilizzata è (∗, ∗, 2, ∗, 3). Qui un asterisco in una tessera implica prendere quella dimensione e combinarla con la successiva dimensione minore. È possibile raggruppare più dimensioni adiacenti in un'unica dimensione. Una dimensione sommata è rappresentata da un valore di riquadro di -1 in quella dimensione del riquadro, che non è altrimenti valido in un riquadro come dimensione di dimensione.

Più precisamente, se la dimensione i della forma viene eliminata tramite un asterisco nella piastrella, prima che venga applicata la definizione precedente di piastrellatura, tale dimensione viene rimossa sia dalla forma piastrellata che dal vettore della piastrella, e qual era la dimensione i-1 della forma ha il limite di matrice aumentato da d i-1 a d i d i-1 . Questo passaggio viene ripetuto per ogni asterisco nel vettore delle tessere.