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

Layout em mosaico

figura 1

A Figura 1 mostra como um array F32 [3,5] é organizado na memória com 2x2 lado a lado. Uma forma com esse layout é escrita como F32 [3,5] {1,0: (2,2)}, em que 1,0 se relaciona à ordem física das dimensões (campo menor_para_maior no Layout) enquanto (2,2) após o cólon indica lado a lado das dimensões físicas por um bloco 2x2.

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

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

Fórmulas de índice linear para ladrilhos com 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 mais pequena) é apresentada por ordem maior para menor na posição:

índice_ linear (e, d)
= índice_ linear ((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 simplificar a notação neste documento, assumimos que um bloco possui o mesmo número de dimensões que o array. Na implementação de lado a lado do XLA, isso é generalizado para inclinações com menos dimensões, mantendo inalteradas as dimensões iniciais mais importantes e aplicando o lado a lado apenas às dimensões menores, para que o lado a lado especificado mencione um sufixo das dimensões físicas do elemento. forma sendo lado a lado.

Quando lado a lado do 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:

índice_ linear com o arquivo (e, d, t)
= índice_ linear ((⌊e / t⌋, e mod t), (⌈d / t⌉, t)) (a aritmética é elemento a elemento, (a, b) é concatenação)
= índice_ linear ((⌊e n / t n ⌋, ..., 1e 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 ))
= índice_ linear ((⌊e n / t n ⌋, ..., 1e 1 / t 1 ⌋), (⌈d n / t n ⌉, ..., 1d 1 / t 1 ⌉)) n t n t n-1 ... t 1 + índice_ linear ((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: (ne n / t n n , ..., 1e 1 / t 1 ⌋), que corresponde a um índice de bloco em uma matriz de blocos de tamanho ((d n / t n ⌉, ..., 1 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. Os ladrilhos e os elementos nos ladrilhos são dispostos recursivamente sem lado a lado.

Para o exemplo na Figura 1, o elemento (2,3) possui índice de bloco (1,1) e índice dentro do bloco (0,1), para um vetor de coordenadas combinadas de (1, 1, 0, 1). Os índices do bloco têm limites (2, 3) e o próprio bloco é (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))
= índice_ linear ((1,1,0,1), (2,3,2,2))
= índice_ linear ((1,1), (2,3)) ∙ 2 ∙ 2 + índice_ linear ((0,1), (2,2))
= (1 ∙ 3 + 1) ∙ 2 ∙ 2 + (0 ∙ 2 + 1)
= 17

Lado a lado como pad-remodelar-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 pequena). Quando é apresentado com lado a lado do tamanho (t n , t n-1 , ..., t 1 ) (t 1 é a dimensão menor), esse lado a lado pode ser descrito em termos de transposição de forma de almofada na seguinte maneira.

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

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

Ladrilhos repetidos

A disposição do XLA se torna ainda mais flexível, aplicando-a repetidamente.

Figura 2

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

Observe que um segundo bloco ou mais recente pode se referir às 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 podem se referir ao bloco transversal principal dimensões do ladrilho anterior.

Combinando cotas usando blocos

A telha 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} primeiro antes de colocar lado a lado com (2,3 ) O bloco usado é (∗, ∗, 2, ∗, 3). Aqui, um asterisco em um bloco implica pegar essa 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 igual a -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, antes que a definição anterior de ladrilho seja aplicada, essa dimensão será removida da forma que está sendo ladrilhada e do vetor do ladrilho e do que era a 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 lado a lado.