2015/12/11

WebGLで暴力的な並列ソートに挑戦する

日照時間足りてますか？やまだです。
KLab Advent Calender 11日目の記事です。

最近WebGLで実装できるちょっと強そうなソートアルゴリズムを知ったので書いてみました。
WebGLよくわかんないやーって人も雰囲気だけでも伝われば幸いです。

強さの秘訣

WebGLは3Dを描くAPIとして著名ですが、2Dはもちろん、工夫次第で汎用計算（GPGPU）にも応用できます。
WebGLでの演算は通常のJavaScript演算とは違い、GPUを使った並列な浮動小数点演算を可能とします。

そして、ソートアルゴリズムの中には計算量こそ地味なものの、並列処理が可能なものがあります。
そのうちの一つがバイトニックソートです。

GPUを使い、物量で計算する暴力的な並列ソートをやっていきましょう。

バイトニックソートとは

前述のとおり、並列処理が可能なソートアルゴリズムです。

配列を小さい領域に分割し、その小さい領域内部で昇順（赤い部分）と降順（青い部分）に分けて整列させます。
揃え終わると次は隣接する領域とマージして、これも昇順と降順に揃えます。
これを繰り返すことにより最終的にソートが完成します。

直列時の最良計算量及び最悪計算量は共に onlogn2 ですが、並列にするとn要素同時に比較することによりnが削除され ologn2 となります。

古典的なGPGPU

続いてWebGLでこのソートをどうやって実現するかです。
WebGLでは頂点位置の計算を行うバーテックスシェーダとピクセル単位の計算を行うフラグメントシェーダの2種類をGLSLというシェーダ用言語で記述します。

バーテックスシェーダは主に頂点位置座標をカメラから見た位置に変換し、フラグメントシェーダはピクセルに色をつけます。

通常、WebGLはCanvasに対して画像をレンダリングします。
ですが、Canvasに対してではなく、テクスチャに対してレンダリングもできます。
これをオフスクリーンレンダリングといいます。
こちらはwgldにとてもわかりやすく解説されています。

古典的なGPGPUのポイントはフラグメントシェーダに対する書き込みです。
テクスチャは通常1画素あたりRGBAの4要素が格納されていますが、必ずしもこれを色情報として扱う必要はありません。
そのためRGBAをXYZWと読み替えれば座標情報のソートができます。

ただし、テクスチャの座標は二次元座標ですので左下から右上にむかってインデックスをふり、一次元にする必要があります。
これにより1画素を配列の1要素とみなしてソートできるようになりました。

texture_index

ソートのために乱数でテクスチャを埋めて見ると次のようになります。

noise

このノイズ画像をシェーダに渡してからソート開始なのですが、WebGLは三角形や点をレンダリングする機能しか持ちません。
そこでCanvasいっぱいに四角形ポリゴンをレンダリングし、ソートする要素の数だけフラグメントシェーダが実行されるようにします。
そのため、バーテックスシェーダはとてもシンプルな実装になる一方で、フラグメントシェーダが複雑になっていきます。

フラグメントシェーダ

t-pot『Bitonic sort』にDirectX用のシェーダ言語であるHLSLのシェーダがあるのでそれを参考にWebGL版を実装しました。
参考にというか変数名に至るまでほぼそのままです、助かりました~~（実はパラメータの求め方よくわかってない）~~。

uniform vec2 resolution;    //画像解像度
uniform vec2 halfDelta;     //テクセル半分

uniform sampler2D texture;  //ソート対象の画像

uniform float stepno;
uniform float offset;
uniform float stage;

void main() {
  vec2 elem2d = floor(gl_FragCoord.xy);
  float elem1d = elem2d.x + elem2d.y * resolution.x;

  float csign = (mod(elem1d, stage) < offset) ? 1.0 : -1.0;

  float cdir  = (mod(floor(elem1d / stepno), 2.0) <= 0.5) ? 1.0 : -1.0;

  vec4 val0 = texture2D(texture, elem2d / resolution + halfDelta);

  float adr1d = csign * offset + elem1d;
  vec2 adr2d = vec2(mod(adr1d, resolution.x), floor(adr1d / resolution.x));

  vec4 val1 = texture2D(texture, adr2d / resolution + halfDelta);

  vec4 cmin = (val0.y < val1.y) ? val0: val1;
  vec4 cmax = (val0.y < val1.y) ? val1: val0;

  vec4 dst = (csign == cdir) ? cmin : cmax;

  gl_FragColor = dst;
}