CUB - CUDA用のお手軽・超高速な並列計算プリミティブライブラリ

CUBとは

GPUを活用したアルゴリズムでは、往々にしてその基本的な構成要素としてReduceやスキャン(Prefix Sum)やソートが使われています。CPU上のアルゴリズムであればとりあえず適当な標準ライブラリの関数(C++ならSTL)でも使ってしまえば、そういった部分に時間を取られることなくアルゴリズムの本質の実装に集中できます。自分で実装するにしても並列化などを考えなければそれほど難しくありません(実装例も探せばいくらでも出てくるでしょう)。一方でGPUの場合、それら構成要素の実装はなかなかに面倒で本質の実装にとりかかるまでに時間がかかってしまいますし、優れた実装ができなければ本来実装したいアルゴリズム全体の処理時間に影響を与えてしまい評価も難しくなります。

CUBはNVIDIAによって作られたCUDA用の並列計算プリミティブを提供するライブラリです。
https://nvlabs.github.io/cub/
ワープ単位(Warp-wide)やブロック単位(Block-wide)、そしてデバイス単位(Device-wide)のReduceやScan、Radix Sortなどの高速な実装が提供されており、これらを使用すればGPU上のアルゴリズムに関する論文などを読んで実装に取り掛かる際に、本質にのみ集中することができます。
上のページではGitHubへのリンクが貼られていますが、CUDA 11.0からCUBはCUDA SDK本体に取り込まれるようになったので、使い始めるにあたって追加でなにかダウンロードする必要はありません。#include <cub/cub.cuh>で使い始められます。(Intellisenseと相性悪かったりするみたいなのでVisual Studioを使っている場合は#if !defined(__INTELLISENSE__) ... #endifで囲ったほうが安全かも。)

Radix SortのGPU実装などは最近でも新たなアルゴリズム、例えば以下に示す通称 "Onesweep Radix Sort"

• A FASTER RADIX SORT IMPLEMENTATION
  https://developer.download.nvidia.com/video/gputechconf/gtc/2020/presentations/s21572-a-faster-radix-sort-implementation.pdf
• Onesweep: A Faster Least Significant Digit Radix Sort for GPUs
  https://arxiv.org/ftp/arxiv/papers/2206/2206.01784.pdf

が提案されていたりしますが、CUBを使っておけば(恐らく)インターフェースは変わることなく高速な実装に乗っかっておくことができるでしょう。またNVIDIA自身が提供しているライブラリなので、GPUのハードウェア・アーキテクチャー観点でも最適化が施されていることが期待できます。

Thrustとの違い

デバイス単位のソートなどの処理を提供するNVIDIAのライブラリとしてはCUBよりもThrustのほうが有名かもしれません。ThrustはSTLに近い使用感で非常に簡単ではあるのですが、例えばソートに必要な作業用GPUメモリの確保などが裏で勝手に走ったり、streamの指定など細かい制御に難があるので、CPU/GPUの同期のオーバーヘッドなどが問題にならない超巨大な問題でもない限り個人的にはCUBのほうがおすすめです。CUBも十分簡単に使用することができます。

使用例

以下に各階層における処理の例を示します。CUDAにおけるGPUメモリの確保などについては示しません。CUBの使い方も文章を読むよりもコードを見ればすぐにわかるでしょう。

デバイス単位のプリミティブ例

C++コード

// Sum
// values, resultの型はint32_t*, uint32_t*, float*などのGPUアドレス。
// resultに総和が格納される。
cub::DeviceReduce::Sum(tempStorage, tempStorageSize,
                       values, result, numElements);

// ArgMin
// values, resultの型はint32_t*, uint32_t*, float*などのGPUアドレス。
// resultには最小値とそのインデックスが格納される。
cub::DeviceReduce::ArgMin(tempStorage, tempStorageSize,
                          values, result, numElements);

// Exclusive Prefix Sum
// values, prefixSumsの型はint32_t*, uint32_t*, float*などのGPUアドレス。
// prefixSumsに結果が格納される。
cub::DeviceScan::ExclusiveSum(tempStorage, tempStorageSize,
                              values, prefixSums, numElements);

// Radix Sort (Key/Value Pairs)
// keys/valuesに元の列を入れておく。
// valuesの型は何でも良い。
cub::DoubleBuffer<uint64_t> keys(keysA, keysB);
cub::DoubleBuffer<uint32_t> values(valuesA, valuesB);
cub::DeviceRadixSort::SortPairs(tempStorage, tempStorageSize,
                                keys, values, numElements);
// この時点でkeys.Current(), values.Current()で取得できるGPUアドレスに結果が格納される。

作業用GPUメモリのサイズ tempStorageSize は以下のように、作業用メモリにnullptrを渡すことで取得できるようになっています。

// tempStorageSizeは参照が渡っており、結果が返ってくる。
cub::DeviceReduce::Sum(nullptr, tempStorageSize,
                       values, result, numElements);

各処理はデフォルトではCUDAのデフォルトストリームを使用するようになっていますが、以下のように明示的にストリームを指定することができます。

// ストリームとしてcuStreamを指定。
cub::DeviceReduce::Sum(tempStorage, tempStorageSize,
                       values, result, numElements,
                       cuStream);

ブロック単位やワープ単位のプリミティブ例

CUDAカーネル

// ブロック単位のReduce (Min)
// valueはブロック内のスレッドごとの値。
// ブロック内スレッド0のreducedValueにブロック単位の最小値が格納される。
using BlockReduce = cub::BlockReduce<float, BlockWidth, cub::BLOCK_REDUCE_RAKING, BlockWidth>;
__shared__ typename BlockReduce::TempStorage s_reduceStorage;
float reducedValue = BlockReduce(s_reduceStorage).Reduce(value, cub::Min());

// ワープ単位のScan (Exclusive Prefix Sum)
// valueはワープ内のスレッドごとの値。
// resultにスレッドに対応する位置のExclusive Prefix Sumの結果が格納される。
using WarpScan = cub::WarpScan<uint32_t>;
__shared__ typename WarpScan::TempStorage s_scanStorage[NumWarpsPerBlock];
int warpIdx = threadIdx.x / 32;
WarpScan(s_scanStorage[warpIdx]).ExclusiveSum(value, result);

CUBd

CUB自体は非常に便利なものの、CUBのヘッダーファイルにはCUDAの予約語が含まれています。これは次のような欠点を意味します。

• CUBのヘッダーをincludeしたファイルはnvcc経由でコンパイル(遅い)する必要がある。
• CUDA Runtime APIを使わなければならない。(Driver APIを使いたい人にとっては欠点)

GPUカーネルはCUBの有無に関わらずそもそもnvccでコンパイルしないといけないため、ワープ単位やブロック単位のプリミティブを使う場合は特に問題にならないのですが、デバイス単位のプリミティブはホストコードから呼び出すため、それらを呼び出したコードもnvcc経由でコンパイルする必要があります。

上記欠点を回避するため、CUBのinclude、デバイス単位プリミティブ呼び出しを隔離したライブラリ(の例)を作りました。
https://github.com/shocker-0x15/CUBd
総和やスキャン、Radixソートなど一部のデバイス単位のプリミティブを典型的な型に対して定義しています。CUB自体はテンプレートライブラリなので、ユーザー定義の型や処理を織り交ぜることが可能でそれらについては当然定義できていませんが、定義済みの型に倣えば比較的簡単に追加できると思います。
cubd.hにはCUBのincludeが含まれていないため、このヘッダーを使用するファイルは通常のC++ソースとしてコンパイルすることが可能です。.h/.cppの直接使用・静的ライブラリ・動的ライブラリの3種類の方法で使えるようにしてあります。

使用例

例えばCUBdにおけるRadixソートの呼び出し方は次のように、名前空間がcubdに変わっているだけでCUBとほぼ同じになっています。

// Radix Sort (Key/Value Pairs)
cubd::DoubleBuffer<uint64_t> keys(keysA, keysB);
cubd::DoubleBuffer<uint32_t> values(valuesA, valuesB);

cubd::DeviceRadixSort::SortPairs(tempStorage, tempStorageSize,
                                 keys, values, numElements);