Kategoria: czas kompilacji: Mosaic Unproven Memory Access Alignment
Ten błąd występuje, gdy kompilator analizuje operację dostępu do pamięci (np. vector.load, vector.store, tpu.load lub tpu.store) i nie może statycznie udowodnić, że dynamiczny indeks używany w określonym wymiarze jest wielokrotnością wymaganego rozmiaru kafelkowania.
Przykładowe komunikaty o błędach:
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>
Backendy XLA: TPU
Przegląd
Gdy jądro wczytuje lub zapisuje wektor, adres pamięci (obliczony na podstawie wskaźnika bazowego plus indeks dynamiczny) musi być zgodny z rozmiarem kafelkowania wektora na sprzęcie. Jeśli na przykład wymiar jest podzielony na 128 elementów, dynamiczny indeks używany do uzyskiwania do niego dostępu musi mieć wartość 0, 128, 256 itd. Pamiętaj, że wiele operacji (takich jak wczytywanie i zapisywanie wektorów) nie ma takich wymagań dotyczących indeksów statycznych.
Kompilator egzekwuje to wymaganie za pomocą analizy statycznej. Śledzi historię zmiennej indeksu wstecz, przez operacje arytmetyczne, które ją utworzyły (np. mnożenia, dodawania). Jeśli kompilator nie może zagwarantować (w czasie kompilacji), że wynikowa wartość będzie zawsze podzielna przez rozmiar kafelka, zgłasza ten błąd.
Kompilator traktuje „udowodnione niezgodności” i „nieznane dopasowania” w identyczny sposób.
Jeśli więc używasz indeksu, który z matematycznego punktu widzenia na pewno będzie nieprawidłowo dopasowany (np.
i * 128 + 32), kompilator zgłosi ten sam błąd.
Ten błąd może więc wystąpić, gdy
- Do uzyskiwania dostępu do pamięci używasz zmiennej środowiska wykonawczego (indeksu dynamicznego).
- Logika obliczania indeksu jest zbyt złożona, aby kompilator mógł ją przeanalizować.
- Indeks jest matematycznie prawidłowy, ale w kodzie nie ma wyraźnego dowodu.
- Analiza statyczna wskazuje „udowodnione niezgodności”.
Debugowanie
Aby rozwiązać ten problem, możesz:
1. Wyraźne potwierdzenie zgodności
Jeśli wiesz, że indeks jest prawidłowy, ale kompilator nie może tego udowodnić, użyj operacji tpu.assume_multiple. Jest to obietnica dla kompilatora, że wartość jest podzielna przez określony czynnik.
2. Używanie opcji Wyrównaj obciążenia i obróć
W sytuacjach, w których przesunięcie jest zamierzone, zamiast wczytywać mały, nieprawidłowo wyrównany segment wektorowy:
- Wczytaj większy, w pełni wyrównany kafelek, a następnie obróć wartości o dynamiczną wartość, aby przesunąć żądane dane na odpowiednią pozycję (ponieważ wycinki wektorowe z dynamicznymi indeksami początkowymi nie są obsługiwane).
- Zmień kształt tensora lub dodaj do niego dopełnienie, aby dane zaczynały się od indeksu 0, a krok między dostępami był zgodny z wyrównaniem sprzętowym.
- Przykład: jeśli iterujesz po fragmentach o rozmiarze 32, zaczynając od przesunięcia 1, przesunięcia to 1, 33, 65 itd. (niezgodne).
- Rozwiązanie: przepakuj dane w nowy tensor, w którym pierwszy fragment znajduje się w pozycji 0, a wymiar jest uzupełniony do 128. Przesunięcia wynoszą 0, 128, 256 itd., co spełnia wymagania dotyczące wyrównania.
Te metody zużywają więcej pamięci, ale często upraszczają logikę jądra i eliminują konieczność ręcznego potwierdzania wyrównania.