Se usó la API de Cloud Translation para traducir esta página.
Switch to English

Disposición en mosaico

Figura 1

La Figura 1 muestra cómo una matriz F32 [3,5] se coloca en la memoria con un mosaico de 2x2. Una forma con este diseño se escribe como F32 [3,5] {1,0: T (2,2)}, donde 1,0 se relaciona con el orden físico de las dimensiones (campo minor_to_major en Layout) mientras que (2,2) después de los dos puntos indica el mosaico de las dimensiones físicas con un mosaico de 2x2.

Intuitivamente, los mosaicos se colocan para cubrir la forma y luego dentro de cada mosaico, los elementos se colocan sin mosaicos, como en el ejemplo anterior, donde la parte derecha del ejemplo muestra el diseño en la memoria, incluidos los elementos de relleno blancos que se agregan. para tener mosaicos completos de 2x2 aunque los límites de la matriz original no sean uniformes.

No es necesario que los elementos adicionales del relleno contengan ningún valor en particular.

Fórmulas de índice lineal para mosaicos con una forma y un mosaico

Sin mosaico, un elemento e = (e n , e n-1 , ..., e 1 ) en una matriz con límites de matriz d = (d n , d n-1 , ..., d 1 ) (d1 es la dimensión más secundaria) se presenta por orden mayor a menor en la posición:

índice_lineal (e, d)
= índice_lineal ((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 la notación en este documento, asumimos que un mosaico tiene el mismo número de dimensiones que la matriz. En la implementación de los mosaicos de XLA, esto se generaliza a los mosaicos con menos dimensiones dejando las dimensiones iniciales más importantes sin cambios y aplicando el mosaico solo a las dimensiones más pequeñas, de modo que el mosaico que se especifica menciona un sufijo de las dimensiones físicas de la forma de mosaico.

Cuando se utiliza un mosaico de tamaño (t n , t n-1 , ..., t 1 ), un elemento de la matriz con índices (e n , e n-1 , ..., e 1 ) se asigna a este posición en el diseño final:

índice_lineal_con_tile (e, d, t)
= índice_línea ((⌊e / t⌋, e mod t), (⌈d / t⌉, t)) (la aritmética es por elementos, (a, b) es concatenación)
= índice_lineal ((⌊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 ))
= índice_lineal ((⌊e n / t n ⌋, ..., ⌊e 1 / t 1 ⌋), (⌈d n / t n ⌉, ..., ⌈d 1 / t 1 ⌉)) ∙ t n t n-1 ... t 1 + índice_lineal ((e n mod t n , ..., e 1 mod t 1 ), (t n , t n-1 , ..., t 1 ))

Se puede pensar que el diseño tiene dos partes: (⌊e n / t n ⌋, ..., ⌊e 1 / t 1 ⌋), que corresponde a un índice de mosaico en una matriz de mosaicos de tamaño (⌈d n / t n ⌉, ..., ⌈d 1 / t 1 ⌉) y (e n mod t n , ..., e 1 mod t 1 ), que corresponde a un índice dentro del mosaico. La función ceil aparece en ⌈d i / t i ⌉ porque si los mosaicos sobrepasan los límites de la matriz más grande, el relleno se inserta como en la Figura 1. Tanto los mosaicos como los elementos dentro de los mosaicos se colocan de forma recursiva sin mosaicos.

Para el ejemplo de la Figura 1, el elemento (2,3) tiene índice de mosaico (1,1) e índice dentro del mosaico (0,1), para un vector de coordenadas combinado de (1, 1, 0, 1). Los índices de mosaicos tienen límites (2, 3) y el mosaico en sí es (2, 2) para un vector combinado de (2, 3, 2, 2). El índice lineal con mosaico para el elemento con índice (2, 3) en la forma lógica es entonces

índice_lineal_con_tile ((2,3), (3,5), (2,2))
= índice_lineal ((1,1,0,1), (2,3,2,2))
= índice_lineal ((1,1), (2,3)) ∙ 2 ∙ 2 + índice_lineal ((0,1), (2,2))
= (1 ∙ 3 + 1) ∙ 2 ∙ 2 + (0 ∙ 2 + 1)
= 17.

Mosaico como pad-remodelar-transponer

El diseño basado en mosaicos funciona de la siguiente manera:
Considere una matriz de dimensiones (d n , d n-1 , ..., d1) (d1 es la dimensión más pequeña). Cuando se presenta con mosaicos de tamaño (t n , t n-1 , ..., t 1 ) (t 1 es la dimensión más pequeña), ese mosaico se puede describir en términos de pad-reshape-transpose en lo siguiente camino.

  1. La matriz se rellena a (⌈d n / t n ⌉ ∙ t n , ..., ⌈d 1 / t 1 ⌉ ∙ t 1 ).
  2. Cada dimensión i se divide en (⌈d i / t i ⌉, t i ), es decir, la matriz se reforma para
    (⌈D n / t n ⌉, t n , ..., ⌈d 1 / t 1 ⌉, t 1 ).
    No hay ningún cambio de diseño físico en esta remodelación por sí misma, por lo que esta remodelación es un bitcast. Si uno no está pensando explícitamente en un mosaico, esta remodelación podría expresar cualquier forma con el mismo número de elementos que la forma acolchada; el ejemplo aquí es cómo expresar un mosaico de esta manera.
  3. Una transposición ocurre moviendo t n , ..., t 1 a las dimensiones más menores mientras se mantiene su orden relativo, de modo que el orden de las dimensiones de la mayor a la menor se convierte en
    (⌈D n / t n ⌉, ..., ⌈d 1 / t 1 ⌉, t n , ..., t 1 ).

La forma final tiene el prefijo
(⌈D n / t n ⌉, ..., ⌈d 1 / t 1 ⌉), que describe el número de mosaicos en cada dimensión. Un elemento de la matriz (e n , ..., e 1 ) se asigna a este elemento en la forma final:
(⌊E n / t n ⌋, ..., ⌊e 0 / t 0 ⌋, e n mod t n , ..., e 1 mod t 1 ). Es fácil ver que el índice lineal del elemento sigue la fórmula anterior como se esperaba.

Azulejos repetidos

El mosaico de XLA se vuelve aún más flexible al aplicarlo repetidamente.

Figura 2

La Figura 2 muestra cómo una matriz de tamaño 4x8 está dividida en mosaicos con dos niveles de mosaico (primero 2x4 y luego 2x1). Representamos este mosaico repetido como (2,4) (2,1). Cada color indica un mosaico de 2x4 y cada cuadro de borde rojo es un mosaico de 2x1. Los números indican el índice lineal en la memoria de ese elemento en el formato de mosaico. Este formato coincide con el formato utilizado para BF16 en TPU, excepto que el mosaico inicial es más grande, es decir, el mosaico es (8,128) (2,1), donde el propósito del segundo mosaico por 2x1 es recopilar dos valores de 16 bits para forman un valor de 32 bits de una manera que se alinea con la arquitectura de una TPU.

Tenga en cuenta que un segundo mosaico o posterior puede referirse a las dimensiones menores dentro del mosaico, que simplemente reorganiza los datos dentro del mosaico, como en este ejemplo con (8,128) (2,1), pero también puede referirse al mosaico transversal principal dimensiones del alicatado anterior.

Combinando dimensiones usando mosaicos

El mosaico de XLA también admite la combinación de dimensiones. Por ejemplo, puede combinar dimensiones en F32 [2,7,8,11,10] {4,3,2,1,0} en F32 [112,110] {1,0} primero antes de colocarlo en mosaico con (2,3 ). El mosaico utilizado es (∗, ∗, 2, ∗, 3). Aquí, un asterisco en un mosaico implica tomar esa dimensión y combinarla con la siguiente dimensión menor. Múltiples dimensiones adyacentes pueden subsumirse juntas en una dimensión. Una dimensión subsumida está representada por un valor de mosaico de -1 en esa dimensión del mosaico, que de otro modo no es válido en un mosaico como tamaño de dimensión.

Más precisamente, si la dimensión i de la forma se elimina mediante un asterisco en el mosaico, entonces antes de aplicar la definición anterior de mosaico, esa dimensión se elimina tanto de la forma que se está colocando en mosaico como del vector de mosaico, y lo que era la dimensión i-1 de la forma tiene su límite de matriz aumentado de d i-1 a d i d i-1 . Este paso se repite para cada asterisco en el vector de mosaico.