スターリンソートをGPUで爆速にする。

# 1. 役に立たないアルゴリズム

YouTubeを物色していると、こんな動画を見つけました。

<iframe width="950" height="534" src="https://www.youtube.com/embed/KGZ_X076xdM" title="YouTube video player" frameborder="0" allow="accelerometer; autoplay; clipboard-write; encrypted-media; gyroscope; picture-in-picture" allowfullscreen></iframe>

これを見て思ったのです。

**役に立たないアルゴリズムも作れないものか**

自分で考案するのが理想なのですが、役に立たないアルゴリズムを作る楽しさを感じたことのない私は、まず既存のものを作ってみようと思ったわけです。

そしていい題材として[スターリンソート](https://qiita.com/Sirloin/items/c9e5c74be3366df65bce)を見つけました。

参考：[話題の粛清ソートアルゴリズム「スターリンソート」をPythonで実装した](https://qiita.com/Sirloin/items/c9e5c74be3366df65bce)

このまま作っても面白くないので、お得意のGPUで高速化して$O(\log n)$の計算量に爆速化してみます。
意外に学びのある結果になったのではと思います。

:::note info

もっと高速化した記事を公開しました。
[スターリンソートをGPUで爆速にする。その2](https://qiita.com/tommyecguitar/items/4cdfa20e87619d7f8668)

:::

# 2. スターリンソートとは

上記の参考記事を見ていただきたいのですが、要はいらないやつを粛清するソートです。
元の配列が以下のようなものだったとき、

| index | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 |
|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|
|element| 4 | 8 | 6 | 8 | 12 | 2 | 12 | 19 |

index = 0から順番にelementを見て行って、それまでの最大値より小さい要素が来たら粛清するというアルゴリズムです。今の例だと6, 2が粛清されて

| index | 0 | 1 | 2 | 3 | 4 | 5 |
|:-:|:-:|:-:|:-:|:-:|:-:|:-:|
|element| 4 | 8 | 8 | 12 | 12 | 19 |

がソート結果となります。

これをGPUで爆速化して、役に立たない度を上げていこうと思います。

# 3. GPUでのアルゴリズム

CPU(正確にはシングルスレッド)なら簡単なことも、GPUではひと工夫必要です。順番に処理を考えてみました。

## 3.1. ある要素がそれ以前の要素より大きいか判定する。

GPU化にあたって一番難しいのがこれです。先ほどの例のindex = 5, element = 2を見てみましょう。

| index | 0 | 1 | 2 | 3 | 4 | *5* | 6 | 7 |
|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|
|element| 4 | 8 | 6 | 8 | 12 | ***2*** | 12 | 19 |

GPUは並列演算を得意とするユニットで、配列を扱う場合1スレッドあたりの担当範囲を決めなければなりません。
今回1スレッドあたり1要素を割り振るとしましょう。スレッドインデックスを`idx`と呼ぶことにすると、そのスレッドは`element[idx]`を主な担当として読み込みます、そしてそれ以前のインデックスの要素のうち、自分よりも大きい数字が存在したら**粛清対象**とします。ざっくりこんなコードになるでしょう。

```cuda
// これはスレッドインデックスを求めるための式
const int idx = blockIdx.x * blockDim.x + threadIdx.x;

for(int i = 0; i < idx; i++){
    if(element[idx] < element[i]){
        // この要素を粛清する。
    }
}
```

ただこれ、あまり効率よくないのです。このやり方だと要素数nの時に、最悪n-1回のループ計算が走ってしまいます。GPUのスレッド計算速度は遅いので、この方法だとCPUよりも遅くなってしまいます。

## 3.2. 累積和をとる処理をうまく使えないか

GPU界隈ではもう説明し尽くされていますが、GPUで配列の累積和を取るときは以下のようにします。

![](https://developer.nvidia.com/sites/all/modules/custom/gpugems/books/GPUGems3/elementLinks/39fig02.jpg)

この処理を利用して、あるインデックスまでの要素の最大値を記録できないか考えてみます。

| 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 |
|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|
| $a_0$ | $a_1$ | $a_2$ | $a_3$ | $a_4$ | $a_5$ | $a_6$ | $a_7$ |
| $a_0$ | max($a_0$,$a_1$) | max($a_1$,$a_2$)  | max($a_2$,$a_3$)  | max($a_3$,$a_4$)  | max($a_4$,$a_5$)  | max($a_5$,$a_6$)  | max($a_6$,$a_7$)  |
| $a_0$ | max($a_0$,$a_1$) | max($a_0\cdots a_2$)  | max($a_0\cdots a_3$)  | max($a_1\cdots a_4$)  | max($a_2\cdots a_5$)  | max($a_3\cdots a_6$)  | max($a_4\cdots a_7$)  |
| $a_0$ | max($a_0$,$a_1$) | max($a_0\cdots a_2$)  | max($a_0\cdots a_3$)  | max($a_0\cdots a_4$)  | max($a_0\cdots a_5$)  | max($a_0\cdots a_6$)  | max($a_0\cdots a_7$)  |

和を取る処理を全て`max`に変えれば、所望の動作ができます。ただこのフローはブロック間の同期が多くなるので、[GPU Gems 3](https://developer.nvidia.com/gpugems/gpugems3/part-vi-gpu-computing/chapter-39-parallel-prefix-sum-scan-cuda)や[Single-pass Parallel Prefix Scan with Decoupled Look-back](https://research.nvidia.com/publication/single-pass-parallel-prefix-scan-decoupled-look-back)を参考に勉強しましょう。今回はCUBを使います。

## 3.3. 粛清判定

これで粛清の判定ができます。

| index | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 |
|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|
|element| 4 | 8 | 6 | 8 | 12 | 2 | 12 | 19 |
|max value| 4 | 8 | 8 | 8 | 12 | 12 | 12 | 19 |
|survive flag| 1 | 1 | 0 | 1 | 1 | 0 | 1 | 1 |

`max_value[index] > element[index]`のものを粛清すれば良いことになります。
後段の都合上、粛清フラグではなくて存命フラグにしておきます。

いきなり粛清されるのではなく、前もって誰が対象かを決めて通告する形なので、なかなか残酷度が高くなったと思います。並列で同時に何人も粛清するにはこうするしかないんです。

# 3.4. 出力配列への書き込み

GPUコンピューティングは並列処理なので、**配列の末尾に順番に追加せよ**みたいなことはできません。あらかじめ出力用の配列を用意しなければいけません。**これにも累積和が役に立ちます。**

`survive flag`をintと解釈して累積を取ります。

**入力系**
| index | 0 | 1 | 2 | 3 | 4 | 5 | 6 | 7 |
|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|:-:|
|element| 4 | 8 | 6 | 8 | 12 | 2 | 12 | 19 |
|max value| 4 | 8 | 8 | 8 | 12 | 12 | 12 | 19 |
|survive flag| 1 | 1 | 0 | 1 | 1 | 0 | 1 | 1 |
|scan survive flag| 1 | 2 | 2 | 3 | 4 | 4 | 5 | 6 |

**出力系**
| index | 0 | 1 | 2 | 3 | 4 | 5 |
|:-:|:-:|:-:|:-:|:-:|:-:|:-:|
|output| 4 | 8 | 8 | 12 | 12 | 19 |

こうすることで`scan_survive_flag`の末尾が出力用素数となり、`survive_flag == 1`のインデックスの`scan_survive_flag`の場所に粛清結果を書き込めばいいことになります。

# 4. GPUコーディングして爆速で粛清してみる。

いよいよコーディングしてみます。

GPU計算の高速化で地味に聞いてくるのが**メモリ確保・メモリコピーの時間**です。GPUプログラミングでは、CPUから元の配列をGPUメモリに送信したり、幾らかの領域を確保するという命令を送ります。配列が大きいほどこの時間が馬鹿にならないほど遅く、ちゃんと作らないとメモリ確保だけでCPU計算に負けます。

できることなら最初にどかっとメモリを確保して、それを使い回すということをしたいのです。今回必要なメモリは以下のものです。

| 配列名 | サイズ |
|:-:|:-:|
| element | 元の配列サイズ |
| max value | 元の配列サイズ |
| survive flag | 元の配列サイズ |
| scan survive flag | 元の配列サイズ |
| output | 粛清後のサイズ |

なるべくならメモリサイズは小さくしたいので、配列への上書きができないか考えてみます。ただし、elementとoutputのメモリは分け、elementはいじらずに保持することとします。(用途によってはここも上書きしてしまえばいい)

## 4.1. メモリは上書きして使いまわそう

この条件下で考えると、elementとmax valueは別の領域にならざるを得ません。
max valueとsurvive flagとscan survive flagの関係が非常に難しいのですが、よく考えるとmax valueってsurvive flagを立てるためにいるんです。survive flagを作り終えたらmax valueはいらないので、max valueにsurvive flagを上書きしてしまえばいいのです。

:::note info
elementやmax valueがfloatとかlong xxだったりしたらめんどくさいので、intしか入ってこないものとします。仮にそういうのが来たとしてもreinterpret_castとかで無理矢理どうにかできるでしょう。
:::

そして最後にsurvive flag == 1の時のscan survive flagを使って並べ替えるわけですが、ここはわざわざsurvive flagを残す必要はなく、`scan_survive_flag[idx] != scan_survive_flag[idx-1]`だったら`survive_flag == 1`とわかるので、survive flagに上書きする形でscan survive flagを作ればいいです。

## 4.3. 静的にメモリ確保しよう

というわけで、元の配列サイズをNとして型をintとすると、確保するべきメモリは`3 * N * sizeof(int)`です。あれ？outputサイズってNじゃないんじゃない？って思った方、正解です。もちろんoutputサイズを粛清サイズにして確保してもいいのですが、元の配列がどうであるかによって動的にサイズが変わってしまいメモリ確保が2回に割れてしまうので、静的に確保するためこうしています。

## 4.4 コード

まずカーネルのコードから

```cuda
// 生存フラグを立てるカーネル
__global__ void SetSurviveFlag(int num, int* org, int* max_element){

    const int idx = blockIdx.x * blockDim.x + threadIdx.x;

    if(idx >= num){
        return;
    }

    max_element[idx] = (int)(org[idx] == max_element[idx]);
}

// 粛清して並べ替えるカーネル
__global__ void Purge(int num, int* scan_survive_flag, int* input, int* output){

    const int idx = blockIdx.x * blockDim.x + threadIdx.x;

    if(idx >= num){
        return;
    }

    const int val = scan_survive_flag[idx];
    const int pre_val = (idx > 0) ? scan_survive_flag[idx - 1] : 0;

    if(val != pre_val){
        const int move_idx = val - 1;
        output[move_idx] = input[idx];
    }
}
```

続いてメイン側のコードです

```cuda
        // scan max value
        cub::DeviceScan::InclusiveScan(scan_storage, scan_storage_byte_size, d_input, stalin_buffer, max_op, num);
        // survive flag
        SetSurviveFlag<<<grid, block>>>(num, d_input, stalin_buffer);
        // scan survive flag
        cub::DeviceScan::InclusiveSum(scan_storage, scan_storage_byte_size, stalin_buffer, stalin_buffer, num);
        // output
        Purge<<<grid, block>>>(num, stalin_buffer, d_input, d_output);
```

### 4.4.1 【追記】DevicePartitionがあるのを忘れてました。

CUBライブラリにはDevicePartitionがあるのをすっかり忘れてました。アルゴリズムとしては上の通りで正しいのですが、より楽で最適化されているであろうCUBのコードを使いまわしてみます。

DevicePartitionとは、上のコードで言う`survive_flag == 1`の要素を配列の前に詰める処理の部分、つまり`scan_survive_flag`を生成してoutputに書き込むところに該当します。

これを踏まえてコードを書き換えると、カーネルの`Purge()`が不要になり、ホストコードがこのようになります。

```cuda
        // scan max value
        cub::DeviceScan::InclusiveScan(scan_storage, scan_storage_byte_size, d_input, stalin_buffer, max_op, num);

        // survive flag
        SetSurviveFlag << <grid, block >> > (num, d_input, stalin_buffer);

        // partition
        cub::DevicePartition::Flagged(scan_storage, scan_storage_byte_size, d_input, stalin_buffer, d_output, d_purge_num, num);
```

以降の測定時間などには結果を並べて記載しておきます。

## 4.5 粛清時間

CPUでもスターリンソートを組んでみて、時間を計測してみました。計測条件は以下の通りです。

| 項目 | 内容 |
|:-:|:-:|
| CUDA ver | 11.6 |
| GPU | NVIDIA RTX 2070 |
| CPU | Intel Core i7-10750H @2.6GHz|
| サンプル数 | 1000回 |

**CPU計測条件**
- 初期値の生成は含まない。
- 初期配列の入力から出力配列を生成するところまでを計測する。
- chronoを使用する。

**GPU計測条件**
- メモリ確保・コピーの時間は含まない。(最初はデバイス初期化が入って余計な時間が入ることと、GPUプログラミングはGPUで完結するものをつくるのでメモリコピーは含まない。)
- 初期配列の入力から出力配列を生成するところまでを計測する。
- cudaEventを使用する。(chronoではGPU同期を待ってくれないため。)

**計測結果**
| 要素数 | CPU粛清時間[ms] | GPU粛清時間[ms] | GPU粛清時間(DevicePartition使用)[ms] |
|:-:|:-:|:-:|:-:|
| 10 | 0.005 | 31.5388 | 16.0113 |
| 100 | 0.077 | 37.4764 | 16.3853 |
| 1000 | 0.573 | 19.886 | 17.8869 |
| 10000 | 10.58 | 26.9168 | 25.0143 |
| 100000 | 87.223 | 27.4412 | 24.2178 |
| 1000000 | 893.693 | 119.263 | 121.326 |
| 10000000 | 6487.83 | 1072.87 | 1046.07 |

![image.png](https://qiita-image-store.s3.ap-northeast-1.amazonaws.com/0/530192/86e20e14-742c-7217-c8e6-52e92a1da3c2.png)

**CPU粛清のコード**

```cpp
h_output_cpu[0] = h_input[0];
output_idx = 1;

int max_val = h_input[0];
for(int i = 1; i < num; i++){
    if(h_input[i] >= max_val){
        h_output_cpu[output_idx] = h_input[i];
        max_val = h_input[i];
        output_idx++;
    }
}
```

# 5. まとめ

結果の考察をします。

**要素数が少ない時はCPUのほうが速い。**
10000要素くらいまではCPUのほうが速いという結果になりました。CPUとGPUのクロック周波数や、アルゴリズム(カーネル同期など)の差だと思われます。

**1000要素を超えるとGPUが突然速くなる**
なぜか1000要素の粛清のほうが10, 100要素のよりも速く殺せるという結果になりました。プロファイルをとればどこが速くなったかわかりますが、そこまでやってないので予想で言うと、CUBのスキャンが要素数によってアルゴリズムを変えることがあり、1block内に全要素収まるような場合は極端に遅いのだと思います。
要素が少ない時は[もっと効率的なスキャン](https://arxiv.org/abs/1811.09736)があるらしいのであながち間違っていないはずです。

**GPUでは1000000要素あたりから粛清時間が比例的?**
GPUの粛清時間は並列演算の効果もあり100000要素あたりまではほぼ変わりませんが、それ以降は比例増加っぽいです。GPUも無限に並列できるわけではないので、今回の条件ではこの辺が並列の限界なのでしょうね。

**cub優秀**
一部の処理をDevicePartitionに置き換えたら小さい要素数の時の速度が改善しました。内部的にブロック数などを調整しているんでしょうか。
巨大な配列の時は自力コードとあまり差がないので、扱うサイズが常に大きい場合はメンテ性も考えると自力で組んでおいてもいいかもしれないです。