Kod błędu: E2002

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 podpasmopasmo. 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:

  1. Podzielność: wymiar bloku musi być wielokrotnością wymiaru kafelka w tensorze bazowym lub
  2. 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 128100 != 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).