LUT и SWAR

Ночные эксперименты с нейросеткой на «ПикоКалке» дали, кстати, неплохой результат — теперь на токен тратится чуть больше секунды. Тут есть вклад оборудования, например, разгона центрального процессора до 270 МГц (выше плата зависает), но есть и чисто программные моменты. Я испытал несколько способов перемножить несколько чисел — в вычислениях любой нейросетки такого много.

Быстрее всего оказался обычный INT8 — прямое целочисленное умножение. Веса у меня четырёхбитные, чтобы сетка поместилась в память: в одном байте хранятся два веса. Для расчёта они распаковываются и по отдельности умножаются на восьмибитные активации.

Произведения складываются в 32-битном аккумуляторе, затем применяется масштаб. По сравнению с вариантом LUT_FLOAT32 здесь меняется и точность активаций: вместо чисел с плавающей точкой используются восьмибитные целые, но на работе нейросети это не сказалось.

Пока с этим возился, познакомился с несколькими интересными методами считать такие суммы, о двух таких алгоритмах хочу рассказать. В моём «ПикоКалке» стоит RP2040 с целочисленным умножителем, но нет FPU — аппаратного блока вычислений с плавающей точкой.

В каждом случае будем считать скалярное произведение двух векторов: веса — 3, 6, 1, 4 и активации — 2, 3, 5, 7.

LUT

Первый называется LUT («Look-Up Table»), я его вчера уже упоминал. Он очень простой для понимания.

Сначала веса раскладываем на степени двойки (то есть по битам):

3 = 1 + 2  6 = 2 + 4  1 = 1  4 = 4

Теперь нам надо получить скалярное произведение вектора весов и вектора активаций:

(1 + 2) × 2 + (2 + 4) × 3 + (1) × 5 + (4) × 7

Если раскрыть скобки и сгруппировать слагаемые по степеням двойки, каждый разряд весов будет выбирать свой набор активаций:

1 × (2 + 5) + 2 × (2 + 3) + 4 × (3 + 7)

Из четырёх активаций получается шестнадцать возможных наборов, включая пустой. Заранее считаем сумму каждого набора и записываем в таблицу LUT.

Чтобы дать наборам адреса, назначаем позициям активаций флажки 1, 2, 4, 8. Адрес набора, или маска, — сумма флажков выбранных позиций. Например, младшие биты весов 3, 6, 1, 4 дают выбор 1, 0, 1, 0: берём первую и третью активации, то есть 2 и 5. Адрес этого набора — 1 + 4 = 5, поэтому LUT[5] = 2 + 5 = 7.

Для разряда ×2 выбираются первая и вторая активации: адрес 1 + 2 = 3, сумма LUT[3] = 2 + 3 = 5. Для разряда ×4 — вторая и четвёртая: адрес 2 + 8 = 10, сумма LUT[10] = 3 + 7 = 10. Нужное скалярное произведение таким образом находится по таблице:

1 × LUT[5] + 2 × LUT[3] + 4 × LUT[10] = 7 + 10 + 40 = 57

Таблицы строятся для текущего вектора активаций и используются для всех строк матрицы весов. У каждой строки свои маски, а таблицы готовых сумм общие. При смене вектора активаций их пересчитывают.

Максимальная ширина входа слоя у меня — 384 активации. Делим их на четвёрки, для каждой храним 16 сумм. Одна сумма в моём варианте LUT_FLOAT32 занимает четыре байта, получается: (384 / 4) × 16 × 4 = 6144 байта. Столько внутренней SRAM занимает весь набор таблиц LUT.

SWAR

Второй метод называется SWAR («SIMD Within A Register»), где SIMD — «Single Instruction, Multiple Data».

Тут мы считаем два произведения одной командой. Проще всего проиллюстрировать принцип в привычной нам десятичной системе. Пусть нам надо подсчитать следующее скалярное произведение:

3 × 2 + 6 × 3

Для этого мы применяем такой любопытный приём:

(6 × 100 + 3) × (2 × 100 + 3) или 603 × 203 = 122409

Если взять средние две цифры, это будет сумма произведений:

3 × 2 + 6 × 3 = 24

Почему это работает? Станет понятно, если по-другому расставить скобки:

12 × 10000 + (6 × 3 + 3 × 2) × 100 + 9

Остаются вопросы — как мы получаем среднюю часть числа и где тут экономия, ведь в раскрытой формуле куча умножений?

Эта формула объясняет результат одной команды. В коде мы не вычисляем её слагаемые по отдельности: упаковываем числа сдвигами и битовыми масками, выполняем одно обычное умножение и извлекаем нужную часть результата.

В настоящем расчёте вместо основания 100 используется 65536: два числа упаковываются в 32-битное слово с двумя 16-битными полями. Команда умножения сохраняет младшие 32 бита произведения, поэтому член с множителем 65536² выпадает сам собой — для него просто нет места. В нашем примере остаётся 24 × 65536 + 9. Сдвиг вправо на 16 бит убирает младшее поле и оставляет 24.

Размеры полей и диапазоны чисел нужно подбирать так, чтобы переносы между полями не испортили результат. Для отрицательных активаций нужны дополнительные поправки; положительный десятичный пример показывает только принцип.