WebGL2でGPGPUバイトニックソート

前回の記事で、JavaScriptでバイトニックソートを実装してみました。
JavaScriptでバイトニックソート - Qiita

アルゴリズムを理解できたので、バイトニックソートの本領が発揮できるWebGLでGPGPUバイトニックソートを実装してみました。

WebGLではサイズが2のn乗のテクスチャを用意して、そのテクスチャにオフスクリーンレンダリングで何度も途中状態を書き込みながら値をソートしていきます。テクスチャは2次元ですが、内部ではテクセルの値をもとに１次元に変換してソートを行っています。

以下はバイトニックソートの進行していく過程のテクスチャの様子となっています。最終的に左下が小さい値(黒)、右上が大きい値(白)になっていきます。

長いのでソースコード全文は掲載しません。この記事では要点だけを解説していきます。
Githubにソースコードを置いておいたので、全文はそちらで確認してください。実際に動いているものも見ることができます。
aadebdeb/WebGL2_BitonicSort: GPGPU Bitonic Sort with WebGL2

まず、テクスチャの初期化を行います。今回はビジュアライズしたときにわかりやすいように、以下のシェーダーで0から1のランダムな値を各テクセルに書き込んでおきます。

#version 300 es

precision highp float;

out float o_value;

uniform vec2 u_randomSeed;

float random(vec2 x){
  return fract(sin(dot(x,vec2(12.9898, 78.233))) * 43758.5453);
}

void main(void) {
  o_value = random(gl_FragCoord.xy * 0.01 + u_randomSeed);
}

JavaScript版バイトニックソートではbitonicsort関数で２重ループを回し、kernel関数でスワップ処理を行っていました。bitnoicsort関数の２重ループはWebGLb版でも以下のように同様に行いますが、kernel関数で行っている処理はフラグメントシェーダーを利用して行います。

for (let i = 0; i < sizeN; i++) {
  for (let j = 0; j <= i; j++) {
    executeKernel(writableFramebuffer.framebuffer, readableFramebuffer.texture, textureSize, i, j);
    swapFramebuffer();
  }
}

executeKernel関数では以下のフラグメントシェーダーを利用して、スワップ処理を行っています。やっていることはJavaScript版バイトニックソートの最後に掲載したコードと同じです。u_blockStepとu_subBlockStepがJavascript版のkernel関数の引数p,とqにそれぞれ対応しています。２次元の座標から１次元のインデックスへの変換とその逆方向への変換を行うためにconvertCoordToIndexとconvertIndexToCoordを定義しています。

#version 300 es

precision highp float;

out float o_value;

uniform sampler2D u_valueTexture;
uniform uint u_size;
uniform uint u_blockStep;
uniform uint u_subBlockStep;

uint convertCoordToIndex(uvec2 coord) {
  return coord.x + coord.y * u_size;
}

uvec2 convertIndexToCoord(uint index) {
  return uvec2(index % u_size, index / u_size);
}

void main(void) {
  uint index = convertCoordToIndex(uvec2(gl_FragCoord.xy));
  uint d = 1u << (u_blockStep - u_subBlockStep);

  bool up = ((index >> u_blockStep) & 2u) == 0u;

  uint targetIndex;
  if ((index & d) == 0u) {
    targetIndex = index | d;
  } else {
    targetIndex = index & ~d;
    up = !up;
  }

  float a = texelFetch(u_valueTexture, ivec2(gl_FragCoord.xy), 0).x;
  float b = texelFetch(u_valueTexture, ivec2(convertIndexToCoord(targetIndex)), 0).x;
  if ((a > b) == up) {
    o_value = b; // swap
  } else {
    o_value = a; // no_swap
  }
}

追記

ここまでに紹介した実装では1つの値しか持たない配列をソートする場合を例にしました。このときには値が同じときにはスワップしてもしなくてもどちらでも問題ありませんでした。しかし、vec2など2つ以上の値を持つ配列をそのうちの1要素でソートする場合には、以下のように同じ値のときに必ずスワップするかスワップしないかを決めておかないといけません。

  vec2 a = texelFetch(u_valueTexture, ivec2(gl_FragCoord.xy), 0).xy;
  vec2 b = texelFetch(u_valueTexture, ivec2(convertIndexToCoord(targetIndex)), 0).xy;
  if ((a.x == b.x || (a.x > b.x) == up) { // x要素でソート, 値が同じときは必ずswapする
    o_value = b; // swap
  } else {
    o_value = a; // no_swap
  }