Kategorie:Kompilierzeit: Nicht nachgewiesene Ausrichtung des Speicherzugriffs in Mosaic
Dieser Fehler tritt auf, wenn der Compiler einen Speicherzugriffsvorgang (z. B. vector.load, vector.store, tpu.load oder tpu.store) analysiert und statisch nicht nachweisen kann, dass der für eine bestimmte Dimension verwendete dynamische Index ein Vielfaches der erforderlichen Kachelgröße ist.
Beispiele für Fehlermeldungen:
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>
XLA-Back-Ends:TPU
Übersicht
Wenn Ihr Kernel einen Vektor lädt oder speichert, muss die Speicheradresse (berechnet aus dem Basiszeiger plus dem dynamischen Index) mit der Kachelgröße des Vektors auf der Hardware übereinstimmen. Wenn eine Dimension beispielsweise in 128 Elemente unterteilt ist, muss der dynamische Index, der für den Zugriff darauf verwendet wird, 0, 128, 256 usw. sein. Beachten Sie, dass für viele Vorgänge (z. B. Vektor-Lade- und ‑Speichervorgänge) keine solchen Anforderungen an statische Indizes gelten.
Der Compiler erzwingt diese Anforderung mithilfe der statischen Analyse. Dabei wird die Historie der Indexvariablen anhand der arithmetischen Operationen, die zu ihr geführt haben (z.B. Multiplikationen, Additionen), nachvollzogen. Wenn der Compiler zur Kompilierzeit nicht garantieren kann, dass der resultierende Wert immer durch die Kachelgröße teilbar ist, wird dieser Fehler ausgegeben.
Der Compiler behandelt „proven misalignment“ (nachgewiesene Abweichung) und „unknown alignment“ (unbekannte Ausrichtung) identisch.
Wenn Sie also einen Index verwenden, der mathematisch garantiert nicht übereinstimmt (z.B.
i * 128 + 32) wird vom Compiler derselbe Fehler ausgegeben.
Dieser Fehler kann also in folgenden Fällen auftreten:
- Sie verwenden eine Laufzeitvariable (dynamischer Index), um auf den Speicher zuzugreifen.
- Die Logik für die Indexberechnung ist zu komplex, als dass der Compiler sie analysieren könnte.
- Der Index ist mathematisch gültig, aber es gibt keinen expliziten Beweis im Code.
- Bei der statischen Analyse wird „nachgewiesene Abweichung“ ermittelt.
Debugging
Sie haben folgende Möglichkeiten, diesen Fehler zu beheben:
1. Übereinstimmung explizit bestätigen
Wenn Sie wissen, dass Ihr Index gültig ist, der Compiler dies aber nicht beweisen kann, verwenden Sie die tpu.assume_multiple-Operation. Dies ist ein Versprechen an den Compiler, dass ein Wert durch einen bestimmten Faktor teilbar ist.
2. Aligned Loads und Rotate verwenden
In Szenarien, in denen die Abweichung beabsichtigt ist, sollten Sie anstatt eines kleinen, nicht ausgerichteten Vektorsegments Folgendes tun:
- Laden Sie eine größere, vollständig ausgerichtete Kachel und drehen Sie die Werte dann um einen dynamischen Betrag, um die gewünschten Daten in die richtige Position zu bringen (da Vektorslices mit dynamischen Startindizes nicht unterstützt werden). oder
- Formen Sie den Tensor um oder füllen Sie ihn auf, sodass die Daten bei Index 0 beginnen und der Abstand zwischen den Zugriffen der Hardwareausrichtung entspricht.
- Beispiel: Wenn Sie Blöcke der Größe 32 ab dem Offset 1 durchlaufen, sind Ihre Offsets 1, 33, 65 usw. (nicht ausgerichtet)
- Lösung: Packen Sie die Daten in einen neuen Tensor um, wobei der erste Chunk bei 0 beginnt und die Dimension auf 128 aufgefüllt wird. Ihre Offsets werden zu 0, 128, 256 … und erfüllen damit die Ausrichtungsanforderung.
Diese Methoden benötigen mehr Arbeitsspeicher, vereinfachen aber oft die Kernellogik und machen manuelle Ausrichtungsassertions überflüssig.