Код ошибки: E2002

Категория: Время компиляции: Мозаичный ввод/вывод, несовпадение блоков и мозаичное расположение

Эта ошибка возникает, когда форма блока ввода или вывода ядра не совпадает со стандартным расположением блоков данных данного типа на используемом оборудовании TPU.

Примеры сообщений об ошибках:

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 Backends: TPU

Обзор

Тензорные ядра (ТЯ) в TPU имеют двумерные векторные регистры. Эти два измерения называются подобластью и полосой . Поскольку вычислительные блоки TPU (например, MXU) работают с детализацией векторных регистров, массивы XLA размещаются в памяти TPU в виде плиток.

Такая разбивка на блоки (например, 8x128 ) минимизирует преобразования данных при подаче их в вычислительные блоки. Точные размеры разбивки (подблоки × блоки) зависят от поколения оборудования и типа данных. Например, распространенная разбивка для большинства типов данных — 8×128 .

На этапе компиляции XLA применяет следующие ограничения к малому и второму малому измерениям каждого входного/выходного параметра ядра:

  1. Делимость: Размерность блока должна быть кратна размерности плитки в базовом тензоре, или
  2. Исключение для полной формы: если размерность блока не делится на две части, она должна быть равна полной размерности этой части в базовом тензоре.

Эта ошибка возникает, когда блок нарушает оба условия. Например, загрузка блока формы (8, 100) из входных данных формы (8, 1024) на аппаратном обеспечении с мозаичным расположением элементов формы (8, 128) завершается неудачей, поскольку 100 не делится на 128 и 100 != 1024 Однако это было бы допустимо, если бы входные данные имели форму (32, 100) .

Отладка

Для устранения этой ошибки убедитесь, что форма блоков вашего ядра соответствует текущей аппаратной разметке. Измените код ядра, чтобы размер блока был кратен требуемой разметке.

  • Пример: Если в сообщении об ошибке указано, что мозаичность — (16, 128) , а форма вашего блока — (8, 128) , измените спецификацию блока так, чтобы форма соответствовала (16, 128) .