Diseño de mosaico

Figura 1

En la Figura 1, se muestra cómo se dispone un array F32[3,5] en la memoria con mosaicos de 2x2. Una forma con este diseño se escribe como F32[3,5]{1,0:T(2,2)}, en la que 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 indican la distribución de las dimensiones físicas con una tarjeta de 2 × 2.

De manera intuitiva, las tarjetas se disponen para cubrir la forma y, dentro de cada tarjeta, se disponen los elementos sin mosaicos, como en el ejemplo anterior, donde la parte derecha del ejemplo muestra el diseño en la memoria, incluidos los elementos de padding blanco que se agregan para tener mosaicos completos de 2 × 2, aunque los límites originales del array no sean uniformes.

No es necesario que los elementos adicionales en el padding contengan ningún valor en particular.

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

Sin mosaicos, un elemento e=(e_n, e_n-1, ... , e₁) en un array con límites de array d=(d_n, d_n-1, ... , d₁) (d1 es la dimensión más secundaria) se dispone de mayor a menor en la posición:

linear_index(e, d)
= linear_index((e_n, e_n-1, ... , e₁), (d_n, d_n-1, ... , d₁))
= e_nd_n-1... d₁ +_n-1...

Para simplificar la notación en este documento, suponemos que una tarjeta tiene la misma cantidad de dimensiones que el array. En la implementación de mosaicos de XLA, esto se generaliza en mosaicos con menos dimensiones. Para ello, no se modifican las dimensiones más importantes iniciales y se aplican las mosaicos solo a las dimensiones más pequeñas, de modo que en las tarjetas que se especifiquen se mencione un sufijo de las dimensiones físicas de la forma que se está mostrando.

Cuando se usan mosaicos de tamaño (t_n, t_n-1, ... , t₁), se asigna un elemento del array con índices (e_n, e_n-1, ..., e₁) a esta posición en el diseño final:

t linear_index_with_tile(e, d, t)
= linear_index(⌊e/t⌋, e mod t), (⌈d/t⌉, t)) (aritmética es elementwise, (a,b) es concatenación)
_nnn
n

Se puede considerar que el diseño tiene dos partes: (⌊e_n/t_n⌋, ... , ⌊e₁/t₁⌋), que corresponde a un índice de mosaicos en un array de mosaicos de tamaño (⌈d_n/t_n⌉, ... , ⌈d ⌈d₁/mode). La función ceil aparece en ⌈d_i/t_i⌉ porque si los mosaicos exceden los límites del array más grande, el padding se inserta como en la Figura 1. Tanto las tarjetas como los elementos dentro de las tarjetas se disponen de forma recursiva sin mosaicos.

Para el ejemplo de la Figura 1, el elemento (2,3) tiene un índice de mosaico (1,1) y un índice dentro de la tarjeta (0,1), para un vector de coordenadas combinado de (1,1,0,1). Los índices de tarjetas tienen límites (2,3) y la propia tarjeta es (2,2) para un vector combinado de (2,3,2,2). Luego, se lleva a cabo el índice lineal con mosaico del elemento que tiene el índice (2,3) en la forma lógica.

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,2)

Mosaico como pad-reshape-transpose

El diseño basado en mosaicos funciona de la siguiente manera:
Considera un array de dimensiones (d_n, d_n-1, ... , d1) (d1 es la dimensión más secundaria). Cuando se presenta con mosaicos de tamaño (t_n, t_n-1, ... , t₁) (t₁ es la dimensión más secundaria), esa representación en mosaicos se puede describir en términos de pad-reshape-transpose de la siguiente manera.

1. El array se rellena a (⌈d_n/t_n⌉∙t_n, ... , ⌈d₁/t₁⌉∙t₁).
2. Cada dimensión i se divide en (⌈d_i/ti⌉, t_i), es decir, el array se cambia de forma a
   (⌈d_n/t_n⌉, t_n, ... , ⌈d₁/t₁⌉, t₁).
   No hay cambios en el diseño físico cuando se trata de este cambio de forma por sí solo, por lo que se trata de un bitcast. Si no se piensa explícitamente en una tarjeta, esta modificación podría expresar cualquier forma con la misma cantidad de elementos que la forma con padding. El ejemplo aquí es cómo expresar una tarjeta de esta manera.
3. Una transposición se produce moviendo t_n, ... , t₁ 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 mayor de las menores sea
   (⌈d_n/t_n⌉, ... , ⌈d₁/t₁⌉, ... t_n).

La forma final tiene el prefijo
(⌈d_n/t_n⌉, ... , ⌈d₁/t₁⌉), que describe la cantidad de mosaicos en cada dimensión. Un elemento del array (e_n, ... , e₁) se asigna a este elemento en la forma final:
(⌊e_n/t_n⌋, ... , ⌊e₀/t₀⌋, e_n mod t_n, ... , e₁). Es fácil ver que el índice lineal del elemento sigue la fórmula anterior como se esperaba.

Mosaicos repetidos

La creación de mosaicos de XLA se vuelve aún más flexible si se la aplica repetidamente.

Figura 2

En la Figura 2, se muestra cómo un array de tamaño de 4 × 8 está dividido en dos niveles de mosaicos (primero 2 × 4, luego 2 × 1). Representamos esta representación en mosaico repetido como (2,4)(2,1). Cada color indica una tarjeta de 2 × 4 y cada cuadro de borde rojo es una tarjeta de 2 × 1. Los números indican el índice lineal en la memoria de ese elemento en el formato de mosaico. Este formato coincide con el formato que se usa para BF16 en la TPU, con la excepción de que el mosaico inicial es más grande, es decir, las mosaicos son (8,128)(2,1), donde el propósito de la segunda división en 2x1 es recopilar dos valores de 16 bits para formar un valor de 32 bits que se alinee con la arquitectura de la TPU.

Ten en cuenta que una segunda tarjeta o una posterior puede referirse a ambas dimensiones menores dentro de la tarjeta, que simplemente reorganiza los datos dentro de la tarjeta, como en este ejemplo con (8,128)(2,1), pero también puede hacer referencia a las dimensiones principales de mosaicos cruzados de la tarjeta anterior.

Cómo combinar dimensiones con mosaicos

La mosaicos 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} antes de crear mosaicos con (2,3). El mosaico utilizado es (∗,∗,2,∗,3). En este caso, un asterisco en una tarjeta implica tomar esa dimensión y combinarla con la siguiente dimensión más secundaria. Múltiples dimensiones adyacentes se pueden sumar juntas en una dimensión. Una dimensión subsumida se representa con un valor de mosaico de -1 en esa dimensión, que, de lo contrario, no es válida en una tarjeta como tamaño de dimensión.

Más precisamente, si se elimina la dimensión i de la forma mediante un asterisco en la tarjeta, antes de que se aplique la definición anterior de mosaicos, esa dimensión se quita tanto de la forma en mosaico como del vector de mosaico, y aumenta el límite del array de la dimensión i-1 de d_i-1 a d_id_i-1. Este paso se repite para cada asterisco en el vector de mosaicos.