Код ошибки: E2003

Категория: Время компиляции: Мозаичное выравнивание доступа к памяти (Unproven Memory Access Alignment)

Эта ошибка возникает, когда компилятор анализирует операцию доступа к памяти (например, vector.load , vector.store , tpu.load или tpu.store ) и не может статически доказать, что динамический индекс, используемый для конкретного измерения, является кратным требуемому размеру тайлинга .

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

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

Обзор

Когда ядро ​​загружает или сохраняет вектор, адрес в памяти (вычисляемый из базового указателя плюс динамический индекс) должен совпадать с размером ячейки вектора на аппаратном уровне. Например, если измерение состоит из 128 элементов, динамический индекс, используемый для доступа к нему, должен быть 0 , 128 , 256 и т. д. Обратите внимание, что для многих операций (таких как загрузка и сохранение векторов) таких требований к статическим индексам нет.

Компилятор обеспечивает выполнение этого требования с помощью статического анализа . Он отслеживает историю переменной индекса вплоть до арифметических операций, которые привели к её созданию (например, умножения, сложения). Если компилятор не может гарантировать (на этапе компиляции), что результирующее значение всегда будет делиться на размер тайлинга, он выдаёт эту ошибку.

Компилятор одинаково обрабатывает «доказанное смещение» и «неизвестное выравнивание». Поэтому, если вы используете индекс, который математически гарантированно смещен (например, i * 128 + 32 ), компилятор выдаст ту же ошибку.

Таким образом, эта ошибка может возникнуть, когда

  1. Для доступа к памяти используется переменная времени выполнения (динамический индекс).
  2. Логика вычисления индекса слишком сложна для анализа компилятором.
  3. Индекс математически корректен, но в коде отсутствует явное доказательство.
  4. Статический анализ определяет «доказанное несоосность».

Отладка

Для устранения этой ошибки у вас есть следующие варианты:

1. Явное подтверждение выравнивания.

Если вы знаете, что ваш индекс действителен, но компилятор не может это доказать, используйте операцию tpu.assume_multiple . Она действует как обещание компилятору, что значение делится на определенный множитель.

2. Используйте выровненные нагрузки и поверните их.

В случаях, когда смещение является преднамеренным, вместо загрузки небольшого невыровненного сегмента вектора:

  • Загрузите более крупный, полностью выровненный фрагмент, а затем поверните значения на динамический угол, чтобы сместить нужные данные в нужное положение (поскольку векторные срезы с динамическими начальными индексами не поддерживаются).
  • Измените форму или дополните тензор таким образом, чтобы данные начинались с индекса 0, а шаг между обращениями соответствовал аппаратному выравниванию.
    • Пример: Если вы будете перебирать фрагменты размером 32, начиная со смещения 1, ваши смещения будут 1, 33, 65... (невыровненные).
    • Исправление: Перепакуйте данные в новый тензор, где первый фрагмент находится по адресу 0, а размерность дополняется до 128. Ваши смещения станут 0, 128, 256..., что удовлетворяет требованию выравнивания.

Эти методы потребляют больше памяти, но часто упрощают логику ядра и устраняют необходимость в ручной проверке выравнивания.