pdep、pdep抜きで - ubiquitinブログ

pdepという命令があります。ビット演算のための命令セットBMI2に属していて、タイミング的にはSSE4.2とAVXの間くらいに登場したものです。2つのuint64_tを受け取ってuint64_tを返す命令で、計算の中身は以下のように書き下せます。

inline uint64_t pdep_naive(uint64_t a, uint64_t mask) {
	uint64_t dst = 0, k = 0;

	for (uint64_t m = 0; m < 64; ++m) {
		if (mask & (1ULL << m)) {
			if (a & (1ULL << k++)) {
				dst += (1ULL << m);
			}
		}
	}

	return dst;
}

引数maskのビットを下から順になめるとします。立っているビットをk回目に見つけたとき、その場所はmaskの下からm番目だったとします。このとき、引数aの下からk番目のビットを、(それが0か1かは問わずに)dstのm番目のビット位置に代入します。

そんなpdep命令ですが、Intel Intrinsics Guideを見るとLatency 3 Throughput 1なのですが、AMDのCPUではメチャクチャ遅いと言われています。

Achilles heel of #AMD Zens: data dependency of PDEP/PEXT instructions. For every BMI2-capable #Intel all values are 1. pic.twitter.com/mtdIXByoj2
— InstLatX64 (@InstLatX64) 2019年12月23日

pdepなどといういかがわしい命令が来てもクラッシュしない寛大で慈悲深い世界最高のAMD製CPUに無制限の絶対忠誠を誓いましょう。こういうのはK8時代のbsfとか色々あって別に珍しくはないのですが、実際じゃあpdepを使わずにpdepと等価な処理を書くにはどうするのがいいのか？　と思いいくつか試してみました。

ちょっと改善したnaive

最初のコードではfor文の中に2個のif文が入っていましたが、頑張ればif文は取り除けます。すると以下のようになります。

inline uint64_t pdep_naive2(uint64_t a, uint64_t mask) {
	uint64_t dst = 0;

	for (uint64_t m = 0; mask; mask /= 2, ++m) {
		const uint64_t flag = mask & 1ULL;
		dst += (flag & a) << m;
		a >>= flag;
	}

	return dst;
}

2次元の表引きで頑張る

これから注目すべきmaskの4bitと、これから代入していくべきaの4bitとが決まれば、dstの4bitに入れるべき値が確定します。なのでmaskとaの4bitを添字として、2次元配列 uint8_t[16][16] にアクセスして、引いた値をdstに突っ込むのを16回繰り返せばいいわけです。具体的には以下のようになります。以下のコードでinit_tables()は最初に1度だけ呼び出すものです。

alignas(64) uint8_t table_16_16[16][16];
alignas(64) uint8_t table_16_16_popcount[16];

inline uint64_t popcount64_naive(uint64_t x) {
	uint64_t answer = 0;
	for (uint64_t i = 0; i < 64; ++i) {
		if (x & (1ULL << i))++answer;
	}
	return answer;
};

void init_tables() {
	for (uint64_t x = 0; x < 16; ++x)for (uint64_t y = 0; y < 16; ++y) {
		const uint64_t p = pdep_naive(x, y);
		assert(p < 16ULL);
		table_16_16[x][y] = (uint8_t)p;
	}
	for (uint64_t x = 0; x < 16; ++x) {
		const uint64_t p = popcount64_naive(x);
		assert(p <= 4ULL);
		table_16_16_popcount[x] = (uint8_t)p;
	}
}

inline uint64_t pdep_table_16_16(uint64_t a, uint64_t mask) {
	uint64_t dst = 0;

	for (uint64_t m = 0; mask; mask /= 16, m += 4) {
		const uint64_t b = mask % 16;
		dst += ((uint64_t)table_16_16[a % 16][b]) << m;
		a >>= table_16_16_popcount[b];
	}

	return dst;
}

いくつかバリエーションが考えられます。

8bit値2つを添字として uint8_t[256][256] を引くと倍速になる説（表がデカくなると厳しい説もある）
添字2つを逆にするほうがいい説（ただ、maskがゼロビットだらけのときaは変化しないのだから、a由来の添字を上位にするほうが良い気はする）
popcountを表引きではなく組み込み命令_mm_popcnt_u64でやるべき説

4bitの表引きをpshufbでやる

おなじみのやつ（？）です。uint8_t[16][16]は__m128i[16]と同一視できるので、pshufbを使って以下のように書けます。

alignas(64) __m128i table_16_pshufb[16];

void init_tables() {
	for (uint64_t x = 0; x < 16; ++x) {
		uint8_t tmp[16];
		for (uint64_t y = 0; y < 16; ++y) {
			const uint64_t p = pdep_naive(x, y);
			assert(p < 16ULL);
			tmp[y] = (uint8_t)p;
		}
		table_16_pshufb[x] = _mm_loadu_si128((__m128i*)tmp);
	}
}

inline uint64_t pdep_pshufb(uint64_t a, uint64_t mask) {

	const __m128i mask_lo = _mm_set_epi64x(0xFFFF'FFFF'FFFF'FFFFULL, (mask & 0x0F0F'0F0F'0F0F'0F0FULL));
	const __m128i mask_hi = _mm_set_epi64x(0xFFFF'FFFF'FFFF'FFFFULL, (mask & 0xF0F0'F0F0'F0F0'F0F0ULL) >> 4);

	__m128i bytemask = _mm_set_epi64x(0, 0xFF);
	__m128i answer = _mm_setzero_si128();

	for (int i = 0; i < 8; ++i) {

		const __m128i x_lo = _mm_shuffle_epi8(table_16_pshufb[a % 16], mask_lo);
		a >>= table_16_16_popcount[mask % 16];
		mask /= 16;
		const __m128i x_hi = _mm_shuffle_epi8(table_16_pshufb[a % 16], mask_hi);
		a >>= table_16_16_popcount[mask % 16];
		mask /= 16;

		answer = _mm_or_si128(answer, _mm_and_si128(bytemask, _mm_or_si128(x_lo, _mm_slli_epi64(x_hi, 4))));
		bytemask = _mm_slli_epi64(bytemask, 8);
	}

	return (uint64_t)_mm_cvtsi128_si64(answer);
}

結果

遺伝研スパコンのEPYC7501（Zen1世代、2.0GHz）とXeonGold6136（SkyLake世代、3.0GHz）で実験しました。
コンパイラはGCC8.2.0で、オプションは g++ -std=c++1y -O3 -flto -march=native で、各CPUの上でコンパイルして走らせました。
aとmaskをxorshift64で生成してpdepするのを $2^{30}$ 回繰り返すのにかかる時間を測定しました。
実験に使ったコードは全てGitHubにあげてあります。
https://github.com/eukaryo/pdep-senpai

種類	EPYC7501での時間(s)	XeonGold6136での時間(s)
pshufb(pop表)	17.856	15.652
intrinsics	62.133	1.452
naive	517.005	323.403
naive2	74.033	52.795
table(16,normal,pop表)	21.773	17.398
table(256,normal,pop表)	12.244	8.389
table(16,inverse,pop表)	22.075	18.575
table(256,inverse,pop表)	12.240	8.710
table(16,normal,pop命令)	22.502	17.343
table(256,normal,pop命令)	12.463	8.387
table(16,inverse,pop命令)	23.288	18.588
table(256,inverse,pop命令)	12.753	8.823

normalとinverseは前述した添字の前後関係のやつで、前述のコード通りのほうがnormalです。

結論

繰り返してるだけなんだからキャッシュにギリギリ載るクソでかい表を引くのが最速、それはそう

追記: pextについてもやってみた

コードはGitHubにあります。詳細な説明は省きますが、同じような感じでやっていき、同じような結果が出ました。クロック周波数が1.5倍違うことを考えるとEPYCのほうが実質速いのではないかとか錯覚しそうになりましたが、intrinsicsのパワーの前には無力です。

種類	EPYC7501での時間(s)	XeonGold6136での時間(s)
pshufb(pop表)	21.169	19.920
intrinsics	59.705	1.452
naive	531.238	325.163
naive2	71.229	52.557
table(16,normal,pop表)	20.731	17.376
table(256,normal,pop表)	11.635	8.284
table(16,inverse,pop表)	21.085	18.566
table(256,inverse,pop表)	11.839	8.708
table(16,normal,pop命令)	21.464	17.429
table(256,normal,pop命令)	12.176	8.405
table(16,inverse,pop命令)	22.165	18.475
table(256,inverse,pop命令)	12.554	14.092