Loop fusion is a compiler transformation in which two adjacent loops are merged into a single loop over the same index range. This transformation is typically applied to reduce loop overhead and improve run-time performance.
— Intel compiler guide (перевод).
Привет, Хабр! Меня зовут Пётр Чекмарёв, я старший инженер компании YADRO, занимаюсь компьютерным зрением на мобильных устройствах и низкоуровневой оптимизацией плотных вычислительных функций.
Оптимизация кода — вечная тема, особенно актуальная в дни триумфального шествия искусственного интеллекта. Оптимально написанные, но изолированные ядра сетей составляются в разные последовательности в зависимости от архитектуры модели. Однако, если дать им информацию друг о друге во время компиляции, сеть удастся заметно ускорить. Выгружать программу для перекомпиляции, будь она движком инференса или СУБД — бессмысленно, поэтому компилировать надо во время работы, Just-In-Time. В предыдущей статье AI-дивизиона YADRO Илья Знаменский рассказывал про JIT на базе Xbyak. Продолжая тему, я расскажу про пет-проект векторной JIT-кодогенерации, который я веду, и покажу, как она может помогать в оптимизации.
Оптимизация, векторизация и JIT
Когда возникает задача оптимизации кода, а все мы хотим, чтобы наше приложение или фреймворк не тормозили, есть несколько классических подходов. Если речь идет о вычислительно емких алгоритмах обработки значительных массивов данных, векторизация, то есть раскручивание циклов и задействование векторных инструкций SIMD, обрабатывающих пачку данных за раз, — это один из самых очевидных и мощных способов ускорения кода, работающего с однородными данными. Например, с данными компьютерного зрения, вычислительной физики, нейронных сетей и даже, как ни странно, БД. Причем на мобильных устройствах альтернатива в виде портирования кода на GPU оказывается довольно дорогой, а иногда может даже проигрывать в скорости и энергоэффективности.
Почему это важно для «железа» заказчиков? В B2B-планшетах цена миллисекунды — это не только скорость интерфейса, но и время работы от батареи, нагрев, стабильность в длительных сессиях и предсказуемость на разных ревизиях устройств. Поэтому задачи оптимизации у нас приходят от полевых сценариев клиентов: распознавание/классификация на устройстве, пост-обработка, локальная аналитика — там, где облако недоступно или запрещено политиками безопасности.
Современные компиляторы могут автоматически векторизовать ваш код, но вы об этом не узнаете, пока не полезете в ассемблер. Несмотря на впечатляющее развитие возможностей автовекторизации, практика показывает, что иногда даже с несложными задачами они справиться не могут. Сложные же ядра автовекторизатор может либо не обработать, либо сделать это так, что из получившегося кода ручная оптимизация может выжать кратное ускорение. Например, программист может хитро изменить схему доступа к данным и порядок их обхода, заточив не только алгоритм, но и структуры данных специально под вектора. Поэтому работа с ассемблером и интринсиками более чем востребована.
Интринсики — это специализированные функции, которые почти в неизменном виде пролетают через пайплайн компилято��а. Они плюс-минус соответствуют одной инструкции ассемблера. По сравнению с ассемблером интринсики добавляют очень ощутимый синтаксический сахар:
вы работаете не с ограниченным числом регистров, а с бесконечным числом переменных;
не нужно создавать ассемблерные вставки: код можно отлаживать вместе с остальным.
Конечно, это не избавляет оптимизатора от необходимости переосмыслить алгоритм в терминах векторных операций и даже периодически почитывать генерируемый ассемблерный код, но работа значительно упрощается. Интринсики — это главный инструмент человека, который занимается ручной векторизацией. Он органично и удобно описывает машинерию векторных расширений.
В качестве простого примера посмотрим на упрощенный код функции сглаживания 3x3 на NEON, составляющей для каждого пикселя среднее из пикселей окрестности:
Код
cv::Mat src, dst;
// Специальный тип регистра с четырьмя float'ми
float32x4_t denominator = vdupq_n_f32(1./9.);
constexpr int NEON_32BIT_LANES = 4;
for(int y = 0; y < hd; y++)
{
float* srcPrevLine = src.ptr<float>(y-1);
float* srcCurLine = src.ptr<float>(y);
float* srcNextLine = src.ptr<float>(y+1);
float* dstCurLine = dst.ptr<float>(y);
// Будем обрабатывать пиксели не по одному, а по четыре.
for(int x = 0; x <= wd - NEON_32BIT_LANES; x+=NEON_32BIT_LANES)
{
float32x4_t sum = vld1q_f32(srcPrevLine + x - 1);
// Эти функции со специфическими названиями и есть
// интринсики, add - суммирование, ld - загрузка из
// памяти, mul - умножение, st - выгрузка в память.
// Немного практики - и все становится довольно просто.
sum = vaddq_f32(sum, vld1q_f32(srcPrevLine + x ));
sum = vaddq_f32(sum, vld1q_f32(srcPrevLine + x + 1));
sum = vaddq_f32(sum, vld1q_f32(srcCurLine + x - 1));
sum = vaddq_f32(sum, vld1q_f32(srcCurLine + x ));
sum = vaddq_f32(sum, vld1q_f32(srcCurLine + x + 1));
sum = vaddq_f32(sum, vld1q_f32(srcNextLine + x - 1));
sum = vaddq_f32(sum, vld1q_f32(srcNextLine + x ));
sum = vaddq_f32(sum, vld1q_f32(srcNextLine + x + 1));
vst1q_f32(dstCurLine + x, vmulq_f32(sum, denominator));
}
}Давайте теперь поговорим про JIT и про то, как он связан с оптимизацией. Классический подход к оптимизации — это использование обычного AOT-компилятора C++ с поддержкой интринсиков. Термин «JIT» наверняка хорошо знаком Java-программистам — он отсылает к Hotspot. Это стандартный JIT-компилятор, который на ходу преобразует вызываемые куски байт-кода в нативный код. Последний п��офилируется по ходу и при необходимости дооптимизируется, то есть компилируется с более агрессивными настройками. Поэтому часто люди, которые слышат «JIT-компилятор», представляют себе недоступный для управления программистом процесс автоматической компиляции байткода в машинный код.
Чтобы избежать путаницы, я буду пользоваться термином «JIT-кодогенерация», которая также осуществляется на лету, но управляется программистом. Это совершенно другая практика, заключающаяся в написании программного кода в процессе работы программы с последующим запуском. В той же Java есть библиотека JIT-кодогенерации Janino, использующаяся, например, в Spark, для написания оптимального кода для фильтров конкретных запросов (насколько вообще может быть оптимальным Java-код).
В YADRO мы занимаемся оптимизацией под конкретные сценарии заказчиков. Часть таких сценариев живет на KVADRA_T — корпоративных планшетах: офлайн-инференс на устройстве, ограничения по энергии и термопакету, требования по безопасности и длительной поддержке. Поэтому JIT-кодогенерация и векторные оптимизации для нас — не абстракция, а способ выжать стабильную производительность в реальных B2B-нагрузках без зависимости от облака.
JIT-кодогенерация полезна в оптимизации — она позволяет на лету специализировать обобщенный код под текущий случай. Посмотрим на несколько примеров.
Пример 1. Удаление условных переходов
Самый простой пример — это удаление условных переходов, когда вы точно знаете, что какое-то условие не будет реализовано. AOT-компилятор с этим не справится, потому что ему неоткуда взять информацию о принципиальной нереализуемости ветвления.
Пример 2. Параметризованные ядра нейронных сетей
Пример, точнее передающий идею метапрограммирования, заложенную в JIT — это параметризованные ядра нейронных сетей, таких как свертки разных размеров, ресайзы или пулинги. Возьмем MaxPool, у которого несколько параметров: размер окна, stride и dilation. По классике сначала пишется общая реализация — медленная, но работающая при любых комбинациях параметров. Затем выбираются несколько наиболее часто употребляемых комбинаций ({3x3, stride 1x1}, {3x3, stride 2x2}, {5x5, stride 1x1},...) и переписываются оптимально. Но таких комбинаций много не напишешь, а их число очень быстро растет, увеличивая размер исходника и трудозатраты на написание и поддержку.
Наличие JIT же позволяет реализовать «метаоперацию MaxPool», которая создаст оптимальный код MaxPool для конкретных параметров, конкретного типа данных (float, uint8, …) и под конкретное железо (например, ARM v8.2).
В качестве PoC библиотеки Loops мы ускоряли ядра из экспериментального движка ficusNN, в частности MaxPool. MaxPool — сепарабельное ядро, то есть вместо O(N2) прохода по пикселям окна, его можно разбить на два O(N) прохода, сначала по вертикалям, потом по горизонталям. Однако, результат первого прохода нужно где-то хранить, обычно для этого используют память, а нам благодаря JIT удалось использовать под хранение промежуточного результата регистры. В виде обобщенной AOT-реализации сделать это в принципе невозможно: регистры хардкодятся в бинарное представление программы, перебирать их в цикле по номерам нельзя. Либо вручную для отдельных комбинаций ядер, либо обобщенно с помощью JIT. В результате на типичных ядрах мы проигрывали ядрам ручной оптимизации небольшие проценты почти в пределах погрешности, а на экзотических, вроде 13 на 13, которые встречаются, например, в YOLO, получали ускорение в десятки раз, что позволяло ощутимо ускорять всю нейросеть.
Пример 3. Loop fusion
Лучший пример — это loop fusion, вынесенный в заголовок и давший название проекту. Два идущих подряд цикла с одним диапазоном и порядком, с независимыми данными или данными, зависимость которых можно уложить в итерацию. Циклы можно объединить в один, что позволит уменьшить накладные расходы на организацию циклов, количество операций чтения и записи, а если данные достаточно большие, то избежать лишней прокачки кэша.
Например, когда мы загружаем в память нейронную сеть и строим граф вычислений, мы можем проанализировать граф и понять, что какие-то операции можно и нужно объединить. В качестве примера идеально подойдет применение функций активации после какой-нибудь сложной по данным операции. Применение функции активации это поэлементная операция, поэтому может быть без труда подклеена к предыдущему слою, например к свертке или к полносвязному слою (матричному умножению).
Применимость JIT
В общем нам как оптимизаторам понятно, что JIT — это хорошо, векторный JIT — очень хорошо. А есть ли недостатки? В теории — нет. А вот чтобы решить, нужен ли вам на практике в вашем проекте JIT, нужно для себя ответить на несколько вопросов:
Скорость компиляции. Все-таки мы занимаемся оптимизацией, поэтому будет странно, если мы полминуты будем собирать код, чтобы ускорить его исполнение на долю секунды.
Добавление зависимостей. Компиляторы — это обычно большие проекты, интегрировать их в свой проект — значит усложнить сборку, число трудозатрат на поддержку этой сборки и заметно увеличить объем приложений.
Простота использования JIT. В идеале нам хотелось бы писать код как обычно мы пишем код с интринсиками, при этом каким-то образом указывая, какие переменные станут в собранном коде константным значением параметра, а какие будут именно переменными.
Кросс-платформенность. Векторизация — это низкоуровневая практика, сильно завязанная на архитектуру. Можно ли как-то добиться независимости наших кодогенераторов от типа машины?
Мы, в составе небольшой команды из Вадима Писаревского и меня, постарались дать хорошие ответы на эти вопросы в нашем open source-проекте.
Мы задаем эти вопросы не в вакууме: в корпоративных планшетах и других персональных устройствах заказчиков любые «дорогие» решения быстро проявляются в эксплуатации — в автономности, температуре, SLA и стоимости владения. Поэтому R&D-задачи по JIT/векторизации у нас формулируются от требований клиентов: ускорить конкретный пайплайн, уложиться в энергобюджет, убрать зависимость от сети/облака, обеспечить воспроизводимость производительности на парке устройств.
Истоки проекта
Одним из источников вдохновения для нас стали universal intrinsics – технология из OpenCV. Фактически это abstraction layer. Такой подход позволяет писать код на векторных интринсиках независимо от платформы. Среди поддерживаемых платформ, например, Intel (AVX2/AVX512), ARM (Neon) и RISC-V (RVV). Благодаря технологии universal intrinsic большая часть векторизованных алгоритмов OpenCV написана разом под все поддерживаемые CPU-архитектуры. Это AOT-технология, поэтому ожидать способности избегать комбинаторного взрыва или ускорения кода с помощью слияния циклов тут не нужно.
Приглашаем в ИИ-команду YADRO, у нас есть задачи практически на любой вкус и для любого уровня:
→ AI Tech Lead (LLM)
→ Senior AQA Engineer (AI)
→ Product Owner AI Solutions
Интересно, что установить полное соответствие между разными архитектурами невозможно: ряда инструкций нет в NEON, а других — в AVX2. Поэтому abstraction layer иногда приходится разворачивать отдельные его инструкции в небольшие последовательности нативных интринсиков конкретной платформы. Рассмотрим тот же пример с blur 3x3 на универсальных интринсиках:
Код
cv::Mat src, dst;
// Специальный тип регистра с четырьмя float'ми
float32x4_t denominator = vdupq_n_f32(1./9.);
constexpr int NEON_32BIT_LANES = 4;
for(int y = 0; y < hd; y++)
{
float* srcPrevLine = src.ptr<float>(y-1);
float* srcCurLine = src.ptr<float>(y);
float* srcNextLine = src.ptr<float>(y+1);
float* dstCurLine = dst.ptr<float>(y);
// Будем обрабатывать пиксели не по одному, а по четыре.
for(int x = 0; x <= wd - NEON_32BIT_LANES; x+=NEON_32BIT_LANES)
{
float32x4_t sum = vld1q_f32(srcPrevLine + x - 1);
// Эти функции со специфическими названиями и есть
// интринсики, add - суммирование, ld - загрузка из
// памяти, mul - умножение, st - выгрузка в память.
// Немного практики - и все становится довольно просто.
sum = vaddq_f32(sum, vld1q_f32(srcPrevLine + x ));
sum = vaddq_f32(sum, vld1q_f32(srcPrevLine + x + 1));
sum = vaddq_f32(sum, vld1q_f32(srcCurLine + x - 1));
sum = vaddq_f32(sum, vld1q_f32(srcCurLine + x ));
sum = vaddq_f32(sum, vld1q_f32(srcCurLine + x + 1));
sum = vaddq_f32(sum, vld1q_f32(srcNextLine + x - 1));
sum = vaddq_f32(sum, vld1q_f32(srcNextLine + x ));
sum = vaddq_f32(sum, vld1q_f32(srcNextLine + x + 1));
vst1q_f32(dstCurLine + x, vmulq_f32(sum, denominator));
}
}Второй и главный исток проекта, его прародитель — это Xbyak, библиотека JIT-кодогенерации от Сигэо Мицунари (Shigeo Mitsunari). Она используется в Intel OpenVINO для ускорения инференса нейронных сетей за счет компиляции параметризованных ядер (слоев нейронных сетей), специализации обобщенного кода и слияния циклов. Есть несколько версий Xbyak для:
Все версии Xbyak либо предоставляют возможность закодировать весь набор ассемблерных инструкций конкретной архитектуры, либо стремятся к этому. И это можно только приветствовать! Но именно поэтому нет особого смысла объединять эти проекты в один: точное следование архитектуре процессора ведет к необходимости писать под разные архитектуры разный пользовательский код.
Работа c Xbyak похожа на работу с ассемблером. Вы условно объявляете функцию (на самом деле создаете буфер в памяти) и начинаете вызывать «инструкции», то есть функции, имена которых повторяют названия ассемблерных инструкций. При вызове этих функций вместо исполнения инструкций происходит их добавление в конец буфера. Когда функция завершена, под нее специальным образом выделяется память, на нее можно получить указатель и запустить. Из-за полного соответствия набора функций ассемблерным командам трансляция происходит мгновенно. Это преимущество. Поскольку это JIT-кодогенерация, открываются все возможности ускорения с помощью специализации и слияния циклов.
Главная головная боль программиста на Xbyak – это распределение регистров. Любой реальный процессор имеет ограниченное количество регистров, которые с некоторой натяжкой можно спроецировать на понятие переменной. Например, у x86_64 — 16 регистров общего назначения и 16 векторных в AVX2/SSE или, если есть AVX512 — 32 векторных. Если отрезок кода использует больше переменных, их надо сохранять на стеке до момента использования (то есть имеет место так называемый register spilling).
В Xbyak процессами распределения регистров, загрузки, выгрузки переменных, отслеживания текущей локации — как в стеке, так и на регистре — управляет программист. Это тяжелая рутинная работа, знакомая всем, кто пишет на ассемблере. При использовании C/C++ или любого другого языка высокого уровня эти вопросы возьмет на себя компилятор, разгрузив голову человека. Но в Xbyak, несмотря на то что мета-функция пишется на C++, мы снова спускаемся на уровень ассемблера. В Loops мы тоже решаем этот вопрос. Это, пожалуй, его главное преимущество перед Xbyak.
Схожее преимущество — это работа с вызовами функций. Даже у одного процессора на разных операционных системах может использоваться разный calling convention – договоренность о том, как функции передаются аргументы. Часть из них размещается в регистрах, часть определенным образом укладывается в стек. Для правильного извлечения аргументов и подготовки фрейма стека у функций есть пролог и эпилог, которые в рамках работы с xbyak также должен писать программист. Поэтому одну и ту же Xbyak-функцию нужно писать по-разному для Windows и Linux. В Loops мы автоматизировали эту работу.
Пример с blur 3x3 выглядит громоздко, но зато мы в рантайме можем менять этот код в зависимости от наших потребностей:
Код
// Этот отрезок кода отличается от кода на NEON-интринсиках
// ровно одним: он запустится на Intel и Risc-V так же,
// как на Arm.
cv::Mat src, dst;
v_float32x4 denominator = v_setall_f32(1./9.);
constexpr int UI128_32BIT_LANES = 4;
for(int y = 0; y < hd; y++)
{
float* srcPrevLine = src.ptr<float>(y-1);
float* srcCurLine = src.ptr<float>(y);
float* srcNextLine = src.ptr<float>(y+1);
float* dstCurLine = dst.ptr<float>(y);
for(int x = 0; x <= wd - UI128_32BIT_LANES; x+=UI128_32BIT_LANES)
{
v_float32x4 sum = v_load(srcPrevLine + x - 1);
sum = v_add(sum, v_load(srcPrevLine + x ));
sum = v_add(sum, v_load(srcPrevLine + x + 1));
sum = v_add(sum, v_load(srcCurLine + x - 1));
sum = v_add(sum, v_load(srcCurLine + x ));
sum = v_add(sum, v_load(srcCurLine + x + 1));
sum = v_add(sum, v_load(srcNextLine + x - 1));
sum = v_add(sum, v_load(srcNextLine + x ));
sum = v_add(sum, v_load(srcNextLine + x + 1));
v_store(dstCurLine + x, v_mul(sum, denominator));
}
}Loops. Определение
Loops – это минималистичная скоростная кросс-платформенная C++-библиотека JIT-кодогенерации с распределением регистров, которая предназначена для написания векторного кода. То есть Loops — это JIT-компилятор, оформленный как библиотека, для удобной интеграции его в пользовательские проекты. Основные особенности — ниже.
Минимализм — не пытаясь повторить набор инструкций современных процессоров во всей их полноте, библиотека предоставляет все необходимые целочисленные скалярные инструкции для организации логики программы: базовых вычислений, циклов, условий, а также векторные инструкции, где в полной мере поддерживаются целочисленные операции и операции с плавающей запятой для всех основных типов данных, от 8-битных целых (int8_t/uint8_t) до 64-битных вещественных (double).
Кросс-платформенность — целочисленные скалярные инструкции реализованы для Intel/AMD x64, ARM aarch64 и RISC-V (тоже 64-бит). Векторные инструкции пока реализованы только для x64 (AVX2) и aarch64 (NEON), поддержка расширения RVV для RISC-V в среднесрочных планах. В отличие от Xbyak, ключевая кросс-платформенная особенность Loops — использование одного и того же имени для семантически одинаковых операций, независимо от платформы. Например, вместо mm256max_ps (интринсик для AVX2) или vmaxq_f32 (аналогичный интринсик для NEON) мы просто пишем loops::max, а тип операции (f32) будет автоматически выведен из типов аргументов.
Поддержка операционных систем:
Windows (x86_64),
Linux (x86_64, aarch64, riscv64),
MacOS (aarch64).
Из кроссплатформенности следует необходимость в платформо-независимом промежуточном представлении, о нем подробнее будет рассказано чуть ниже.
Скорость — Loops является очень быстрым многопроходным компилятором. Каждая функция компилируется независимо, и каждый проход имеет теоретическую сложность не больше чем O(N2) от числа команд N в функции, да и те в-среднем работают за O(N) или O(N*log(N)). На данный момент проект еще не профилировался, но предполагается, что можно дополнительно ускорить компиляцию за счет правильного подбора std-контейнеров, мелкой алгоритмической оптимизации и так далее. Loops сложно назвать оптимизирующим компилятором, предполагается, что он будет использоваться людьми с опытом векторизации с помощью интринсиков. Это очевидный и главный недостаток при сравнении его с большими компиляторными проектами вроде LLVM.
Распределение регистров — из-за желания добиться максимальной скорости компиляции вместо тяжеловесных алгоритмов раскраски графа используется усовершенствованный вариант простого алгоритма linear scan. Опыт показывает, что для наших задач этого вполне достаточно для получения качественного кода, особенно на современных платформах вроде ARM или RISC-V с большим количеством регистров: 32. Качество определяется количеством промежуточных сохранений регистров на стек с последующим восстановлением (так называемых «спиллов» / register spill). Чем их меньше, тем лучше. Важнее всего минимизировать количество «спиллов» во внутренних циклах, чего наша реализация linear scan и пытается добиться.
Простота сборки — хотя Loops, в отличие от Xbyak, не является суперлегковесным проектом, состоящим всего из нескольких заголовочных файлов, но довольно к этому близок — он состоит из нескольких заголовочных и исходных файлов на С++. Собирается с помощью CMake и не имеет внешних зависимостей.
Как упростить процесс отладки в CMake с помощью встроенного отладчика и профилировщика — читайте в статье «Как победить CMake: отладка CMake-скриптов».
Внутренняя машинерия
Разберемся с архитектурными решениями проекта.
Кросс-платформенность неизбежно требует введения IR. Исторически, первой архитектурой, которую мы поддержали в Loops, была ARM aarch64 с NEON, поэтому во многом IR повторяет их систему команд, в целом, довольно полную и удобную. Loops позволяет для каждой сгенерированной функции распечатать IR. Также можно распечатать в виде ассемблера получившийся из IR машинный код.
Исходя из минимализма, мы решили сделать структуры данных максимально простыми. IR-инструкции укладываются в простую планарную последовательность, без деревьев вроде AST. Это помешало бы написанию сложных оптимизирующих проходов, но для наших задач в самый раз — просто и быстро. Впрочем, на самом раннем этапе, при сборе кода в буфер, деревья все-таки используются, но это касается в основном интерфейса, а основная работа происходит над планарным представлением.
Идеальной простоты Xbyak, где псевдоинструкция сразу транслирует код и добавляет его в буфер, нам достигнуть не удастся, да и не хочется. Тут есть проходы по функциям как массивам IR-инструкций, уложенные в пайплайн, выполняющий компиляторную работу. Ниже мы разберем проходы по отдельности.
Интерфейс Loops имитирует привычные действия с кодом, как это делает Xbyak, но при этом содержит гораздо больше синтаксического сахара. Здесь вы не встретите команду mov — вместо нее нужно присваивать переменные и делать арифметические операции, что достигается благодаря агрессивной перегрузке операторов. Есть конструкции для цикла while и ветвлений. Одна из странностей — отсутствие оператора goto, что компенсируется другой странностью — операторами break(x) и continue(x), которые позволяют выходить не из самого вложенного цикла, а из нескольких вложений.
В течение работы проходы анализируют инструкции, опрашивая объект бэкэнда для выявления различных свойств той или иной инструкции. Например, самый первый проход после сбора буфера функции — это CP_IMMEDIATE_IMPLANTATION. Он определяет, можно ли вписать константный аргумент прямо в тело инструкции, или придется добавить отдельный mov. На Arm и RISC-V с их фиксированным размером инструкций такие вопросы возникают особенно часто. Для этого достаточно посмотреть, сколько места в бинарном представлении инструкции отведено под операнд, то есть рассмотреть энкодинг, который формируется в самом конце пайплайна и сильно зависит от целевой платформы. Проход работает с платформо-независимым IR, но при этом вынужден опираться на свойства инструкций нативной архитектуры, как бы подсматривать, что будет с этой IR-инструкцией в конце пайплайна. Такие опросы о свойствах инструкции распространены во всем пайплайне и обобщенно называются концепцией peeking. Она делает проходы универсальными, а всю необходимую информацию об инструкциях компактно собирает в одном месте: в энкодинге.
Список компиляторных проходов
Рассмотрим проходы по порядку. Префикс CP_ означает compiler pass.
CP_COLLECTING — это метапроход, обозначающий состояние до начала основной обработки, в котором происходит сбор инструкций для функции, которую мы будем компилировать. Интерфейс при этом все же выполняет некоторую работу: например, разворачивает деревья выражений для присваиваний и использования в усл��виях циклов и ветвлений. На этом этапе циклы и ветвления помечаются специальными аннотационными инструкциями, которые позже исчезнут из кода. Например, для циклов есть три такие аннотации: annotation:whilecstart, annotation:whilecend и annotation:endwhile. Это компенсирует отсутствие древовидной структуры кода и фактически позволяет воссоздавать ее по этим аннотациям в тех проходах, где она нужна: анализе времени жизни и аллокации регистров.
CP_IMMEDIATE_IMPLANTATION — уже упомянутый проход, распределяющий константные значения в телах инструкций, делая peeking энкодинга для определения ширины поля в битах.
CP_ELIF_ELIMINATION — удаляет elif-инструкции, превращая их в систему вложенных if-else-endif.
CP_<ARCHNAME>_BRA_SNIPPETS. BRA_SNIPPETS расшифровывается как Before Register Allocation Snippets и решает задачу, близкую к Universal Intrinsics в OpenCV: разворачивает инструкцию в последовательность инструкций, если на текущей архитектуре нет ее прямого аналога. <ARCHNAME> выделено потому, что на всех машинах имя отличается. В Loops есть система платформо-зависимых проходов, которые добавляются в двух точках: до аллокации регистров и после. Это задел под систему пользовательских проходов.
CP_LIVENESS_ANALYSIS — один из двух наиболее алгоритмически сложных проходов. Он выполняет подготовительную работу для прохода распределения регистров и тесно с ним связан. На выходе у него не только и не столько преобразованный код, сколько карта времен жизни переменных. Кроме того, проход разбивает время жизни переменных на минимально возможные отрезки, что улучшает качество работы аллокатора регистров. Интересно, что в общем случае эта задача решается за O(N^4), однако отказ от произвольных goto позволяет снизить сложность до O(N^2), а в большинстве случаев — до O(M*N), где N — число инструкций, а M — число базовых блоков. Именно поэтому в Loops нет произвольных goto, но есть break(x) и continue(x). На этом этапе выполняется peeking информации о том, какие аргументы инструкции являются входными, а какие — выходными.
CP_REGISTER_ALLOCATION — сердце Loops, собирающее все ниточки в одном месте. Здесь происходит переход от бесконечного числа скалярных и векторных переменных к конечному числу регистров, доступных на целевой машине. Если регистра под переменную не хватает, место под нее выделяется на стеке и добавляются инструкции spill и unspill. Вместо тяжеловесных алгоритмов раскраски графа используется простой, быстрый и, на практике, вполне эффективный алгоритм линейного сканирования. [Poletto, M. Linear scan register allocation / M. Poletto, V. Sarkar // ACM Transactions on Programming Languages and Systems (TOPLAS). —1999. — Т. 21, № 5. — С. 895—913.] Удивительно, но и этот, и предыдущий проход удалось сделать платформо-независимыми, ограничившись peeking энкодинга. Здесь же добавляются пролог и эпилог функции — они тесно связаны с набором использованных регистров.
CP_CONTROLFLOW_TO_JUMPS — простой проход, который удаляет annotation-инструкции и заменяет их обычными командами условного и безусловного перехода.
CP_<ARCHNAME>_ARA_SNIPPETS — это технические ARA-сниппеты (After Register Allocation), которые зависят не только от платформы, но и от ее регистров, поэтому выполняются после распределения регистров. Например, операция скалярного целочисленного деления на Intel требует, чтобы делимое и результат деления находились в регистре rax. Поэтому rax сохраняется на стеке, используется для деления, а затем восстанавливается. Одна из самых важных задач, которые решают ARA-сниппеты — вызов функций из сгенерированного кода.
CP_IR_TO_ASSEMBLY — преобразование общего IR в ассемблерное представление. К этому этапу почти все инструкции переводятся простым отображением «команда в команду», «аргумент в аргумент». Иногда нужно поменять порядок аргументов или добавить дополнительные, но это несложная работа. Кроме того, на этом этапе появляется информация о размере кода, а значит о сдвигах, поэтому здесь переходы на метки превращаются в переходы на адреса, для вычисления этих сдвигов peeking запрашивает размеры команд.
CP_ASSEMBLY_TO_HEX — перевод ассемблерного представления в бинарный вид и формирование тела функции в специально выделенном буфере. В основном это просто склеивание бинарных полей, но иногда бывают сложные случаи: например, с большими константами на Arm, где в ряде инструкций нужно обрабатывать как маленькие immediate, так и большие, из-за чего способ кодирования весьма нетривиальный. Энкодинг на Intel с его бесконечными слоями legacy самый громоздкий и запутанный. В то же время у минималистичного RISC-V есть всего четыре-пять заранее определенных форматов команд, работать с ним одно удовольствие.
Примеры использования Loops
Пример 1. Сложение
Для начала создадим какой-нибудь совсем простой пример и посмотрим, как он скомпилируется на разных машинах. Hello world на ассемблере писать не так-то просто, поэтому начнем с функции, которая складывает два входных числа:
#include <iostream>
#include <loops/loops.hpp>
int main()
{
loops::Context CTX; // Создаем контекст
USE_CONTEXT_(CTX); // Обеспечиваем корректную работу макросов
loops::IReg a, b; // Создаем переменные-аргументы функции
STARTFUNC_("sum", &a, &b) // Объявляем новую функцию с именем и аргументами
{
loops::IReg c = a + b; // Объявили переменную и присвоили ей результат сложения
RETURN_(c); // Макрос для возврата
}
loops::Func sfunc = CTX.getFunc("sum"); // Достали по имени функцию из контекста
std::cout << "=========IR LISTING=========\n";
sfunc.printIR(std::cout); // Печать IR(самый поздний проход)
std::cout << "\n========ASM LISTING=========\n";
sfunc.printAssembly(std::cout); // Печать ассемблерного кода
typedef int64_t (*sum_t)(int64_t a, int64_t b);
sum_t func = (sum_t)sfunc.ptr(); // Компилируем и достаем указатель на функцию
std::cout << "\n\n5 + 4 = " << func(5, 4) << "\n"; // Вызываем наш первый JIT!
return 0;
}
Итак, по пунктам:
Сначала нужно создать контекст (loops::Context). Это контейнер функций, с помощью него они создаются и удаляются, доступ к функциям по именам.
Если вы собираетесь что-то писать в буфер текущей функции, то нужно инициализировать макросы с помощью USE_CONTEXT_. Если только брать уже написанные функции, печатать их код или запускать, USE_CONTEXT_ не нужен.
Благодаря USE_CONTEXT_ хорошо сработает следующий макрос — STARTFUNC_, который «объявляет» функцию. В фигурных скобках внутри можно писать весь нужный вам код. Само собой, описываемые вами действия — не выполняются, вы только добавляете в буфер их последовательность, прямо во время выполнения программы решая, какую последовательность действий туда записать.
Loops::IReg — абстракция для регистра общего назначения, в более широком смысле — это переменная. Это всегда 64-битное знаковое целое, но вы можете загружать туда и выгружать оттуда типы меньших размеров — поддержаны соответствующие инструкции.
Loops::IReg должен объявляться только внутри блока объявленной функции, иначе будут ошибки. Единственное исключение — это объявление аргументов функций. Они должны объявляться по умолчанию и подаваться в макрос STARTFUNC_ для обозначения аргументов генерируемой нами функции.
Внутри функции вы можете пользоваться loops::IReg, как обычной переменной: складывать, присваивать и передавать в функции. Она будет вести себя примерно так же, как и обычная переменная. Однако, на первых порах не рекомендую объявлять переменные по умолчанию, потому что им неоткуда будет взять информацию о контексте.
Loops::Func описывает функцию, из нее можно получить указатель void* на скомпилированный код. Также позволяет распечатать функцию, причем на разных стадиях обработки: можно указать имя прохода, и printIR выведет код после применения этого прохода.
Перед использованием этот указатель надо преобразовать в соответствии с сигнатурой, и можно выполнять свежескомпилированный код.
Что выведет эта программа на Windows+x86_64:
=========IR LISTING=========
sum(i0, i1)
0 : add i1, i1, i2
1 : mov i0, i1
2 : ret
========ASM LISTING=========
sum(i0, i1)
0 : add rcx, rdx ; 48 01 d1
1 : mov rax, rcx ; 48 89 c8
2 : ret ; c3
5 + 4 = 9Отмечу, что ASM listing будет выглядеть иначе, даже если просто перейти с одной операционной системы на другую, в нашем случае на Linux. Это связано с тем, что в Windows и Linux используются разные соглашения о том, через какие регистры передаются аргументы:
========ASM LISTING=========
sum(i0, i1)
0 : add rdi, rsi ; 48 01 f7
1 : mov rax, rdi ; 48 89 f8
2 : ret ; c3
Тот же листинг на Arm:
========ASM LISTING=========
sum(i0, i1)
0 : add x0, x0, x1 ; 00 00 01 8b
1 : ret x30 ; c0 03 5f d6И на RISC-V:
========ASM LISTING=========
sum(i0, i1)
0 : add a0, a0, a1 ; 33 05 b5 00
1 : ret ; 67 80 00 00Пример 2. Сортировка
Первый пример был довольно простым, поэтому напишем программу посложнее, чтобы посмотреть, как организуются циклы и ветвления.
Далее я привожу примеры на Arm, чтобы не тратить время на объяснения архитектурных различий.
Рассмотрим простую сортировку целых 32-битных знаковых чисел:
#include <iostream>
#include <loops/loops.hpp>
int main()
{
loops::Context CTX;
USE_CONTEXT_(CTX);
loops::IReg ptr32, size;
STARTFUNC_("sort", &ptr32, &size)
{
loops::IReg curnum = CONST_(0); // Начальное значение переменной задается с помощью макроса CONST_
WHILE_(curnum < size - 1) // Цикл, как цикл. Но генерирует код, а не гоняет итерации.
{
loops::IReg minnum = curnum;
loops::IReg searchnum = curnum + 1;
WHILE_(searchnum < size - 1)
{
// Обращение к элементам массива - задача чуть более сложная, чем в Си. И обратите
// внимание, что offset домножается на размер int32_t. Без типизации
// адресная арифметика становится сложнее.
loops::IReg minval = loops::load_<int32_t>(ptr32, minnum * sizeof(int32_t));
loops::IReg searchval = loops::load_<int32_t>(ptr32, searchnum * sizeof(int32_t));
IF_(searchval < minval)
minnum = searchnum;
searchnum = searchnum + 1;
}
loops::IReg curval = loops::load_<int32_t>(ptr32, curnum * sizeof(int32_t));
loops::IReg minval = loops::load_<int32_t>(ptr32, minnum * sizeof(int32_t));
loops::store_<int32_t>(ptr32, curnum * sizeof(int32_t), minval);
loops::store_<int32_t>(ptr32, minnum * sizeof(int32_t), curval);
curnum = curnum + 1;
}
}
loops::Func sfunc = CTX.getFunc("sort");
std::cout << "=========IR LISTING=========\n";
sfunc.printIR(std::cout);
std::cout << "\n========ASM LISTING=========\n";
sfunc.printAssembly(std::cout);
typedef void (*sort_t)(int32_t* ptr32, int64_t size);
sort_t func = (sort_t)sfunc.ptr();
std::vector<int32_t> toSort = {5, 2, 15, -4, 10};
func(toSort.data(), (int64_t)toSort.size());
std::cout << "\n========SORTED VALUES=========\n";
for(int i = 0; i < (int)toSort.size(); i++)
std::cout << toSort[i] << std::endl;
return 0;
}
Вывод
=========IR LISTING=========
sort(i0, i1)
0 : mov i2, 0
1 : __loops_label_0:
2 : sub i3, i1, 1
3 : cmp i2, i3
4 : jmp_ge __loops_label_2
5 : mov i3, i2
6 : add i4, i2, 1
7 : __loops_label_3:
8 : sub i5, i1, 1
9 : cmp i4, i5
10 : jmp_ge __loops_label_5
11 : mov i5, 4
12 : mul i5, i3, i5
13 : load.i32 i5, i0, i5
14 : mov i6, 4
15 : mul i6, i4, i6
16 : load.i32 i6, i0, i6
17 : cmp i6, i5
18 : jmp_ge __loops_label_7
19 : mov i3, i4
20 : __loops_label_7:
21 : add i4, i4, 1
22 : jmp 3
23 : __loops_label_5:
24 : mov i4, 4
25 : mul i4, i2, i4
26 : load.i32 i4, i0, i4
27 : mov i5, 4
28 : mul i5, i3, i5
29 : load.i32 i5, i0, i5
30 : mov i6, 4
31 : mul i6, i2, i6
32 : store.i32 i0, i6, i5
33 : mov i5, 4
34 : mul i3, i3, i5
35 : store.i32 i0, i3, i4
36 : add i2, i2, 1
37 : jmp 0
38 : __loops_label_2:
39 : ret
========ASM LISTING=========
sort(i0, i1)
0 : eor x2, x2, x2 ; 42 00 02 ca
1 : __loops_label_0: ;
2 : sub x3, x1, #0x01 ; 23 04 00 d1
3 : cmp x2, x3 ; 5f 00 03 eb
4 : b.ge __loops_label_2 ; ea 03 00 54
5 : mov x3, x2 ; e3 03 02 aa
6 : add x4, x2, #0x01 ; 44 04 00 91
7 : __loops_label_3: ;
8 : sub x5, x1, #0x01 ; 25 04 00 d1
9 : cmp x4, x5 ; 9f 00 05 eb
10 : b.ge __loops_label_5 ; 8a 01 00 54
11 : mov x5, #0x04 ; 85 00 80 d2
12 : mul x5, x3, x5 ; 65 7c 05 9b
13 : ldrsw x5, [x0, x5] ; 05 68 a5 b8
14 : mov x6, #0x04 ; 86 00 80 d2
15 : mul x6, x4, x6 ; 86 7c 06 9b
16 : ldrsw x6, [x0, x6] ; 06 68 a6 b8
17 : cmp x6, x5 ; df 00 05 eb
18 : b.ge __loops_label_7 ; 4a 00 00 54
19 : mov x3, x4 ; e3 03 04 aa
20 : __loops_label_7: ;
21 : add x4, x4, #0x01 ; 84 04 00 91
22 : b __loops_label_3 ; f3 ff ff 17
23 : __loops_label_5: ;
24 : mov x4, #0x04 ; 84 00 80 d2
25 : mul x4, x2, x4 ; 44 7c 04 9b
26 : ldrsw x4, [x0, x4] ; 04 68 a4 b8
27 : mov x5, #0x04 ; 85 00 80 d2
28 : mul x5, x3, x5 ; 65 7c 05 9b
29 : ldrsw x5, [x0, x5] ; 05 68 a5 b8
30 : mov x6, #0x04 ; 86 00 80 d2
31 : mul x6, x2, x6 ; 46 7c 06 9b
32 : str w5, [x0, x6] ; 05 68 26 b8
33 : mov x5, #0x04 ; 85 00 80 d2
34 : mul x3, x3, x5 ; 63 7c 05 9b
35 : str w4, [x0, x3] ; 04 68 23 b8
36 : add x2, x2, #0x01 ; 42 04 00 91
37 : b __loops_label_0 ; e0 ff ff 17
38 : __loops_label_2: ;
39 : ret x30 ; c0 03 5f d6
========SORTED VALUES=========
-4
2
5
15
10Посмотрим, на что в этом примере стоит обратить внимание.
В коде можно встретить макрос CONST_ при объявлении переменных. Дело в том, что Loops должен помещать операции в буфер программы, а для этого ему нужно знать в какой именно буфер. Эта информация обычно хранится в операндах — уже объявленных IReg. Это сделано ради синтаксического сахара и лаконичности выражений. Однако только конструирующемуся IReg взять эту информацию неоткуда. CONST_ порождает аргумент уже с ней, подключая IReg к буферу. Также можно объявить переменную без инициализации, с помощью макроса DEF_().
В Loops переменные описывают обобщенные регистры, что характерно для низкоуровневого кода. Память при этом рассматривается как внешняя среда, взаимодействие с которой происходит через специальные операции ввода/вывода (load/store). Также нужно обратить внимание на особенности адресной арифметики. Типизация указателей не используется, поэтому указатели и их смещения измеряются в байтах. Возможно, позже мы это улучшим.
Пример 3. Векторы и вызовы
Концептуально loops выстроен вокруг SIMD-регистров. Давайте поработаем с ними и напишем простой одномерный maxpool: в каждом элементе целевого массива — максимальное значение из трех подряд идущих элементов оригинала.
Код упрощенный, размеры входных массивов нужно проверять снаружи. Чтобы показать еще что-нибудь интересное, добавим отладо��ный вызов, который будет выводить некоторые значения:
// Аргументы вызываемых функций - только int64_t
// Возврат - либо void, либо int64_t
void debug_print(int64_t a)
{
float* a_f32 = (float*)(void*)(&a);
printf("Loops debug: %f\n", *a_f32);
}
//...
STARTFUNC_("maxpool1d_3", &ptrin_f32, &ptrout_f32, &sizeout)
{
loops::IReg curpos = CONST_(0);
WHILE_(curpos < sizeout)
{
// Векторные регистры типизированы, благодаря чему под арифметические и иные действия
// автоматически подбираются корректные инструкции
loops::VReg<float> in0 = loops::loadvec<float>(ptrin_f32, curpos * sizeof(float));
loops::VReg<float> in1 = loops::loadvec<float>(ptrin_f32, (curpos + 1) * sizeof(float));
loops::VReg<float> in2 = loops::loadvec<float>(ptrin_f32, (curpos + 2) * sizeof(float));
loops::VReg<float> out = loops::max(loops::max(in0, in1), in2);
VOID_CALL_(debug_print, loops::getlane(out, 1)); // Вызов функции с одним аргументом
loops::storevec(ptrout_f32, curpos * sizeof(float), out);
curpos = curpos + CTX.vlanes<float>();// Количество lane'ов в векторе на разных машинах разное
}
}
Полученный код
=========IR LISTING=========
maxpool1d_3(i0, i1, i2)
0 : sub i31, i31, 416
1 : arm_stp i31, 400, i29, i30
2 : mov i29, i31
3 : mov i3, 0
4 : __loops_label_0:
5 : cmp i3, i2
6 : jmp_ge __loops_label_2
7 : mov i4, 4
8 : mul i4, i3, i4
9 : vld.fp32 v0, i0, i4
10 : add i4, i3, 1
11 : mov i5, 4
12 : mul i4, i4, i5
13 : vld.fp32 v1, i0, i4
14 : add i4, i3, 2
15 : mov i5, 4
16 : mul i4, i4, i5
17 : vld.fp32 v2, i0, i4
18 : max.fp32 v0, v0, v1
19 : max.fp32 v0, v0, v2
20 : getlane.fp32 i4, v0, 1
21 : mov i5, 25416
22 : arm_movk i5, 47206, 16
23 : arm_movk i5, 43690, 32
24 : arm_stp i31, 0, i0, i1
25 : arm_stp i31, 16, i2, i3
26 : arm_stp i31, 32, i4, i5
27 : arm_stp i31, 48, i6, i7
28 : arm_stp i31, 64, i8, i9
29 : arm_stp i31, 80, i10, i11
30 : arm_stp i31, 96, i12, i13
31 : arm_stp i31, 112, i14, i15
32 : arm_stp i31, 128, i16, i17
33 : mov i0, i4
34 : add i10, i31, 144
35 : vst_lane.u64 i10, v0
36 : vst_lane.u64 i10, v4
37 : vst_lane.u64 i10, v8
38 : vst_lane.u64 i10, v12
39 : call_noret [i5]()
40 : add i10, i31, 144
41 : vld_lane.u64 v0, i10
42 : vld_lane.u64 v4, i10
43 : vld_lane.u64 v8, i10
44 : vld_lane.u64 v12, i10
45 : arm_ldp i0, i1, i31, 0
46 : arm_ldp i2, i3, i31, 16
47 : arm_ldp i4, i5, i31, 32
48 : arm_ldp i6, i7, i31, 48
49 : arm_ldp i8, i9, i31, 64
50 : arm_ldp i10, i11, i31, 80
51 : arm_ldp i12, i13, i31, 96
52 : arm_ldp i14, i15, i31, 112
53 : arm_ldp i16, i17, i31, 128
54 : mov i4, 4
55 : mul i4, i3, i4
56 : vst.fp32 i1, i4, v0
57 : add i3, i3, 4
58 : jmp 0
59 : __loops_label_2:
60 : arm_ldp i29, i30, i31, 400
61 : add i31, i31, 416
62 : ret
========ASM LISTING=========
maxpool1d_3(i0, i1, i2)
0 : sub sp, sp, #0x1a0 ; ff 83 06 d1
1 : stp x29, x30, [sp, #0x190] ; fd 7b 19 a9
2 : mov x29, sp ; fd 03 1f aa
3 : eor x3, x3, x3 ; 63 00 03 ca
4 : __loops_label_0: ;
5 : cmp x3, x2 ; 7f 00 02 eb
6 : b.ge __loops_label_2 ; aa 06 00 54
7 : mov x4, #0x04 ; 84 00 80 d2
8 : mul x4, x3, x4 ; 64 7c 04 9b
9 : ldr q0, [x0, x4] ; 00 68 e4 3c
10 : add x4, x3, #0x01 ; 64 04 00 91
11 : mov x5, #0x04 ; 85 00 80 d2
12 : mul x4, x4, x5 ; 84 7c 05 9b
13 : ldr q1, [x0, x4] ; 01 68 e4 3c
14 : add x4, x3, #0x02 ; 64 08 00 91
15 : mov x5, #0x04 ; 85 00 80 d2
16 : mul x4, x4, x5 ; 84 7c 05 9b
17 : ldr q2, [x0, x4] ; 02 68 e4 3c
18 : fmax v0.4s, v0.4s, v1.4s ; 00 f4 21 4e
19 : fmax v0.4s, v0.4s, v2.4s ; 00 f4 22 4e
20 : umov w4, v0.s[1] ; 04 3c 0c 0e
21 : mov x5, #0x6348 ; 05 69 8c d2
22 : movk x5, #0xb866, lsl #16 ; c5 0c b7 f2
23 : movk x5, #0xaaaa, lsl #32 ; 45 55 d5 f2
24 : stp x0, x1, [sp, #0] ; e0 07 00 a9
25 : stp x2, x3, [sp, #0x10] ; e2 0f 01 a9
26 : stp x4, x5, [sp, #0x20] ; e4 17 02 a9
27 : stp x6, x7, [sp, #0x30] ; e6 1f 03 a9
28 : stp x8, x9, [sp, #0x40] ; e8 27 04 a9
29 : stp x10, x11, [sp, #0x50] ; ea 2f 05 a9
30 : stp x12, x13, [sp, #0x60] ; ec 37 06 a9
31 : stp x14, x15, [sp, #0x70] ; ee 3f 07 a9
32 : stp x16, x17, [sp, #0x80] ; f0 47 08 a9
33 : mov x0, x4 ; e0 03 04 aa
34 : add x10, sp, #0x90 ; ea 43 02 91
35 : st1 {v0.2d, v1.2d, v2.2d, v3.2d}, [x10], #0x40 ; 40 2d 9f 4c
36 : st1 {v4.2d, v5.2d, v6.2d, v7.2d}, [x10], #0x40 ; 44 2d 9f 4c
37 : st1 {v8.2d, v9.2d, v10.2d, v11.2d}, [x10], #0x40 ; 48 2d 9f 4c
38 : st1 {v12.2d, v13.2d, v14.2d, v15.2d}, [x10], #0x40 ; 4c 2d 9f 4c
39 : blr x5 ; a0 00 3f d6
40 : add x10, sp, #0x90 ; ea 43 02 91
41 : ld1 {v0.2d, v1.2d, v2.2d, v3.2d}, [x10], #0x40 ; 40 2d df 4c
42 : ld1 {v4.2d, v5.2d, v6.2d, v7.2d}, [x10], #0x40 ; 44 2d df 4c
43 : ld1 {v8.2d, v9.2d, v10.2d, v11.2d}, [x10], #0x40 ; 48 2d df 4c
44 : ld1 {v12.2d, v13.2d, v14.2d, v15.2d}, [x10], #0x40 ; 4c 2d df 4c
45 : ldp x0, x1, [sp, #0] ; e0 07 40 a9
46 : ldp x2, x3, [sp, #0x10] ; e2 0f 41 a9
47 : ldp x4, x5, [sp, #0x20] ; e4 17 42 a9
48 : ldp x6, x7, [sp, #0x30] ; e6 1f 43 a9
49 : ldp x8, x9, [sp, #0x40] ; e8 27 44 a9
50 : ldp x10, x11, [sp, #0x50] ; ea 2f 45 a9
51 : ldp x12, x13, [sp, #0x60] ; ec 37 46 a9
52 : ldp x14, x15, [sp, #0x70] ; ee 3f 47 a9
53 : ldp x16, x17, [sp, #0x80] ; f0 47 48 a9
54 : mov x4, #0x04 ; 84 00 80 d2
55 : mul x4, x3, x4 ; 64 7c 04 9b
56 : str q0, [x1, x4] ; 20 68 a4 3c
57 : add x3, x3, #0x04 ; 63 10 00 91
58 : b __loops_label_0 ; cb ff ff 17
59 : __loops_label_2: ;
60 : ldp x29, x30, [sp, #0x190] ; fd 7b 59 a9
61 : add sp, sp, #0x1a0 ; ff 83 06 91
62 : ret x30 ; c0 03 5f d6Код довольно наглядный, все же добавлю пару пояснений
CTX.vlanes<float> возвращает количество элементов данного типа, которые умещаются в векторном регистре. Если мы говорим про float, то на Intel (AVX2) их будет восемь, а на Arm (NEON) — четыре. Понятно, что если мы ведем обработку векторами, то «шагать» по массиву надо именно на такое количество.
Отладка JIT-кода — задача не из простых. В лучшем случае приходится «ползать» по ассемблеру. Даже простая отладочная печать не всегда доступна — нельзя просто вставить вызов printf в код.
С появлением поддержки вызовов функций ситуация немного упростилась. Сигнатура вызываемых функций сильно ограничена. Аргументы могут быть только типа int64_t. Их количество зависит от платформы: на Intel — до четырех аргументов на Windows и до шести на Linux, на ARM и RISC-V — до восьми. То есть максимальное количество аргументов совпадает с количеством аргументов, которые передаются через регистры согласно соглашению о вызове функций (calling convention). В коде можно увидеть много инструкций spill. Это префикс и постфикс вызова. Другими словами, вызовы хоть и есть, реализованы не очень хорошо и пока стоит их избегать, по крайней мере в нагруженных участках кода.
Пример 4. Параметрический поэлементный полином
Приведем, наконец, пример, с которым может справиться только JIT — параметрическая компиляция. Loops очень хорошо подходит для поэлементных или попиксельных операций, поэтому рассмотрим функцию, которая поэлементно высчитывает полином. Параметрами будут степень и коэффициенты этого полинома. Не будем все это сильно оптимизировать, напишем проще и посмотрим код, который получается в зависимости от входных данных:
loops::Func create_polynomic_function(loops::Context& CTX, const std::vector<float> ratios)
{
USE_CONTEXT_(CTX);
loops::IReg ptrin_f32, ptrout_f32, sizeout;
STARTFUNC_("poly", &ptrin_f32, &ptrout_f32, &sizeout)
{
loops::IReg curpos = CONST_(0);
std::vector<loops::VReg<float> > ratiosV(ratios.size());
for(int rNum = 0; rNum < ratios.size(); rNum++)
if(rNum == 0 || (ratios[rNum] != 0.f && ratios[rNum] != 1.f))
{// Отложенная инициализация регистра необходима при работе с массивом VReg.
ratiosV[rNum].copyidx(VCONST_(float, ratios[rNum]));
}
WHILE_(curpos < sizeout)
{
loops::VReg<float> result = ratiosV[0];
loops::VReg<float> x = loops::loadvec<float>(ptrin_f32, curpos * sizeof(float));
// res = ratios[0] + ratios[1] * x * ratios[2] * x^2 + ...;
loops::VReg<float> power = x;
// Цикл, добавляющий инструкции в зависимости от числа входных коэффициентов
for(size_t pNum = 1; pNum < ratiosV.size(); pNum++)
{// Нулевые компоненты не добавляем, если коэффициент единица - плюсуем степень без умножения.
if(ratios[pNum] == 1.f) // Главное не путать if'ы и IF_'ы
result += power;
else if(ratios[pNum] != 0.f)
result += power * ratiosV[pNum];
if(pNum + 1 < ratiosV.size())
power *= x;
}
loops::storevec(ptrout_f32, curpos * sizeof(float), result);
curpos = curpos + CTX.vlanes<float>();
}
}
return CTX.getFunc("poly");
}
int main()
{
loops::Context CTX;
USE_CONTEXT_(CTX);
loops::IReg ptrin_f32, sizeout, ptrout_f32;
loops::Func sfunc = create_polynomic_function(CTX, {1});
std::cout << "=========IR LISTING=========\n";
// Мы можем указать проход после которого мы хотим напечатать код
sfunc.printIR(std::cout, 3, "CP_COLLECTING");
std::cout << "\n========ASM LISTING=========\n";
sfunc.printAssembly(std::cout);
return 0;
}
Посмотрим на код в том виде, в котором он появляется в буфере до всех проходов. Для этого достаточно в функцию printIR передать название нужного пасса, в данном случае “CP_COLLECTING”.
Для вырожденного полинома({1}):
poly(i0, i1, i2)
0 : mov i3, 0
1 : mov v0, 1065353216
2 : annotation:whilecstart 0
3 : cmp i3, i2
4 : jmp_ge __loops_label_2
5 : annotation:whilecend
6 : mov v1, v0
7 : mul i4, i3, 4
8 : vld.fp32 v2, i0, i4
9 : mov v3, v2
10 : mul i5, i3, 4
11 : vst.fp32 i1, i5, v1
12 : add i3, i3, 8
13 : annotation:endwhile 0, 2
Для квадратичного случая {1,2,1} :
poly(i0, i1, i2)
0 : mov i3, 0
1 : mov v0, 1065353216
2 : mov v1, 1073741824
3 : annotation:whilecstart 0
4 : cmp i3, i2
5 : jmp_ge __loops_label_2
6 : annotation:whilecend
7 : mov v2, v0
8 : mul i4, i3, 4
9 : vld.fp32 v3, i0, i4
10 : mov v4, v3
11 : mul.fp32 v5, v4, v1
12 : add.fp32 v2, v2, v5
13 : mul.fp32 v4, v4, v3
14 : add.fp32 v2, v2, v4
15 : mul i5, i3, 4
16 : vst.fp32 i1, i5, v2
17 : add i3, i3, 8
18 : annotation:endwhile 0, 2Для кубического случая {0.5, 0, 0, 2}:
poly(i0, i1, i2)
0 : mov i3, 0
1 : mov v0, 1056964608
2 : mov v1, 1073741824
3 : annotation:whilecstart 0
4 : cmp i3, i2
5 : jmp_ge __loops_label_2
6 : annotation:whilecend
7 : mov v2, v0
8 : mul i4, i3, 4
9 : vld.fp32 v3, i0, i4
10 : mov v4, v3
11 : mul.fp32 v4, v4, v3
12 : mul.fp32 v4, v4, v3
13 : mul.fp32 v5, v4, v1
14 : add.fp32 v2, v2, v5
15 : mul i5, i3, 4
16 : vst.fp32 i1, i5, v2
17 : add i3, i3, 8
18 : annotation:endwhile 0, 2
Если не считать граничных случаев, которые можно аккуратно обработать в генераторе, получается оптимальный код.
Посмотрим, на что в этом примере стоит обратить внимание.
Отложенная инициализация регистров — важный момент. Думаю, уже понятно, что у синтаксического сахара, который позволяет маскировать добавление действий в буфер под сами эти действия, есть своя цена. В частности — сложности с инициализацией.
Золотое правило работы с Loops-регистрами — они должны инициализироваться сразу. Однако это не всегда возможно. Например, если нужно работать с целым контейнером регистров, то инициализацию приходится откладывать. В таких случаях на помощь приходит функция copyidx, которая позволяет создать несколько ссылок на один регистр, то есть скопировать его индекс.
Цикл, который увеличивает регистр с результатом, показывает главный принцип: в генерирующем коде мы видим понятный цикл, а в сгенерированном — фиксированную последовательность действий, без итераций. Когда код становится сложным, можно начать путать циклы генератора и циклы, которые мы создаем. Генерирующие и генерируемые управляющие конструкции путаются, это довольно странное ощущение, но со временем к этому привыкаешь.
Коэффициенты мы перевели в векторный вид заранее, предполагая, что полином будет не слишком большой и нам хватит векторных регистров. Часть регистров теперь зарезервирована под хранение коэффициентов.
Под катом — Neon-код последнего случая {0.5, 0, 0, 2}:
Код
========ASM LISTING=========
poly(i0, i1, i2)
0 : eor x3, x3, x3 ; 63 00 03 ca
1 : eor x4, x4, x4 ; 84 00 04 ca
2 : movk x4, #0x3f00, lsl #16 ; 04 e0 a7 f2
3 : dup v0.4s, w4 ; 80 0c 04 4e
4 : eor x4, x4, x4 ; 84 00 04 ca
5 : movk x4, #0x4000, lsl #16 ; 04 00 a8 f2
6 : dup v1.4s, w4 ; 81 0c 04 4e
7 : __loops_label_0: ;
8 : cmp x3, x2 ; 7f 00 02 eb
9 : b.ge __loops_label_2 ; ea 01 00 54
10 : mov v2.4s, v0.4s ; 02 1c a0 4e
11 : mov x4, #0x04 ; 84 00 80 d2
12 : mul x4, x3, x4 ; 64 7c 04 9b
13 : ldr q3, [x0, x4] ; 03 68 e4 3c
14 : mov v4.4s, v3.4s ; 64 1c a3 4e
15 : fmul v4.4s, v4.4s, v3.4s ; 84 dc 23 6e
16 : fmul v3.4s, v4.4s, v3.4s ; 83 dc 23 6e
17 : fmul v3.4s, v3.4s, v1.4s ; 63 dc 21 6e
18 : fadd v2.4s, v2.4s, v3.4s ; 42 d4 23 4e
19 : mov x4, #0x04 ; 84 00 80 d2
20 : mul x4, x3, x4 ; 64 7c 04 9b
21 : str q2, [x1, x4] ; 22 68 a4 3c
22 : add x3, x3, #0x04 ; 63 10 00 91
23 : b __loops_label_0 ; f1 ff ff 17
24 : __loops_label_2: ;
25 : ret x30 ; c0 03 5f d6
Замеры
Приведу небольшое сравнение скоростей обобщенного AOT-кода с кодом, сгенерированным в Loops. В качестве PoC мы написали два оптимальных типичных ядра: depthwise convolution и maxpool и посмотрели, как они будут работать отдельно и в составе сети. Заметим, что ускорение достигается за счет параметрической компиляции, а не других техник, вроде loops fusion. Конфигурация для замера: Apple M1, 8 Gb, Debian, 4 потока. В YOLO есть большие и очень редкие maxpool, в efficientnet есть depthwise convolution. Замеряемая статистическая характеристика: минимум времени в 100 повторениях.
Case |
AOT |
JIT |
Speedup |
Maxpool 13x13(fp16) |
13.9 ms |
0.39 ms |
35.64x |
YOLOv4(fp16) |
173.33 ms |
162.2 ms |
1.07x |
Depthwise convolution 5x5 (fp32) |
3.48 ms |
1.97 ms |
1.77x |
efficientnet-lite4-11(fp32) |
12.61 ms |
11.07 ms |
1.14x |
Сравнение с аналогами
Для полноты картины нужно сравнить Loops c похожими проектами и обозначить уникальную нишу нашего проекта.
Про Xbyak мы написали уже достаточно.
В любом разговоре о компиляторах сложно обойти стороной LLVM — пожалуй, самый главный компиляторный проект на сегодня. Он умеет все, включая JIT-компиляцию. Проблема в его размере, сложности использования, скорости компиляции и в том, что поддерживать зависимость от него не так-то просто. О one-header простоте в случае с LLVM, конечно, речи не идет.
Также упомяну очень интересный проект Владимира Макарова, MIR. Насколько я понимаю, на данный момент проект мигрирует в сторону полноценного легковесного и быстрого C-компилятора, с оптимизирующими проходами, но использоваться в тех же целях он может. Одна проблема — в нем нет векторных регистров.
Сравним упомянутые выше проекты:
Категория |
Xbyak |
LLVM |
MIR |
Loops |
Минималистичный |
+ |
- |
+ |
+ |
Кросс-платформенный |
- |
+ |
+ |
+ |
Оптимизирующий |
- |
++ |
+- |
- |
Распределение регистров |
- |
+ |
+ |
+ |
Векторные инструкции |
+ |
+ |
- |
+- |
Синтаксический сахар |
- |
- |
- |
+ |
Вопросы, наверное, может вызвать только пункт «синтаксический сахар». Имеется в виду заполнение буфера с помощью имитации исполнения. Это неоднозначный подход, и для сложных задач можно использовать альтернативный интерфейс. Однако у него есть важное преимущество — он значительно снижает порог входа.
Статус и перспективы
У Loops пока нет продуктового качества, релиз первой версии планируется в середине-конце 2026 года. Нужно улучшить аллокатор регистров, провести профилировку для ускорения компиляции кода, добавить векторные расширения для RISC-V, а также подумать над SVE2 и AVX512/AVX10, провести хорошее большое сравнение с аналогами в цифрах.
Тем не менее, библиотека работает и уже хорошо себя показала:
В проекте есть PoC-библиотека генераторов loopslayers, включающая экспериментальные мета-функции для слоев нейронных сетей. Мы успешно использовали ее для ускорения инференса GoogLeNet и YOLO.
В рамках уже рабочего, а не PoC проекта мы использовали проект для ускорения SQL движка, когда при поступления SQL запроса к базе на лету генерировался машинный код для его вычисления.
Библиотека хорошо подходит для любых тяжелых вычислений: нейронных сетей, матфизики, БД, графики. Кроме этого у нее есть еще одно важное применение— с ее помощью можно прототипировать инструкции при работе над расширениями процессоров. Добавить новую инструкцию в общем случае займет около получаса.
Так что скачивайте, пробуйте и задавайте вопросы в комментариях — буду рад ответить.