Categoría: Tiempo de compilación: Bloque y mosaico de entrada/salida desalineados
Este error se produce cuando la forma de bloque de una entrada o salida del kernel no se alinea con la segmentación predeterminada del tipo de datos en el hardware de TPU específico que se usa.
Ejemplos de mensajes de error:
UNIMPLEMENTED: Mosaic failed to compile TPU kernel: Failed to set window params
for input 0: Operand of shape (..., 256, 8192) has tiling (16, 128), but its
block shape (..., 8, 8192) is not divisible by tiling evenly nor matches the
full shape.
Back-ends de XLA: TPU
Descripción general
Los TensorCores (TC) en las TPU tienen registros de vectores bidimensionales. Las dos dimensiones se llaman carril secundario y carril. Dado que las unidades de procesamiento de TPU (p.ej., MXU) operan con la granularidad de los registros de vectores, los arreglos de XLA se disponen en la memoria de la TPU en mosaicos.
Esta división en mosaicos (p.ej., 8x128) minimiza las transformaciones de datos cuando se alimentan las unidades de procesamiento. Las dimensiones exactas de la división en mosaicos (carriles secundarios × carriles) dependen de la generación de hardware y el tipo de datos. Por ejemplo, un mosaico común para la mayoría de los tipos es 8 × 128.
En tiempo de compilación, XLA aplica las siguientes restricciones para las dimensiones secundaria y secundaria 2 de cada entrada o salida del kernel:
- Divisibilidad: La dimensión del bloque debe ser un múltiplo de la dimensión de la segmentación en el tensor subyacente.
- Excepción de forma completa: Si la dimensión del bloque no es divisible, debe ser igual al tamaño completo de esa dimensión en el tensor subyacente.
Este error se activa cuando un bloque incumple ambas condiciones. Por ejemplo, cargar un bloque con forma (8, 100) desde una entrada con forma (8, 1024) en hardware con segmentación en forma (8, 128) falla porque 100 no es divisible por 128 y 100 != 1024. Sin embargo, se permitiría si la forma de entrada fuera (32, 100).
Depuración
Para resolver este error, asegúrate de que las formas de los bloques del kernel se alineen con la segmentación de hardware actual. Modifica el código del kernel para alinear el tamaño del bloque de modo que sea un múltiplo de la segmentación requerida.
- Ejemplo: Si el error indica que el mosaico es
(16, 128), pero la forma de tu bloque es(8, 128), cambia la especificación del bloque para que la forma coincida con(16, 128).