Hatena::Grouptopcoder

nodchipのTopCoder日記 このページをアンテナに追加 RSSフィード

 | 

2014-12-02

Competitive Programming Advent Calendar 2014 3日目 自動ベクトル化で O(N^2) N=60,000 を通す 00:00 Competitive Programming Advent Calendar 2014 3日目 自動ベクトル化で O(N^2) N=60,000 を通す - nodchipのTopCoder日記 を含むブックマーク はてなブックマーク - Competitive Programming Advent Calendar 2014 3日目 自動ベクトル化で O(N^2) N=60,000 を通す - nodchipのTopCoder日記 Competitive Programming Advent Calendar 2014 3日目 自動ベクトル化で O(N^2) N=60,000 を通す - nodchipのTopCoder日記 のブックマークコメント

この記事は Competitive Programming Advent Calendar 2014 3日目の記事です。

http://partake.in/events/9b53847a-3a97-4aac-b754-5e681c3c7197

はじめに

最近のコンパイラには自動ベクトル化という最適化機能が実装されています。これはループ処理をSIMD命令を使って高速化するというものです。詳しくはWikipedia等を参照してください。

この最適化機能によってどれくらい処理速度が向上するのか、競技プログラミングの問題を使って測ってみました。使用したコンパイラのバージョンと最適化オプションは以下の通りです。

g++ (GCC) 4.8.2
Copyright (C) 2013 Free Software Foundation, Inc.
This is free software; see the source for copying conditions.  There is NO
warranty; not even for MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.

-O3 -ftree-vectorizer-verbose=2 -mavx

また、実行時間の計測に使用した環境は以下のとおりです。

Inspiron 17R Special Edition
CPU: Intel(R) Core(TM) i7-3610QM CPU @ 2.30GHz
メモリ: 16GB

現時点においてはこれらの環境は競技プログラミングコンテスト本番で使用することはできません。思考実験の一種と考えていただければ幸いです。

TopCoder Single Round Match 436 Div1 Hard 1000 CircularShifts

問題

長さNの配列XとYが与えられる。一方を循環シフトし、他方との内積を取る。内積の値の最大値を求めよ。

  •  1 \le N \le 60,000
  •  0 \le X\[i\], Y\[i\] < 100 for all 0 \le i < N

解法

まず初めにO(N^2)のナイーブなコードを書きます。N \le 60,000ですので、このままでは通りません。

#include <cassert>
#define REP(i,n) for(int i=0;i<(int)n;++i)
template<class T> void MAX_UPDATE(T& a, const T& b) { if (a < b) a = b; }

int X[64 * 1024];
int Y[64 * 1024];
long long Z[128 * 1024];
class CircularShifts {
public:
	int maxScore(int N, int Z0, int A, int B, int M) {
		Z[0] = Z0 % M;
		for (int i = 1; i < N * 2; ++i) {
			Z[i] = (Z[i - 1] * A + B) % M;
		}
		REP(i, N) {
			X[i] = Z[i] % 100;
			Y[i] = Z[i + N] % 100;
		}

		int result = 0;
		for (int offset = 0; offset < N; ++offset) {
			int temp = 0;
			for (int i = 0; i < N; ++i) {
				temp += X[i] * Y[(i + offset) % N];
			}
			MAX_UPDATE(result, temp);
		}

		return result;
	}

	void run_test(int Case) {
		int Arg0 = 60000;
		int Arg1 = 96478;
		int Arg2 = 24834;
		int Arg3 = 74860;
		int Arg4 = 92112;
		int Arg5 = 188580624;

		assert(Arg5 == maxScore(Arg0, Arg1, Arg2, Arg3, Arg4));
	}
};

int main() {
	CircularShifts ___test;
	___test.run_test(-1);
}

ローカルでの実行時間は以下のようになりました。

./a.exe  6.55s user 0.01s system 99% cpu 6.589 total

続いてコンパイラが自動ベクトル化をしやすいように書き換えてみます。コンパイラのログに

CircularShifts.old.cpp:24: note: not consecutive access _32 = Y[_31];

という行があります。これはおそらく配列へのアクセスが連続していないため、ベクトル化ができなかったということなのだと思います。これを踏まえ、配列へのアクセスが連続になるよう、配列Yの長さを2倍化し、"% N"を取り除きます。

#include <cassert>
#define REP(i,n) for(int i=0;i<(int)n;++i)
template<class T> void MAX_UPDATE(T& a, const T& b) { if (a < b) a = b; }

int X[64 * 1024];
int Y[128 * 1024];
long long Z[128 * 1024];
class CircularShifts {
public:
	int maxScore(int N, int Z0, int A, int B, int M) {
		Z[0] = Z0 % M;
		for (int i = 1; i < N * 2; ++i) {
			Z[i] = (Z[i - 1] * A + B) % M;
		}
		REP(i, N) {
			X[i] = Z[i] % 100;
			Y[i] = Y[i + N] = Z[i + N] % 100;
		}

		int result = 0;
		for (int offset = 0; offset < N; ++offset) {
			int temp = 0;
			for (int i = 0; i < N; ++i) {
				temp += X[i] * Y[i + offset];
			}
			MAX_UPDATE(result, temp);
		}

		return result;
	}

	void run_test(int Case) {
		int Arg0 = 60000;
		int Arg1 = 96478;
		int Arg2 = 24834;
		int Arg3 = 74860;
		int Arg4 = 92112;
		int Arg5 = 188580624;

		assert(Arg5 == maxScore(Arg0, Arg1, Arg2, Arg3, Arg4));
	}
};

int main() {
	CircularShifts ___test;
	___test.run_test(-1);
}

実行結果は以下のようになりました。

./a.exe  0.75s user 0.01s system 97% cpu 0.786 total

TopCoderのサーバーはローカルより若干遅いようですが、これくらいの実行時間なら通ると思います。

出力されたアセンブラから重要な部分を抜粋します。

	movl	$_Y, %edx			// Y + offset
	vpxor	%xmm3, %xmm3, %xmm3		// result = 0; 256bitレジスタに32bit変数を8つ格納し、外側のループ8個分のresultを同時に計算する
L5:
	xorl	%eax, %eax			// i = 0;
	vpxor	%xmm0, %xmm0, %xmm0		// temp = 0; 256bitレジスタに32bit変数を8つ格納し、外側のループ8個分のtempを同時に計算する
	.p2align 4,,7
L9:
	vbroadcastss	_X(,%eax,4), %xmm2	// X[i] の値を256bitレジスタの各32bit領域にロードする
	vmovdqu	(%edx,%eax,4), %xmm1		// Yのi + offset番目から8個分の値をレジスタにロードする
	addl	$1, %eax			// ++i
	cmpl	$60000, %eax			// i < 60000
	vpmulld	%xmm1, %xmm2, %xmm1		// X[i] * Y[i + offset]
	vpaddd	%xmm1, %xmm0, %xmm0		// temp += X[i] * Y[i + offset]
	jne	L9				// 内側のループの終わり
	addl	$16, %edx			// ++offset
	vpmaxsd	%xmm0, %xmm3, %xmm3		// result = max(result, temp) を要素ごとに行う
	cmpl	$_Y+240000, %edx		// offset < N
	jne	L5				// 外側のループの終わり
	vpsrldq	$8, %xmm3, %xmm0		// ここからreductionと呼ばれる処理 resultを8バイトシフト
	vpmaxsd	%xmm3, %xmm0, %xmm3		// 上記とresultのmaxを取る
	vpsrldq	$4, %xmm3, %xmm0		// resultを4バイトシフト
	vpmaxsd	%xmm3, %xmm0, %xmm3		// 上記とresultのmaxを取る
	vpextrd	$0, %xmm3, %eax
	cmpl	$188580624, %eax		// result == 188580624
	jne	L14

temp 8個分を同時に計算しているのが驚きでした。最後のreductionも美しいですね。

UTPC 2009 Problem F: 天使の階段

問題

N段の階段のそれぞれのステップに12種類の音階T_{i}が設定されている。またM個の音階からなる譜面S_{i}が与えられる。階段の各ステップに乗ると以下の様な音がなる。

  • k+1 段目に降りた -> T_{k+1} の音が鳴る。
  • k+2 段目に降りた -> T_{k+2} を半音上げた音が鳴る。
  • k-1 段目に戻った -> T_{k-1} を半音下げた音が鳴る。

譜面を正しく演奏しつつ階段を降り終えられるかどうか判定せよ。

解法

まずはじめにO(N^2)のナイーブなコードを書きます。N \le 50,000ですのでこのままでは通りません。

#include <iostream>
#include <map>
using namespace std;

int T[64 * 1024];
int S[64 * 1024];
int dp[2][64 * 1024];
int main() {
	std::ios::sync_with_stdio(false);

	map<string, int> noteToIndex = {
		{ "C", 0 },
		{ "C#", 1 },
		{ "D", 2 },
		{ "D#", 3 },
		{ "E", 4 },
		{ "F", 5 },
		{ "F#", 6 },
		{ "G", 7 },
		{ "G#", 8 },
		{ "A", 9 },
		{ "A#", 10 },
		{ "B", 11 },
	};

	int N, M;
	cin >> N >> M;
	for (int n = 0; n < N; ++n) {
		string t;
		cin >> t;
		T[n + 1] = noteToIndex[t];
	}

	for (int m = 0; m < M; ++m) {
		string s;
		cin >> s;
		S[m] = noteToIndex[s];
	}

	int front = 0;
	int back = 1;
	dp[back][0] = true;
	for (int m = 0; m < M; ++m) {
		for (int n = 0; n <= N; ++n) {
			dp[front][n] = 0;
		}

		// k + 1 段目に降りた
		for (int n = 0; n + 1 <= N; ++n) {
			dp[front][n + 1] |= dp[back][n] && T[n + 1] == S[m];
		}

		// k + 2 段目に降りた
		for (int n = 0; n + 2 <= N; ++n) {
			dp[front][n + 2] |= dp[back][n] && (T[n + 2] + 1) % 12 == S[m];
		}

		// k - 1 段目に戻った
		for (int n = 1; n <= N; ++n) {
			dp[front][n - 1] |= dp[back][n] && (T[n - 1] + 11) % 12 == S[m];
		}

		swap(front, back);
	}

	cout << ((dp[back][N - 1] || dp[back][N]) ? "Yes" : "No") << endl;
}

実行時間は以下の通りになりました。

./a < in.txt  9.34s user 0.00s system 99% cpu 9.345 total

入力データの生成器は以下の通りです。

#include <iostream>
#include <cstdlib>
using namespace std;

int main() {
	std::ios::sync_with_stdio(false);

	const char* notes[] = { "C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B" };
	int n = 50000;
	int m = 50000;
	cout << n << " " << m << endl;

	for (int i = 0; i < n; ++i) {
		cout << notes[rand() % 12] << " ";
	}
	cout << endl;

	for (int i = 0; i < m; ++i) {
		cout << notes[rand() % 12] << " ";
	}
	cout << endl;
}

元のソースコードをコンパイラが自動ベクトル化しやすいように書き換えてみます。書き換えた箇所は以下の通りです。

  • 1本のレジスタに多くの要素を入れられるようにするため、dp配列の型をboolに変更する。
  • 階段の音階と曲の音階が一致しているかどうかの判定で使う==を消すため、階段の情報の持ち方を変更する。

ポイントはループの中を"a |= b & c"と、コンパイラに理解しやすい形にした点です。

#include <iostream>
#include <map>
#include <string>
using namespace std;

int S[64 * 1024];
bool dp[2][64 * 1024];
bool T[12][64 * 1024];
int main() {
	std::ios::sync_with_stdio(false);

	map<string, int> noteToIndex = {
		{ "C", 0 },
		{ "C#", 1 },
		{ "D", 2 },
		{ "D#", 3 },
		{ "E", 4 },
		{ "F", 5 },
		{ "F#", 6 },
		{ "G", 7 },
		{ "G#", 8 },
		{ "A", 9 },
		{ "A#", 10 },
		{ "B", 11 },
	};

	int N, M;
	cin >> N >> M;
	for (int n = 0; n < N; ++n) {
		string t;
		cin >> t;
		T[noteToIndex[t]][n + 1] = 1;
	}

	for (int m = 0; m < M; ++m) {
		string s;
		cin >> s;
		S[m] = noteToIndex[s];
	}

	int front = 0;
	int back = 1;
	dp[back][0] = true;
	for (int m = 0; m < M; ++m) {
		for (int n = 0; n <= N; ++n) {
			dp[front][n] = 0;
		}

		// k + 1 段目に降りた
		for (int n = 0; n < N; ++n) {
			dp[front][n + 1] |= dp[back][n] & T[S[m]][n + 1];
		}

		// k + 2 段目に降りた
		for (int n = 0; n < N - 1; ++n) {
			dp[front][n + 2] |= dp[back][n] & T[(S[m] + 11) % 12][n + 2];
		}

		// k - 1 段目に戻った
		for (int n = 1; n <= N; ++n) {
			dp[front][n - 1] |= dp[back][n] & T[(S[m] + 1) % 12][n - 1];
		}

		swap(front, back);
	}

	cout << ((dp[back][N - 1] || dp[back][N]) ? "Yes" : "No") << endl;
}

実行時間は以下の通りになりました。

./a < angel.4.in.txt  0.89s user 0.03s system 97% cpu 0.942 total

これなら通りそうです。

出力されたアセンブラから、3つ並んだループの1つ目の部分を抜粋します。

	vmovdqu	(%ebx,%edx), %xmm1
	vmovdqu	(%ecx,%edx), %xmm0
	vinsertf128	$0x1, 16(%ebx,%edx), %ymm1, %ymm1
	vinsertf128	$0x1, 16(%ecx,%edx), %ymm0, %ymm0
	vandps	(%edi,%edx), %ymm1, %ymm1
	addl	$1, %esi
	vorps	%ymm0, %ymm1, %ymm0
	vmovdqu	%xmm0, (%ecx,%edx)
	vextractf128	$0x1, %ymm0, 16(%ecx,%edx)

レジスタにロード、and、or、メモリにストアというストレートフォワードなアセンブラが出力されました。ymmレジスタへのロードを2回に分けて行っているところが気になりますが、これはコンパイラの内部でレジスタ幅が128bitに設定されているためだと思います。

まとめ

コンパイラの自動ベクトル化を使ってO(N^2) (N=60,000 or 50,000)の問題を定数倍高速化し、ローカルでの実行時間を時間制限以内に収めることが出来ました。将来、問題作成者が自動ベクトル化のことを考慮したデータサイズの設定をしなければならない時代が来るのかもしれませんね。

明日12/4の発表者はsh19910711さんです。

トラックバック - http://topcoder.g.hatena.ne.jp/nodchip/20141202
 |