Kategorie:Kompilierungszeit: Falsch ausgerichteter Block und falsche Kachelung bei der Mosaikeingabe/-ausgabe
Dieser Fehler tritt auf, wenn die Blockform einer Kerneleingabe oder -ausgabe nicht mit der Standardaufteilung des Datentyps auf der verwendeten TPU-Hardware übereinstimmt.
Beispiele für Fehlermeldungen:
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.
XLA-Back-Ends:TPU
Übersicht
TensorCores (TCs) in TPUs haben zweidimensionale Vektorregister. Die beiden Dimensionen heißen sublane und lane. Da TPU-Recheneinheiten (z.B. MXU) arbeiten mit der Granularität von Vektorregistern. XLA-Arrays werden im TPU-Arbeitsspeicher in Kacheln angeordnet.
Diese Kachelung (z.B. 8x128) minimiert Datentransformationen beim Zuführen der Recheneinheiten. Die genauen Kachelungsdimensionen (Sublanes × Lanes) hängen von der Hardwaregeneration und dem Datentyp ab. Ein gängiges Kachelmuster für die meisten Typen ist beispielsweise 8×128.
Zur Kompilierzeit erzwingt XLA die folgenden Einschränkungen für die untergeordneten und zweituntergeordneten Dimensionen jeder Kernel-Ein-/Ausgabe:
- Teilbarkeit:Die Blockdimension muss ein Vielfaches der Kacheldimension im zugrunde liegenden Tensor sein.
- Ausnahme für die vollständige Form:Wenn die Blockdimension nicht teilbar ist, muss sie der vollen Größe dieser Dimension im zugrunde liegenden Tensor entsprechen.
Dieser Fehler wird ausgelöst, wenn ein Block gegen beide Bedingungen verstößt. Wenn Sie beispielsweise einen Block mit der Form (8, 100) aus einer Eingabe mit der Form (8, 1024) auf Hardware mit der Form (8, 128) laden, schlägt das Kacheln fehl, da 100 nicht durch 128 und 100 != 1024 teilbar ist. Das wäre jedoch zulässig, wenn die Eingabeform (32, 100) wäre.
Debugging
Achten Sie darauf, dass die Blockformen Ihres Kernels mit der aktuellen Hardware-Tiling-Konfiguration übereinstimmen, um diesen Fehler zu beheben. Ändern Sie Ihren Kernel-Code so, dass die Blockgröße ein Vielfaches der erforderlichen Kachelung ist.
- Beispiel:Wenn im Fehler angegeben wird, dass die Kachelung
(16, 128)ist, Ihr Block aber die Form(8, 128)hat, ändern Sie die Blockspezifikation so, dass die Form mit(16, 128)übereinstimmt.