Forme e layout

Struttura di un'operazione XLA

Considera un esempio di HLO:

add.936 = bf16[8,1,1280,16384]{3,2,0,1:T(8,128)(2,1)}
          add(exponential.183, broadcast.3115)

È costituito dai seguenti componenti:

• Nome operazione: add.936
  • Il nome univoco dell'operazione.
• Forma: bf16[8,1,1280,16384]
  • Questa è la forma di output dell'operazione. Qui il dtype è bf16 e la forma è [8,1,1280,16384].
• Layout (con affiancamento): 3,2,0,1:T(8,128)(2,1)
  • Descrive come l'array viene memorizzato in memoria. 3,2,0,1 indica l'ordine degli assi in memoria (ad es. colonna principale, riga principale e così via) e T(8,128)(2,1) indica il tiling e il padding utilizzati.
  • Il layout è facoltativo. Se non specificato, non è presente alcun tiling e si presume che le dimensioni siano ordinate dalla più importante alla meno importante.
• Operazione: add
  • L'operazione in corso. In questo caso, è Add, che viene menzionato anche nel nome dell'operazione.
• Argomenti: exponential.183, broadcast.3115
  • Questa operazione accetta due argomenti, specificati con i rispettivi nomi univoci.

Vediamo un altro esempio, un'operazione di fusione:

%fusion.3 = bf16[32,32,4096]{2,1,0:T(8,128)(2,1)S(1)}
            fusion(bf16[32,32,8192]{2,1,0:T(8,128)(2,1)S(1)} %fusion.32),
            kind=kCustom, calls=%all-reduce-scatter.3

Oltre ai componenti descritti in precedenza, è costituito da:

• Attributi: kind e calls
  • Questi forniscono maggiori informazioni sull'operazione eseguita, in questo caso: fusione.
• Posizione in memoria (identificatore dello spazio di memoria): S(1)
  • Indica lo spazio/la posizione di memoria in cui è archiviato l'array. S(1) qui indica che questo array si trova in VMEM (su una TPU).
• Dettagli su forma e layout per l'argomento di input %fusion.32

Le sezioni seguenti descrivono le forme, il layout e gli identificatori dello spazio di memoria. Puoi scoprire di più sull'affiancamento in Layout affiancato.

Forme

Il proto XLA ShapeProto (xla_data.proto) descrive il numero di dimensioni, la dimensione e il tipo di dati di un array N-dimensionale (array in breve).

Terminologia, notazione e convenzioni

• Il numero effettivo di dimensioni di un array è il numero di dimensioni con una dimensione maggiore di 1.

• Le dimensioni sono numerate da 0 fino a N-1 per un array dimensionale N. La dimensione di una dimensione è un numero intero non negativo. In particolare, la dimensione 0 è valida. I numeri delle dimensioni sono etichette arbitrarie per comodità. L'ordine di questi numeri di dimensione non implica un ordine particolare minore/maggiore nel layout della forma. Il layout è determinato dal proto LayoutProto.

• Per convenzione, le dimensioni sono elencate in ordine crescente in base al numero della dimensione. Ad esempio, per una matrice tridimensionale di dimensioni [A x B x C], la dimensione 0 ha dimensione A, la dimensione 1 ha dimensione B e la dimensione 2 ha dimensione C.

  Alcune utilità in XLA supportano anche l'indicizzazione negativa in stile Python: Dimension -1 è l'ultima dimensione (equivalente a N-1 per un array di dimensioni N). Ad esempio, per l'array tridimensionale descritto sopra, la dimensione -1 ha dimensione C, la dimensione -2 ha dimensione B e così via.

• Gli array bidimensionali, tridimensionali e quadridimensionali spesso hanno lettere specifiche associate alle dimensioni. Ad esempio, per un array 2D:

  • dimensione 0: y
  • dimensione 1: x

  Per un array 3D:

  • dimensione 0: z
  • dimensione 1: y
  • dimensione 2: x

  Per un array 4D:

  • dimensione 0: p
  • dimensione 1: z
  • dimensione 2: y
  • Dimensione 3: x

• Le funzioni nell'API XLA che accettano dimensioni lo fanno in ordine crescente di numero di dimensione. Corrisponde all'ordinamento utilizzato quando si passano le dimensioni come initializer_list, ad es.

  ShapeUtil::MakeShape(F32, {A, B, C, D})

  creerà una forma il cui array di dimensioni è costituito dalla sequenza [A, B, C, D].

Layout

Il proto LayoutProto descrive come viene rappresentato un array in memoria. Include i seguenti campi:

message LayoutProto {
  repeated int64 minor_to_major;
  int64 tail_padding_alignment_in_elements;
  ...
}

Ordinamento delle dimensioni da minore a maggiore

L'unico campo obbligatorio è minor_to_major. Questo campo descrive l'ordine dalla dimensione secondaria a quella principale all'interno di una forma. I valori in minor_to_major sono un ordinamento delle dimensioni dell'array (da 0 a N-1 per un array di dimensioni N), con il primo valore che è la dimensione meno importante fino all'ultimo valore, che è la dimensione più importante. La dimensione meno importante è quella che cambia più rapidamente quando si scorrono gli elementi dell'array disposti nella memoria lineare.

Ad esempio, considera la seguente matrice bidimensionale di dimensioni [2 x 3]:

a b c
d e f

Qui la dimensione 0 è la taglia 2 e la dimensione 1 è la taglia 3. Se il campo minor_to_major nel layout è [0, 1], la dimensione 0 è la dimensione meno importante e la dimensione 1 è la dimensione più importante. Questo corrisponde al seguente layout nella memoria lineare:

a d b e c f

Questo ordine delle dimensioni da minore a maggiore, da 0 a N-1, è simile all'ordine column-major (per le dimensioni bidimensionali). Supponendo un ordinamento monotono delle dimensioni, un altro modo in cui potremmo fare riferimento a questo layout nel codice è semplicemente "dim 0 is minor".

D'altra parte, se il campo minor_to_major nel layout è [1, 0], il layout nella memoria lineare è:

a b c d e f

Un ordine delle dimensioni da minore a maggiore, da N-1 a 0 per un array N dimensionale, è simile a row-major (per array bidimensionali). Supponendo un ordinamento monotonico delle dimensioni, un altro modo in cui potremmo fare riferimento a questo layout nel codice è semplicemente "dim 0 è principale".

Ordinamento predefinito dal minore al maggiore

Il layout predefinito per le forme appena create è "L'ordine delle dimensioni è da maggiore a minore" (ovvero [N-1, ..., 0]).

Spaziatura interna

Il campo tail_padding_alignment_in_elements definisce l'allineamento dell'array affiancato in termini di numero di elementi. Dopo l'applicazione dell'affiancamento, gli elementi con spaziatura interna verranno aggiunti alla fine del layout finché il numero totale di elementi non sarà un multiplo di questo valore.

Indicizzazione negli array

La classe IndexUtil in index_util.h fornisce utilità per la conversione tra indici multidimensionali e indici lineari dato una forma e un layout. Gli indici multidimensionali includono un indice int64 per ogni dimensione. Gli indici lineari sono un singolo valore int64 che viene utilizzato per indicizzare il buffer contenente l'array. Consulta shape_util.h e layout_util.h nella stessa directory per le utilità che semplificano la creazione e la manipolazione di forme e layout.

Identificatori dello spazio di memoria

In HLO, ogni array può essere annotato con un identificatore di spazio di memoria, scritto come S(n).

• S(0) (spesso omesso) indica la memoria ad alta larghezza di banda (HBM) del dispositivo.
• S(1) rappresenta la memoria virtuale (VMEM) sul dispositivo.
• S(2), S(3) e così via corrispondono a spazi di memoria aggiuntivi specifici per il dispositivo.
• S(5) indica la memoria dell'host.