Компьютерное зрение на микроконтроллере за $2: SIMD на ESP32-S3 ускорил FAST в 2,2 раза

Я написал ускоренный детектор особых точек FAST для ESP32-S3 на его 128-битных SIMD-инструкциях. Он работает примерно вдвое быстрее эталонной реализации и обрабатывает кадр QVGA (320×240) примерно за 6 мс. Исходный код лежит здесь.

За свои деньги ESP32-S3 — очень мощный микроконтроллер. В невзрачном пластиковом корпусе прячется двухъядерный процессор с частотой до 240 МГц и куча периферии, включая радиомодули Wi-Fi и Bluetooth Low Energy. Копаясь в его техническом справочнике, я обнаружил, что чип поддерживает ограниченный набор SIMD-инструкций. Для кремния, который стоит дешевле чашки кофе, это здорово. Посмотрим, что с ним можно сделать!

Пока я писал детектор углов FAST на SIMD, я освоил основы ассемблера ESP32-S3, научился обходить его странные ограничения и в итоге написал собственный простой распределитель регистров basm.

Что уже было

Поискав ещё, я понял, что эти инструкции вовсе не секретные, просто малоизвестные. Мне хотелось поэкспериментировать с ними самому и написать что-то сложнее того, что можно найти в сети. Компьютерное зрение меня в последнее время очень увлекает, поэтому я решил сделать SIMD-версию детектора особых точек FAST для ESP32-S3.

Начал я с поиска примеров кода с SIMD-инструкциями. На момент написания статьи я нашёл всего три источника:

• блог Ларри Бэнка, где он разбирает эту возможность и приводит минимальный пример;

• библиотеку Espressif esp-dsp для задач в духе DSP (свёртки, БПФ и т. п.);

• библиотеку Espressif esp-dl для ускорения ИИ прямо на чипе.

Просмотрев их, я с облегчением понял, что ассемблер на ESP32-S3 не так уж страшен. В итоге я остановился на такой общей структуре:

    .section .rodata# read-only data here    .section .text    .align 4    .global add    .type   add,@function# a simple function that adds two numbers# int32_t add(int32_t a, int32_t b)add:    entry           sp, 48 # sp (stack pointer) is an alias for a1    #define r_a a2    #define r_b a3# code hereadd a2, r_a, r_b # a2 = a + b    retw.n # return, with a2 containing the return value.

Первое ограничение: невыровненное чтение

Дальше я взялся за SIMD-версию предварительной проверки угла в FAST. По сути, она берёт по одному пикселю в каждой из четырёх сторон света и проверяет, есть ли среди них хотя бы три «экстремальных». Экстремальным здесь считается пиксель, который отличается от центрального по модулю больше заданного порога.

Я решил написать функцию, которая считает число экстремумов сразу для шестнадцати пикселей: в каждый векторный регистр помещается шестнадцать восьмибитных значений. Для этого нужно прочитать четыре «блока» пикселей — по одному с каждой стороны — и сравнить их с соответствующими центральными пикселями.

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

               [ north ]    [ west ]   [ center ]   [ east ]               [ south ]

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

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

Второе ограничение: сравнение только со знаком

Так, данные пикселей есть. Пора сравнивать, да? А вот и нет. Оказалось, что ESP32-S3 умеет сравнивать только знаковые восьмибитные числа. Это усложнило дело, ведь пиксели у меня хранились как беззнаковые восьмибитные значения. Если просто переинтерпретировать беззнаковые данные как знаковые, получится странное:

>>> np.int8(np.uint8(range(0, 256, 8)))array([   0,    8,   16,   24,   32,   40,   48,   56,   64,   72,   80,         88,   96,  104,  112,  120, -128, -120, -112, -104,  -96,  -88,        -80,  -72,  -64,  -56,  -48,  -40,  -32,  -24,  -16,   -8],      dtype=int8)

Диапазон [128, 255] нелинейно отображается в [-128, -1], и все сравнения ломаются. Чтобы это заработало, нужно было сдвинуть значения из [0, 255] в [-128, 127]. Казалось бы, обычное вычитание, но и тут помешал набор инструкций: на ESP32-S3 вся SIMD-арифметика с насыщением. Значит, все значения пикселей выше 127 вместо того, чтобы перейти через границу и получить правильный результат, упёрлись бы в -128. Пришлось как-то вычесть 128 из каждого числа параллельно, не пользуясь инструкцией вычитания. Сначала задача показалась мне неподъёмной, но в итоге решение нашлось.

Вычитание через XOR

Пусть x — уменьшаемое. Вот веса его разрядов:

-128 64 32 16 8 4 2 1

Замечаете что-нибудь интересное? Вычесть 128 — то же самое, что прибавить -128. Поэтому, чтобы вычислить x - 128, достаточно «прибавить единицу» к старшему разряду x, вес которого равен -128. А это просто инверсия старшего бита x. XOR любого бита с единицей равносилен побитовому НЕ, значит, старший бит x можно точечно инвертировать, сделав XOR с маской 0b10000000, то есть 0x80. Следовательно, x - 128 == x ^ 0x80.

Теперь каждый раз, когда загружаются беззнаковые восьмибитные значения, к ним применяется XOR с 0x80. Так они переходят в подходящий (линейный) диапазон, где арифметика и сравнения снова имеют смысл.

С этими двумя приёмами мне удалось обойти ограничения ESP32-S3 и написать SIMD-версию предварительной проверки угла в FAST. Ещё я написал на SIMD функцию оценки углов, где используется похожий набор операций, но о ней здесь рассказывать не буду.

Результат

В итоге в моих тестах производительность детектора FAST выросла примерно до 220% от исходной: с 5,1 до 11,2 мегапикселя в секунду. Этого с запасом хватает для компьютерного зрения в реальном времени: ESP32-S3 спокойно справляется с потоком VGA при 30 кадрах в секунду. Неплохо для чипа за два доллара!