malkovsky 10 часов назад

Ранг-селект словари

Сложный

9 мин

686

Программирование * Алгоритмы * Математика * C++ *

Пример битвектора и результаты операций ранг/селект над ним, справа реализации ранга/селекта с помощью специальных процессорных инструкций.

Это первая статья из планируемой серии про succinct data structures - это класс очень компактных структур данных. Канонический пример такой структуры - это представление дерева в виде правильной скобочной последовательности, дерево извершин таким образом представляется с помощьюбит в то время как типичная динамическая реализация требовала бы как два указателя по 64-бит на каждый узел (разумеется можно немного сократить простыми оптимизациями, но даже близко 2 бита не получить). Фундамент подобных структур - это rank-select словарь, представляющий собой битовый вектор и дополнительную структуру для выполнению двух операций ранг и селект. В указанном примере с деревом с помощью ранга и селекта можно сделать базовую навигацию: найти номера потомков/родителей, узнать размер поддерева. В статье расскажу как делать эти операции быстро используя при этом всего 3,6% дополнительной памяти.

Введение

Rank-select dictionary (RSD) - это структура данных, построенная поверх некоторой битовой строки, позволяющая эффективно отвечать на два запроса:

- количество единичных бит на позициях $0, \ldots k-1$ , т.е.

$Rank(k)=\sum_{i=0}^{k-1}B[i]$

- позиция-ого единичного бита, т.е.

$Select_1(k) = \min\{ \ell~|~Rank_1(\ell) = k \}$

Например для строки (договоримся, что в статье работаем со строками, которые следует читать слева-направо)

$Select_1(k) = \min\{ \ell~|~Rank_1(\ell) = k \}$
$Select_1 = \{1, 4, 6, 8, 9, 10\}$

Аналогичным образом определяются

- количество нулевых бит на позициях $0, \ldots, k-1$

$Rank_0(k)=\sum_{i=0}^{k-1}(1-B[i])$

- позиция-ого нулевого бита

$Select_0(k)=\min\{\ell~|~Rank_0(\ell)=k\}$

Результаты операций ранга/селекта на битвекторе

LOUDS Представление дерева

RSD не является самостоятельной структурой полезной для широкого класса задач как например хэш-таблицы или бинарные деревья поиска, но на ней строятся все известные succinct структуры. Подробно про них рассказывать я планирую в других статьях, а пока небольшой пример одной из таких структур - Level-Ordered Unary Degree Sequence (LOUDS) представление дерева

Несколько вариантов компактного представления дерева

LOUDS строится следующим образом: начинаем со строки, делаем обход в ширину дерева, при посещении вершины с детьми дописываем в конец строки единиц и один ноль. Таким образом на дерево из вершин будет записано нулей и единиц. Если предположить, что вершины пронумерованы в соответствии с обходом в ширину, то следующие операции по дереву выражаются через Rank/Select

Степень вершины

Номер-ого ребёнка-ого узла.

Номер родителя-ого узла

Обычно этого достаточно для того, чтобы делать обходы в глубину/ширину или например сделать поиск в бинарном дереве поиска, но при этом объем данных, необходимый для хранения структурной информации — 2 бита на узел против 2-х 64-битных указателей в стандартной реализации. Не очень сложными махинациями можно добиться $2\log_2n$ бит на указатель, но не более того.

Наивная реализация RSD

Самый просто способ реализации Rank/Select поверх битвектора длины выглядит как-то так: заводим массив длины и храним в частичные суммы

$R[i]=\sum_{j=0}^{i-1}B[j]$

Тогда, а для подсчёта нужно сделать бинарный поиск по. Таким образом ранг вы умеем считать за $\mathcal{O}(1)$ , селект за $\mathcal{O}(\log n)$ . Значения в пробегают в худшем случае от до, в результате чего нужно $n\log_2n$ бит.

Succinct rank

В 1998 Якобсон представил структуру, с помощью которой Rank/Select можно выполнять за $\mathcal{O}(1)$ используя $\mathcal{O}(n)$ памяти. Первичная идея схожа с методом четырёх русских: для всех возможных последовательностей битовых строк длины не больше $L=\frac{1}{2}\lceil\log_2n\rceil$ количество единичных бит в этой строке можно предподсчитать. Таких строк не больше $2\times 2^{\frac{1}{2}\lceil\log2n\rceil}=\mathcal{O}(\sqrt{n})$ . Пусть- ранг первого элемента-ого блока длины, т.е.

$R_L[i]=\sum_{j=0}^{iL-1}B[j]$

тогда ранг можно разбить на две части — ранг ближайшего блока и остаток длины не больше.

$_Rank(k)=R_L\left[k/L\right]+\sum_{i=L\times[k/L]}^{k-1}B[i]$

Теперь же для вычисления ранга мы используем $\mathcal{O}(\log n)$ операций с предподсчетом $\mathcal{O}(\sqrt{n})$ для блоков длины, массив имеет длину, каждый элемент занимает $\mathcal{O}(\log n)$ бит — в итоге мы используем $\mathcal{O}(n)$ памяти, то есть пока еще не succinct. Чтобы добиться succinct можно использовать дельта кодирование или кодирование Элиаса-Фано для сжатия монотонной последовательности. Якобсон предложил следующую явную структуру. К блокам длины добавляем суперблоки длины $K=\lceil\log_2^2n\rceil$ . Теперь, ранг мы храним только для каждого суперблока, а для обычных блоков храним локальный ранг, т.е. ранг по отношению к ближайшему суперблоку. Для локального ранга нужно уметь хранить $\lceil\log_2^2n\rceil$ различных значений, для этого нужно $\lceil 2\log_2\log_2n \rceil$ бит. В итоге получаем

$\mathcal{O}\left(\frac{n}{\log n}\right)$ для хранения рангов суперблоков
$\mathcal{O}(\sqrt{n})$ для хранения таблиц базовых блоков
$\mathcal{O}\left(\frac{n\log\log n}{\log n}\right)$ для хранения локальных рангов

с результирующей формулой

$Rank(k)=R_K[k/K]+R_L\left[k/L\right]+\sum_{i=L\times[n/L]}^{n-1}B[i]$

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

Succinct select

Опишу упрощённую версию версию селекта, с помощью которой можно получить асимптотическое время запроса $\mathcal{O}((\log\log n)^2)$ , на практике попытка добиться $\mathcal{O}(1)$ оказывается контрпродуктивной. Итак, 3 основные идеи для селект

Как и в случае с рангом для блоков длины не больше $L=\frac{1}{2}\log_2 n$ можно предподсчитать таблицы ответов для всех возможных блоков такой длины запросом ранга в диапазоне $0\ldots L-1$ , из чего в результате получаем таблицы размера $\mathcal{O}(\sqrt{n}\log_2 n)$ .
Как уже отмечалось, можно использовать посчитанные таблицы ранга для использования бинарного поиска для нахождения селекта. Тут стоит отметить, что стоит аккуратно оценивать сложность бинарного поиска так как операция сравнения стоит $\mathcal{O}(S)$ для чисел длины. Бинарный поиск на рангах суперблоков стоит $\mathcal{O}(\log n)$ на итерацию и делает $\mathcal{O}(\log n)$ итераций. Бинарный поиск на базовых блоках стоит $\mathcal{O}(\log\log n)$ на итерацию и делает $\mathcal{O}(\log \log n)$ итераций.
Бинарный поиск на суперблоках дорогой, но от него можно избавиться с помощью предподсчёта позиции каждого-ого единичного бита. Если обозначить, то

$S_K(\lfloor i/K \rfloor)\leq Select(i)<S_K(\lfloor i/K\rfloor + 1)$

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

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

Практические аспекты ранга/селекта

Для начала стоит отметить общие моменты по железу, ориентируемся на CPU

Для всех разумных применений $n\leq 2^{64}$ и, соответственно, $\log_2n\leq 64$ . Операция над $\log_2 n$ битами такая же по сложности как операция над 1 битом.
Таблицы размера $\sqrt{n}$ достаточно большие чтобы точно не вмещаться в L1 кэш, а может даже и в L2/L3. Из-за этого Broadword programming/SWAR (см. например статью) подходы предпочтительней табличных.
Еще про кэши: мы хотим делать быстрые запросы используя 2-3 таблицы поверх большого объема данных, количество кэш промахов оказывается в этой ситуации боттлнеком и это первое, что стоит оптимизировать при выборе параметров архитектуры. В частности размер базового блока стоит выбирать по размеру кэш линии, т.е. 64 байт/512 бит для большинства CPU. Есть одна оптимизация, использующая в большинстве реализаций - interleaving, заключается в том, чтобы хранить исходный вектор вперемешку с информацией о базовых блоках или же информацию о базовых и суперблоках.
Большинство процессоров предоставляют инструкции POPCNT, TZCNT и PDEP и их SIMD версии. Первая — подсчёт единичных битов в числе, вторая - подсчёт первых нулевых бит, третья немного сложнее, опишу позже. Ранг вплоть до 512 бит можно сделать с помощью POPCNT, селект также вплоть до 512 бит можно проделать с помощью PDEP+TZCNT. POPCNT и TZCNT являются частью стандарта С++20, а вот c PDEP придётся возиться вручную. Также стоит отметить, что на момент написания статьи есть некоторые опасения по поводу AVX-512, кто хорошо подкован в вопросе, напишите пожалуйста в комментариях.

Вычисления ранга/селекта на 512 битах

SIMD within a register(SWAR)/broadword programming — термины для параллельных алгоритмов на данных размера машинного слова. Канонический пример — вычисление количества единичных битов в целом числе (как раз то, что нам нужно для ранга)

uint64_t popcount64(uint64_t x) {
    x = x - ((x >> 1) & 0x5555555555555555ULL);               
    x = (x & 0x3333333333333333ULL) + ((x >> 2) & 0x3333333333333333ULL);
    x = (x + (x >> 4)) & 0x0F0F0F0F0F0F0F0FULL;               
    x = x * 0x0101010101010101ULL;                            
    return x >> 56;                                           
}

Сейчас этот алгоритм не имеет большого практического значения из-за того, что современные процессоры предоставляют инструкцию POPCNT, вычисление с помощью которой будет быстрее, но если вдруг по какой-то причине она недоступна, то этот код наиболее производителен из оставшихся альтернатив. C SIMD вариантом есть пара нюансов. Чтобы сделать popcount первых count бит можно сделать что-то такое:

uint64_t rank_64(uint64_t x, uint64_t count) {
  // не забудьте, что при count >= 64 1ull << count -- это UB
  // в этом контексте проверка дорогая и остаётся на усмотрение
  // разработчика
  return std::popcount(x & ((1ull << count) - 1));
}

SIMD вариант — это инструкции mmXpopcnt_epi64для параллельного вычисления для нескольких 64-битных чисел + mmXreduce_add_epi64для суммирования результата. И тут внезапно оказывается, что подсчет маски ((1 << count) - 1))для count ≥ 64одной инструкцией сделать не получится. Одной инструкцией можно сделать для count=8k, но для оставшихся 8 бит приходится что-то придумывать, лучшее что я нашел - инструкция shldv(a, b, k), которая берёт две последовательности чисел a, b и составляет из них последовательность с, каждый элемент которой - это k верхних бит b и (x-k) младших бит a, в итоге получаем вот такое

uint64_t rank_512(uint64_t *x, uint64_t count) {
  __m512i a = _mm512_maskz_set1_epi64((1ull << ((count >> 6))) - 1,
                                      std::numeric_limits<uint64_t>::max());
  __m512i b = _mm512_maskz_set1_epi64((1ull << ((count >> 6) + 1)) - 1,
                                      std::numeric_limits<uint64_t>::max());
  __m512i mask = _mm512_shldv_epi64(a, b, _mm512_set1_epi64(count % 64));

  __m512i res = _mm512_load_epi64(x);
  res = _mm512_and_epi64(res, mask);
  __m512i cnt = _mm512_popcnt_epi64(res);
  return _mm512_reduce_add_epi64(cnt);
}

Для подсчёта селекта есть SWAR алгоритм на основе ранга и параллельного бинарного поиска, реализация есть вот тут. По сути мы просто делаем бинарный поиск по рангу, его можно посчитать через popcount, но можно и проще (но не факт, что быстрее). В SWAR алгоритме для подсчёта единичных битов на промежуточных шагах строятся суммы каждого блока размера, если сохранить эту информацию, то на-ой итерации бинарного поиска нам достаточно информации по блокам размера $2^{6-i}$ (для 64-битного числа). В итоге получаем алгоритм из 10 итераций, на каждой итерации несколько арифметических действий, сами итерации к сожалению не параллелятся.

Как оказывается, селект можно сделать через 2 процессорные инструкции:

Parallel bit deposit (PDEP): принимает два аргумента x, mask. Допустим, единичные биты mask находятся на позициях $i_1, i_2, \ldots, i_k$ , PDEP берёт первые бит xи располагает их на этих позициях, является частью BMI2
Count trailing zeros (TZCNT): это довольно стандартная операция для подсчёта младших нулевых бит или, что тоже самое, позиции первого единичного бита. Начиная с С++20 это std::countr_zero.

Селект выражается очень просто через эти две операции (и битовый сдвиг)

size_t select64(uint64_t x, size_t rank) {
  return _tzcnt_u64(_pdep_u64(1ull << rank, x));
}

Parallel bit deposit берет биты одного числа и раскладывает их по единичным битам другого числа

Поиск нужного блока

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

Бинарный поиск: хорошо в теории, на практике не очень из-за того, что итераций не так много, а скакать по разным участкам памяти не выгодно по локальности
Линейный поиск: плохо в теории, но хорошо на практике, хорошо ускоряется SIMD инструкциями
Интерполяционный поиск: если мы знаем позицию-ого бита и-ого бита, а распределение единиц почти линейное, тогда позицию $i\leq k \leq j$ можно приблизить как.

Как отмечают в статьях если запросы случайные, то ключевым фактором является количество кэш промахов. Для большого битвектора можно ожидать, что каждый запрос к таблице — это кэш промах, на большинстве процессоров кэш линия 512 бит, поэтому есть смысл просматривать данные в целой кэш линии. Из-за этого линейныq поиск довольно хорош на практике. Для интерполяционного поиска нужно сделать чуть больше вычислений, но он помогает уменьшить количество кэш промахов.

Что в итоге

Две референсных библиотеки:

Spider - довольно свежая статья, по заявлениям у них очень эффективный селект, в котором как раз используется интерполяционный поиск с дополнительными ухищрениями. В целом к сожалению не рекомендую, потому что код написан довольно кустарно "по-академически".
pasta-toolbox - тоже относительно свежая реализация, линейный поиск +SIMD для селекта, эта либа написана довольно хорошо и компактно.

С незначительными отличиями структура ранг-селект битвектора выглядит так:

Базовый блок 512 бит
Супер блок $2^{16}$ бит
64-битные глобальные ранги суперблоков 0.98% от размера битвектора, 16-битные локальные ранги базовых блоков 3.125% от размера битвектора.
Позиции каждого 8192-ого единичного бита 0.78% от размера битвектора.

Ранг вычисляется как и было описано раньше: глобальный ранг + локальный ранг + popcount на базовом блоке. Селект: определяем позицию ближайшего сэмпла, а дальше линейный поиск сначала по суперблокам, потом по базовым, в конце pdep(tzcnt). С помощью 512 регистров можно просматривать либо 8 64-битных глобальных рангов за раз либо 32 16-битных локальных рангов, для локального ранга код выглядит примерно так

for (size_t i = 0; i < 4; ++i) {
  auto batch =
      _mm512_loadu_epi16(&basic_block_rank[128 * s_block + 32 * i]);
  auto cmp = _mm512_cmplt_epu16_mask(batch, rank_32);
  auto count = std::popcount(cmp);
  if (count < 32) {
    return kBlocksPerSuperBlock * s_block + pos + count - 1;
  }
  pos += 32;
}

Свою собственную реализацию я собрал тут

Друзья и коллеги! С удовольствием хотел бы прорекламировать CS Space — открытый научный клуб по CS-related темам; идейных последователей питерского Computer Science Club (и CS Center), расформировавшегося после известных событий. Ребята организуют крутые лекции и курсы по CS от профессионалов своего дела, да еще и помогают мне с написанием научно-популярных статей!

Сайт сообщества: csspace.io
Telegram-канал: t.me/csspace

Если вам понравилась статья — поставьте плюс, автору всегда приятно когда его работу ценят. Возможно вас также заинтересует мой канал А зачем это нужно? где я рассказываю о том, что математику и алгоритмы придумали не только для собеседований в бигтехи.

Хабы: