Segmented Prefix Scan

GPGPU Advent Calendar22日目です！（ほんとは15日に書く予定だったのを急遽代わってもらいました。

とは言うもののだいぶネタが尽きてきているので、Segmented Prefix Scanの紹介でもします。
Segmented Prefix Scanとは、例えば

input.cpp

int groupIndeces = {0, 0, 0, 1, 1, 2, 2, 2, 2, 3};
int values       = {0, 1, 2, 3, 4, 5, 6, 7, 8, 9};

というデータがあった時に、グループごとのPrefix Scanすなわち部分和を計算して

answer.cpp

int results = {3, 7, 26, 9};

という結果を計算する手法です（※和じゃなくても加法群であればいいようです？）。
グループ番号を行番号に見立てるなどして、行列計算する時に使ったりします。

通常のPrefix Scan

まず復習で、普通にPrefix Scanをしましょう（以降、exclusive scanとします）。
色々最適化はありますが、素朴にやるとこんなかんじになるはずです。

ループ回数	#0	#1	#2	#3	#4	#5	#6	#7	#8	#9	説明
0	0	1	2	3	4	5	6	7	8	9	初期状態
1	0	1	3	5	7	9	11	13	15	17	1つ左を足す
2	0	1	3	6	10	14	18	22	26	30	2つ左を足す
3	0	1	3	6	10	15	21	28	36	44	4つ左を足す
4	0	1	3	6	10	15	21	28	36	45	8つ左を足す

繰り返し回数はlog2(全要素数)です。

Segmented Prefix Scan

これにグループごとにする場合は、なんと基本的なアルゴリズムは普通のPrefix Scanそのままでちょっと加えればできちゃいます。方法は、MITの資料か、NVIDIAの資料に書いてあります。

つまり、最初に、グループの先頭をtrueあとをfalseとするフラグを生成します

flags.cpp

int groupIndeces = {0, 0, 0, 1, 1, 2, 2, 2, 2, 3};
int flags        = {T, F, F, T, F, T, F, F, F, T};

次に、こんな演算Add2を定義します。

Add2.cpp

void Add2(int i, int j)
{
	values[j] = flags[j] ? values[j] : (values[i] + values[j]);
	flags[j] |= flags[i];
}

つまり、フラグは常に論理和を取り、自分のフラグがfalseだった場合のみ値を足し合わせるという演算です。こんな演算を使って普通のPrefix Scanをすれば

ループ回数	#0	#1	#2	#3	#4	#5	#6	#7	#8	#9	説明
	0	1	2	3	4	5	6	7	8	9	初期状態
0	T	F	F	T	F	T	F	F	F	T	フラグ
	0	1	3	3	7	5	11	13	15	9
1	T	T	F	T	T	T	T	F	F	T	1つ左を足す
	0	1	3	3	7	5	11	18	26	9
2	T	T	T	T	T	T	T	T	T	T	2つ左を足す

となって結果が求められます。
ね、簡単でしょう（？）

注意点

ここで重要なのは（当たり前ですが）Segmented Prefix Scanにすると、条件分岐が入ってくることと、一方で繰り返し回数は最大でlog2(グループ内の最大要素数)に減ることです。

まとめ

というわけで、疎行列とベクトルの積とかをこれで簡単に計算できるようになります。が、行列・ベクトル演算なら、素直にライブラリ（cuSPARSEとかMAGMAとか）使ったほうがよいと思います。

というわけで、次回（明後日）はcuSPARSEの話をします。

その前に、明日は@tanakmuraさんです！