Code d'erreur : E2002

Catégorie : Temps de compilation : bloc et mosaïque d'entrée/sortie non alignés

Cette erreur se produit lorsque la forme de bloc d'une entrée ou d'une sortie de noyau ne correspond pas au tiling par défaut du type de données sur le matériel TPU spécifique utilisé.

Exemples de messages d'erreur :

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.

Backends XLA : TPU

Présentation

Les TensorCores (TC) des TPU disposent de registres vectoriels bidimensionnels. Les deux dimensions sont appelées sous-voie et voie. Étant donné que les unités de calcul TPU (par exemple, MXU) fonctionnent avec la précision des registres vectoriels, les tableaux XLA sont disposés dans la mémoire TPU en mosaïques.

Ce pavage (par exemple, 8x128) minimise les transformations de données lors de l'alimentation des unités de calcul. Les dimensions exactes du tiling (sous-canaux × canaux) dépendent de la génération du matériel et du type de données. Par exemple, une mosaïque courante pour la plupart des types est 8x128.

Au moment de la compilation, XLA applique les contraintes suivantes aux dimensions mineures et de deuxième niveau de chaque entrée/sortie du noyau :

  1. Divisibilité : la dimension du bloc doit être un multiple de la dimension du bloc dans le Tensor sous-jacent, ou
  2. Exception de forme complète : si la dimension du bloc n'est pas divisible, elle doit être égale à la taille complète de cette dimension dans le Tensor sous-jacent.

Cette erreur se produit lorsqu'un bloc ne respecte aucune des deux conditions. Par exemple, le chargement d'un bloc de forme (8, 100) à partir d'une entrée de forme (8, 1024) sur un matériel avec une forme de pavage (8, 128) échoue, car 100 n'est pas divisible par 128 et 100 != 1024. Toutefois, il serait autorisé si la forme d'entrée était (32, 100).

Débogage

Pour résoudre cette erreur, assurez-vous que les formes de bloc de votre noyau sont alignées sur le tiling matériel actuel. Modifiez le code de votre noyau pour aligner la taille des blocs de sorte qu'elle soit un multiple du tiling requis.

  • Exemple : Si l'erreur indique que le pavage est (16, 128), mais que la forme de votre bloc est (8, 128), modifiez la spécification du bloc pour que la forme corresponde à (16, 128).