バイトニックソート(Bitonic Sort)の概要

バイトニックソート(Bitonic Sort)は主にGPU等の並列計算器でソートを実装しようとするときに使われるソートである。
計算量のオーダーはO(n log^2 n)であり、クイックソートのO(n log n)には負けるものの並列化による高速化が勝るという感じなのでいろんなところに使われている。

対象読者

キーワード「GPU」「バイトニックソート」で検索してこの記事にたどり着いただろう方が対象。
この記事では
・バイトニックソートでなぜソートできるか
・どうやったら高速化できるか
という点について重点的に書いている。

高速化については
OpenCLでバイトニックソートを実装している海外サイト
パクリ参考にした。
これのComputeShader移植+αかつ日本語資料として、一応価値はあるかと思うのでまとめてみた。

時間がない方は完成版コードの
BitonicSort_ComputeShader

をご参照を。

ソートはソートでも、ただ配列をソートしているのではなく、配列のもとのindexを保持したままソートしている実装なので可読性より実用性重視になっている。

アルゴリズム

wikinoyatuとりあえず英語のwikipediaの画像がわかりやすいので転載してみた。
16要素のソートの場合のアルゴリズムを図示したもののようだ。
左から見ていく。
最初はソートされてない数字が縦に並んでいて、縦の矢印のところで比較、swapが行なわれるという意味らしい。
矢印の方向によって大きい値が矢印の先に移動するといった具合でやると、最終的に一番画面右の時点で上から小さい順に数字が並んでいる状態になる。

なぜソートできるのか

4要素の時

あまり数学的な説明はできないが、自分の思考過程をまとめてみた。
sikou0
まず4要素のソートで考える。(2だと自明すぎるので)
最終的にこのような形にするためには、という考え方をすると後ろから見ていくとよさそうだ。

sikou1

この時点で、上2つは1と3が、下2つは4と6が入ってないと絶対いけない。

つまりこの図の操作(矢印の数でいうと4回)後には、上に小さい方が値2つ、下に大きい方の値2つが振り分けられてないといけないことがわかった。

sikou2

この操作だけでは、1,4,3,6と入力したときに「1,4」「3,6」となり振り分けが不十分だ。

sikou3
やっぱり最初の操作をしないと「1,3」「4,6」の振り分けがうまくいかないようだ。

要素数がもっと大きいとき

それではもっと要素数が大きくなった場合に、この「振り分け」が常にきちんとできているのか気になる。
例えば

sikou4
このように2n個の数字を入力すると、大きい方n個、小さい方n個に分けられて出力されないといけない。
と、ここで入力側は真ん中を境に2方向にソートされていたことを思い出す。

sikou5



Webp.net-gifmaker

このように入力がソートされていれば、16入力のうち小さいものと大きいものを8個ずつに分けることができる。

例えば小さい方の8個というのは、赤の数値を小さい順にみて集めていき、赤と青のお互いの上下がひっくりかえったところで今度は青の数値を集めていけば、それが必ず小さい8個の集合になる。

256入力になったときの図
gifmaker2

この例を見ても確かにうまくふるい分けられている。

Bitonic sequence

では今の現状を整理する。
sikou6

真ん中でソートされているn+nの入力ならば、出力側で上下分離できていることが分かった。
しかし先ほどのgifの出力の「へ」の字みたいなのは必ずしも真ん中でソートされているわけではない。それが次の入力になっているが、はたして大丈夫なのか。

実際にやってみた。

gifmaker3

このようにいびつな形でも上下分離はできるようだ。
今度はルート記号のような形がでてきた。こんなんでも実は上下分離できる。
gifmaker4

実はこの形のことをbitonic sequenceと呼ぶ。
https://en.wikipedia.org/wiki/Bitonic_sorter
wikiitibu

単調増加→単調減少となっている数列で、かつそれが循環シフトしていてもよい。
上の画像のルート記号を縦にしたような形は、よくみると赤の一番上の位置と青の一番下の位置が繋がっている。循環シフトすれば「ヘ」の字のように単調増加→単調減少の形になるはずだ。

結局のところ

バイトニックソートの内部構造を改めてみてみる。wikinoyatu


結局のところバイトニックソートは、bitonic sequenceを入力すると上下分離されたbitonic sequenceが出力され、それが次の入力になって・・・という形が続いているだけである。
なのである範囲内で集合の上下関係を確定させることが可能で、その結果最後はソートされた数列になっているというわけ。

改めて見るとなんだか美しい図に見えてくるなぁ


(超超余談ですが上のgifを作るのためにHSPでプログラムを書いて絵を出力しました。案外1-2時間程度で簡便に作れたので手短にかけるHSPは楽で良いです)

バイトニックソートのGPU実装

GPU実装。まずどこを並列化できるかを探す。
gpuimpliment0
このようにStep分けすると、各Step内はほぼ同じような処理をしているためここを並列化できそうだ。
Step7を抽出してみると

gpuimpliment1

このように1threadが2点の比較をするように組むとして、あとは自threadがどの2点にアクセスするかの計算コードを書けば良い。

Unityで実装

Unityを起動し
空のGameObjectを作成
下記のMain.csをアタッチしBitonicSort.computeを紐づけ
saisyo


C#(Main.cs)
using System;
using System.Collections;
using System.Collections.Generic;
using System.Runtime.InteropServices;
using UnityEngine;

//keyを配列index付きでソートする
public class Main : MonoBehaviour
{
    struct mystruct
    {
        public float key;//ソートしたい値
        public uint index;//一緒にソートされるindex
    }

    public ComputeShader shader;
    ComputeBuffer gpu_data;
    int kernel_ParallelBitonic;
    mystruct[] host_data;
    const int THREADNUM_X = 64;

    //配列の要素数
    int N = 1 << 19;//下限 1<<7,上限 1<<27

    void kernelfindStart()
    {
        kernel_ParallelBitonic = shader.FindKernel("ParallelBitonic");
    }

    void host_Init()
    {
        for (uint i = 0; i < N; i++)
        {
            host_data[i].key = UnityEngine.Random.Range(0.0f, 1.0f);
            host_data[i].index = i;
        }
    }

    void Start()
    {
        host_data = new mystruct[N];//ソートしたいデータ
        gpu_data = new ComputeBuffer(host_data.Length, Marshal.SizeOf(host_data[0]));

        //カーネル初期設定
        kernelfindStart();

        //CPU側で初期値代入
        host_Init();

        // host to device
        gpu_data.SetData(host_data);

        //ソート実装部分
        BitonicSort(gpu_data);
        gpu_data.GetData(host_data);

        //結果表示
        resultDebug();
    }

    void BitonicSort(ComputeBuffer gpu_data)
    {
        int n = gpu_data.count;
        //引数をセット
        shader.SetBuffer(kernel_ParallelBitonic, "data", gpu_data);

        int nlog = (int)(Mathf.Log(n, 2));
        int lninc, inc;

        for (int i = 0; i < nlog; i++)
        {
            inc = 1 << i;
            for (int j = 0; j < i + 1; j++)
            {
                shader.SetInt("inc", inc);
                shader.SetInt("dir", 2 << i);
                shader.Dispatch(kernel_ParallelBitonic, n / 2 / THREADNUM_X, 1, 1);
                inc /= 2;
            }
        }
    }

    void resultDebug()
    {
        // device to host
        gpu_data.GetData(host_data);

        Debug.Log("GPU上でソートした結果");
        for (int i = 0; i < Mathf.Min(1024, gpu_data.count); i++)
        {
            Debug.Log("index=" + host_data[i].index + " key=" + host_data[i].key);
        }
    }

    private void OnDestroy()
    {
        //解放
        gpu_data.Release();
    }
}


ComputeShader(BitonicSort.compute)
#pragma kernel ParallelBitonic

//構造体と比較を自前で定義する必要あり。
struct Mystruct
{
	float key;
	uint index;
};
#define data_t Mystruct
#define COMPARISON(a,b) ( a.key > b.key )
//ここまで

RWStructuredBuffer<data_t> data;
int inc;
int dir;

[numthreads(64, 1, 1)]
void ParallelBitonic(uint threadid : SV_DispatchThreadID)
{
	int t = threadid; // thread index
	int low = t & (inc - 1); // low order bits (below INC)
	int i = (t << 1) - low; // insert 0 at position INC
	bool reverse = ((dir & i) == 0); // asc/desc order

	// Load
	data_t x0 = data[i];
	data_t x1 = data[inc + i];

	// Sort
	bool swap = reverse ^ COMPARISON(x0, x1);
	data_t auxa = x0;
	data_t auxb = x1;
	if (swap) { x0 = auxb; x1 = auxa; }

	// Store
	data[i] = x0;
	data[inc + i] = x1;
}

結果

kekka

Key(float型)が大きい順にソートされている。

ソートする前のindexの情報も含んだ構造体でソートしているので、indexも一緒にソートできている。

この構造体はmystructで定義しているので一応自由に変えられる。
C#のmystructを変えたらComputeShader側のmystructも変える必要あり。

また、使用の注意点としてはソースコードにもあるよう、2のべき乗の要素数でしかソートできない。もし端数なら、-∞や+∞で数値を埋めておいてソート結果に影響しないようにしておく必要がある。
(追記、Apache License 2.0なのでお好きに使ってください。)

速度計測

単位はms
要素数/カーネル名B2
2560
5120
10240
20480
40960
81920
163840
327680
655361
1310721
2621442
5242884
104857610
209715219
419430439
838860883
16777216179
33554432382
67108864818
1342177281744
元の記事を参考にしているのでこのB2というのはカーネル名。
(1threadあたり2点読み込んでいるためと思われる。)

高速化Tips

これでも十分速いけど高速化はロマン!

コアレスアクセス

UnityでGPGPUその2
でさんざん書いてきたよう、コアレスアクセスは速い。上のコードではすでにコアレスアクセスはできているので特段変える必要はない。
corelessacs


改めて図示するとこんな感じに、同じwarp(wavefront)が実行するカーネルはなるべく連続したメモリアドレスにアクセスするようする。

Shared Memoryを使う

今までのコードは、1threadがグローバルメモリの2か所を読み取って比較し、適宜swapした後またメモリに書き戻す、という作業をしていた。このグローバルメモリはいくらコアレスアクセスができてもなお遅いので可能な限りグローバルメモリを介さず計算したいと考える。
そこでShared Memoryを使えないかと考える。1Gruopあたり16threadだとするとその16threadはshared memoryを共有できる。

バイトニックソートのアルゴリズム図をもっと遠くからみた図
shared0

各Stepで結果をグローバルメモリにいちいち書き戻さなくても、shared memoryに書き込んで使いまわしていくことで、グローバルメモリにアクセスするのはカーネルの最初と最後だけにすることができる。

shared1

この青のところが全部shared memoryで高速化できるところだ!
ここでは図示しやすいように1Group=16threadで書いているが、実際は1Group=64threadや1Group=128threadで実装している。

これをやることで
要素数/カーネル名B2B2C2
25600
51200
102400
204800
409600
819200
1638400
3276800
6553610
13107211
26214421
52428843
1048576108
20971521911
41943043924
83886088351
16777216179109
33554432382236
67108864818508
13421772817441095

1.5-6倍速くなった!


2Step同時計算

上図のピンク領域はShared Memoryにおさまりきらないので、従来のグローバルメモリに読み書きする方式でやっている。今度はそこを高速化しようという話。
従来はthread0~thread7がそれぞれ2点ずつ読み込んで、という処理だったが

gpuimpliment2

今度はthread0~thread3がそれぞれ4点のメモリを読み込んで2Step分まとめて計算しようという発想だ。

従来の方法だと
Step7を計算するのに必要なメモリアクセス回数が、Load16回+Store16回、そしてStep8を計算するのにLoad16回+Store16回で合計64回の読み書きが発生していた。
しかし2Step同時計算をすることでLoad16回+Store16回で合計32回の読み書きで済む。

こうしてグローバルメモリのアクセス回数を減らしていくことで高速化をしていく。
要素数/カーネル名B2B2C2B2B4C2
256000
512000
1024000
2048000
4096000
8192000
16384000
32768000
65536100
131072110
262144211
524288432
10485761086
209715219119
4194304392419
8388608835138
1677721617910977
33554432382236163
67108864818508345
13421772817441095735

最初から2倍以上速くなった!

n Step同時計算

ではStepを3つや4つまとめて計算すればもっと良いのではという発想が出るのは自然なことだ。
しかし4つまとめて、とやると1threadあたりの作業量が16点になり、5つで32点・・・と指数的に増えていってしまう。さすがにレジスタの容量も限られているので4Stepまとめくらいが限度だろう。
OpenCLのサイトの実装も4Stepまでであった。
要素数/カーネル名B2B2C2B2B4C2B2B4B8C2B2B4B8B16C2
25600000
51200000
102400000
204800000
409600000
819200000
1638400000
3276800000
6553610000
13107211000
26214421111
52428843222
1048576108665
209715219119118
41943043924191616
83886088351383333
16777216179109776863
33554432382236163139127
67108864818508345292271
13421772817441095735614556

効果はかなりあるようだ。


ここまでの高速化Tipsを全て踏まえ、さらにShared memory内の計算にも2Step同時計算を適応したコードを適応すると
要素数/カーネル名B2B2B4B8B16C2B2B4B8B16C2C4
256000
512000
1024000
2048000
4096000
8192000
16384000
32768000
65536100
131072100
262144211
524288422
10485761055
20971521987
4194304391614
8388608833329
167772161796360
33554432382127127
67108864818271264
1342177281744556552

これはあまり速くならなかったが気持ち少し速くなっている!

ともあれ最初と比べ3.16倍になった。高速化は大成功!!

実行メモリバンド幅測定

このベンチマークをとったのはRTX2080Ti。このGPUのメモリのバンド幅は616GB/s。
B2B4B8B16C2C4カーネルで要素数最大の134217728でやった場合のメモリアクセス量は82回*2回×
134217728*8byte=176.1GB
これを1秒あたりに直して319GB/sの実行メモリバンド幅とわかる。

約51.8%の性能が出せている。個人的には十分だと思う。

行列転置で全範囲Sared memory化

ここからは自分オリジナルの案になるが、結論から言うと速くならなかった。
発想としてはFFTのブロック化の案(参考pdf)を使うことで、メモリアクセスの局所化をはかる。

sikou7

これを

sikou8
こうする。
なお、書いてある数字はメモリの中身ではなくメモリアドレスだと思ってほしい。
このアルゴリズム図はメモリレイアウトが変わっただけで、今までと全く同じことをしている。
違うのは各Stepを計算するとき、メモリアクセスが4領域内にまとめられているということ。
画像でいうと最初の処理で(0,4,8,12)、(1,5,9,13)、(2,6,10,14)、(3,7,11,15)にまとめてアクセスしており、後半処理では(0,1,2,3)、(4,5,6,7)、(8,9,10,11)、(12,13,14,15)にまとめてアクセスしている。この「まとまっている」がコアレスアクセスに相当する。
最初の処理では(0,4,8,12)と飛び飛びだからコアレスアクセスできてないと思うかもしれないが、事前に行列転置をしておくことで可能になる。
GPUを使った行列転置は今後記事にする予定だが、最初0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15だったのを0,4,8,12,1,5,9,13,2,6,10,14,3,7,11,15に効率的に配置できる。同じように
0,4,8,12,1,5,9,13,2,6,10,14,3,7,11,15から0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15へも効率的に配置できる(これは画像の通り)。
この行列転置はShared memoryを使うことで読み込み書き込みともにコアレスアクセスにできる。

これで何が嬉しいかと言うと、今までShared memoryに入りきらなかったStepの計算もすべて
Shared memoryを使いつつコアレスアクセスにできることである。
イメージとしてはこうなる。

shared2

Shared memoryを使った計算では8Stepまとめて計算できるので、今までの4Step+4Stepの処理が1つにまとまる。しかしその分行列転置の処理が入り込むので結局メモリアクセス量はたいして変わらずかむしろ増えたようで、Totalの速度ではむしろ1.2倍くらい遅くなってしまった。


なお
Shared memoryの容量的にまだ余裕があったので10Stepまとめて計算するコードに変えてみたが意外と速くならなかった。なかなか高速化は難しい・・・。

Shared memoryだけ使わないと

環境によってはShared memoryが使えないという状況もあると思い追加で速度測定

B2B4B8B16というやつ

要素数/カーネル名B2B2C2B2B4B8B16C2C4B2B4B8B16
2560000
5120000
10240000
20480000
40960000
81920000
163840000
327680000
655361000
1310721100
2621442111
5242884321
104857610854
2097152191179
419430439241418
838860883512936
167772161791096073
33554432382236127154
67108864818508264322
13421772817441095552670

1.2倍ほど遅くなったがこれでも十分すぎるほど速いのではないか。

参考文献



【Qiita】Compute Shaderでのソートアルゴリズムの処理時間比較

【Qiita】ComputeShaderで近傍探索の高速化