Недавно я выпустил StringZilla v4 — первый релиз с поддержкой CUDA моей библиотеки для обработки строк. нацеленной в первую очередь на SIMD. Это означает, что теперь она стала быстрой не только на CPU, но и на GPU!

• Я хотел добавить ускорение ROCm для GPU AMD

• Я хотел добавить параллельный мультипаттерновый алгоритм поиска

• Я хотел опубликовать всё это ещё в декабре 2024 года

Итак, не всё пошло по плану, но StringZilla 4 CUDA наконец-то здесь, и она добавляет 500 с лишним GigaCUPS вычислений редакторского расстояния; при этом пакет можно установить через pip install. Также в ней есть некоторые другие трюки, предназначенные для крупномасштабных систем извлечения данных, баз данных и озёр данных, а также биоинформационных задач. И всё это под разрешительной опенсорсной лицензией Apache 2.0, позволяющей свободно использовать библиотеку в коммерческих целях. В этом посте я рассмотрю самые интересные части релиза, и в том числе:

• Быструю оценку алгоритмов динамического программирования на GPU,

• Хэширование CRC32, MurMurHash, xxHash, aHash и не только, а также

• Фингерпринтинг биологических последовательностей 52-битными целыми числами

Предыстория и источники вдохновения

Изначально StringZilla родилась, как материал для доклада с конференции в конце 2010-х, демонстрирующий мощь AVX-512 и тонкости векторизации нагрузок с непараллельными данными. За последующие годы я превратил её из нескольких ядер поиска подстрок в монстра, соревнующегося с GLibC за звание самого быстрого memcpy (да, я знаю, что о подобном заявляют многие). Позже в неё была добавлена поддержка платформ little- и big-endian; множество поколения расширений AVX-512 на x86, Arm NEON, SVE и SVE2; динамическая диспетчеризация; прямые привязки к Python, Rust, JavaScript и даже к Swift, транслирующие внутреннюю реализацию C99.

Теперь StringZilla v4 ещё больше расширила функциональность:

• Ещё одна некриптографическая хэш-функция, строковые PRNG для AES и другие портопараллельные команды SIMD.

• Новые алгоритмы нахождения пересечения и сортировки для больших наборов строк в JOIN и ORDER BY DBMS.

• Ядра нахождения схожести строк с ускорением GPU и CPU, обрабатывающие расстояния Левенштейна, алгоритмы Нидлмана — Вунша и Смита — Ватермана с расширениями Гото для «аффинных пробелов», необходимых в биоинформатике.

• Ядра фингерпринтинга MinHashing с ускорением GPU и CPU для крупномасштабного извлечения информации и избавления от дубликатов.

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

Традиционные способы вычисления схожести строк

Динамическое программирование и порядок вычисления расстояний Левенштейна

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

Здесь ​ — это редакторское расстояние между первыми символами строки  и первыми символами строки . Три члена в минимуме обозначают, соответственно, затраты на удаление, вставку и замену, где  равно 1, если символы отличаются, и 0, если они совпадают.

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

При параллелизации таких алгоритмов можно поступать умнее: изучить цепочки зависимостей, препятствующие векторизации. В случае таких расчётов схожести строк, в том числе и расстояния Левенштейна, всё просто — нужно вместо строк вычислять диагонали! Мы будем хранить три диагонали вместо двух строк, и каждая последующая диагональ будет вычисляться из двух предыдущих. Затраты на подстановку будут браться из первой диагонали, а затраты на вставку и удаление — из второй.

Помогло ли это? Производительность таких алгоритмов вычисляется в обновлениях ячеек в секунду (Cell Updates Per Second, CUPS). При сравнении строк длиной ≅ 1000 байт результаты были такими:

• rapidfuzz::levenshtein: 14316 MCUPS на ядре Intel Sapphire Rapids,

• stringzillas::LevenshteinDistances: 13084 MCUPS на том же одном ядре,

• stringzillas::LevenshteinDistances: 624730 MCUPS на GPU Nvidia H100.

Одна из самых активно используемых библиотек Python, NLTK, с более чем миллиардом скачиваний, обеспечивает приблизительно ≅ 2 MCUPS. Разумеется, обёрнутая в Python RapidFuzz обеспечивает гораздо большую производительность, но основная её часть всё равно теряется в слое привязок. Привязки между C и Python в StringZilla — одни из самых тонких в GitHub, поэтому по сравнению с другими пакетами деградация минимальна, в том числе и по сравнению с CuDF компании Nvidia. Даже при предварительной подготовке объектов cudf.Series производительность для разных длин строк в MCUPS выглядит следующим образом:

? обозначает привязки Python, ? обозначает Rust.

CuDF вызывает nvtext::levenshtein. Она хорошо работает с коллекциями небольших входных данных, но не крупных. Вероятнее всего, она не распределяет большие входные данные по ядрам GPU так же агрессивно, как StringZilla. И какими же получаются результаты?

• Рост производительности в 46 раз для строк длиной ≅ 1000 байт.

• Рост производительности в 109 раз для строк длиной ≅ 10000 байт.

Если честно, cudf — это не совсем библиотека для вычисления схожести строк, и у Nvidia есть отдельный коммерческий проект Clara Parabricks, выполняющий именно эту задачу. Тем не менее, производительность в 109 раз больше, чем у CuDF — это отличное начало, которое будет идеальным фундаментом для дальнейшего совершенствования!

Затраты на подстановки и штрафы аффинных пробелов

В биоинформатике расстояние Левенштейна используется для сравнения строк ДНК. Вставка и удаление символов в парах строк сигнализируют о физических разрывах в очень длинных молекулах. Присутствие таких пробелов — это более весомый фактор, чем определённая длина пробелов, поэтому оно оценивается иначе — при помощи «штрафов аффинных пробелов» (affine gap penalties). В этом случае хранения одной матрицы уже недостаточно: нам нужно хранить три матрицы, в том числе две новые, позволяющие различать продолжение пробелов и их начало.

Аналогично, алгоритм Нидлмана — Вунша обобщает расстояние Левенштейна для учёта переменных затрат на подстановки. Он удобен при работе с последовательностями белков. Эти последовательности представлены в виде строк с гораздо меньшим алфавитом из 20 аминокислот: ACDEFGHIKLMNPQRSTVWY, поэтому нам нужно хранить матрицу подстановок 20x20 или большего размера, если мы хотим учитывать неоднозначные аминокислоты.

StringZilla уже использует команды Nvidia DP4A и DPX для SIMD-обработки малых целочисленных значений в таких задачах динамического программирования. С другой стороны, текущая реализация StringZilla ошибочно использует для хранения матрицы подстановок память констант; это узкое место и в будущем ситуация будет улучшена. Тем не менее, результаты всё равно не так уж плохи по сравнению с другими решениями, которые можно установить через pip install. Для последовательностей аминокислот длиной ≅ 1000 мы получаем следующее:

• biopython обеспечивает ≅ 303 MCUPS,

• stringzillas-cpus обеспечивает ≅ 276 MCUPS,

• stringzillas-cuda обеспечивает ≅ 10098 MCUPS.

Ещё две хэш-функции?

Цели разработки

В мире очень много хэш-функций! Некоторые из них очень хороши.

• Например, XXH3 Янна Колле — это интересная кодовая база, которую можно использовать в продакшене; в ней может найти что-то интересное даже очень опытный разработчик.

• MurMurHash — это классика, выпущенная Остином Эплбаем в 2008 году вместе с фреймворком SMHasher.

• Также есть очень качественный пакет ahash для Rust разработчика Тома Кейтчака, популяризирующий команды AES, ранее предложенные Эндрю Роджерсом в AquaHash.

В целом, мне нужна была хэш-функция, обладающая следующими свойствами:

• Быстрая и для коротких (скорость), и для длинных строк (производительность).

• Поддержка инкрементного (потокового) хэширования, когда данные поступают по блокам.

• Поддержка порождающих значений (seed) для хэшей, влияющих на все части выходных данных.

• Обеспечение динамической диспетчеризации для разных архитектур с целью упрощения развёртывания.

• Использование не только AVX2 & NEON SIMD, но и маскированного AVX-512 и предикатного SVE2.

• Документация логики и генерация одинакового результата на всех платформах.

• Вывод 64-битных или 128-битных хэшей и прохождение тестов SMHasher --extra.

AES и параллелизм портов

Применение команд AES для хэширования — не новая, но очень логичная идея. Посмотрите на объём смешения сигналов, происходящий в одном раунде AES:

На каждой итерации AES выполняет четыре операции с матрицей 4∗4=16 байт:

1. SubBytes — нелинейный этап подстановки, на котором с каждым байтом сопоставляется другой.

2. ShiftRows — перестановка, при которой последние три строки циклично смещаются.

3. MixColumns — линейная операция смешения, работающая со столбцами состояния.

4. AddRoundKey — простая операция XOR с ключом раунда.

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

• VAESENC (ZMM, ZMM, ZMM) и VAESDEC (ZMM, ZMM, ZMM):

  • на Intel Ice Lake: 5 тактов на порте 0.

  • на AMD Zen4: 4 тактов на портах 0 или 1.

Как это часто происходит с целочисленными нагрузками на процессорах Intel, команда SIMD всегда маршрутизируется только на один порт исполнения, обычно это порт 0 или 5. AES выполняет маршрутизацию на 0, поэтому чтобы смешивать параллельно больше данных, нам нужно сочетать это с малозатратными командами отображения на другие порты, например:

• VPSHUFB_Z (ZMM, K, ZMM, ZMM)

  • на Intel Ice Lake: 3 такта на порте 5.

  • на AMD Zen4: 2 такта на портах 1 или 2.

• VPADDQ (ZMM, ZMM, ZMM):

  • на Intel Ice Lake: 1 такт на портах 0 или 5.

  • на AMD Zen4: 1 такт на портах 0, 1, 2, 3.

Это надёжное решение, идеалогически близкое к ahash, но имеющее некоторые тонкости:

• Для входных данных больше 64 байт используются большее состояние и больший размер блоков, что даёт выигрыш благодаря более широким регистрам современных CPU. Как и у многих других хэш-функций, состояние инициализируется при помощи seed и множества констант пи. В отличие от других функций, здесь мы берём больше битов пи (1024), но только 64 бита seed, чтобы API не был слишком безумным.

• Длина входных данных не примешивается в блок AES в начале, чтобы обеспечить возможность инкрементной генерации, когда окончательная длина заранее неизвестна.

• Векторные нагрузки не чередуются, то есть каждый байт входных данных имеет одинаковый вес в хэше. При реализации на старых платформах это требует дополнительного перемешивания, а на новых платформах выполняется при помощи «маскированных» нагрузок в AVX-512 и «предикатных» команд в SVE2.

Любопытный факт: команды AES на x86 и Arm выполняются немного по-разному, но если это компенсировать, производительность всё равно останется отличной. Бенчмарки StringWa.rs на ядре Intel Sapphire Rapids дают следующие результаты:

¹ Портируемость означает доступность на нескольких других языках программирования, например, C, C++, Python, Java, Go, JavaScript и так далее.

Генерация случайных строк

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

1. CTR (Counter Mode)

2. OFB (Output Feedback Mode)

3. CFB (Cipher Feedback Mode)

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

Часть логики StringZilla превосходит по скорости memcpy благодаря использованию невременных хранилищ, обходящих кэши CPU и выполняющих запись данных непосредственно в ОЗУ. Это снижает степень загрязнения кэшей и повышает энергоэффективность нагрузок, активно использующих ввод-вывод. Поэтому для PRNG это крайне логично, но если вы хотите, чтобы PRNG мог использовать seed и был воспроизводимым, то повышение сложности из-за обработки невыравненных операций записи того не стоит! Даже без этих трюков результаты по сравнению с другими Rust-вариантами оказываются впечатляющими:

52-битные вычисления для биоинформатики?

Однако есть и тип хэширования, в котором я по-прежнему вместо примитивов AES использую целочисленную арифметику по модулю, но с неожиданной особенностью!

Для идентификации элементов мой поисковый движок USearch использует 40-битные целые числа — нестандартный размер выбран, чтобы уравновесить эффективность памяти с адресным пространством, достаточным для размещения одного триллиона векторов на одной машине. Продолжая эту тенденцию использования необычных целочисленных размеров, StringZilla v4 использует 52-битные integer для вычисления MinHash, или «Min-wise independent permutations Locality Sensitive Hashing».

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

Из текстового документа длиной байт мы можем извлечь  N-грамм длиной . Например, строка «hello» содержит следующие 3-граммы: «hel», «ell» и «llo». MinHash определяет хэш-функций , , … . Для каждой хэш-функции ​ мы вычисляем хэшей — по одному для каждой N-граммы — и берём минимальное значение:

Эти  минимальных значений хэша образуют MinHash-сигнатуру документа, -мерный вектор. С увеличением каждый из этих хэшей становится всё более вычислительно затратным. То есть общая сложность вычисления MinHash-сигнатуры составляет . Это много!

Для снижения сложности до можно перейти к скользящим хэшам в стиле Рабина, что было реализовано и в StringZilla v3, но этого оказалось недостаточно! Для вычисления высококачественных сигнатур требуются достаточно хорошие промежуточные хэши, то есть:

1. простая схема смешения с большими 64-битными представлениями состояний хэшера, или

2. более простая схема смешения с меньшими 32-битными представлениями состояний хэшера.

Но всегда лучше третий вариант, правда? Для многих из нас, в том числе для читателей этого блога и Less Slow C++, уже привычно использовать популярный трюк — float для выполнения грязной работы с небольшими целыми числами. Но многие забывают, что существуют и double, обладающие 53-битной точностью (52 хранимых битов плюс 1 неявный). То есть если мы сможем уместить состояние хэшера в 52 бит, то нам удастся выполнять всё смешение и арифметику по модулю с использованием double, а затем сохранять окончательный результат, как 32-битный integer.

Современные векторизированные CPU с 512-битными блоками FMA справляются с этим на удивление хорошо, однако GPU — ещё лучше! Мне потребовалось какое-то время, чтобы избавиться от всех пограничных случаев и чтобы ядра CPU и GPU генерировали одинаковые фингерпринты. Итоги:

• однопоточный код на Rust: 0,5 МиБ/с.

• код H100 CUDA: 392,37 МиБ/с.

В 2023 году Nvidia также добавила MinHash в свой пакет nvtext::. Он использует в качестве фундамента MurmurHash3_32, и по архитектуре ближе к пакету Rust probabilistic_collections. Сравнивать эти решения исключительно по одной производительности было бы нечестно, потому что важно качество генерируемых сигнатур, а оно сильно различается.

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

Сортировка и пакетные операции

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

bool is_less(char const* a, size_t a_len, char const* b, size_t b_len) noexcept {
    char const* a_comparison_end = a + std::min(a_len, b_len);
    for (; a < a_comparison_end; ++a, ++b) // ! Цикл, которого нужно избегать!
        if (*a != *b) return *a < *b;
    return a_len < b_len;
}

В предыдущих версиях StringZilla я показал, что существенного ускорения можно достичь при помощи тривиальной оптимизации: сортировки целочисленных префиксов перед остальной частью строки. Однако реализация оставалась proof of concept; теперь она гораздо более удобна и результаты её гораздо лучше:

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

Попробуйте сами

Кроме алгоритмических инноваций, сложной задачей при разработке стал выпуск библиотеки для различных платформ и языков:

• Для одних только привязок Python я выпускал на каждый релиз в PyPI больше специфичных для платформ wheel, чем NumPy…

• И, в отличие от NumPy, мои библиотеки проставляются с собственными ядрами SIMD и GPGPU, которые нужно перекомпилировать для каждой платформы вместо использования интерфейсов C библиотек BLAS…

• И, в отличие от NumPy, StringZilla также выпущена для C, C++, CUDA, Rust, JavaScript, Go и Swift, а у каждого из них есть свои тонкости…

… и всё это из CI одного репозитория на бесплатном тарифе GitHub Actions. Так что не удивлюсь, если всё ещё будут обнаруживаться флаги сборки, распространяющие странные проблемы и вызывающие странную поддержку SIMD на некоторых платформах. Чтобы убедиться, что у вас всё работает так, как задумано, для Python выполните следующее:

pip install stringzilla         # для последовательных алгоритмов
pip install stringzillas-cpus   # для параллельных бэкендов с несколькими CPU
pip install stringzillas-cuda   # для параллельного бэкенда GPU Nvidia

Для проверки обнаруженных возможностей:

python -c "import stringzilla as sz; print(sz.__version__, sz.__capabilities__)"       # для последовательных алгоритмов
python -c "import stringzillas as szs; print(szs.__version__, szs.__capabilities__)"   # для параллельных алгоритмов

Информацию по другим языкам программирования см. в README.md, а если что-то не работает, отправляйте свои issue в тему в GitHub. Впрочем, для Rust всё должно работать нормально, потому что я активно тестировал библиотеку для StringWa.rs, так что вы можете воспроизвести всё на собственном оборудовании:

RUSTFLAGS="-C target-cpu=native" \
    STRINGWARS_DATASET=your-favorite-dataset \
    STRINGWARS_TOKENS=lines \
    cargo criterion --features bench_hash bench_hash --jobs $(nproc)

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

Особая благодарность Nebius — моему самому часто используемому облаку за надёжные мощности GPU для подобных личных экспериментов, а также опенсорсному Unum.