Catégorie : Temps de compilation : alignement non prouvé de l'accès à la mémoire Mosaic
Cette erreur se produit lorsque le compilateur analyse une opération d'accès à la mémoire (telle que vector.load, vector.store, tpu.load ou tpu.store) et ne peut pas prouver de manière statique que l'index dynamique utilisé pour une dimension spécifique est un multiple de la taille de mosaïque requise.
Exemples de messages d'erreur :
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>
Backends XLA : TPU
Présentation
Lorsque votre noyau charge ou stocke un vecteur, l'adresse mémoire (calculée à partir du pointeur de base plus l'index dynamique) doit correspondre à la taille de mosaïque du vecteur sur le matériel. Par exemple, si une dimension est mosaïquée par 128 éléments, l'index dynamique utilisé pour y accéder doit être 0, 128, 256, etc. Notez que de nombreuses opérations (comme les chargements et les stockages de vecteurs) n'ont pas de telles exigences pour les index statiques.
Le compilateur applique cette exigence à l'aide de l'analyse statique. Il retrace l'historique de la variable d'index à travers les opérations arithmétiques qui l'ont produite (par exemple, les multiplications et les additions). Si le compilateur ne peut pas garantir (au moment de la compilation) que la valeur résultante sera toujours divisible par la taille du pavé, il génère cette erreur.
Le compilateur traite "désalignement avéré" et "alignement inconnu" de la même manière.
Par conséquent, si vous utilisez un index dont le désalignement est mathématiquement garanti (par exemple,
i * 128 + 32), le compilateur générera la même erreur.
Cette erreur peut donc se produire dans les cas suivants :
- Vous utilisez une variable d'exécution (index dynamique) pour accéder à la mémoire.
- La logique de calcul de l'index est trop complexe pour que le compilateur puisse l'analyser.
- L'index est mathématiquement valide, mais il manque une preuve explicite dans le code.
- L'analyse statique détermine le "désalignement avéré".
Débogage
Pour résoudre cette erreur, vous avez les options suivantes :
1. Affirmer explicitement l'alignement
Si vous savez que votre index est valide, mais que le compilateur ne peut pas le prouver, utilisez l'opération tpu.assume_multiple. Cela sert de promesse au compilateur qu'une valeur est divisible par un facteur spécifique.
2. Utiliser "Charges alignées" et "Rotation"
Dans les cas où le désalignement est intentionnel, au lieu de charger un petit segment de vecteur non aligné :
- Chargez une carte plus grande et entièrement alignée, puis faites pivoter les valeurs d'un montant dynamique pour déplacer les données souhaitées (car les tranches vectorielles avec des indices de début dynamiques ne sont pas prises en charge).
- Remodelez ou complétez le Tensor pour que les données commencent à l'index 0 et que le stride entre les accès corresponde à l'alignement matériel.
- Exemple : Si vous itérez sur des blocs de taille 32 en commençant à l'offset 1, vos offsets sont 1, 33, 65, etc. (non aligné).
- Correction : reconditionnez les données dans un nouveau Tensor où le premier bloc est à 0 et la dimension est complétée à 128. Vos décalages deviennent 0, 128, 256, etc., ce qui répond à l'exigence d'alignement.
Ces méthodes consomment plus de mémoire, mais simplifient souvent la logique du noyau et éliminent le besoin d'assertions d'alignement manuel.