Kategoria: czas kompilacji: niezgodność bloku i kafelkowania danych wejściowych/wyjściowych mozaiki
Ten błąd występuje, gdy kształt bloku danych wejściowych lub wyjściowych jądra nie jest zgodny z domyślnym podziałem na kafelki typu danych na konkretnym używanym sprzęcie TPU.
Przykładowe komunikaty o błędach:
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.
Backendy XLA: TPU
Przegląd
Rdzenie Tensor Core (TC) w TPU mają dwuwymiarowe rejestry wektorowe. Te 2 wymiary to podpasmo i pasmo. Ponieważ jednostki obliczeniowe TPU (np. MXU) działają na poziomie rejestrów wektorowych, a tablice XLA są rozmieszczone w pamięci TPU w postaci kafelków.
To kafelkowanie (np. 8x128) minimalizuje przekształcenia danych podczas przekazywania ich do jednostek obliczeniowych. Dokładne wymiary kafelków (podścieżki × ścieżki) zależą od generacji sprzętu i typu danych. Na przykład w przypadku większości typów powszechnym podziałem jest 8 × 128.
W czasie kompilacji XLA wymusza następujące ograniczenia dotyczące wymiarów podrzędnych i drugorzędnych każdego wejścia/wyjścia jądra:
- Podzielność: wymiar bloku musi być wielokrotnością wymiaru kafelka w tensorze bazowym lub
- Wyjątek dotyczący pełnego kształtu: jeśli wymiar bloku nie jest podzielny, musi być równy pełnemu rozmiarowi tego wymiaru w tensorze bazowym.
Ten błąd jest wywoływany, gdy blok narusza oba warunki. Na przykład wczytywanie bloku o kształcie (8, 100) z danych wejściowych o kształcie (8, 1024) na sprzęcie z kafelkowaniem o kształcie (8, 128) nie powiedzie się, ponieważ 100 nie jest podzielne przez 128 i 100 != 1024. Byłoby to jednak dozwolone, gdyby kształt danych wejściowych był równy (32, 100).
Debugowanie
Aby naprawić ten błąd, sprawdź, czy kształty bloków jądra są zgodne z obecnym kafelkowaniem sprzętowym. Zmodyfikuj kod jądra, aby dostosować rozmiar bloku tak, aby był wielokrotnością wymaganego kafelkowania.
- Przykład: jeśli błąd wskazuje, że kafelkowanie ma kształt
(16, 128), ale blok ma kształt(8, 128), zmień specyfikację bloku tak, aby kształt był zgodny z(16, 128).