Categoría: Tiempo de compilación: Alineación no probada del acceso a la memoria de Mosaic
Este error ocurre cuando el compilador analiza una operación de acceso a la memoria (como vector.load, vector.store, tpu.load o tpu.store) y no puede demostrar de forma estática que el índice dinámico que se usa para una dimensión específica es un múltiplo del tamaño de segmentación requerido.
Ejemplos de mensajes de error:
INTERNAL: Mosaic failed to compile TPU kernel: cannot statically prove that index in dimension 1 is a multiple of 128
at location: ...
The MLIR operation involved:
%14372 = "vector.load"(%14371, %93, %14363) : (memref<4x256xf32, #tpu.memory_space<vmem>>, index, index) -> vector<1x32xf32>
Back-ends de XLA: TPU
Descripción general
Cuando tu kernel carga o almacena un vector, la dirección de memoria (calculada a partir del puntero base más el índice dinámico) debe alinearse con el tamaño de segmentación del vector en el hardware. Por ejemplo, si una dimensión se segmenta en 128 elementos, el índice dinámico que se usa para acceder a ella debe ser 0, 128, 256, etcétera. Ten en cuenta que muchas operaciones (como las cargas y los almacenamientos de vectores) no tienen tales requisitos para los índices estáticos.
El compilador aplica este requisito con el análisis estático. Realiza un seguimiento del historial de la variable de índice a través de las operaciones aritméticas que la produjeron (p.ej., multiplicaciones, sumas). Si el compilador no puede garantizar (en tiempo de compilación) que el valor resultante siempre será divisible por el tamaño de segmentación, genera este error.
El compilador trata la "desalineación comprobada" y la "alineación desconocida" de la misma manera.
Por lo tanto, si usas un índice que está garantizado matemáticamente que no está alineado (p.ej.,
i * 128 + 32), el compilador generará el mismo error.
Por lo tanto, este error puede ocurrir en los siguientes casos:
- Usas una variable de tiempo de ejecución (índice dinámico) para acceder a la memoria.
- La lógica de cálculo del índice es demasiado compleja para que la analice el compilador.
- El índice es matemáticamente válido, pero no tiene una prueba explícita en el código.
- El análisis estático determina la "falta de alineación comprobada".
Depuración
Para resolver este error, tienes las siguientes opciones:
1. Cómo afirmar la alineación de forma explícita
Si sabes que tu índice es válido, pero el compilador no puede probarlo, usa la operación tpu.assume_multiple. Esto actúa como una promesa para el compilador de que un valor es divisible por un factor específico.
2. Cómo usar las funciones Cargas alineadas y Rotar
En situaciones en las que la desalineación es intencional, en lugar de cargar un segmento de vector pequeño y desalineado, haz lo siguiente:
- Carga una tarjeta más grande y completamente alineada, y, luego, rota los valores en una cantidad dinámica para desplazar los datos deseados a la posición correcta (ya que no se admiten las segmentaciones de vectores con índices de inicio dinámicos).
- Cambia la forma del tensor o agrégale relleno para que los datos comiencen en el índice 0 y la distancia entre los accesos coincida con la alineación del hardware.
- Ejemplo: Si iteras sobre fragmentos de tamaño 32 a partir del desplazamiento 1, tus desplazamientos serán 1, 33, 65… (sin alineación).
- Corrección: Vuelve a empaquetar los datos en un tensor nuevo en el que el primer fragmento se encuentra en 0 y la dimensión se completa hasta 128. Tus desplazamientos se convierten en 0, 128, 256… lo que satisface el requisito de alineación.
Estos métodos consumen más memoria, pero suelen simplificar la lógica del kernel y eliminan la necesidad de aserciones de alineación manuales.