TadaoYamaokaの日記

山岡忠夫Homeで公開しているプログラムの開発ネタを中心に書いていきます。

将棋AIの実験ノート:AVX対応

コンピュータチェスのCeresでは、PUCTによるノード選択の処理をAVXを使って高速化している。
これは、Ceres独自の「parallelized descent algorithm」(並列降下アルゴリズム)と合わせて使用することで、効果を発揮するもののようだ。

Ceresで実際にどれくらいNPSがでるか確認したところ、15ブロック、192フィルタのモデルを使用して、RTX30901枚で、初期局面のNPSは、142,685であった。

go nodes 1000000
Loaded network weights: 750181: 15x192 WDL MLH  from H:\src\Ceres\lc0\weights_run3_750181.pb.gz
info depth 1 seldepth 2 time 1016 nodes 1 score cp -0 tbhits 0 nps 1 pv g2g3  string M= 157
(略)
info depth 10 seldepth 22 time 7005 nodes 999471 score cp 8 tbhits 0 nps 142685 pv c2c4 e7e6 g2g3 d7d5 f1g2 d5d4 g1f3 c7c5 e2e3 b8c6 e3d4 c5d4 e1g1 g8f6 f1e1 f8d6 a2a3 a7a5 d2d3 e8g8 b1d2  string M= 157
bestmove c2c4

一方、dlshogiは、10ブロック、192フィルタのモデルで、初期局面のNPSは48,245である。
チェスと将棋の盤面サイズ、終端ノードでの詰み探索のありなどの条件の違いはあるが、Ceresはdlshogiより大きいモデルで約2.96倍のNPSが出ている。

dlshogiの探索部もまだまだ高速化の余地があることが分かる。

Ceresの並列アルゴリズムをdlshogiにも導入したいと考えているが、とりあえず、並列アルゴリズムは現在のままで、ノード選択処理のAVX化を行ってみた。

Ceresは並列アルゴリズムと組み合わせてAVX化しているが、単にノードを選択する際に、複数の子ノードのUCBを計算して、UCBが最大のノードを選ぶ処理をAVX化した。

AVX化前の処理

	for (int i = 0; i < child_num; i++) {
		const WinType win = uct_child[i].win;
		const int move_count = uct_child[i].move_count;

		if (move_count == 0) {
			// 未探索のノードの価値に、親ノードの価値を使用する
			q = parent_q;
			u = init_u;
		}
		else {
			q = (float)(win / move_count);
			u = sqrt_sum / (1 + move_count);
		}

		const float rate = uct_child[i].nnrate;
		const float ucb_value = q + c * u * rate;
		if (ucb_value > max_value) {
			max_value = ucb_value;
			max_child = i;
		}
	}

AVX化後の処理

const __m256i m256i_zero{};
const __m256i m256i_one = _mm256_set1_epi32(1);
const __m256 m256_one = _mm256_set1_ps(1);
const __m256i m256i_eight = _mm256_set1_epi32(8);

	__m256 m256_c = _mm256_broadcast_ss(&c);
	__m256 m256_parent_q = _mm256_broadcast_ss(&parent_q);
	__m256 m256_init_u = _mm256_broadcast_ss(&init_u);
	__m256 m256_sqrt_sum = _mm256_broadcast_ss(&sqrt_sum);

	__m256 vmaxvalue = _mm256_set1_ps(-FLT_MAX);
	__m256i vnowposition = _mm256_set_epi32(7, 6, 5, 4, 3, 2, 1, 0);
	__m256i vmaxposition = vnowposition;
	for (size_t i = 0; i < child_num; i += 8) {
		if (i + 8 > child_num) {
			// 残り8未満
			__m256i mask_rest;
			switch (child_num - i) {
			case 1:
				mask_rest = _mm256_set_epi32(0, 0, 0, 0, 0, 0, 0, -1);
				break;
			case 2:
				mask_rest = _mm256_set_epi32(0, 0, 0, 0, 0, 0, -1, -1);
				break;
			case 3:
				mask_rest = _mm256_set_epi32(0, 0, 0, 0, 0, -1, -1, -1);
				break;
			case 4:
				mask_rest = _mm256_set_epi32(0, 0, 0, 0, -1, -1, -1, -1);
				break;
			case 5:
				mask_rest = _mm256_set_epi32(0, 0, 0, -1, -1, -1, -1, -1);
				break;
			case 6:
				mask_rest = _mm256_set_epi32(0, 0, -1, -1, -1, -1, -1, -1);
				break;
			case 7:
				mask_rest = _mm256_set_epi32(0, -1, -1, -1, -1, -1, -1, -1);
				break;
			default:
				// unreachable
				mask_rest = _mm256_set1_epi32(0);
				break;
			}
			__m256i m256i_move_count = _mm256_maskload_epi32(move_count + i, mask_rest);
			__m256i mask = _mm256_cmpgt_epi32(m256i_move_count, m256i_zero);

			//	q = (float)(win / move_count);
			__m256 m256_win = _mm256_maskload_ps(win + i, mask_rest);
			__m256 m256_move_count = _mm256_cvtepi32_ps(m256i_move_count);
			__m256 m256_q_tmp = _mm256_div_ps(m256_win, m256_move_count);
			__m256 m256_q = _mm256_blendv_ps(m256_parent_q, m256_q_tmp, _mm256_castsi256_ps(mask));

			//	u = sqrt_sum / (1 + move_count);
			__m256 m256_move_count_plus1 = _mm256_add_ps(m256_move_count, m256_one);
			__m256 m256_u_tmp = _mm256_div_ps(m256_sqrt_sum, m256_move_count_plus1);
			__m256 m256_u = _mm256_blendv_ps(m256_init_u, m256_u_tmp, _mm256_castsi256_ps(mask));
			__m256 m256_rate = _mm256_maskload_ps(nnrate + i, mask_rest);

			//const float ucb_value = q + c * u * rate;
			__m256 m256_ucb_value = _mm256_mul_ps(m256_c, m256_u);
			m256_ucb_value = _mm256_mul_ps(m256_ucb_value, m256_rate);
			m256_ucb_value = _mm256_add_ps(m256_q, m256_ucb_value);

			// mask
			m256_ucb_value = _mm256_and_ps(m256_ucb_value, _mm256_castsi256_ps(mask_rest));

			// find max
			__m256 vcmp = _mm256_cmp_ps(m256_ucb_value, vmaxvalue, _CMP_GT_OS);
			vmaxvalue = _mm256_max_ps(m256_ucb_value, vmaxvalue);
			vmaxposition = _mm256_blendv_epi8(vmaxposition, vnowposition, _mm256_castps_si256(vcmp));
			vnowposition = _mm256_add_epi32(vnowposition, m256i_eight);

			break;
		}

		//if (move_count == 0) {
		__m256i m256i_move_count = _mm256_load_si256((__m256i*)(move_count + i));
		__m256i mask = _mm256_cmpgt_epi32(m256i_move_count, m256i_zero);

		//	// 未探索のノードの価値に、親ノードの価値を使用する
		//	q = parent_q;
		//	u = init_u;
		//  --> 下記のelseの計算結果と合わせて_mm256_blendv_psで設定する
		//}
		//else {
		//	q = (float)(win / move_count);
		__m256 m256_win = _mm256_load_ps(win + i);
		__m256 m256_move_count = _mm256_cvtepi32_ps(m256i_move_count);
		__m256 m256_q_tmp = _mm256_div_ps(m256_win, m256_move_count);
		__m256 m256_q = _mm256_blendv_ps(m256_parent_q, m256_q_tmp, _mm256_castsi256_ps(mask));

		//	u = sqrt_sum / (1 + move_count);
		__m256 m256_move_count_plus1 = _mm256_add_ps(m256_move_count, m256_one);
		__m256 m256_u_tmp = _mm256_div_ps(m256_sqrt_sum, m256_move_count_plus1);
		__m256 m256_u = _mm256_blendv_ps(m256_init_u, m256_u_tmp, _mm256_castsi256_ps(mask));
		//}

		//const float rate = uct_child[i].nnrate;
		__m256 m256_rate = _mm256_load_ps(nnrate + i);

		//const float ucb_value = q + c * u * rate;
		__m256 m256_ucb_value = _mm256_mul_ps(m256_c, m256_u);
		m256_ucb_value = _mm256_mul_ps(m256_ucb_value, m256_rate);
		m256_ucb_value = _mm256_add_ps(m256_q, m256_ucb_value);

		// find max
		__m256 vcmp = _mm256_cmp_ps(m256_ucb_value, vmaxvalue, _CMP_GT_OS);
		vmaxvalue = _mm256_max_ps(m256_ucb_value, vmaxvalue);
		vmaxposition = _mm256_blendv_epi8(vmaxposition, vnowposition, _mm256_castps_si256(vcmp));
		vnowposition = _mm256_add_epi32(vnowposition, m256i_eight);
	}
	// find max
	const int* maxposition = (int*)&vmaxposition;
	__m256 vallmax = _mm256_max_ps(vmaxvalue, _mm256_shuffle_ps(vmaxvalue, vmaxvalue, 0xb1));
	vallmax = _mm256_max_ps(vallmax, _mm256_shuffle_ps(vallmax, vallmax, 0x4e));
	vallmax = _mm256_max_ps(vallmax, _mm256_permute2f128_ps(vallmax, vallmax, 0x01));
	__m256 vcmp = _mm256_cmp_ps(vallmax, vmaxvalue, _CMP_EQ_US);
	int mask = _mm256_movemask_ps(vcmp);
	max_child = maxposition[__builtin_ctz(mask)];

パッと見には何をやっているか分からない(;'∀')
8未満を処理するためにswichで分けている部分が美しくない。

最大値のインデックスを探す処理は、はじめ思いつかなかったが、Discordで教えを乞うたところ、やねうらお氏から参考情報を教えていただき、それを元に実装したコードを、@wain_CGP氏に添削していただいた(最終的にはもらったコードほぼそのまま)。

測定

上記のAVX化した処理を探索部に組み込んで、どれくら速くなるか測定を行った。
V100 8枚で、floodgateからサンプリングした100局面(1局面は詰みなので除外)で測定した。
DeepLearningShogi/benchmark.py at master · TadaoYamaoka/DeepLearningShogi · GitHub

測定結果
AVXなし AVXあり
平均 229621 239431 1.04
中央値 235244 242866 1.03
最大 264526 392199 1.48
最小 148579 152466 1.03

平均で、4%程NPSが上昇した。
最大では、1.48倍の局面がある。
探索する局面の子ノードの数によって効果が変わってくる。

コード

探索部に組み込んだ処理はこちら。
DeepLearningShogi/UctSearch.cpp at 63d17b043c4ca0a37d15ea4e417a63cdffb29cc6 · TadaoYamaoka/DeepLearningShogi · GitHub
詰みのAND-OR木探索も含んでいるので、さらに美しくなくなっている。

まとめ

Ceresにインスパイアされてノード選択処理のAVX化を行った。
結果、平均で4%程高速化することができた。

AVX化を行ってもCeresのNPSには遠く及んでいない。
Ceresは、AVX化よりも衝突を回避した並列アルゴリズムの方がNPSへの寄与が大きいと思われる。
そのうち、Ceresの並列アルゴリズムをdlshogiにも導入したいと考えている。

2021/4/4 追記

AVX化の前後で強さを確認を測定した。
持ち時間3分、1手1秒加算

   # PLAYER          :  RATING  ERROR  POINTS  PLAYED   (%)  CFS(%)    W    D    L  D(%)
   1 dlshogi_avx2    :     0.2    9.0   500.5    1000    50      52  370  261  369    26
   2 dlshogi         :    -0.2    9.0   499.5    1000    50     ---  369  261  370    26

AVX化しても強くなっていない。