rank/selectのselect1でpdepを使う

以前、pdep命令がRyzenでめっちゃ遅いことについて調べました。
pdep、pdep抜きで - ubiquitinブログ

その記事のベンチマークではpdepの使いみちについては特に書きませんでした。圧縮全文索引の話で出てくるrank/selectデータ構造のselect1関数をpdepで高速に計算できるらしいということを最近知ったので、確かめてみました。

FM-Indexのことが懐かしくなったので実装してみた - ubiquitinブログ

select1関数の実装

select1関数は、整数xのうち下位から数えてcount番目に立っているビットの位置を返す関数です。ナイーブには以下のように書けます。（引数countは0-originだとします）

int select1_naive(uint64_t x, int count) {

	assert(0 <= count && count < 64);
	assert(_mm_popcnt_u64(x) > count);

	for (int n = 0, i = 0; i < 64; ++i) {
		if (x & (1ULL << i)) {
			if (count == n++)return i;
		}
	}

	assert(0);
	return -1;
}

popcnt命令で二分探索することもできます。多分遅いと思いますが。（以降、最初と最後のassertとかは省略します）

int select1_popcnt_binarysearch(uint64_t x, int count) {

	int lb = 0, ub = 64;
	while (lb + 1 < ub) {
		const int mid = (lb + ub) / 2;
		const int pop = _mm_popcnt_u64(x & ((1ULL << mid) - 1));
		if (pop <= count) lb = mid;
		else ub = mid;
	}
	return lb;
}

整数xの「最下位のビットを取り除く」処理をcount回やってからbsf命令で答えを得る方法もあります。

int select1_bsf(uint64_t x, int count) {

	for (int i = 0; i < count; ++i)x &= x - 1;
	uint32_t index = 0;
	bitscan_forward64(x, &index);
	return int(index);
}

pdep命令の第一引数にを入れて、mask（第二引数）にxを入れることで、『「最下位のビットを取り除く」処理をcount回やる』がpdep命令一発で済みます。

int select1_pdep(uint64_t x, int count) {

	const uint64_t answer_bit = _pdep_u64(1ULL << count, x);
	uint32_t index = 0;
	bitscan_forward64(answer_bit, &index);
	return int(index);
}

ベンチマーク結果

遺伝研スパコンのEPYC7702（Zen2世代、2.0GHz）とXeonGold6136（SkyLake世代、3.0GHz）で実験しました。
xをxorshift64で生成してselect1するのを回繰り返すのにかかる時間を測定しました。
実験に使ったコードは全てGitHubにあげてあります。

pdep-senpai/select1 at master · eukaryo/pdep-senpai · GitHub

種類	EPYC7702での時間(s)	XeonGold6136での時間(s)
naive	71.241	75.074
popcnt	18.311	18.594
bsf	7.713	8.397
pdep	45.809	1.769

pdep/pext命令は一見意味不明ですが独特の中毒性があって、慣れてくると様々な場所で使いたくなってきてしまいます。Zen3世代で改善されているか気になります。