Мы в Smart Engines активно занимаемся низкоуровневой оптимизацией нейронных сетей. Наши библиотеки работают на центральных процессорах, поэтому мы активно используем Single Instruction Multiple Data (SIMD) расширения. SSE, AVX и NEON уже стали привычными, но в свежих версиях архитектуры ARM были представлены новые расширения: SVE, SVE2, а также SME. Вот о них мы и поговорим в этой статье.

Но сначала объясним, почему это важная тема. В последнее время каждая компания завела свой собственный нейроускоритель или сопроцессор: Google TPU, Intel Movidius, Apple Matrix, Huawei Ascend и т.д. Они здорово ускоряют вычисления в глубоких нейронных сетях, но для разработчика кроссплатформенного программного обеспечения общего назначения они совсем непрактичны.

Очевидные причины:

• заводить ветку кода под каждый нейроускоритель крайне трудозатратно;

• к каждому ускорителю (в лучшем случае) прилагаются свои библиотеки, работающие на ограниченном наборе операционных систем;

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

Неочевидные причины

Зайдем издалека: результат умножения даже двух вещественных матриц на разных устройствах может немного отличаться из-за 1) разного порядка вычислений и 2) разной аппаратной реализации. Первая проблема известна давно и существуют специальные алгоритмы действительно точных вычислений с числами с плавающей точкой (см., например, тут). Однако для исполнения нейронных сетей они не применяются, так как там в первую очередь важна скорость. За иллюстрацией второй проблемы также далеко ходить не надо: примером тут могут служить fused multiply-add (FMA) устройства на ARM: компилятор может объединить умножение и сложение в одну инструкцию, которая будет исполняться быстрее и без дополнительного округления. Результат станет немного точнее, но он разойдется с результатом на х86!

Пока у нас есть только один тип данных, мы можем смириться с тем, что на разных платформах парой распознанных примеров будет больше — парой меньше. Однако нейронные сети для действительно быстрого исполнения должны использовать малобитные типы данных. Вот только для одного ускорителя это будет fp16, для другого bfloat16, INT8 или INT4, которые также могут быть имплементированы по-разному. Помимо этого, существует ещё и разные схемы квантования для этих типов данных!

В результате, для достижения максимальной производительности придется обучать свою модель для каждого из устройств. Разумеется, эти модели будут иметь немного разное качество, что будет огорчать пользователей и доводить до нервного тика ML-инженеров. 

Таким образом, эти ускорители могут помочь разработчикам программно-аппаратных платформ, которые точно знают свое железо, ОС и доступный набор библиотек. Мы же к ним не относимся и работаем на разных платформах с центральными процессорами x86, ARM и Эльбрус. На MIPS тоже, когда это кому-нибудь нужно. Поэтому мы используем SIMD – стандартизованные расширения для каждой архитектуры процессора.

Scalable Vector Extensions

В 2016 году в ARMv8-A в качестве опционального расширения было представлено Scalable Vector Extension (SVE). Важное нововведение: длина SVE регистра должна быть степенью двойки от 128 до 2048 бит и определяется производителем процессора. Таким образом производитель может выбирать оптимальную длину векторных регистров с учётом задач процессора, желаемого энергопотребления и других факторов. Таких регистров 32, как и в NEON, причем они являются их непосредственным расширением. Также есть 16 предикатных регистров, с помощью которых и осуществляется обработка данных блоками неизвестного заранее размера. Результирующий код может быть перенесен между любыми устройствами с SVE без ограничений. В 2019 году появился первый процессор с SVE – суперкомпьютерный Fujitsu A64FX c 512-битным SVE. Еще через год появились процессоры Neoverse c 128- или 256-битным SVE.

Затем в 2019 году поддерживаемых инструкций стало больше: было анонсировано расширение SVE2. Оно является надмножеством команд SVE.

В процессорах архитектуры ARMv9-A наличие имплементации SVE и SVE2 обязательно. 

Кроме того, в 2021 году было анонсировано опциональное расширение ARMv9 для матричных вычислений: Scalable Matrix Extension (SME). Оно включается только в специальном режиме SVE Streaming Mode и позволяет эффективно выполнять внешнее умножение вектора на вектор и транспонирование. При этом становится доступным дополнительное хранилище для аккумуляторов, новые инструкции, а длина регистров может возрастать.

На этом описательная часть закончена, перейдем к экспериментальной. На дворе настал 2024 год, процессоры с ARMv9.2 вышли на рынок. 

Мы купили и дождались из Китая телефон с Snapdragon 8 gen 3. Популярный процессор для флагманских моделей телефонов на ближайшие время с одним ядром Cortex-X4, 5 Cortex-A720 и 2 Cortex-A520. В описании процессора ничего не написано про поддерживаемые SIMD-расширения, но это нас не насторожило. 

Каково же было наше удивление, когда getauxval(AT_HWCAP) & HWCAP_SVE показал отсутствие SVE! SVE2 и уж тем более SME тоже не было. Мы пошли дальше и запустили код, содержащий SVE инструкции и получили падение с illegal instruction.

При этом в Arm Cortex-X4 Technical Reference Manual имплементация SVE вполне фигурирует.  Тогда мы решили погуглить и обнаружили несколько сообщений от пользователей форумов AnandTech, HackerNews и твиттера Х, в которых они писали о том, что именно Qualcomm блокирует доступ к SVE на уровне гипервизора.

Мы решили проверить эту информацию и протестировали ряд мобильных процессоров, которые нашли в офисе. Результаты представлены в таблице 1.

Таблица 1. Векторные расширения в мобильных процессорах ARM (согласно флагам процессора/фактическое наличие)

Обратите внимание на левый нижний угол: при наличии SVE2 и фактическом наличии SVE, процессоры это отрицают!

В целом выводы оказались следующие: Qualcomm'овские Snapdragon правда крайне популярные. Однако возможности использовать SVE, SVE2 и SME в пользовательском софте там нет. Последние процессоры MediaTek (Dimensity 9000 и выше) и Exynos (2200 и 2400) поддерживают SVE и SVE2. Длина регистра при этом 128 бит, то есть кратного ускорения за счёт обработки большего количества данных не ожидается, только за счет увеличения пропускной способности (возможно) и наличия новых инструкций.

При подготовке данного материала мы также нашли вот этот проект. Результаты автора совпадают с нашими за исключением Snapdragon 8 gen 1, на котором мы SVE не обнаружили.

Таким образом, мы не опоздали с освоением новых векторных расширений, а даже немного поторопились: сейчас SVE и SVE2 есть лишь на небольшом количестве телефонов, процессоры которых и так достаточно мощные и поддерживают 128-битный NEON. 

P.S. Все процессоры Apple имеют архитектуру ARMv8 и SVE на них отсутствует. Там есть векторно-матричный сопроцессор AMX, пользоваться которым разработчики могут через фреймворк Accelerate, а это уже совсем другая история.