Mira conferencias magistrales, sesiones de productos, talleres y más de Google I / O Ver lista de reproducción

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 menor_a_mayor en Diseño) mientras que (2,2) después de los dos puntos indica el mosaico de las dimensiones físicas mediante un mosaico de 2x2.

Los mosaicos se colocan intuitivamente 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 el mosaico 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 del 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 mosaico.

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 menor), 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 piensa 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 formar un valor de 32 bits de una manera que se alinee 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 las dimensiones en F32 [2,7,8,11,10] {4,3,2,1,0} en F32 [112,110] {1,0} primero antes de colocarlas 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 se pueden subsumir 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 cuál era la dimensión i-1 de la forma tiene su cota de matriz aumentada de d i-1 a d i d i-1 . Este paso se repite para cada asterisco en el vector de mosaico.