Esta página foi traduzida pela API Cloud Translation.
Switch to English

Layout lado a lado

figura 1

A Figura 1 mostra como um array F32 [3,5] é disposto na memória com 2x2 tiling. Uma forma com este layout é escrita como F32 [3,5] {1,0: T (2,2)}, onde 1,0 se relaciona à ordem física das dimensões (campo minor_to_major no Layout) enquanto (2,2) após os dois pontos indica o ladrilho das dimensões físicas por um ladrilho 2x2.

Os ladrilhos são dispostos intuitivamente para cobrir a forma e, em seguida, dentro de cada ladrilho, os elementos são então dispostos sem ladrilhos, como no exemplo acima, onde a parte direita do exemplo mostra o layout na memória, incluindo os elementos de preenchimento branco que são adicionados para ter blocos 2x2 completos, mesmo que os limites da matriz original não sejam iguais.

Os elementos extras no preenchimento não precisam conter nenhum valor específico.

Fórmulas de índice linear para ladrilhos dados uma forma e um ladrilho

Sem lado a lado, um elemento e = (e n , e n-1 , ..., e 1 ) em uma matriz com limites de matriz d = (d n , d n-1 , ..., d 1 ) (d1 é a dimensão menor) é estabelecida por ordem maior para menor na posição:

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

Para simplicidade de notação neste documento, presumimos que um ladrilho tem o mesmo número de dimensões que a matriz. Na implementação de ladrilhos do XLA, isso é generalizado para ladrilhos com menos dimensões, deixando as dimensões principais iniciais inalteradas e aplicando o ladrilho apenas às dimensões menores, de modo que o ladrilho que é especificado mencione um sufixo das dimensões físicas do forma sendo lado a lado.

Quando o tiling de tamanho (t n , t n-1 , ..., t 1 ) é usado, um elemento na matriz com índices (e n , e n-1 , ..., e 1 ) é mapeado para este posição no layout final:

linear_index_with_tile (e, d, t)
= linear_index ((⌊e / t⌋, e mod t), (⌈d / t⌉, t)) (a aritmética é elemento a elemento, (a, b) é concatenação)
= 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 ))

O layout pode ser considerado como tendo duas partes: (⌊e n / t n ⌋, ..., ⌊e 1 / t 1 ⌋), que corresponde a um índice de ladrilhos em uma matriz de ladrilhos de tamanho (⌈d n / t n ⌉, ..., ⌈d 1 / t 1 ⌉), e (e n mod t n , ..., e 1 mod t 1 ), que corresponde a um índice dentro do bloco. A função ceil aparece em ⌈d i / t i ⌉ porque se os ladrilhos ultrapassam os limites da matriz maior, o preenchimento é inserido como na Figura 1. Tanto os ladrilhos quanto os elementos dentro dos ladrilhos são dispostos recursivamente sem ladrilhos.

Para o exemplo na Figura 1, o elemento (2,3) tem índice de bloco (1,1) e índice dentro de bloco (0,1), para um vetor de coordenadas combinado de (1, 1, 0, 1). Os índices dos blocos têm limites (2, 3) e o bloco em si é (2, 2) para um vetor combinado de (2, 3, 2, 2). O índice linear com bloco para o elemento com índice (2, 3) na forma lógica é então

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.

Tiling como pad-reshape-transpose

O layout baseado em ladrilhos funciona da seguinte maneira:
Considere uma matriz de dimensões (d n , d n-1 , ..., d1) (d1 é a dimensão mais secundária). Quando é disposto com ladrilhos de tamanho (t n , t n-1 , ..., t 1 ) (t 1 é a dimensão menor), esse ladrilho pode ser descrito em termos de pad-reshape-transpose a seguir maneira.

  1. A matriz é preenchida com (⌈d n / t n ⌉ ∙ t n , ..., ⌈d 1 / t 1 ⌉ ∙ t 1 ).
  2. Cada dimensão i é dividida em (⌈d i / t i ⌉, t i ), ou seja, a matriz é remodelada para
    (⌈D n / t n ⌉, t n , ..., ⌈d 1 / t 1 ⌉, t 1 ).
    Não há alteração de layout físico nesta remodelação por si só, portanto, esta remodelação é um bitcast. Se alguém não estiver pensando explicitamente em um ladrilho, essa reformulação pode expressar qualquer forma com o mesmo número de elementos que a forma acolchoada - o exemplo aqui é de como expressar um ladrilho dessa forma.
  3. Uma transposição acontece movendo t n , ..., t 1 para as dimensões menores, mantendo sua ordem relativa, de modo que a ordem das dimensões da maior para a menor torna-se
    (⌈D n / t n ⌉, ..., ⌈d 1 / t 1 ⌉, t n , ..., t 1 ).

A forma final tem o prefixo
(⌈D n / t n ⌉, ..., ⌈d 1 / t 1 ⌉), que descreve o número de blocos em cada dimensão. Um elemento na matriz (e n , ..., e 1 ) é mapeado para este elemento na forma final:
(⌊E n / t n ⌋, ..., ⌊e 0 / t 0 ⌋, e n mod t n , ..., e 1 mod t 1 ). É fácil perceber que o índice linear do elemento segue a fórmula acima conforme o esperado.

Mosaico repetido

O mosaico do XLA torna-se ainda mais flexível ao aplicá-lo repetidamente.

Figura 2

A Figura 2 mostra como uma matriz de tamanho 4x8 é ladrilhada por dois níveis de ladrilho (primeiro 2x4 e depois 2x1). Representamos esse mosaico repetido como (2,4) (2,1). Cada cor indica um ladrilho 2x4 e cada caixa de borda vermelha é um ladrilho 2x1. Os números indicam o índice linear na memória desse elemento no formato lado a lado. Este formato corresponde ao formato usado para BF16 na TPU, exceto que o bloco inicial é maior, ou seja, o bloco é (8.128) (2,1), onde o objetivo do segundo bloco por 2x1 é reunir dois valores de 16 bits para formar um valor de 32 bits de uma maneira que se alinhe à arquitetura de uma TPU.

Observe que um segundo bloco ou posterior pode se referir a ambas as dimensões menores dentro do bloco, que apenas reorganizam os dados dentro do bloco, como neste exemplo com (8.128) (2,1), mas também pode se referir ao bloco cruzado principal dimensões da telha anterior.

Combinando dimensões usando blocos

O ladrilho do XLA também suporta a combinação de dimensões. Por exemplo, ele pode combinar dimensões em F32 [2,7,8,11,10] {4,3,2,1,0} em F32 [112,110] {1,0} antes de colocá-lo lado a lado com (2,3 ) O bloco usado é (∗, ∗, 2, ∗, 3). Aqui, um asterisco em um ladrilho implica em pegar aquela dimensão e combiná-la com a próxima dimensão menor. Várias dimensões adjacentes podem ser agrupadas em uma dimensão. Uma dimensão subsumida é representada por um valor de bloco de -1 nessa dimensão do bloco, que de outra forma não é válido em um bloco como um tamanho de dimensão.

Mais precisamente, se a dimensão i da forma for eliminada por meio de um asterisco no ladrilho, então antes que a definição anterior de ladrilho seja aplicada, essa dimensão é removida tanto da forma que está sendo ladrilhada e do vetor do ladrilho, e o que era dimensão i-1 da forma tem seu limite de matriz aumentado de d i-1 para d i d i-1 . Esta etapa é repetida para cada asterisco no vetor do bloco.