Привет, Хабр! Меня зовут Илья Знаменский, я ведущий инженер в группе оптимизации алгоритмов искусственного интеллекта в AI-дивизионе YADRO. 

Популярность RISC-V растет стремительными темпами, и на рынке появляется все больше новых отладочных плат. Моей команде поставили задачу: узнать, как эти платы будут справляться с простыми AI-нагрузками (задачи запуска LLM-on-device не стояло). В процессе работы мы внесли вклад в развитие собственного тензорного компилятора и создали библиотеку математических ядер, которая позволила существенно увеличить производительность инференса моделей на RISC-V. С какими трудностями мы столкнулись и что в итоге из всего этого получилось — читайте в статье.

Архитектурный бум: почему RISC-V?

Архитектура RISC-V уверенно развивается: формируются новые рабочие группы, активно работают комитеты, появляются стартапы. Ее популярность во многом обусловлена открытостью и масштабируемостью. Так она проектировалась изначально: есть компактный базовый набор инструкций и широкий спектр модульных расширений, что позволяет масштабироваться от простых встроенных микроконтроллеров до высокопроизводительных ядер и специализированных ускорителей. Любой разработчик или компания могут использовать и расширять базовую ISA RISC-V без лицензий и бесплатно, а это снижает порог входа и стимулирует инновации. 

Консалтинговое агентство Omdia прогнозирует, что объемы поставок процессоров на базе RISC-V будут в среднем расти на 50% ежегодно до 2030 года и к концу этого периода достигнут 17 млрд штук в год.

Еще в 2023 году рынок СнК на архитектуре RISC-V превысил отметку в 6 млрд долларов — это в четыре раза больше уровня 2022 года. Архитектура внедряется технологическими гигантами, такими как Google, NVIDIA, Bosch, Infineon, Nordic, а также сообществом разработчиков Qualcomm. Они видят в RISC-V потенциал, поэтому немало инвестируют в ее развитие и распространение.

Как архитектура успешно нашла свое применение в мобильных устройствах— читайте в нашей статье. 

Научное сообщество также активно включилось в развитие RISC-V: создаются исследовательские центры и лаборатории, многие вузы сотрудничают с представителями индустрии. Новейшие разработки быстро находят отражение в IP-блоках и отладочных платах.

AI-сегмент не остается в стороне: ожидается уверенный рост до ~550 миллионов поставок к 2031 году. И уже сейчас процессоры с RISC-V находят применение в различных устройствах, связанных с работой систем искусственного интеллекта. Так, например, SiFive удалось успешно запустить LLM на RISC-V за счет компиляторного подхода. Компания доказывает, что RISC-V с векторными расширениями — это серьезная альтернатива в AI-нагрузках таким традиционным архитектурам как x86 и ARM.

Сначала были страусы

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

Несколько лет назад существовали определенные трудности с отладкой на платах с RISC-V: их было мало, многие из них были достаточно «сырыми», поддерживали они лишь набор самых стандартных расширений RISC-V (IMAFDC). На тот момент ни о какой поддержке векторных расширений речи не шло. QEMU же «из коробки» предоставлял полностью функциональный RISC-V процессор, пусть и виртуальный, но с поддержкой RVV и широк��ми возможностями отладки через GDB. Интеграция QEMU в CI/CD позволила нам начать предварительные эксперименты с оптимизациями под RVV до того момента, как мы заказали и получили платы с поддержкой векторного расширения. 

Функциональная отладка на QEMU работает без нареканий, однако нет возможности оценить, какой будет производительность на самом деле. Различные симуляторы (например, gem5) тоже не решают проблему. Они медленные: реальная тактовая частота отладочной платы в любом случае будет выше, чем у симулятора. К тому же, еще и неточные: симулятор может неплохо справляться с простыми сценариями, но при оценке производительности сложного многопоточного кода может дать погрешность в десятки процентов.

«Не попробуешь на реальной плате — не узнаешь», — поняли мы. И, когда на рынке появились RISC-V платы с поддержкой векторного расширения RVV, приступили к тестированию.

Disclaimer. RISC-V развивается на высоких скоростях: каждый месяц приносит нам новости о новых разработках и продуктах на открытой архитектуре. Пока мы занимались исследованием и писали статью, появились более современные платы, которые в этот обзор не вошли, — например, Milk-V Jupiter и Milk-V Megrez. Если статья вызовет интерес у комьюнити, мы бы выпустили продолжение с обзором этих и других новых плат. Оставляйте комментарии и stay tuned!

Обзор отладочных плат: различные фрукты, Kendryte и бракованный «пионер»

Mango Pi MQ Pro

Эта плата совсем небольшая, примерно 6,5 на 3 см. Она основана на процессорном ядре Xuantie — T-Head C906. Единственное ядро имеет небольшую тактовую частоту в 1 ГГц, обусловленную проектными нормами 22 нм. Поэтому чудес в плане производительности от Mango Pi ожидать не стоит. Но зато она отлично подойдет в качестве недорогого базового инструмента для желающих опробовать RISC-V вживую. 

Mango Pi работает под управлением операционной системы Armbian. Есть базовый набор интерфейсов, чтобы использовать Mango Pi в качестве одноплатного компьютера. Но это скорее удел энтузиастов, так как у платы не слишком высокая производительность. Нас же она интересовала наличием векторного расширения промежуточной версии RVV 0.7.1. 

RVV реализован в C906 через кастомное расширение XTheadVector. Векторные регистры в C906 имеют ширину 128 бит, что при объединении регистров в группы (LMUL) дает возможность работать с большим числом операндов в рамках одной векторной инструкции. Более подробное описание характеристик VPU процессора C906 можно найти на странице с результатами бенчмарков платы. 

Мы запустили на этой плате инференс более чем десяти популярных моделей нейросетей: классификация, сегментация и другие. Производительность кода была увеличена примерно в два раза за счет оптимизации под RVV, но чуда не случилось: одно ядро на невысоких частотах даже с векторизацией кода работало значительно медленнее, чем хотелось бы. Заявленные 4 Гфлопс пиковой производительности на реальных compute-bound нагрузках нам достигнуть не удалось. 

Мы интегрировали плату в CI/CD, и там она успешно работала, пока на рынок не вышла ее «старшая сестра» — Lichee Pi 4A.

Lichee Pi 4A

Sipeed, производитель Lichee Pi 4A, заявляет, что эта плата с базой из четырех ядер Xuantie T-Head C910 — аналог Raspberry Pi 4 с улучшенной производительностью и расширенными функциональными возможностями. Lichee Pi — модульное устройство, состоящее из расширительной платы и вычислительного модуля LM4A. Наш опыт подтверждает, что у Sipeed вышел достаточно качественный продукт: очевидных проблем с надежностью мы не заметили. Lichee Pi 4A действительно способна обогнать RPI4, если затратить определенные усилия на оптимизацию ПО. 

Плата оснащена 16 ГБ памяти LPDDR4x, отдельными GPU и NPU IP-блоками Imagination. Модульная архитектура, GPU с пиковыми ~50 Гфлопс и современный набор интерфейсов позволяют использовать Lichee как основу для разных устройств — есть, например, мини-ноутбук Lichee Console 4A и планшет Lichee Pad 4A. Игровой производительности от платы ожидать не стоит, но для веб-серфинга и видео в среднем качестве возможностей GPU достаточно.

СнК T-Head TH1520 изготовлена по проектным нормам 12 нм. Четыре ядра C910 (до 1,85 ГГц) имеют L1i / L1d кеши объемом 64 КБ каждый и общий кеш L2 1 МБ. Архитектура соответствует RV64GCV и, как Mango Pi, поддерживает кастомное векторное расширение XTheadVector (RVV 0.7.1). Векторный блок быстрый — задержки 3–5 тактов с пропускной способностью до двух инструкций за такт. Но есть особенности: операции «вектор + скаляр» иногда заметно медленнее, чем «вектор + вектор». На других архитектурах такие инструкции, напротив, дают дополнительное ускорение за счет лучшей утилизации векторных регистров.

Ограничения подсистемы памяти и кешей ощутимы: на типичных задачах ИИ реальный пик производительности достигает лишь ~50% от заявленного. Наши тесты также показали, что пропускная способность памяти у Lichee Pi заметно ниже, чем у х86-процессоров Intel и AMD.

NPU — это чип Imagination AX3386 NNA c заявленной пиковой производительностью 4 TOPs @ int8. Для инференса нейросетей используется фреймворк HHB на базе Apache TVM. Поддерживаются популярные форматы моделей. Но, когда мы тестировали фреймворк, он был сыроват: без доработок запускались только базовые модели, и это не позволило в полной мере оценить возможности NPU.

Несмотря на ограничения, Lichee Pi 4A оказалась для нас лучшим вариантом по сочетанию стабильности, производительности и функциональности. Мы отказались от Mango Pi в CI, полностью перейдя на Lichee, и добавили несколько плат для удаленного доступа разработчикам.

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

→ AI Tech Lead (LLM)

→ Senior Full-Stack Developer (JavaScript + Python) 

Kendryte CanMV-K230

Мы ждали Kendryte CanMV-K230 как первую доступную плату с официальной поддержкой RVV ратифицированной версии 1.0. Она построена на двух ядрах Xuantie C908 с 8-стадийным in-order конвейером и интересна кастомным ускорителем для квантованных нейросетей — расширением XTheadVdot.

В базовой спецификации RVV нет поддержки int8-инференса с накоплением в 32-битных аккумуляторах, что приводит к переполнению результата. Приходится прибегать к промежуточным 16-битным регистровым буферам. Этого удается избежать в x86 за счет инструкции VPDPBUSD и в ARM за счет инструкций SDOT, UDOT, SUDOT — опционально в ARMv8.2/v8.3, обязательно начиная с v8.4. Эти инструкции позволяют значительно ускорить инференс квантованных нейросетей и применяются в популярных библиотеках oneDNN и XNNPACK.

Из-за отсутствия аналогов в RVV T-Head создала расширение XTheadVdot, и C908 его поддерживает. Мы подтвердили его работоспособность, но не успели оценить эффективность. Для сравнения: применение SDOT на ARM и VPDPBUSD на x86 дает ускорение в 1,5–2 раза на int8-свертках.

Banana Pi BPI-F3

Эта плата пришла к нам одной из последних. Она разработана компанией SpacemiT и содержит восемь процессорных ядер (SpacemiT K1 Octa-core X60), соответствующих спецификации RV64GCVB. Заявлена полная поддержка профиля RVA22. CPU-ядра объединены в два кластера по четыре ядра в каждом. Плата поддерживает версию RVV 1.0 и имеет широкие 256-битные векторные регистры, что в два раза больше, чем у Lichee Pi.

Одно из ключевых преимуществ платы — поддержка кастомного матричного расширения, реализованного на одном из процессорных кластеров. По классификации оно относится к классу интегрированных: SpacemiT IME переиспользует векторные регистры, соответствующие контрольные регистры и регистры состояния, обеспечивая за счет этого практически полную бинарную совместимость с RVV. Есть поддержка большого числа форматов данных низкой разрядности: от int4 до int16, а также плавающей точки от 4-разрядной до 16-разрядной и в блочном формате bf16. Пиковая производительность процессорного кластера, поддерживающего IME, составляет более 2 TOPS (по ~0,5 TOPS на одно CPU-ядро) при умножении 8-разрядных целочисленных операндов с накоплением в 32-разрядных аккумуляторах.

Пиковая производительность SpacemiT K1

В бенчмарке Pigirons есть поддержка RVV 1.0 и SpacemiT IME. Ниже на картинках представлены результаты замеров пиковой производительности SpacemiT K1 в разных конфигурациях из: 

• двух кластеров (восьми ядер), FP16/FP32/FP64, инструкции vfmacc.{vf, vv};

• одного процессорного ядра с поддержкой IME, S8/U8/S32, FP16/FP32/FP64, инструкции vmadot{u, us, su, slide}, vfmacc.{vf, vv};

• одного кластера (четыре ядра) с поддержкой IME.

Для оценки пиковой производительности операции доступа к памяти не используются, а также минимизируются накладные расходы на организацию циклов, ядро которых — многократные повторения операции «умножения с накоплением».

Больше про виды и применение матричных расширений читайте в статье Валерии Пузиковой, эксперта группы разработки математических библиотек в YADRO.

Мы не успели потестировать IME, но идея реализации матричного расширения на основе уже имеющегося векторного нам показалась интересной. Так можно реализовать нечто похожее на NPU без интеграции внешнего IP-блока. Если у вас такой опыт был — поделитесь им в комментариях! 

Milk-V Pioneer

Отдельным пунктом в графе «интересы» у нас стояла серверная история на RISC-V. Плату Milk-V Pioneer на базе СнК Sophon SG2042 можно считать первым решением на базе RISC-V серверного класса. У нее 64 ядра C920, основным отличием которых от C910 является увеличенная до ~2 ГГц тактовая частота. Организация — шестнадцать кластеров по четыре CPU-ядра в каждом, индивидуальный кеш L1 (64 КБ + 64 КБ) и общий L2 размером 1 МБ на кластер. Также присутствует системный кеш L3 размером 64 МБ. На плате установлено 128 ГБ DDR4 RAM и 1 ТБ SSD.

Мы планировали установить обе платы в корпуса, подключить периферию. Для этого приобрели и подготовили все необходимые комплектующие. Даже создали какие-то задачи по тестированию серверных нагрузок, чтобы ускорить дальнейший процесс, пока платы будут корпусировать и настраивать наши DevOps-инженеры. Но, к сожалению, серверная история с RISC-V закончилась, так и не начавшись: на платах присутствовали следы некачественной пайки и перегрева, и попытки оживить их до уровня вывода чего-то осознанного в консоль не дали успеха. 

Вполне возможно, что нам просто не повезло — обе платы мы вернули по гарантии. Если у вас получилось успешно запустить Milk-V Pioneer, опишите свои кейсы в комментариях. Нам будет интересно почитать про ваш опыт. 

RISC-V инференс, или зачем нам эти платы

Тензорный компилятор YADRO

Разработку собственного тензорного компилятора в YADRO начал Владислав Виноградов: он подробно описал процесс в своей статье на Хабре. 

Тензорный компилятор — это отдельный проект или компонент в фреймворке, который занимается оптимизацией и переводом модели глубокого обучения в исполняемый формат под конкретное устройство. Этим устройством может быть CPU, GPU или специализированные AI-акселераторы — NPU или TPU.

Сейчас мы активно развиваем собственный компилятор, который вырос в отдельный фреймворк для работы с моделями глубокого обучения. В его состав входят такие компоненты, как Inference Engine, сам компилятор, а также рантайм и утилиты. Дополнительно идут плагины, делегаты и компиляторные фронтенды для OpenVINO и TFLite.

Компилятор поддерживает фреймворки OpenVINO и TFLite, работает с операционными системами Linux и Android и является кросс-платформенным: обеспечена совместимость с х86-64, ARM-64 и RISC-V-64. Также в компилятор интегрирована экспериментальная поддержка GPU и NPU (RK3588) для реализации гетерогенного режима работы. Бэкенд оптимизирован за счет использования известных библиотек для ускорения инференса: oneDNN, XNNPACK, ACL, а также собственной проприетарной библиотеки.

Изначально наша библиотека была заточена только под RISC-V. В дальнейшем мы подстроились под требования бизнеса и добавили оптимизированные имплементации математических ядер для ARM64 и x86.

Что касается RISC-V, то в библиотеке мы оптимизировали и покрыли тестами порядка 30 наиболее частотных слоев нейросетей: Convolution, Pooling, Reorder, Split и многие другие. Помимо этого, была добавлена интеграция с must-have инференс-библиотеками oneDNN и XNNPACK. 

Практически все оптимизации под RISC-V были выполнены для формата данных FP32, так как он поддерживался на уровне спецификации RVV. Чего нельзя сказать про FP16: поддержка этого формата появилась значительно позже с расширением Zvfh. Так как в системе инструкций RVV нет нужных команд для эффективного int8-инференса, FP32 виделся единственным адекватным, простым и дающим максимально точные результаты решением. 

В библиотеке мы реализовали экспериментальную поддержку int8-ядер для ускорения примитивов Convolution и Fully Connected с использованием RVV.  На Lichee Pi удалось получить примерно двукратный прирост производительности по сравнению с FP32. Чтобы избежать переполнения 16-битных аккумуляторов, анализируются значения весов и определяется, сколько итераций умножения с накоплением можно сделать перед прибавлением промежуточного 16-разрядного результата к значению 32-разрядного. Схема получается сложной и уменьшает число доступных векторных регистров из-за буферизации в 16-битных аккумуляторах, но все же дает заметный выигрыш в производительности.

В целом, основную часть тестов по оптимизации инференса и сам компиляторный подход мы прогоняли именно н�� плате Lichee Pi c RVV 0.7.1, так как первые платы с поддержкой RVV 1.0 — Kendryte CanMV-K230 и Banana Pi BPI-F3 — пришлось ждать довольно долго. Их возможности мы успели оценить только поверхностно, но в будущем протестировали бы их в полной мере. 

Наши библиотека и компилятор поддерживают тулчейны Xuantie и Синтакор, работают как с решениями на базе T-Head через Xuantie-тулчейн, так и с релизной версией RVV 1.0. Код можно собирать с RVV и без него, например, для запуска на RV64GC платах. Мы не оптимизировали скорость скалярного кода, сосредоточившись на работе с RVV.

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

Оптимизации RISC-V кода с использованием RVV

При оптимизациях инференса на RISC-V мы старались руководствоваться законом Парето: 20% усилий дают 80% результата. Компиляторные оптимизации и векторные интринсики приносят быстрый результат при достаточной производительности. Если полученных результатов недостаточно, то, например, для наиболее интенсивных вычислительных примитивов типа Convolution или GEMM мы прибегали к «тяжелой артиллерии»: RISC-V Assembler и более продвинутая ступенька — RISC-V JIT Assembler на базе Xbyak. Они как раз и давали нам необходимый performance boost. Соответственно, чем больше усилий мы затрачивали, тем лучше была производительность. Обо всем по порядку.

RVV-интринсики

Компиляторные интринсики — простой и распространенный способ ускорить код за счет векторизации, и в RISC-V также есть свой набор интринсиков. После обновления тулчейна Синтакор мы столкнулись с изменениями в синтаксисе интринсиков. Поэтому мы добавили в код специальные обертки, скрывающие детали реализации, тем самым упростив жизнь программиста. Это позволило собирать код под интринсики разных версий. Такое решение обратно совместимо со старыми версиями тулчейна Синтакор, а также поддерживает интринсики для версий RVV 0.7.1 и 1.0.

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

• нужно получить хороший прирост производительности без глубоких доработок,

• нет задачи получить максимальную производительность на конкретном слое нейросети.

Но для самых нагруженных операций, например Convolution и Fully Connected, интринсиков было недостаточно. И здесь нам пришлось прибегнуть именно к нему. К ассемблеру.

RISC-V ассемблер

Осваивать RISC-V ассемблер мы начинали со вставок, поэтому в библиотеке есть несколько демонстрационных ядер вроде простого eltwise_add — поэлементного сложения двух тензоров. Для полноценной оптимизации свертки пришлось написать ее целиком на RISC-V ассемблере с поддержкой RVV. Больше всего времени ушло на разбор нетривиальной передачи параметров в ассемблерное ядро (calling convention) для соответствия RISC-V ABI, затем — на реализацию самого алгоритма свертки и ее объединения (fusing) с типичными пост-операциями (ReLU, Add, Sum). Чтобы упростить разработку, мы использовали макросы.

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

JIT-ассемблер на базе Xbyak

Мы выжали из RISC-V ассемблера все, что могли, но достигли границы его возможностей: дальше требовалось гибкое решение, которое сохранило бы скорость низкоуровневого кода и позволило бы конфигурировать его под конкретные параметры примитивов. Таким решением стал JIT-ассемблер.

Технология JIT (Just-In-Time) — это возможность компиляции кода непосредственно во время выполнения программы, что дает полный контроль при оптимизации и конфигурируемость генерируемого кода. Существуют разные библиотеки JIT-генерации для C++: Xbyak, Biscuit, Asmjit, MIR. Мы изучили несколько из них и выбрали Xbyak — в первую очередь из-за его тесной интеграции с библиотекой oneDNN, которую мы уже активно использовали.

Xbyak — это header-only библиотека на C++. Вклад в ее портирование на RISC-V во многом внес наш коллега Павел Замелин. В библиотеку была добавлена поддержка основных расширений RISC-V ISA (IMAFDC) и RVV 0.7.1 и 1.0.

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

Для сравнения: пример кода на RISC-V Assembler —

# Multiplication and accumulation (for aligned and unaligned kernels):
# sources: ft0 - ft7 (packed 8 input channels for 1 output width value)
# weights: v0 - v15 (firstly packed 8 output channels, secondly 8 input channels)
# accumulator: v16 - v31 (packed 8 output channels for 1 output width value)
.macro fma_calculations vreg_acc
  vfmacc.vf \vreg_acc, ft0, v0
  vfmacc.vf \vreg_acc, ft1, v2
  vfmacc.vf \vreg_acc, ft2, v4
  vfmacc.vf \vreg_acc, ft3, v6
  vfmacc.vf \vreg_acc, ft4, v8
  vfmacc.vf \vreg_acc, ft5, v10
  vfmacc.vf \vreg_acc, ft6, v12
  vfmacc.vf \vreg_acc, ft7, v14
.endm

# Multiplication and accumulation (for depthwise kernel):
# sources: v0 (packed 8 input channels for 1 output width value)
# weights: v2 (packed 8 groups for 1 output width value)
# accumulator: v16 - v31 (packed 8 output channels for 1 output width value)
.macro fma_calculations_depthwise vreg_acc
  vfmacc.vv \vreg_acc, v0, v2
.endm

# Multiplication and accumulation (for the very first convolution kernel):
# sources: ft0 - ft2 (3 input channels for 1 output width value)
# weights: v0 - v5 (firstly packed 8 output channels, secondly 3 input channels)
# accumulator: v16 - v31 (packed 8 output channels for 1 output width value)
.macro fma_calculations_first_conv vreg_acc
  vfmacc.vf \vreg_acc, ft0, v0
  vfmacc.vf \vreg_acc, ft1, v2
  vfmacc.vf \vreg_acc, ft2, v4
.endm

Как тот же код выглядит на Xbyak:

// Multiplication and accumulation the result
void jit_rvv_conv_kernel::fma_calculations(const Xbyak_riscv::VReg &vreg_acc) {
  const int ic_size = jcp.partial_weights_hint
    ? jcp.ic_subblock_size
    : jcp.ic_block_size;
  if (!jcp.is_dw_convolution) {
    for (int ic = 0; ic < ic_size; ++ic) {
      for (int ow = 0; ow < jcp.ow_subblock_size; ++ow) {
        vfmacc_vf(vReg(vreg_acc.getIdx() + ow * jcp.vreg_group_size), FReg(jcp.freg_src_begin + ow * jcp.ic_block_size + ic), VReg(jcp.vreg_weights_begin + jcp.vreg_group_size * ic)); 
      }
    }
  } else {
    vfmacc_vv(vreg_acc, VReg(jcp.vreg_src_begin), VReg(jcp.vreg_weights_begin));
  }
}

Добавляя производительности, Xbyak также упрощает разработку: код пишем на C++, поэтому он компактнее, читается и отлаживается проще. В отличие от «жестко пришитых» ассемблерных реализаций, JIT дает гибкость без потери скорости. Можно, например, обращаться к регистрам по индексам и использовать это в циклах, а в Compile-Time — убирать ветвления, зависящие от константных значений.

Попробуем представить оптимизацию свертки на «чистом» RISC-V ассемблере. Это «проклятие размерности»: миллионы возможных комбинаций разных размеров ядра, strides, dilations, pads, активационных функций, которые желательно «склеивать» со сверткой. В итоге либо приходится писать десятки отдельных реализаций под разные комбинации, либо ограничиваться несколькими ключевыми вариантами и мириться с потерей производительности.

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

Мы реализовали на Xbyak RISC-V несколько самых частотных и тяжелых ядер с оптимизациями под RVV 0.7.1 и 1.0: convolution/deconvolution, pooling, inner product. Был также написан механизм поддержки различных активационных функций в виде пост-операций к примитивам. 

Что мы получили

Тут логично спросить: почему не взять готовое решение — те же TFLite, OpenVINO, ONNX Runtime, — к нему «прикрутить» оптимизированный бэкенд вроде XNNPACK или oneDNN и радоваться качественному инференсу на RISC-V?

Основных причин две: 

• Развитие RISC-V железа опережает развитие фреймворков глубокого обучения и сопутствующих библиотек. Сейчас ситуация уже лучше, но еще пару лет назад в том же XNNPACK почти не было полноценного векторного кода для RISC-V — только хорошо оптимизированные скалярные ядра.

• Быстрое развитие RISC-V привносит некоторый хаос по части тулчейнов и расширений ISA. Крупнейшие вендоры (T-Head, SiFive) не ждут, пока появятся ратицифированные и стандартизированные ISA. Вместо этого они по-быстрому создают свои собственные системы инструкций и тулчейны, чтобы не проиграть на рынке и вывести на него свои чипы раньше конкурентов. На выходе мы имеем несколько версий векторных расширений (как стандартизированных, так и кастомных), калейдоскоп матричных расширений (SiFive VCIX, SpacemiT IME, а также стандартные RISC-V IME / AME) и целый набор тулчейнов и версий интринсиков со своими уникальными особенностями.

Поэтому свой тензорный компилятор при поддержке JIT-ассемблера выглядят достаточно мощной связкой, помогающей не запутаться во всем этом «зоопарке» и получить возможность быстрее проводить эксперименты с разными ISA на новых RISC-V платах.

Запуск нескольких моделей на Lichee Pi через собственный компилятор также впечатлил приростом производительности: примерно в два раза по сравнению с TFLite + XNNPACK. Причем итоговая производительность платы соответствовала уровню Raspberry Pi 4, а там «под капотом» ARM64 с NEON, который оптимизирован через XNNPACK.

Мы тестировали в основном легковесные модели, которые способны работать на мобильных устройствах без NPU, — часть из них используется в типичных задачах обработки изображений в YADRO. На таких задачах Lichee Pi продемонстрировала производительность примерно уровня Cortex-A72 и A73.

Выводы

Мы провели эксперименты по оптимизациям инференса на RISC-V, а также внедрили поддержку этой архитектуры в собственный инференс-фреймворк. Тесты показали: современные отладочные платы с векторными и матричными расширениями можно использовать для запуска легковесных AI-нагрузок на RISC-V-ядрах. Mango Pi была полезна для первых шагов, Banana Pi удивила наличием матричного расширения, а Lichee Pi стала нашим фаворитом по части стабильности и производительности.

Наиболее популярные фреймворки и библиотеки с оптимизированным инференс-бэкендом все еще слабо поддерживают RISC-V. Есть проблемы в дизайнах: может вылезти неприятная errata или illegal instruction, а пропускная способность подсистемы памяти не выдерживает конкуренции с решениями x86 и ARM. Все это оставляет много интересных тем для будущих материалов. Если вам хочется увидеть продолжение экспериментов с новыми платами — дайте знать в комментариях!