8x8 bitboardへの対称な変換を全部やるのにAVX2を使う

符号なし64bit整数を8x8 bitboardとみなして、(1)横方向に鏡映 (2)縦方向に鏡映 (3)行列転置 という3つの処理をこの順番でやってもやらなくてもよいとしたとき、結果得られる最大8通りのbitboardのうち、整数としての値が最小のものを得たいという問題を考えます。

C++で基本命令だけで書くと以下のようになります。get_unique_naive関数が上記の問題に答える関数です。

uint64_t vertical_mirror(uint64_t b) {
	b = ((b >> 8) & 0x00FF00FF00FF00FFULL) | ((b << 8) & 0xFF00FF00FF00FF00ULL);
	b = ((b >> 16) & 0x0000FFFF0000FFFFULL) | ((b << 16) & 0xFFFF0000FFFF0000ULL);
	b = ((b >> 32) & 0x00000000FFFFFFFFULL) | ((b << 32) & 0xFFFFFFFF00000000ULL);
	return b;
}
uint64_t horizontal_mirror(uint64_t b) {
	b = ((b >> 1) & 0x5555555555555555ULL) | ((b << 1) & 0xAAAAAAAAAAAAAAAAULL);
	b = ((b >> 2) & 0x3333333333333333ULL) | ((b << 2) & 0xCCCCCCCCCCCCCCCCULL);
	b = ((b >> 4) & 0x0F0F0F0F0F0F0F0FULL) | ((b << 4) & 0xF0F0F0F0F0F0F0F0ULL);
	return b;
}
uint64_t transpose(uint64_t b) {
	uint64_t t;
	t = (b ^ (b >> 7)) & 0x00AA00AA00AA00AAULL;
	b = b ^ t ^ (t << 7);
	t = (b ^ (b >> 14)) & 0x0000CCCC0000CCCCULL;
	b = b ^ t ^ (t << 14);
	t = (b ^ (b >> 28)) & 0x00000000F0F0F0F0ULL;
	b = b ^ t ^ (t << 28);
	return b;
}
uint64_t symmetry_naive(const uint32_t s, uint64_t b) {
	uint64_t answer = b;
	if (s & 1)answer = horizontal_mirror(answer);
	if (s & 2)answer = vertical_mirror(answer);
	if (s & 4)answer = transpose(answer);
	return answer;
}
uint64_t get_unique_naive(const uint64_t b) {
	uint64_t answer = b;
	for (uint32_t i = 1; i <= 7; ++i) {
		const uint64_t new_code = symmetry_naive(i, b);
		if (answer > new_code) {
			answer = new_code;
		}
	}
	return answer;
}

これをAVX2で書くことを考えます。まず、vertical_mirror関数とhorizontal_mirror関数は、命令列が完全に同じでオペランドの定数が異なるだけなので、SIMDならば同時に計算できます。例えば以下の関数のように書けます。これの返り値は、(64bit整数4要素の配列とみなすと)[1]が縦方向への鏡映、[2]が横方向への鏡映、[0]と[3]が引数そのままとなります。

__m256i func1(const uint64_t b) {
	constexpr uint64_t FULL = 0xFFFF'FFFF'FFFF'FFFFULL;

	const __m256i bb0 = _mm256_set1_epi64x(int64_t(b));

	const __m256i tt1lo = _mm256_and_si256(_mm256_srlv_epi64(bb0, _mm256_set_epi64x(0, 8, 1, 0)), _mm256_set_epi64x(FULL, 0x00FF00FF00FF00FFLL, 0x5555555555555555LL, FULL));
	const __m256i tt1hi = _mm256_and_si256(_mm256_sllv_epi64(bb0, _mm256_set_epi64x(0, 8, 1, 0)), _mm256_set_epi64x(FULL, 0xFF00FF00FF00FF00LL, 0xAAAAAAAAAAAAAAAALL, FULL));
	const __m256i tt1 = _mm256_or_si256(tt1lo, tt1hi);

	const __m256i tt2lo = _mm256_and_si256(_mm256_srlv_epi64(tt1, _mm256_set_epi64x(0, 16, 2, 0)), _mm256_set_epi64x(FULL, 0x0000FFFF0000FFFFLL, 0x3333333333333333LL, FULL));
	const __m256i tt2hi = _mm256_and_si256(_mm256_sllv_epi64(tt1, _mm256_set_epi64x(0, 16, 2, 0)), _mm256_set_epi64x(FULL, 0xFFFF0000FFFF0000LL, 0xCCCCCCCCCCCCCCCCLL, FULL));
	const __m256i tt2 = _mm256_or_si256(tt2lo, tt2hi);

	const __m256i tt3lo = _mm256_and_si256(_mm256_srlv_epi64(tt2, _mm256_set_epi64x(0, 32, 4, 0)), _mm256_set_epi64x(FULL, 0x00000000FFFFFFFFLL, 0x0F0F0F0F0F0F0F0FLL, FULL));
	const __m256i tt3hi = _mm256_and_si256(_mm256_sllv_epi64(tt2, _mm256_set_epi64x(0, 32, 4, 0)), _mm256_set_epi64x(FULL, 0xFFFFFFFF00000000LL, 0xF0F0F0F0F0F0F0F0LL, FULL));
	const __m256i tt3 = _mm256_or_si256(tt3lo, tt3hi);
	return tt3;
}

ところで、縦方向と横方向の両方への鏡映変換を行うのは、ビットの配置を逆転させることと等価です。64bit整数のビットの配置を逆転させる処理をAVX2で以下のように書けます。

__m256i func2(const uint64_t b) {
	constexpr auto f = [](const uint8_t i) {
		return uint8_t(((i & 1) << 3) + ((i & 2) << 1) + ((i & 4) >> 1) + ((i & 8) >> 3));
	};

	const __m128i r1 = _mm_set_epi8(f(15), f(14), f(13), f(12), f(11), f(10), f(9), f(8), f(7), f(6), f(5), f(4), f(3), f(2), f(1), f(0));

	const __m128i a1 = _mm_set_epi64x((b >> 4) & 0x0F0F'0F0F'0F0F'0F0FULL, b & 0x0F0F'0F0F'0F0F'0F0FULL);
	const __m128i a2 = _mm_shuffle_epi8(r1, a1);
	const __m128i a3 = _mm_shuffle_epi8(a2, _mm_set_epi8(8, 9, 10, 11, 12, 13, 14, 15, 0, 1, 2, 3, 4, 5, 6, 7));
	const __m128i a4 = _mm_shuffle_epi32(a3, 0b00001110);
	const __m128i a5 = _mm_add_epi32(a4, _mm_slli_epi64(a3, 4));
	return _mm256_zextsi128_si256(a5);
}

__m128i型を8bit整数16個とみなして、そこに引数を4bitずつ刻んで入れます。すると、pshufb命令で「4bitを添字とする8bitの配列への表引き」を全部一気に行えますから、各4bitについてビットの配置を逆転させられます。あとは、8bit単位で配置を逆転させて、刻んだのをもとに戻せばOKです。返り値の最下位64bitに欲しい値が入っていることになります。

このfunc1とfunc2は依存関係なく実行できるのでレイテンシを隠蔽できます。返り値2つをblendすることで、「(縦・横)方向への鏡映を(やる・やらない)」の4通りの結果を__m256i型変数1個に格納できます。

あとは、最初に書いた行列転置の関数を単純にSIMDで書き下して入力すれば、所望の8通りの結果を__m256型変数2個に格納できます。最小値を得る処理を含めると、最終的な関数は以下のようになります。

__m256i get_unique_avx2(const uint64_t b) {
	constexpr uint64_t FULL = 0xFFFF'FFFF'FFFF'FFFFULL;

	const __m256i b1 = func1(b);
	const __m256i b2 = func2(b);

	const __m256i tt = _mm256_blend_epi32(b1, b2, 0b00000011);

	const __m256i x1 = _mm256_and_si256(_mm256_xor_si256(tt, _mm256_srli_epi64(tt, 7)), _mm256_set1_epi64x(0x00AA00AA00AA00AALL));
	const __m256i y1 = _mm256_xor_si256(tt, _mm256_xor_si256(x1, _mm256_slli_epi64(x1, 7)));
	const __m256i x2 = _mm256_and_si256(_mm256_xor_si256(y1, _mm256_srli_epi64(y1, 14)), _mm256_set1_epi64x(0x0000CCCC0000CCCCLL));
	const __m256i y2 = _mm256_xor_si256(y1, _mm256_xor_si256(x2, _mm256_slli_epi64(x2, 14)));
	const __m256i x3 = _mm256_and_si256(_mm256_xor_si256(y2, _mm256_srli_epi64(y2, 28)), _mm256_set1_epi64x(0x00000000F0F0F0F0LL));
	const __m256i zz = _mm256_xor_si256(y2, _mm256_xor_si256(x3, _mm256_slli_epi64(x3, 28)));

	const __m256i a1 = _mm256_sub_epi64(zz, tt);
	const __m256i a2 = _mm256_srai_epi32(a1, 32);
	const __m256i a3 = _mm256_shuffle_epi32(a2, 0b11110101);
	const __m256i a4 = _mm256_blendv_epi8(tt, zz, a3);

	alignas(32) uint64_t result[4] = {};
	_mm256_storeu_si256((__m256i*)result, a4);

	result[0] = std::min(result[0], result[1]);
	result[0] = std::min(result[0], result[2]);
	result[0] = std::min(result[0], result[3]);
	return result[0];
}

単体テストベンチマークテストのコードは以下のページにあげてあります。ベンチマーク結果は省略しますが多少速くなりました。

github.com