понедельник, 1 августа 2016 г.

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

Итак, ACM почти закончился для меня. Лето я провожу бездарно. Поэтому решил что-нибудь написать в блог, чтобы принести хоть какую-то пользу.
Вообще, мне давно были интересно разобраться с SSE/AVX инструкциями. Пару раз я встречал на Codeforces очень сложные кодины (например http://codeforces.com/contest/573/submission/12760770), которые выглядели(да и до сих пор выглядят) для меня как черная магия. Я попытался разобраться с этой магией. Попытаюсь рассказать о том, с чем я разбирался, хоть и похож я больше на шарлатана, чем на мага.


1. Что такое SSE и AVX?

Вообще, SSE и AVX - это специальные наборы инструкций для процессоров Intel, которые позволяют выполнять SIMD (Single Instruction Multiple Data) вычисления.

Концепт SIMD операции
Таким образом, с помощью одной SIMD инструкции мы можем произвести некоторую операцию сразу над группой данных. Реализуется это с помощью специальных регистров (XMM или YMM, в зависимости от того, какое SIMD расширение нужно), а также специальных SIMD ALU, которые содержат в себе обычные, скалярные ALU для выполнения групповых операций. Собственно, SSE представляет собой набор SIMD инструкций, оперирующий над 128-битными блоками данных (XMM регистры), AVX вводит дополнительные 256-битные регистры, но набор операций для них достаточно скуден. Полноценно использовать их можно только с расширением AVX2

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

2. "Все понятно, показывай уже коды"

К сожалению, SSE/AVX инструкции могут быть достаточно капризными к данным, которые они обрабатывают. Поэтому нужно еще обратить внимание на выравнивание данных. Все знают, что процессор читает и записывает данные блоками, размер которых равен нескольким байтам. Причем данные блоки выровнены по некоторой границе (сейчас мы сталкиваемся либо с блоками в 32 бита и соответствующей границей, либо с 64-битным выравниванием). Обычно, работа с данными по невыровненному адресу может просто замедлить исполнение программы (причем обычно не очень значительно). В случае же с SIMD инструкциями время выполнения может увеличиться значительно (хотя вроде бы это не совсем правда). Либо же наша программа может просто экстренно завершиться во время выполнения.

Чтобы избежать такие неприятностей нужно стараться выравнивать все данные, с которыми работают SIMD команды или пользоваться специальными SIMD инструкциями для работы с невыровненными данными.

3. Подготавливаемся

Для начала нужно разобраться, что инклюдить, какие функции/классы использовать, как узнать доступные на сервере расширения.

   • Итак, чтобы проверить, какие SIMD расширения поддерживает процессор, в gcc существует функция __bulitin_cpu_supports. Примеры ее использования:
AVX2 среди популярных сервисов доступен только на CsAcademy, поэтому в основном мы будем пользоваться 128-битными расширениями.

    • Использовать SIMD расширения мы будем в основном с помощью специальных встроенных функций, очень похожих на ассемблерные команды. Есть удобный путеводитель по встроенным функциям для процессоров Intel (там указаны описание и псевдокод функции, задержки выполнения конкретной операции).

    • Чтобы g++ смог скомпилировать intrinsics в соответствующие ассемблерные команды, нужно либо указать специальные ключи компиляции(что мы не можем сделать), либо можно воспользоваться специальными атрибутами g++, а именно атрибутом target (пример использования атрибута совмещен с простеньким кодом, считающим сумму массива)
   • Чтобы использовать необходимые расширения (SSE4.2 и более старые, AVX) нужно подключить заголовочный файл "immintrin.h".
   • И наконец, нужно научиться выравнивать данные. Вообще, в C++11 появились специальные конструкции вида alignas(...)/alignof(...). Первая из них обеспечивает выравнивание определенного объекта по заданной границе. Вторая - возвращает максимальное значение выравнивания, которым обладает данный адрес в памяти:
К сожалению, доступные на Codeforces и Timus компиляторы MSVC не поддерживают данные конструкции, поэтому имеет смысл всегда ориентироваться на компилятор с помощью макросов:
Также тут определен макрос ATTR, который нужно писать перед каждой функцией, непосредственно использующей некоторые SSE intrinsincs. Теперь можно начать. Далее я разберу несколько задач и попытаюсь рассказать, как их можно "упихать" с помощью SIMD расширений. Возможно будет продолжение статьи с мэшапчиком на кф, чтобы можно было в режиме контеста подурачиться.

A. Подпалиндромы (Timus)

Краткое условие: есть строка s. Поступают запросы двух типов:
  1. s[i] = c
  2. Проверить, что подстрока s[l..r] - палиндром
Ограничения: количество запросов и длина строки до 1e5.
Мне давно хотелось сдать эту задачу за O(nm). Давайте сделаем это с помощью SIMD. Поступим следующим образом: будем хранить одновременно строку и перевернутую строку. Запрос изменения символа мы все еще делаем за константу, а чтобы проверить на палиндромность - нам нужно сравнить две подстроки. Я выбрал такой подход по двум причинами:
  1. Почему-то именно это пришло мне в голову первым
  2. В данном случае мы реализуем свою функцию memcmp, быстрее той, которая вызывается на сервере
Ну что же, напишем для начала функцию memcmp_aligned, на вход которой подаются два выровненных по границе в 16 байт указателя и длина сравниваемых подстрок делится на 16. Получится не очень сложный код:
Окей, эта функция может сравнить почти всю длину. К сожалению, мы не можем гарантировать выровненность обоих указателей по нужной границе, поэтому дальше придется использовать вариацию предыдущей функции - memcmp_unaligned - которая совсем чуть-чуть отличается от оригинальной:
Теперь можно использовать эту функцию, чтобы сравнить почти всю строку, а хвост обработать тупо, без векторизации (это в принципе общий подход - начало и хвост обрабатываются наивно, а середина - с помощью SSE).
Вот и все. Это функция сравнивает лексикографически две строки и делает это в 2.5 раза быстрее, чем функция, вызывающаяся на acm.timus.ru (у меня локально данная функция работает медленнее стандартного memcmp; может у меня компилятор новее, или еще от чего зависит. Измерил отношение на тимусе я вот так). С помощью memcmp_fast уже можно сдать задачу, хотя нужно очень аккуратно работать с вводом/выводом, так как в этой задаче достаточно большие входные данные. У меня получилось сдать задачу с этой функцией за 0.421/0.5 сек. Можно еще немного ускорить наш код, развернув цикл. В случае векторных оптимизаций это достаточно полезный прием, с помощью которого удается сэкономить несколько дорогих SIMD инструкций (в данном случае _mm_testz_si128 выполняется дольше чем _mm_or_si128). Кому очень интересно - можно посмотреть реализацию по ссылке. Код, использующий развернутую версию memcmp_unaligned, у меня заходит уже за 0.312/0.5 сек. Вообще, можно избежать работы с невыровненными адресами с помощью функции _mm_alignr_epi8, но одна операция загрузки данных в векторный регистр оказывается быстрее, чем загрузка по выровненному адресу + _mm_alignr_epi8. Я не буду здесь приводить код, полностью все функции вместе с memset_alignr можно посмотреть здесь
Подведем небольшой итог после разбора этой задачи:

B. В поисках истины (Timus)  


Краткое условие: Есть полоска длины n. И есть додекаэдр, который находится в одной из n позиций. Мы хотим угадать где он. Мы сделали m попыток. Каждая попытка - одна позиция p_i. Если додекаэдр в этот момент был в этой позиции - мы выиграли, если нет, то продолжаем играть, а додекаэдр из своей истинной позиции сдвигается либо вправо, либо влево (если там не край поля). Выиграем ли мы?
Ограничения: n, m до 1e5.
В принципе вместо этой задачи можно было бы взять любую другую, где нужно использовать bitset-ы. Наивное решение за квадрат опять же очевидно - поддерживаем все возможные положения Додекаэда, а после каждого хода исключаем из множества одну позицию и считаем побитовый OR двух сдвигов - на один бит влево и вправо. Обычная реализация битсетов заходит под студией за 0.655/1.0 сек, а под g++ получает TL. Ну что же, давайте посмотрим, насколько её ускорит SSE. Для начала можно сделать самую тривиальную оптимизацию - написать OR битсетов на SIMD расширениях. Это очень просто и логично:
Под студией данная оптимизация не дает ускорения, зато под g++ мы уже можем сдать задачу за 0.9/1.0 сек. Можно подметить, что в данном случае разворачивать цикл не имеет особого смысла, так как этим мы не добиваемся экономии дорогих инструкций.
Осталось только реализовать битовый сдвиг с помощью SSE. Это уже не так тривиально, потому что сдвиг не является независимой операцией по компонентам, а SSE/AVX не предоставляет готовой инструкции для такой операции. Однако это можно сделать с помощью комбинации инструкций, которые предоставляются SSE следующим образом:


Собственно реализация выглядит следующим образом:
Аналогичным образом можно написать сдвиг в другую сторону. К сожалению, насколько я понимаю, реализация сдвига на произвольное число бит - нетривиальное занятие. Например по этой ссылке автор разбирает оптимальные комбинаций SSE инструкций для каждого из 128 возможных значений сдвига (они конечно группируются, но все равно ему пришлось заняться кодогенерацией).
Так что же с нашей оптимизацией? Использование SSE в OR и сдвигах позволяет сдать задачу под g++ за 0.468/1.0 сек, а под msvc за 0.39/1.0 сек, что показывает примерно двукратное (или около того) ускорение по сравнению с наивной реализацией битсета.

4. Перерыв

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

Ссылки

[1] Не очень новая статья про выравнивание. Сейчас все уже не так плохо, насколько я понял. Например их тесты, на моем компьютере, почти никак между собой не различаются. Можно заметить падение производительность на 3% при обращении к невыровненному long long-у и все.
[2] Удобный путеводитель по интеловским интринсикам.
[3] Все функции memcmp, которые я имел наглость написать

Комментариев нет:

Отправить комментарий