More than 1 year has passed since last update.

概要

Fermi以降のGPUで使えるwarp vote関数について説明する。

動機

Keplerで新しく導入されたwarp shuffle関数の記事は多くあるのだが、warp vote関数の記事のまとめがあまり見当たらなかったのでここにまとめることにした。あまり使われない機能なのかもしれない。
詳細はCUDA C Programming guideのwarp vote functionsにあるが、サンプルプログラムを書きながらどういった挙動を示すのかまとめることにした。
サンプルコードはgithubにて公開している。

warp vote関数一覧

__any

int __any(int predicate);

関数の説明:predicate$\neq$0となるスレッドがwarp内に一つでも存在する場合には戻り値が1になり、そうでない場合0になる。

__global__ void test_any(int* a,
                         int* b) {
  const auto tid = threadIdx.x + blockIdx.x * blockDim.x;
  a[tid] = __any(threadIdx.x == 16);
  b[tid] = __any(threadIdx.x == 128);
}

...

test_any<<<2, 64>>>(a, b);

この場合、threadIdx.x == 16となるスレッドがwarp内に一つでも存在すれば、そのwarp内すべてのスレッドについて戻り値は1になる。それ以外は0を返す。
つまり、aの配列の0から31、64から95番目には1が入りそれ以外には0がはいる。(以下のコードと同じ)

for (int i = 0; i < 128; i++) {
  if (i >= 0 && i < 32)
    a[i] = 1;
  else if (i >= 64 && i < 96)
    a[i] = 1;
  else
    a[i] = 0;
}

一方でbの配列はすべて0で埋められる。thread blockのサイズを64にしてkernelを起動しているのでthreadIdx.x == 128を満たすwarpが存在しないためである。

__all

int __all(int predicate);

関数の説明:warp内すべてのスレッドについてpredicate$\neq$0となる場合に戻り値が1になり、そうでない場合0になる。

__global__ void test_all(int* a,
                         int* b) {
  const auto tid = threadIdx.x + blockIdx.x * blockDim.x;
  a[tid] = __all(threadIdx.x == 16);
  b[tid] = __all(b[tid] == -1);
}

...

test_all<<<2, 64>>>(a, b);

(ここでbの配列にはkernelを起動する前に-1が埋められているとする。)
aの配列はすべて0で埋められる。warp内でthreadIdx.x == 16になっていないスレッドが31個存在するためである。
一方でbの配列は-1で埋められているのでb[tid] == -1はwarp内のすべてのスレッドについて満たされる。allの戻り値は1になる。

__ballot

unsigned int __ballot(int predicate);

関数の説明:predicate$\neq$0のスレッドの位置に対応するbitが1に、それ以外のbitが0になったunsigned int型の戻り値を返す。

__global__ void test_ballot(int* a,
                            int* b) {
  const auto tid = threadIdx.x + blockIdx.x * blockDim.x;
  a[tid] = __ballot(threadIdx.x == 16);
  b[tid] = __ballot(threadIdx.x % 2 == 0);
}

...

test_ballot<<<2, 64>>>(a, b);

となっているとaとbの配列のデータはそれぞれ

tid a                                b
...
30  00000000000000010000000000000000 01010101010101010101010101010101
31  00000000000000010000000000000000 01010101010101010101010101010101
32  00000000000000000000000000000000 01010101010101010101010101010101
33  00000000000000000000000000000000 01010101010101010101010101010101
...

となる。戻り値の各bitがwarp内のスレッドのpredicateの値と1対1対応になっており、predicate$\neq$0となったスレッドに対応するbitが1になっている。

条件分岐があり、warp divergenceが起きる場合

上記であげた例では条件分岐はなく、すべてのスレッドがany/all/ballotの関数を呼び出すケースのみを考えた。ifがあり、warp divergenceが起きる場合には注意が必要である。
例えば __anyで

__global__ void test_any(int* c) {
  const auto tid = threadIdx.x + blockIdx.x * blockDim.x;
  if (threadIdx.x == 10) c[tid] = __any(threadIdx.x == 16);
}

...

test_any<<<2, 64>>>(c);

としたとする。
この時threadIdx.x == 10の属するwarpにはthreadIdx.x == 16となるスレッドが存在するが、anyの戻り値は0になる。

上記のコードを

__global__ void test_any(int* c) {
  const auto tid = threadIdx.x + blockIdx.x * blockDim.x;
  if (threadIdx.x == 10 || threadIdx.x == 16) c[tid] = __any(threadIdx.x == 16);
}

...

test_any<<<2, 64>>>(c);

のように書き換えてthreadIdx.x == 16を満たすスレッドがanyを呼ぶように変更したらどうなるか。
この場合にはthreadIdx.x == 10のanyの戻り値は0になり、threadIdx.x == 16のanyの戻り値は1になる。
期待通りの振る舞いにはならない。

ballotの場合にはinactiveなスレッドに対応したbitが0になった戻り値が得られる。

余談

ちなみに

__global__ void test_any(int* c) {
  const auto tid = threadIdx.x + blockIdx.x * blockDim.x;
  if (threadIdx.x == 15 || threadIdx.x == 16) c[tid] = __any(threadIdx.x == 16);
}

のようにするとthreadIdx.x == 15とthreadIdx.x == 16に対応するanyの戻り値が両方とも1になる。
隣接したスレッドに対してanyを呼ぶようにすればwarp divergence起きている場合でも正しい振る舞いになる模様。
ただ、ドキュメントにはこんなこと書いていないのであまり使うべきではないと思う。

まとめ

warp vote関数の機能と注意点についてまとめた。any all ballotの三つが使用可能であり、それぞれwarp内スレッドの状態をbroadcastするのに役に立つ。