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

Análisis de indexación

En este documento, se describe el análisis de indexación de HLO, que te permite calcular simbólicamente mapas de indexación para las operaciones de HLO. El mapa de indexación es una función que asigna los índices de un tensor a los índices de otro, p.ej., los índices de una salida de instrucciones de HLO a los índices de las entradas de instrucciones de HLO, o viceversa.

Ejemplo

Para una transmisión de tensor<20xf32> a tensor<10x20x30xf32>

p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}

el mapa de indexación de la salida a la entrada es $(i, j, k) \mapsto (j)$ para $i \in [0, 10]$, $j \in [0, 20]$ and $k \in [0, 30]$.

Motivación

La GPU de XLA usa varias soluciones personalizadas para razonar sobre la fusión, la utilización de operandos y los esquemas de mosaicos (más detalles a continuación). El objetivo del análisis de indexación es proporcionar un componente reutilizable para esos casos de uso. El análisis de indexación se basa en la infraestructura del mapa Affine Map de MLIR y agrega semántica de HLO.

Fusión

El razonamiento sobre la fusión de memoria se vuelve factible en casos no triviales, cuando sabemos qué elementos o porciones de las entradas se leen para calcular un elemento de la salida.

Utilización de operandos

El uso de operandos en XLA indica cuánto se usa cada entrada de una instrucción suponiendo que su salida se usa por completo. Actualmente, el uso tampoco se calcula para un caso genérico. El análisis de indexación permite calcular el uso con precisión.

Mosaico

Un mosaico o porción es un subconjunto hiperrectangular de un tensor parametrizado por desplazamientos, tamaños y segmentos. La propagación de mosaicos es una forma de calcular los parámetros de mosaicos del productor o consumidor de la operación mediante los parámetros de mosaicos de la op en sí. Ya existe una biblioteca que lo hace para softmax y punto. La propagación de mosaicos puede ser más genérica y sólida si se expresa a través de mapas de indexación.

Función y dominio

El mapa de indexación es una función $\boldsymbol{f}(\boldsymbol{d}, \boldsymbol{s})$que asigna un índice múltiple $\boldsymbol{d}$ de un tensor $A$ a elementos o rangos del tensor $B$. El parámetro $\boldsymbol{s}$ hace referencia a los rangos de índices de las dimensiones que están presentes en el tensor $B$, pero no en el tensor $A$.

Por ejemplo, si tenemos una reducción de tensor<2x4x8x16xf32> a tensor<4x8xf32>, el mapa de indexación de la salida en 2D a la entrada 4D es$(d_0, d_1) \mapsto (s_0, d_0, d_1, s_1)$, donde $d_i$ son los parámetros de dimensión que corresponden a los índices del tensor de salida. Los parámetros $s_j$ codifican varios valores, es decir, para calcular un elemento $(d_0, d_1)$ de la salida, necesitamos elementos $(s_0, d_0, d_1, s_1)$ de la entrada, donde $s_0 \in [0, 2)$ y $s_1 \in [0, 16)$.

Esta asignación se puede crear a partir de los atributos de las instrucciones de HLO o las asignaciones de instrucciones no fusionadas se pueden componer a fin de obtener la indexación para una fusión. La asignación también tiene un dominio que especifica qué elementos del tensor existe en la asignación.

\[ \begin{eqnarray} \boldsymbol{f}(\boldsymbol{d}, \boldsymbol{s})\; &s.t.& \\ \boldsymbol{lb}_d &\leq& \boldsymbol{d} \leq \boldsymbol{ub}_d \\ \boldsymbol{lb}_s &\leq& \boldsymbol{s} \leq \boldsymbol{ub}_s \\ \boldsymbol{lb}_g &\leq& \boldsymbol{g}(\boldsymbol{d}, \boldsymbol{s}) \leq \boldsymbol{ub}_g \end{eqnarray} \]

Como queremos minimizar la recomputación, necesitamos una biblioteca para cálculos simbólicos. XLA ya depende de MLIR, por lo que usamos mlir::AffineMap en lugar de escribir una biblioteca aritmética simbólica.

Un AffineMap típico es así

(d0)[s0, s1] -> (s0 + 5, d0 * 2, s1 * 3 + 50)

AffineMap tiene dos tipos de parámetros de manera conveniente: dimensiones y símbolos, que podemos usar para $\boldsymbol d$ y $\boldsymbol s$ , respectivamente. AffineMap no contiene metadatos sobre los rangos de las dimensiones, por lo que debemos proporcionar estos datos nosotros mismos.

struct Range {
 int64_t lower_bound;
 int64_t upper_bound;
};

struct IndexingMap {
 mlir::AffineMap affine_map;
 std::vector<Range> dimension_ranges;
 std::vector<Range> symbol_ranges;
 llvm::DenseMap<mlir::AffineExpr, Range> expr_ranges;
};

dim_ranges codifica las restricciones del cuadro inclusive para los parámetros de dimensión $\boldsymbol{d}$ del mapa de indexación, que suelen coincidir con la forma del tensor de salida para operaciones como transposición, reducción, elemento a nivel y punto, pero hay algunas excepciones como HloConcatenateInstruction.

symbol_ranges codifica valores posibles que pueden tomar los parámetros $\boldsymbol {s}$ .

Analicemos un ejemplo para comprender qué significa realmente todo lo anterior.

Indexa mapas para operaciones no combinadas

A nivel de los elementos

Para una operación a nivel de elementos, la asignación de indexación es una identidad.

  p0 = f32[10, 20] parameter(0)
  p1 = f32[10, 20] parameter(1)
  add = f32[10, 20] add(p0, p1)

El resultado que se asignará a la entrada es el siguiente:

salida -> input_0: $(d_0, d_1) \mapsto (d_0, d_1)$ para $\boldsymbol{d} \in [0,9] \times [0, 19]$, i.e. $\boldsymbol{d} \in {\rm Dom}(salida)$
output -> input_1: $(d_0, d_1) \mapsto (d_0, d_1)$ para $\boldsymbol{d} \in {\rm Dom} (salida)$

Se asigna la entrada a la salida

input_i -> output: $(d_0, d_1) \mapsto (d_0, d_1)$ para $\boldsymbol{d} \in {\rm Dom}(salida)$

Señales de aire

Transmisión significa que algunas de las dimensiones se quitarán cuando asignemos un resultado a la entrada y se agregarán cuando asignemos una entrada a la salida.

p0 = f32[20] parameter(0)
bc0 = f32[10, 20, 30] broadcast(p0), dimensions={1}

El resultado para la asignación de entrada es el siguiente:

output -> input: $(d_0, d_1, d_2) \mapsto (d_1)$ for $\boldsymbol{d} \in {\rm Dom}(salida)$

Asignación de las entradas a las salidas

entrada -> salida: $(d_0) \mapsto (s_0, d_1, s_1)$ para $\boldsymbol{d} \in {\rm Dom}(salida)$ and $\boldsymbol{s} \in [0, 9] \times [0, 29]$.

Ten en cuenta que ahora tenemos $\boldsymbol s$ en el lado derecho para la asignación de entrada a salida. Esos son los símbolos que representan rangos de valores. Por ejemplo, en este caso particular, cada elemento de entrada con el índice $d_0$ se asigna a una porción de 10 x 1 x 30 de la salida.

Iota y Constant

Convenientemente, no tienen ningún parámetro de entrada, por lo que no hay nada para el cual calcular la indexación.

Transposición

El mapa de indexación para la transposición es una permutación de dimensiones de entrada y salida.

p0 = f32[3, 12288, 6, 128] parameter(0)
transpose = f32[3, 6, 128, 12288] transpose(p0), dimensions={0, 2, 3, 1}

El resultado para la asignación de entrada es el siguiente:

output -> input: $(d_0, d_1, d_2, d_3) \mapsto (d_0, d_3, d_1, d_2)$ para $\boldsymbol{d} \in {\rm Dom}(output)$

La entrada para la asignación de salida es la siguiente:

entrada -> salida: $(d_0, d_1, d_2, d_3) \mapsto (d_0, d_2, d_3, d_1)$ para$\boldsymbol{d} \in {\rm Dom}(input)$

Reverso

Mapa de indexación para revertir los cambios de las dimensiones revertidas a $upper_bound(d_i) - d_i$:

p0 = f32[1, 17, 9, 9] parameter(0)
reverse = f32[1, 17, 9, 9] reverse(p0), dimensions={1, 2}

El resultado para la asignación de entrada es el siguiente:

salida -> entrada: $(d_0, d_1, d_2, d_3) \mapsto (d_0, -d_1 + 16, -d_2 + 8, d_3)$ for $\boldsymbol{d} \in {\rm Dom}(output)$

La entrada para la asignación de salida es la siguiente:

entrada -> salida: $(d_0, d_1, d_2, d_3) \mapsto (d_0, -d_1 + 16, -d_2 + 8, d_3)$ for $\boldsymbol{d} \in {\rm Dom}(input)$

(Variadic)Reducir

La reducción variable tiene varias entradas y varias entradas. El mapa de la salida a la entrada agrega las dimensiones reducidas. Por lo tanto, en cierto sentido, se comporta como lo contrario a una transmisión.

p0 = f32[256,10] parameter(0)
p0_init = f32[] constant(-inf)
p1 = s32[256,10] parameter(1)
p1_init = s32[] constant(0)
reduce = (f32[10], s32[10]) reduce(p0, p1, p0_init, p1_init),
  dimensions={0}, to_apply=min

El resultado que se asignará a la entrada es el siguiente:

resultado -> input_j: $(d_0) \mapsto (s_0, d_0)$ para $\boldsymbol{d} \in {\rm Dom}(output)$ and $\boldsymbol{s} \in [0, 9]$
resultado -> init_j: $(d_0) \mapsto ()$ para $\boldsymbol{d} \in {\rm Dom}(salida)$

La entrada a la salida se asigna de la siguiente manera:

input_i -> output_j: $(d_0, d_1) \mapsto (d_1)$ para $\boldsymbol{d} \in {\rm Dom}(entrada)$
init_i -> output_j: $() \mapsto (s_0)$ para $\boldsymbol{s} \in [0, 9]$

para $i, j = 0, \ldots, INPUT\\_COUNT$.

Porción

La indexación desde la salida hasta la entrada para la porción da como resultado un mapa de indexación estirado que es válido para todos los elementos de la salida. La asignación de la entrada a la salida se restringe a un rango estirado de los elementos de la entrada.

p0 = f32[10, 20, 50] parameter(0)
slice = f32[5, 3, 25] slice(f32[10, 20, 50] p0),
  slice={[5:10:1], [3:20:7], [0:50:2]}

El resultado para la asignación de entrada es el siguiente:

output -> input: $(d_0, d_1, d_2) \mapsto (d_0 + 5, 7d_1 + 3, 2d_2)$ para $\boldsymbol{d} \in {\rm Dom}(output)$

La entrada para la asignación de salida es la siguiente:

input -> output: $(d_0, d_1, d_2) \mapsto (d_0, d_1 / 7, d_2 / 2)$ for$\boldsymbol{d} \in [5, 9] \times [3, 19] \times [0, 49]$ con segmentos $[1, 7, 2]$.

Por definir: Indexación de entrada a salida

Reshape

Hay diferentes tipos de cambios.

Contraer forma

Esta es una reforma "linealiza" de N-D a 1D.

p0 = f32[4,8] parameter(0)
reshape = f32[32] reshape(p0)

El resultado para la asignación de entrada es el siguiente:

output -> input: $(d_0) \mapsto (d_0 / 8, d_0 \mod 8)$ para $\boldsymbol{d} \in {\rm Dom}(salida)$

La entrada para la asignación de salida es la siguiente:

entrada -> output: $(d_0, d_1) \mapsto (8 d_0 + d_1)$ for $\boldsymbol{d} \in {\rm Dom}(input)$.

Expandir forma

Esta es una op de “forma de contracción” inversa, que transforma una entrada 1D en una salida N-D.

p0 = f32[32] parameter(0)
reshape = f32[4, 8] reshape(p0)

El resultado para la asignación de entrada es el siguiente:

output -> input: $(d_0, d_1) \mapsto (8 d_0 + d_1)$ for $\boldsymbol{d} \in {\rm Dom}(salida)$

La entrada para la asignación de salida es la siguiente:

entrada -> resultado: $(d_0) \mapsto (d_0 / 8, d_0 \mod 8)$ for $\boldsymbol{d} \in {\rm Dom}(input)$.

Reforma genérica

Estas son las operaciones de cambio de forma que no se pueden representar como una sola forma de expansión o contracción. Solo se pueden representar como una composición de 2 o más formas de expansión o contracción.

Ejemplo 1: Linealización-delineación

p0 = f32[4,8] parameter(0)
reshape = f32[2, 4, 4] reshape(p0)

Esta modificación de forma se puede representar como una composición de la forma de contracción de tensor<4x8xf32> a tensor<32xf32> y, luego, una expansión de la forma a tensor<2x4x4xf32>.

El resultado para la asignación de entrada es el siguiente:

salida -> entrada: $(d_0, d_1, d_2) \mapsto (2d_0 + (4d_1 + d_2) / 8, 4d_1 + d_2) \mod 8)$

durante $\boldsymbol{d} \in {\rm Dom}(output)$

La entrada para la asignación de salida es la siguiente:

entrada -> salida: $(d_0, d_1) \mapsto ((8d_0 + d_1) / 16, ((8d_0 + d_1) \mod 16) / 4, d_1 \mod 4)$

para $\boldsymbol{d} \in {\rm Dom}(input)$.

Ejemplo 2: Subformas expandidas y contraídas

p0 = f32[4, 8, 12] parameter(0)
reshape = f32[32, 3, 4] reshape(p0)

Esta modificación se puede representar como una composición de dos. El primero contrae las dimensiones más externas tensor<4x8x12xf32> a tensor<32x12xf32> y el segundo expande la dimensión más interna tensor<32x12xf32> en tensor<32x3x4xf32>.

El resultado para la asignación de entrada es el siguiente:

output -> input: $(d_0, d_1, d_2) \mapsto (d_0 / 8, d_0 \mod 8, 4d_1 + d_2)$ para $\boldsymbol{d} \in {\rm Dom}(output)$

La entrada para la asignación de salida es la siguiente:

input -> output: $(d_0, d_1, d_2) \mapsto (8d_0 + d_1, d_2 / 4, d_2 \mod 4)$para $\boldsymbol{d} \in {\rm Dom}(input)$.

Bitcast

Una op de transmisión de bits se puede representar como una secuencia de transposición, reforma y transposición. Por lo tanto, sus mapas de indexación son solo una composición de mapas de indexación para esta secuencia.

Concatenate

La asignación de salida a entrada para concat se define para todas las entradas, pero con dominios que no se superponen, es decir, solo se usará una de las entradas a la vez.

p0 = f32[3,50] parameter(0)
p1 = f32[3,30] parameter(1)
concat = f32[3,80] concatenate(f32[3,50] p0, f32[3,30] p1),
  dimensions={1}

El resultado para la asignación de entrada es el siguiente:

salida -> entrada 1:

$(d_0, d_1) \mapsto (d_0, d_1)$ para $\boldsymbol{d} \in [0, 2] \times [0, 49]$

salida -> entrada 2:

$(d_0, d_1) \mapsto (d_0, d_1 - 50)$ para $\boldSymbol{d} \in [0, 2] \times [50, 79]$

Las entradas para la salida se asignan:

entrada 1 -> resultado: $(d_0, d_1) \mapsto (d_0, d_1)$ para $\boldsymbol{d} \in {\rm Dom}(input_1)$.
entrada 2 -> resultado: $(d_0, d_1) \mapsto (d_0, d_1 + 50)$ for $\boldsymbol{d} \in {\rm Dom}(input_2)$.

Punto (de salida a entrada implementada)

Los mapas de indexación para punto son muy similares a los de reducción.

p0 = f32[4, 128, 256] parameter(0)
p1 = f32[4, 256, 64] parameter(1)
dot = f32[4, 128, 64] dot(p0, p1),
  lhs_batch_dims={0}, rhs_batch_dims={0},
  lhs_contracting_dims={2}, rhs_contracting_dims={1}

El resultado a las entradas asigna lo siguiente:

output -> input_1: $(d_0, d_1, d_2) \mapsto (d_0, d_1, s_0)$ para $\boldsymbol{d} \in {\rm Dom}(output)$ y $\boldsymbol{s} \in [0, 255]$
output -> input_2: $(d_0, d_1, d_2) \mapsto (d_0, s_0, d_2)$ para $\boldsymbol{d} \in {\rm Dom}(output)$ y $\boldsymbol{s} \in [0, 255]$

Las entradas de salida se asignan de la siguiente manera:

input_1 -> output: $(d_0, d_1, d_2) \mapsto (d_0, d_1, s_0)$ para $\boldsymbol{d} \in {\rm Dom}(input_1)$ y $\boldsymbol{s} \in [0, 63]$
input_2 -> output: $(d_0, d_1, d_2) \mapsto (d_0, s_0, d_1)$ para $\boldsymbol{d} \in {\rm Dom}(input_2)$ y $\boldsymbol{s} \in [0, 127]$

Reducir período (por definir)

Pad (por definir)

Indexación de mapas para Fusion

El mapa de indexación para las operaciones de fusión es una composición de mapas de indexación para cada operación del clúster. Puede suceder que algunas entradas se lean varias veces con diferentes patrones de acceso.

Una entrada, varios mapas de indexación

Aquí hay un ejemplo para $p_0 + p_0^T$

f {
  p0 = f32[1000, 1000] parameter(0)
  transpose_p0 = f32[1000, 1000]{0, 1} transpose(p0), dimensions={1, 0}
  ROOT a0 = f32[1000, 1000] add(p0, transpose_p0)
}

Los mapas de indexación de salida a entrada para p0 serán $(d_0, d_1) \mapsto (d_0, d_1)$ and $(d_0, d_1) \mapsto (d_1, d_0)$. Esto significa que, para calcular un elemento de la salida, es posible que necesitemos leer el parámetro de entrada dos veces.

Una entrada, mapa de indexación con anulación de duplicación

En algunos casos, los mapas de indexación son iguales, aunque no sea evidente de inmediato.

f {
  p0 = f32[20, 10, 50] parameter(0)
  lhs_transpose_1 = f32[10, 20, 50] transpose(p0), dimensions={1, 0, 2}
  lhs_e = f32[10, 20, 50] exponential(lhs_transpose_1)
  lhs_transpose_2 = f32[10, 50, 20] transpose(lhs_e), dimensions={0, 2, 1}
  rhs_transpose_1 = f32[50, 10, 20] transpose(p0), dimensions={2, 1, 0}
  rhs_log = f32[50, 10, 20] exponential(rhs_transpose_1)
  rhs_transpose_2 = f32[10, 50, 20] transpose(rhs_log), dimensions={1, 0, 2}
  ROOT add = f32[10, 50, 20] add(lhs_transpose_2, rhs_transpose_2)
}

En este caso, el mapa de indexación de salida a entrada de p0 es solo $(d_0, d_1, d_2) \mapsto (d_2, d_0, d_1)$.

Softmax

La indexación de salida a entrada se asigna a parameter 0 para softmax:

$(d_0, d_1, d_2) \mapsto (d_0, d_1, d_2)$
$(d_0, d_1, d_2)[s_0] \mapsto (d_0, d_1, s_0)$

para $\boldsymbol{d} \in {\rm Dom}(output)$ y $\boldsymbol{s} \in [0, 124]$se refiere a la dimensión más interna de la entrada.

Simulador de mapas de indexación

El simplificador predeterminado para mlir::AffineMap en sentido ascendente no puede hacer ninguna suposición sobre los rangos de dimensiones o símbolos. Por lo tanto, no puede simplificar las expresiones con mod y div de manera eficiente.

Podemos aprovechar el conocimiento sobre los límites inferior y superior de las subexpresiones en los mapas afines para simplificarlas aún más.

El simplificador puede reescribir las siguientes expresiones.

$(d_0, d_1) \mapsto (d_0 + d1 / 16, d1 \mod 16)$ para $\boldSymbol{d} \in [0, 6] \times [0, 14]$ becomes $(d_0, d_1) \mapsto (d_0, d_1)$
$(d_0, d_1, d_2) \mapsto ((100d_0 + 10d_1 + d_2) /100, ((100d_0 + 10d_1 + d_2) \mod 100) / 10, d_2 \d_mod, d_0, d_1, d_10)$ for $d_i.$ becomes $
$(d_0, d_1, d_2) \mapsto ((16d_0 + 4d_1 + d_2) /8, (16d_0 + 4d_1 + d_2) \mod 8)$ for $d_i \in [0, 9]$ becomes $(d_0, d_1, d_2) d_2) d_2
$(d_0, d_1) \mapsto (-(-11d_0 - d_1 + 109) / 11 + 9)$ para $\boldSymbol{d} \in [0, 9] \times [0, 10]$ becomes $(d_0, d_1) \mapsto (d_0)$.

El simplificador de mapas de indexación nos permite comprender que algunas de las reformas en cadena de HLO se cancelan entre sí.

p0 = f32[10, 10, 10] parameter(0)
reshape1 = f32[50, 20] reshape(p0)
reshape2 = f32[10, 10, 10] reshape(reshape1)

Después de componer la indexación de mapas y su simplificación, obtendremos

$(d_0, d_1, d_2) \mapsto (d_0, d_1, d_2)$.