sse.md

Векторизация

Рассмотрим следующую программу, в которой считается сумма одномерного целочисленного массива:

#pragma GCC optimize("O3")
// ^ включает самый "агрессивный" уровень оптимизации
// то же самое, что добавить флаг "-O3" при компиляции из консоли

#include <iostream>
using namespace std;

const int n = 1e5;
int a[n], s = 0;

int main() {
    for (int t = 0; t < 100000; t++)
        for (int i = 0; i < n; i++)
            s += a[i];

    return 0;
}

Если скомпилировать этот код под GCC без всяких дополнительных настроек и запустить, он отработает за 2.43 секунды.

Добавим теперь следующую магическую директиву в самое начало программы:

#pragma GCC target("avx2")
// ...остальное в точности как было

Скомпилировавшись и запустившись при тех же условиях, программа завершается уже за 1.24 секунды. Это почти в два раза быстрее, при том, что сам код и уровень оптимизации мы не меняли.

Чтобы понять, что здесь происходит, нам нужно сначала разъяснить некоторые особенности работы современных компьютеров. (Знающие ассемблер могут отмотать примерно до ⅓ статьи.)

Complex Instruction Set Computing

Раньше, во времена, когда компьютеры назывались ЭВМ-ами и занимали целую комнату, увеличение производительности происходило в основном за счёт увеличения тактовой частоты. Тактовая частота условно равна количеству инструкций, выполняемому процессором за единицу времени. (На современных процессорах это не так — разные инструкции занимают разное время, которое ещё и может зависеть от разных обстоятельств.)

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

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

• RISC (англ. reduced instruction set computer), в которых длина кода (идентификатора) самой инструкции ограничена, а значит ограничено и само количество инструкций. Самые первые компьютеры относились к этому типу и могли не иметь даже отдельных инструкций для умножения и деления. Такие процессоры требуют меньше транзисторов, и как следствие сами меньше, дешевле и потребляют меньше энергии. Самое популярное семейство архитектур называется arm и используется почти на всех современных мобильных устройствах.

• CISC (англ. complex instruction set computer), к которому относят всё, что не RISC — в них длина команды не фиксирована, что позволяет поддерживать практически произовльное количество инструкций. Самое популярное семейство архитектур называется x86 и используется почти на всех современных десктопах и серверах.

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

• В обычных CPU довольно быстро добавили инструкцию, которая принимает числа $(x, y)$ и загружает в регистр данные по адресу $a \cdot x + y$. Это полезно при индексации массивов — не нужно отдельно индекс считать.

• На графических сопроцессорах появилась отдельная инструкция, которую называют «saxpy» (сокращенно от выражения s += a * x + y), которая полезна, например, при перемножении матриц.

• В последние GPU от Nvidia добавили «tensor core» — отдельную схему, которая перемножает две матрицы $4 \times 4$ и прибавляет к третьей, как бы производя $4 \times 4 \times 4 = 64$ умножений и $4 \times 4 = 16$ сложений за раз, что сильно ускоряет алгоритмы блочного матричного умножения.

В этой статье мы сфокусируемся на отдельном виде инструкций, которые позволяют выполнять одну и ту же операцию сразу на какой-то последовательности данных. Эта концепция называется SIMD-параллелизмом (англ. single instruction, multiple data).

Streaming SIMD Extensions

SSE — это обобщённое название всех SIMD-инструкций для x86.

Работают они следующим образом. Помимо обычных регистров (самых близких к процессору ячеек памяти, с которыми он непосредственно работает), есть дополнительные, вмещающие не 64, а 128, 256 или даже 512 бит — в зависимости от поддерживаемой версии SSE. В эти регистры загружается последовательные блоки из памяти, над ним производится какая-то последовательность операций, и итоговый результат записывается обратно в память. Сами операции обычно логически разбивают эту булеву последовательность на блоки, например, по 32 бит, и работают уже с ними, причём одновременно.

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

Например, какое-нибудь сложение двух int-овых массивов удаётся таким образом соптимизировать в $\frac{512}{32} = 16$ раз, если процессор поддерживает AVX512, а операции битсета — в 512 раз (реализация из STL, по всей видимости, SSE не использует, поэтому даже c = a | b работает примерно в три раза дольше, чем просто пройтись for-ом по массиву int-ов).

Очень часто SSE используют для работы с действительными числами, и в этой ситуации возникает прямой trade-off между точностью вычислений и скоростью работы: например, вместо double можно использовать float, и тогда в один и тот же регистр поместится в два раза больше чисел. По этой причине в последнее время стали развиваться различные методы квантизации: перевода исходных данных в какой-то более дискретизированный формат на входе какой-нибудь процедуры (например, матричного умножения) и восстановления в исходный формат на выходе.

Конкретный набор инструкций и размеры регистров зависят от вендора и поколения архитектуры. На данный момент (лето 2019 года) большинство процессоров архитектуры x86 производит Intel, поэтому мы сконцентрируемся именно на их наборе инструкций.

Поддержка SIMD-инструкций добавлялись постепенно, сохраняя обратную совместимость. Если третий пентиум в 1999-м году умел работать с регистрами размера 128, то в самых современных i7 есть 512-битные регистры. Автор не является специалистом в проектировании микропроцессоров, но предполагает, то регистры больше 64 байт (512 бит) появятся не скоро, потому что это уже больше размера кэш-линии

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

В компиляторе GCC есть встроенная функция __builtin_cpu_supports, которая берёт строчку-название набора инструкций ("sse", "avx2", "avx512f" и т. п.) и возвращает целое число — ноль или какую-то степень двойки. Эта функция работает так: входная строка во время компиляции переводится в нужную степень двойки, которая в рантайме просто AND-ится с маской из cpuid и возвращается — всё ради эффективности.

#include <iostream>
using namespace std;

int main() {
    cout << __builtin_cpu_supports("sse") << endl;
    cout << __builtin_cpu_supports("sse2") << endl;
    cout << __builtin_cpu_supports("avx") << endl;
    cout << __builtin_cpu_supports("avx2") << endl;
    cout << __builtin_cpu_supports("avx512f") << endl;

    return 0;
}

Экономя время читателю: сервера CodeForces и большинство онлайн джаджей на момент написания статьи поддерживают AVX2, то есть умеют полноценно работать с 256-битными регистрами.

C++ intrinsics

SSE это те же чистые ассемблерные инструкции. Языки с каким-либо более верхним уровнем абстракции напрямую работать с ними уже не могут. Однако не обязательно писать на чистом ассемблере, чтобы их использовать — разработчики компиляторов уже позаботились об этом за вас и сделали встроенные функции-обёртки, которые называют интринзиками (англ. intrinsic — «внутренний»).

Чтобы их подключить, нужно указать include на соответствующий заголовочный файл, а также сказать компилятору о том, что мы хотим использовать конкретный набор или наборы инструкций. В примере из начала статьи мы сделали именно это, прописав target("avx2") — компилятор получил доступ к более широким регистрам и продвинутым инструкциям для них, и смог соптимизировать программу примерно в два раза (по умолчанию включены 128-битные sse и sse2, поэтому в 2, а не в $\frac{256}{32} = 8$).

По аналогии с <bits/stdc++.h>, в GCC есть такой же заголовочный файл <x86intrin.h>, включающий в себя сразу все SSE-интринзики. Шаблон любителя засоренных неймспейсов и избыточно долгой компиляции может начинаться так:

#pragma GCC target("avx2")
#pragma GCC optimize("O3")

#include <x86intrin.h>
#include <bits/stdc++.h>

using namespace std;

Пример. Простой цикл, в котором складывают два массива 64-битных действительных чисел, на SSE-интринзиках будет выглядеть так:

double a[100], b[100], c[100];

for (int i = 0; i < 100; i += 4) {
    // загрузим два отрезка по 256 бит в свои регистры
    __m256d x = _mm256_loadu_pd(&a[i]);
    __m256d y = _mm256_loadu_pd(&b[i]);
    // - 256 означает размер регистров
    // - d означает "double"
    // - pd означает "packed double"

    // просуммируем числа и положим результат в другой регистр:
    __m256d z = _mm256_add_pd(x, y);
    // запишем содержимое регистра в память:
    _mm256_storeu_pd(&c[i], z);
}

Если размер массива не кратен размеру регистра, то программа может отработать некорректно. В этом случае можно поступить одним из двух способов:

1. Добавить в конец массива «нейтральные» элементы, дополнив его до удобной длины.

2. Обработать через SSE столько элементов, сколько получается, а оставшиеся обработать отдельно.

Например, вот так можно считать сумму на массиве произвольного размера:

int sum(int a[], int n) {
    int res = 0;
    
    // создадим регистр, в котором будем хранить 8 текущих сумм
    __m256i x = _mm256_setzero_si256();
    for (int i = 0; i + 8 <= n; i += 8) {
        __m256i y = _mm256_loadu_si256((__m256i*) &a[i]);
        x = _mm256_add_epi32(x, y);
    }

    // сложим все 8 чисел в регистре в одно
    int *b = (int*) &x;
    for (int i = 0; i < 8; i++)
        res += b[i];
    
    // добавим остаток массива
    for (int i = (n / 8) * 8; i < n; i++)
        res += a[i];

    return res;
}

Имена команд. Большинство команд кодируются как _mm<размерность>_<действие>_<тип>.

Несколько других примеров:

• _mm_add_epi16 — складывает две группы 16-битных extended packed integer, проще говоря $\frac{128}{16} = 8$ short-ов (в инструкциях, где размер регистра не указан, он равен 128).

• _mm256_acos_pd — принимает один регистр, содержащий 4 double-ов, и возвращает их арк-косинусы.

• _mm256_broadcast_sd — бродкастит (копирует) double из памяти во все четыре слота в регистре.

• _mm256_ceil_pd — округляет double к ближайшему int-у вверх.

• _mm256_cmpeq_epi32 — сравнивает запакованные int-ы и возвращает вектор-маску, в которой для полностью совпавших элементов будет по 32 единицы.

• _mm256_blendv_ps — по заданной маске берёт значения либо из первого массива, либо из второго. Часто применяется для замены if-а.

Комбинаторно получается огромное количество различных функций. Полная документация по ним — Intel Intrinsics Guide — лежит в закладках браузера у каждого уважающего себя performance engineer-а.

Выравнивание. Отдельно стоит отметить одну делать: операции чтения и записи имеют по две версии — load / loadu и store / storeu. Буква «u» здесь означает «unaligned» (англ. невыровненный). Первые корректно работают только тогда, когда весь считываемый блок помещается на одну кэш-линию (если это не так, то в рантаеме вызвется segfault), в то время как unaligned версия работает всегда и везде.

Иногда, особенно когда операция «лёгкая», это отличие имеет большое значение — если не получается «выровнять» память, то производительность может резко упасть (хотя бы потому, что нужно загружать две кэш-линии вместо одной).

Например, так складывать два массива:

void aplusb_unaligned() {
    for (int i = 3; i + 7 < n; i += 8) {
        __m256i x = _mm256_loadu_si256((__m256i*) &a[i]);
        __m256i y = _mm256_loadu_si256((__m256i*) &b[i]);
        __m256i z = _mm256_add_epi32(x, y);
        _mm256_storeu_si256((__m256i*) &c[i], z);
    }
}

...будет на 30% медленнее, чем так:

void aplusb_aligned() {
    for (int i = 0; i < n; i += 8) {
        __m256i x = _mm256_load_si256((__m256i*) &a[i]);
        __m256i y = _mm256_load_si256((__m256i*) &b[i]);
        __m256i z = _mm256_add_epi32(x, y);
        _mm256_store_si256((__m256i*) &c[i], z);
    }
}

Если предположить, что в первом варианте начало массива совпадает с началом кэш-линии, а её размер 64 байта, то примерно половина loadu и storeu будут «плохими».

Вручную «выровнять» память для последовательного чтения через load в случае с массивами можно так:

alignas(32) float a[n];

for (int i = 0; i < n; i += 8) {
    __m256 x = _mm256_load_ps(&a[i]);
    // ...
}

Указатель на начало массива теперь будет кратен 32 байтам, то есть размеру sse-блока. Теперь любое чтение и запись гарантированно будет внутри кэш-линии.

Типизация. Вообще, загружать и сохранять данные интринзиками и вообще использовать __m-типы на самом деле не обязательно — всё можно сделать и обычным reinterpret_cast-ом. Все данные хранятся в одном и том же формате, и разные типы нужны просто для проверки типов и избежания связанных ошибок.

Для каждой размерности регистра есть 3 типа данных. На примере AVX: __m256 для float, __m256d для double и __m256i для разных int-ов.

Некоторые операции есть только для какого-то одного типа (например, тот же _mm256_blendv_ps не имеет аналога для 32-битных int-ов), однако будут работать с другими типами абсолютно так же. Поэтому, to make compiler happy, нужно применять к ним преобразования типов, которые не будут стоить дополнительных инструкций в рантайме. Они все имеют такой формат: _mm<размерность>_cast<откуда>_<куда>.

Loop unrolling. Добавление флага unroll-loops (или прагмы: #pragma GCC optimize("unroll-loops")) позволяет компилятору делать «размотку» циклов, то есть преобразовывать код вида:

for (int i = 1; i < n; i++)
    a[i] = (i % b[i]);

...во что-то такое:

int i;
for (i = 1; i < n - 3; i += 4) {
    a[i] = (i % b[i]);
    a[i + 1] = ((i + 1) % b[i + 1]);
    a[i + 2] = ((i + 2) % b[i + 2]);
    a[i + 3] = ((i + 3) % b[i + 3]);
}

for (; i < n; i++)
    a[i] = (i % b[i]);

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

Нетривиальный пример

Пусть нам зачем-то понадобилось возвести $10^8$ чисел в какие-то степени.

Сравним два решения: обычное и векторизованное. Тестировать подходы будем таким кодом:

#pragma GCC optimize("O3")
#pragma GCC target("avx2")

#include <x86intrin.h>
#include <bits/stdc++.h>

using namespace std;

typedef unsigned long long ull;
typedef __m256i reg;

const int n = 1e8;
alignas(32) unsigned bases[n], results[n], powers[n];

void timeit(void (*f)()) {
    // запускает другую функцию и меряет время её исполнения
    clock_t start = clock();
    f();
    cout << double(clock() - start) / CLOCKS_PER_SEC << endl;

    for (int i = 0; i < 10; i++)
        cout << results[i] << " ";
    cout << endl;
}

int main() {
    for (int i = 0; i < n; i++) {
        bases[i] = rand();
        powers[i] = rand();
    }

    // timeit(binpow_simple);
    // timeit(binpow_sse);

    return 0;
}

В SSE весьма сложно делить int-ы (см. примечания ниже), поэтому будем считать всё по модулю $2^{32}$, то есть просто переполняя естественным образом unsigned int.

Напишем стандартное итеративное бинарное возведение в степень:

void binpow_simple() {
    for (int i = 0; i < n; i++) {
        unsigned a = bases[i], p = powers[i];

        unsigned res = 1;
        while (p > 0) {
            if (p & 1)
                res = (res * a);
            a = (a * a);
            p >>= 1;
        }

        results[i] = res;
    }
}

Этот код работает за 9.47 секунды.

Теперь попробуем векторизованную версию:

void binpow_sse() {
    const reg ones = _mm256_set_epi32(1, 1, 1, 1, 1, 1, 1, 1);
    for (int i = 0; i < n; i += 8) {
        reg a = _mm256_load_si256((__m256i*) &bases[i]);
        reg p = _mm256_load_si256((__m256i*) &powers[i]);
        reg res = ones;

        // на самом деле, здесь никакого цикла не будет
        // -- компилятор это развернёт в 32 отдельных блока операций
        for (int i = 0; i < 32; i++) {
            // чтобы не писать if, посчитаем для каждого элемента его множитель:
            // это либо единица, либо a, в зависимости от значения нижнего бита p
            // маски элементов, которые нужно домножать на a:
            reg mask = _mm256_cmpeq_epi32(_mm256_and_si256(p, ones), ones);
            // смешаем два вектора по маске:
            reg mul = _mm256_blendv_epi8(ones, a, mask);
            // res *= mul:
            res = _mm256_mullo_epi32(res, mul);
            // a *= a:
            a = _mm256_mullo_epi32(a, a);
            // p >>= 1:
            p = _mm256_srli_epi32(p, 1);
        }

        _mm256_store_si256((__m256i*) &results[i], res);
    }
}

Эта реализация уже работает за 0.7 секунды — в 13.5 раз быстрее. При этом там и дальше есть, что оптимизировать.

Трудности автовекторизации

В самом начале статьи мы приводили пример кода, в котором уже оптимизированный бинарник получается без каких-либо изменений, кроме подключения нужного таргета компиляции. Зачем тогда вообще программисту делать что-либо ещё?

Дело в том, что иногда — очень редко — программист всё-таки умнее компилятора, потому что знает про задачу чуть больше.

Рассмотрим этот же пример, убрав из него всё лишнее:

void sum(int a[], int b[], int c[], int n) {
    for (int i = 0; i < n; i++)
        c[i] = a[i] + b[i];
}

Почему эту функцию нельзя заменить на векторизованный-вариант автоматически?

Во-первых, потому что это не всегда корректно. Предположим, что a[] и c[] пересекаются, причём так, что указатели на начало массивов отличаются на 1-2 позиции. Ну, мало ли — может, мы такой изощрённой свёрткой хотели посчитать последовательность Фибоначчи. Тогда в simd-блоках данные будут пересекаться, и наблюдаемое поведение будет совсем не то, которое мы хотели.

Во-вторых, мы ничего не знаем про выравнивание этих массивов, и можем потерять производительность здесь (для больших циклов неактуально — компилятор оба «края» обрабатывает отдельными циклами).

На самом деле, когда компилятор подозревает, что функция будет использована для больших циклов, то на высоких уровнях оптимизации он сам вставит runtime-проверки на эти случаи и сгенерирует два разных варианта: через SSE и «безопасный».

Но выполнять эти проверки в рантайме не хочется, поэтому можно сообщить компилятору, что мы уверены, что ничего не сломается:

#pragma GCC ivdep
for (int i = 0; i < n; i++)
    // ...

Здесь «ivdep» означает ignore vector dependencies — данные внутри цикла ни от чего не зависят.

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

Разное

С++ в ассемблер. Посмотреть на генерируемые инструкции можно так:

g++ -S program.cpp -o program.s

Это позволяет понять, векторизует ли уже компилятор код или нет (названия векторных инструкций начинаются с буквы v). Во многих IDE есть удобные плагины, позволяющие выяснять это для конкретных функций.

Если указать флаг -fopt-info-vec-optimized, то компилятор прямо укажет на операции, которые он смог векторизовать:

g++ -fopt-info-vec-optimized program.cpp -o run

Можно поменять optimized на missed или all, чтобы посмотреть причины, почему не получилось векторизовать другие.

Распечатать вектор. Для дебага помогает такой код:

template<typename T>
void print(T var) {
    unsigned *val = (unsigned*) &var;
    for (int i = 0; i < 4; i++)
        cout << bitset<32>(val[i]) << " ";
    cout << endl;
}

В данном случае он выводит 4 группы по 32 бита из 128-битного вектора.

Деление. В SSE нет операции деления int-ов, но есть для float-ов и производных. Также нет взятия остатка от деления, что осложняет вычисления в комбинаторике.

Для деления 32-битных целых чисел их можно аккуратно скастовать к даблу, поделить так, и скастовать обратно — точности хватит, хоть это и будет медленно.

Умножение работает в несколько раз быстрее деления, и поэтому для ускорения деления float-ов на известную константу $d$ есть следующий трюк: заменить выражение $x / d$ на $x \cdot \frac{1}{d}$, и при этом $\frac{1}{d}$ посчитать во время компиляции.

Для целочисленных типов такое сделать немного сложнее — нужно заменить деление на умножение и битовый сдвиг. Для этого нужно приблизить $\frac{1}{d} \approx \frac{m}{2^s}$, подобрав «магическое» число $m$ и степень двойки $s$, такие что что x / d == (x * m) >> s для всех x.

Можно показать, что такая пара чисел всегда существует, и компилятор сам оптимизирует деление на константу подобным образом. Вот, например, сгенерированные инструкции для деления unsigned long long на $10^9 + 7$:

movq    %rdi, %rax
movabsq $-8543223828751151131, %rdx ; загружает магическую константу в регистр
mulq    %rdx                        ; делает умножение
movq    %rdx, %rax
shrq    $29, %rax                   ; делает битовый сдвиг результата

Здесь для умножения используется «mixed precision» инструкция mulq, которая берёт два 64-битных числа и записывает 128-битный результат их умножения в два 64-битных регистра (lo, hi).

Для деления long-ов на SSE такой способ пока что не работает: аналогичная инструкция добавилась только в AVX512.