Всем привет! Меня зовут Дарья Ильянова, и сегодня я хочу поделиться с вами опытом создания гетерогенной вычислительной системы для нейросетей, над которой мы работали в команде студентов Инженерно-математической школы НИУ ВШЭ и VK. Вместе с Филиппом Баулиным и Артемом Минеевым мы спроектировали прототип аппаратно-программного ускорителя для машинного обучения.

Что такое гетерогенные вычисления?

Гетерогенные вычисления — это не просто объединение различных типов процессоров, таких как CPU, GPU и NPU, а их совместное использование с максимальной эффективностью. Они позволяют оптимально распределять задачи между архитектурами, снижать энергопотребление и обеспечивать масштабируемость системы.

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

Проектирование аппаратной части

Архитектура ускорителя

Основой нашего ускорителя стал аппаратно-программный блок NPU (Neural Processing Unit), совместимый с любой системой-хостом, построенной на базе архитектуры x86. Его вычислительное ядро базируется на систолическом массиве размером 16x16, что делает его эффективным для матричных операций. Данный блок поддерживает тип данных int8, что упрощает выполнение операций с квантованными моделями нейросетей.

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

Сервисное ядро и управление системой

Для управления системой мы использовали открытое решение — NEORV32, процессор RISC-V архитектуры, написанный на VHDL. Это ядро отвечает за выполнение функций активации, управление данными и координацию работы вычислительных блоков.

On-chip DMA: Для перемещения данных между памятью и вычислительными блоками мы внедрили контроллер прямого доступа к памяти (DMA). Это позволило процессору отдавать команды на копирование данных без его непосредственного участия, что снизило задержки и повысило эффективность работы системы.

Взаимодействие с хостом через PCIe

Для взаимодействия ускорителя с хост-системой используется интерфейс PCI Express 2.0. Этот интерфейс не только обеспечивает передачу данных для входных тензоров и результатов вычислений, но и позволяет загружать прошивки и обмениваться командами. Реализация PCIe основана на IP-блоке Xilinx, подключенном к общей AXI-шине системы.

Архитектура памяти

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

Программная архитектура

Подготовка графов и интеграция с TensorFlow Lite

Для интеграции ускорителя с моделями машинного обучения мы использовали TensorFlow Lite, предоставляющий гибкий интерфейс пользовательских делегатов. Делегаты позволили нам оптимизировать выполнение специфических операций, таких как матричное умножение и активации, непосредственно на нашем NPU.

Использование библиотеки FlatBuffers упростило сериализацию моделей, а благодаря поддержке Python мы смогли анализировать и модифицировать графы алгоритмов на этапе подготовки. Мы объединили последовательности операций в единый пользовательский оператор, что снизило затраты на коммуникацию между хостом и ускорителем.

MLIR и оптимизация нейросетей

Мы активно использовали фреймворк MLIR (Multi-Layer Intermediate Representation) для преобразования графов нейросетей в оптимизированный вид. Благодаря диалектам MLIR, таким как TOSA и linalg, нам удалось разложить операции матричного умножения на блоки, которые эффективно обрабатываются систолическим массивом.

Например, мы разбивали матрицы размером 16x1024 и 1024x16 на подблоки, которые загружались в BRAM (Block RAM) перед началом вычислений. Управление этими процессами осуществлялось через 32-битный вектор параметров, задающий размеры матриц и флаги управления.

Оптимизация и вызовы

Одной из сложностей при разработке системы стало эффективное планирование операций нейронных сетей. Чтобы максимально использовать ресурсы ускорителя, мы разработали алгоритмы динамического распределения задач между вычислительными блоками. Это позволило снизить энергопотребление и минимизировать задержки.

Результаты и перспективы

На данный момент наш прототип NPU поддерживает следующие операции: MatMul, Relu, Softmax и Reshape. Он способен эффективно работать с нейросетями небольших размеров, демонстрируя высокую производительность и низкое энергопотребление — пиковая производительность достигла 40 GOPS на частоте 100 МГц (при расчётной 51,2 GOPS), а энергопотребление составило менее 5 Вт.

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

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